openvinotoolkit/stable-diffusion-webui
1
1
Fork 0
mirror of https://github.com/openvinotoolkit/stable-diffusion-webui.git synced 2024-12-14 22:53:25 +03:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
AUTOMATIC 2023-05-19 09:09:00 +03:00
parent 379fd6204d
commit df6fffb054
7 changed files with 12 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
extensions-builtin/LDSR/scripts/ldsr_model.py
View File

@ -45,9 +45,9 @@ class UpscalerLDSR(Upscaler):
if local_safetensors_path is not None and os.path.exists(local_safetensors_path): if local_safetensors_path is not None and os.path.exists(local_safetensors_path):
model = local_safetensors_path model = local_safetensors_path
else: 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: try:
return LDSR(model, yaml) return LDSR(model, yaml)

3
extensions-builtin/ScuNET/scripts/scunet_model.py
View File

@ -121,8 +121,7 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
def load_model(self, path: str): def load_model(self, path: str):
device = devices.get_device_for('scunet') device = devices.get_device_for('scunet')
if "http" in path: if "http" in path:
filename = load_file_from_url(url=self.model_url, model_dir=self.model_path, file_name="%s.pth" % self.name, filename = load_file_from_url(url=self.model_url, model_dir=self.model_download_path, file_name="%s.pth" % self.name, progress=True)
progress=True)
else: else:
filename = path filename = path
if not os.path.exists(os.path.join(self.model_path, filename)) or filename is None: if not os.path.exists(os.path.join(self.model_path, filename)) or filename is None:

2
extensions-builtin/SwinIR/scripts/swinir_model.py
View File

@ -51,7 +51,7 @@ class UpscalerSwinIR(Upscaler):
def load_model(self, path, scale=4): def load_model(self, path, scale=4):
if "http" in path: if "http" in path:
dl_name = "%s%s" % (self.model_name.replace(" ", "_"), ".pth") 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: else:
filename = path filename = path
if filename is None or not os.path.exists(filename): if filename is None or not os.path.exists(filename):

2
modules/esrgan_model.py
View File

@ -154,7 +154,7 @@ class UpscalerESRGAN(Upscaler):
if "http" in path: if "http" in path:
filename = load_file_from_url( filename = load_file_from_url(
url=self.model_url, url=self.model_url,
model_dir=self.model_path, model_dir=self.model_download_path,
file_name=f"{self.model_name}.pth", file_name=f"{self.model_name}.pth",
progress=True, progress=True,
) )

7
modules/modelloader.py
View File

@ -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 model_url is not None and len(output) == 0:
if download_name is not None: if download_name is not None:
from basicsr.utils.download_util import load_file_from_url 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) output.append(dl)
else: else:
output.append(model_url) output.append(model_url)
@ -144,7 +144,10 @@ def load_upscalers():
for cls in reversed(used_classes.values()): for cls in reversed(used_classes.values()):
name = cls.__name__ name = cls.__name__
cmd_name = f"{name.lower().replace('upscaler', '')}_models_path" 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 datas += scaler.scalers
shared.sd_upscalers = sorted( shared.sd_upscalers = sorted(

2
modules/realesrgan_model.py
View File

@ -73,7 +73,7 @@ class UpscalerRealESRGAN(Upscaler):
return None return None
if info.local_data_path.startswith("http"): 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 return info
except Exception as e: except Exception as e:

1
modules/upscaler.py
View File

@ -34,6 +34,7 @@ class Upscaler:
self.half = not modules.shared.cmd_opts.no_half self.half = not modules.shared.cmd_opts.no_half
self.pre_pad = 0 self.pre_pad = 0
self.mod_scale = None self.mod_scale = None
self.model_download_path = None
if self.model_path is None and self.name: if self.model_path is None and self.name:
self.model_path = os.path.join(shared.models_path, self.name) self.model_path = os.path.join(shared.models_path, self.name)