From 6a0a61411cbf34bfe8601cf57ef1f62d7c733012 Mon Sep 17 00:00:00 2001
From: Rafael Irgolic <hello@irgolic.com>
Date: Thu, 6 Oct 2022 15:36:19 +0200
Subject: [PATCH] modeldownload: Use webui_streamlit.yaml to update models

---
 scripts/modeldownload.py | 74 +++++++++++++++++++---------------------
 scripts/sd_utils.py      |  4 ++-
 2 files changed, 39 insertions(+), 39 deletions(-)

diff --git a/scripts/modeldownload.py b/scripts/modeldownload.py
index acaa1d319..ff6261094 100644
--- a/scripts/modeldownload.py
+++ b/scripts/modeldownload.py
@@ -14,41 +14,39 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>. 
 import os
-import os.path as op
-
-def updateModels():
-    if op.exists('models/ldm/stable-diffusion-v1/model.ckpt'):
-        pass
-    else:
-        os.system('wget https://www.googleapis.com/storage/v1/b/aai-blog-files/o/sd-v1-4.ckpt?alt=media -o models/ldm/stable-diffusion-v1/model.ckpt')
-        # os.rename('models/ldm/stable-diffusion-v1/sd-v1-4.ckpt?alt=media','models/ldm/stable-diffusion-v1/model.ckpt')
-        # os.system('wget https://cdn-lfs.huggingface.co/repos/ab/41/ab41ccb635cd5bd124c8eac1b5796b4f64049c9453c4e50d51819468ca69ceb8/14749efc0ae8ef0329391ad4436feb781b402f4fece4883c7ad8d10556d8a36a?response-content-disposition=attachment%3B%20filename%3D%22modelfull.ckpt%22 -o models/ldm/stable-diffusion-v1/model.ckpt')
-        # os.rename('models/ldm/stable-diffusion-v1/modelfull.ckpt','models/ldm/stable-diffusion-v1/model.ckpt')
-    
-    if op.exists('models/realesrgan/RealESRGAN_x4plus.pth') and op.exists('models/realesrgan/RealESRGAN_x4plus_anime_6B.pth'):
-        pass
-    else:
-        os.system('wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P models/realesrgan')
-        os.system('wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth -P models/realesrgan')
-    
-    if op.exists('models/gfpgan/GFPGANv1.3.pth'):
-        pass 
-    else:
-        os.system('wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth -P models/gfpgan')
-    
-    if op.exists('models/ldsr'):
-        pass
-    else:
-        os.system('git clone https://github.com/devilismyfriend/latent-diffusion.git')
-        os.system('mv latent-diffusion models/ldsr')
-    
-    if op.exists('models/ldsr/model.ckpt'):
-        pass
-    else:
-        os.mkdir('models/ldsr/experiments')
-        os.mkdir('models/ldsr')
-        os.system('wget https://heibox.uni-heidelberg.de/f/31a76b13ea27482981b4/?dl=1 -o models/ldsr/project.yaml')
-        # os.rename('models/ldsr/index.html?dl=1', 'models/ldsr/project.yaml')
-        os.system('wget https://heibox.uni-heidelberg.de/f/578df07c8fc04ffbadf3/?dl=1 -o models/ldsr/model.ckpt')
-        # os.rename('models/ldsr/index.html?dl=1', 'models/ldsr/model.ckpt')
-    
\ No newline at end of file
+from omegaconf import OmegaConf
+import urllib.request
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def download_file(download_link, folder, file_name):
+    model_path = os.path.join(folder, file_name)
+    if os.path.exists(model_path):
+        logger.info(file_name + " already exists")
+        return
+    os.makedirs(folder, exist_ok=True)
+    logger.info('Downloading ' + file_name + '...')
+    try:
+        urllib.request.urlretrieve(download_link, model_path)
+        logger.info("Downloaded " + file_name)
+    except:
+        logger.error("Failed to download " + file_name + " from " + download_link)
+        raise
+
+
+def update_models(models=None):
+    if models is None:
+        models = OmegaConf.load("configs/webui/webui_streamlit.yaml").model_manager.models
+
+    for model_name, model in models.items():
+        model_folder = model.save_location
+        logger.info("Preparing model: " + model_name)
+        for file_ in model.files.values():
+            download_file(file_.download_link, model_folder, file_.file_name)
+        logger.info("Model " + model_name + " is ready.")
+
+
+if __name__ == "__main__":
+    update_models()
diff --git a/scripts/sd_utils.py b/scripts/sd_utils.py
index 2b12425eb..abb6185ca 100644
--- a/scripts/sd_utils.py
+++ b/scripts/sd_utils.py
@@ -115,7 +115,9 @@
 if st.session_state["defaults"].daisi_app.running_on_daisi_io:
     if os.path.exists("scripts/modeldownload.py"):
         import modeldownload
-        modeldownload.updateModels()
+        modeldownload.update_models(
+            models=st.session_state["defaults"].model_manager.models,
+        )
 
 #
 #app = st.HydraApp(title='Stable Diffusion WebUI', favicon="", sidebar_state="expanded",