mrq/stable-diffusion-webui
1
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

implement support for the following command line arguments: --gfpgan-models-path, --codeformer-models-path, --esrgan-models-path, --realesrgan-models-path, --scunet-models-path, --swinir-models-path, --ldsr-models-path

Browse Source
This commit is contained in:
Gabriel Hautclocq 2023-02-02 10:59:02 +01:00
parent 226d840e84
commit 9aa1d2db1a
3 changed files with 22 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
modules/codeformer_model.py
View File

@ -9,12 +9,15 @@ import modules.face_restoration
import modules.shared
from modules import shared, devices, modelloader
from modules.paths import models_path
from modules.shared import cmd_opts
# codeformer people made a choice to include modified basicsr library to their project which makes
# it utterly impossible to use it alongside with other libraries that also use basicsr, like GFPGAN.
# I am making a choice to include some files from codeformer to work around this issue.
model_dir = "Codeformer"
model_path = os.path.join(models_path, model_dir)
if cmd_opts.codeformer_models_path is not None and os.path.isdir(cmd_opts.codeformer_models_path):
model_path = cmd_opts.codeformer_models_path
model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
have_codeformer = False

3
modules/gfpgan_model.py
View File

@ -7,10 +7,13 @@ import gfpgan
import modules.face_restoration
from modules import paths, shared, devices, modelloader
from modules.shared import cmd_opts
model_dir = "GFPGAN"
user_path = None
model_path = os.path.join(paths.models_path, model_dir)
if cmd_opts.gfpgan_models_path is not None and os.path.isdir(cmd_opts.gfpgan_models_path):
model_path = cmd_opts.gfpgan_models_path
model_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
have_gfpgan = False
loaded_gfpgan_model = None

16
modules/upscaler.py
View File

@ -39,6 +39,22 @@ class Upscaler:
if self.model_path is None and self.name:
self.model_path = os.path.join(shared.models_path, self.name)
match self.name:
case "ESRGAN":
if modules.shared.cmd_opts.esrgan_models_path is not None and os.path.isdir(modules.shared.cmd_opts.esrgan_models_path):
self.model_path = modules.shared.cmd_opts.esrgan_models_path
case "LDSR":
if modules.shared.cmd_opts.ldsr_models_path is not None and os.path.isdir(modules.shared.cmd_opts.ldsr_models_path):
self.model_path = modules.shared.cmd_opts.ldsr_models_path
case "RealESRGAN":
if modules.shared.cmd_opts.realesrgan_models_path is not None and os.path.isdir(modules.shared.cmd_opts.realesrgan_models_path):
self.model_path = modules.shared.cmd_opts.realesrgan_models_path
case "ScuNET":
if modules.shared.cmd_opts.scunet_models_path is not None and os.path.isdir(modules.shared.cmd_opts.scunet_models_path):
self.model_path = modules.shared.cmd_opts.scunet_models_path
case "SwinIR":
if modules.shared.cmd_opts.swinir_models_path is not None and os.path.isdir(modules.shared.cmd_opts.swinir_models_path):
self.model_path = modules.shared.cmd_opts.swinir_models_path
if self.model_path and create_dirs:
os.makedirs(self.model_path, exist_ok=True)