change upscalers to download models into user-specified directory (from commandline args) rather than the default models/<...>

7 files changed, 12 insertions, 9 deletions

diff --git a/extensions-builtin/LDSR/scripts/ldsr_model.py b/extensions-builtin/LDSR/scripts/ldsr_model.py
index fbbe9005..c4da79f3 100644
--- a/extensions-builtin/LDSR/scripts/ldsr_model.py
+++ b/extensions-builtin/LDSR/scripts/ldsr_model.py
@@ -45,9 +45,9 @@ class UpscalerLDSR(Upscaler):
         if local_safetensors_path is not None and os.path.exists(local_safetensors_path):
             model = local_safetensors_path
         else:
-            model = local_ckpt_path if local_ckpt_path is not None else load_file_from_url(url=self.model_url, model_dir=self.model_path, file_name="model.ckpt", progress=True)
+            model = local_ckpt_path if local_ckpt_path is not None else load_file_from_url(url=self.model_url, model_dir=self.model_download_path, file_name="model.ckpt", progress=True)
 
-        yaml = local_yaml_path if local_yaml_path is not None else load_file_from_url(url=self.yaml_url, model_dir=self.model_path, file_name="project.yaml", progress=True)
+        yaml = local_yaml_path if local_yaml_path is not None else load_file_from_url(url=self.yaml_url, model_dir=self.model_download_path, file_name="project.yaml", progress=True)
 
         try:
             return LDSR(model, yaml)
diff --git a/extensions-builtin/ScuNET/scripts/scunet_model.py b/extensions-builtin/ScuNET/scripts/scunet_model.py
index cc2cbc6a..45d9297b 100644
--- a/extensions-builtin/ScuNET/scripts/scunet_model.py
+++ b/extensions-builtin/ScuNET/scripts/scunet_model.py
@@ -121,8 +121,7 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
     def load_model(self, path: str):
         device = devices.get_device_for('scunet')
         if "http" in path:
-            filename = load_file_from_url(url=self.model_url, model_dir=self.model_path, file_name="%s.pth" % self.name,
-                                          progress=True)
+            filename = load_file_from_url(url=self.model_url, model_dir=self.model_download_path, file_name="%s.pth" % self.name, progress=True)
         else:
             filename = path
         if not os.path.exists(os.path.join(self.model_path, filename)) or filename is None:
diff --git a/extensions-builtin/SwinIR/scripts/swinir_model.py b/extensions-builtin/SwinIR/scripts/swinir_model.py
index 0ba50487..1c7bf325 100644
--- a/extensions-builtin/SwinIR/scripts/swinir_model.py
+++ b/extensions-builtin/SwinIR/scripts/swinir_model.py
@@ -51,7 +51,7 @@ class UpscalerSwinIR(Upscaler):
     def load_model(self, path, scale=4):
         if "http" in path:
             dl_name = "%s%s" % (self.model_name.replace(" ", "_"), ".pth")
-            filename = load_file_from_url(url=path, model_dir=self.model_path, file_name=dl_name, progress=True)
+            filename = load_file_from_url(url=path, model_dir=self.model_download_path, file_name=dl_name, progress=True)
         else:
             filename = path
         if filename is None or not os.path.exists(filename):
diff --git a/modules/esrgan_model.py b/modules/esrgan_model.py
index a009eb42..2fced999 100644
--- a/modules/esrgan_model.py
+++ b/modules/esrgan_model.py
@@ -154,7 +154,7 @@ class UpscalerESRGAN(Upscaler):
         if "http" in path:

             filename = load_file_from_url(

                 url=self.model_url,

-                model_dir=self.model_path,

+                model_dir=self.model_download_path,

                 file_name=f"{self.model_name}.pth",

                 progress=True,

             )

diff --git a/modules/modelloader.py b/modules/modelloader.py
index 2a479bcb..be23071a 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -47,7 +47,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
         if model_url is not None and len(output) == 0:
             if download_name is not None:
                 from basicsr.utils.download_util import load_file_from_url
-                dl = load_file_from_url(model_url, model_path, True, download_name)
+                dl = load_file_from_url(model_url, places[0], True, download_name)
                 output.append(dl)
             else:
                 output.append(model_url)
@@ -144,7 +144,10 @@ def load_upscalers():
     for cls in reversed(used_classes.values()):
         name = cls.__name__
         cmd_name = f"{name.lower().replace('upscaler', '')}_models_path"
-        scaler = cls(commandline_options.get(cmd_name, None))
+        commandline_model_path = commandline_options.get(cmd_name, None)
+        scaler = cls(commandline_model_path)
+        scaler.user_path = commandline_model_path
+        scaler.model_download_path = commandline_model_path or scaler.model_path
         datas += scaler.scalers
 
     shared.sd_upscalers = sorted(
diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py
index c24d8dbb..99983678 100644
--- a/modules/realesrgan_model.py
+++ b/modules/realesrgan_model.py
@@ -73,7 +73,7 @@ class UpscalerRealESRGAN(Upscaler):
                 return None

 

             if info.local_data_path.startswith("http"):

-                info.local_data_path = load_file_from_url(url=info.data_path, model_dir=self.model_path, progress=True)

+                info.local_data_path = load_file_from_url(url=info.data_path, model_dir=self.model_download_path, progress=True)

 

             return info

         except Exception as e:

diff --git a/modules/upscaler.py b/modules/upscaler.py
index 8acb6e96..7b1046d6 100644
--- a/modules/upscaler.py
+++ b/modules/upscaler.py
@@ -34,6 +34,7 @@ class Upscaler:
         self.half = not modules.shared.cmd_opts.no_half
         self.pre_pad = 0
         self.mod_scale = None
+        self.model_download_path = None
 
         if self.model_path is None and self.name:
             self.model_path = os.path.join(shared.models_path, self.name)