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
This commit is contained in:
parent
226d840e84
commit
9aa1d2db1a
|
@ -9,12 +9,15 @@ import modules.face_restoration
|
||||||
import modules.shared
|
import modules.shared
|
||||||
from modules import shared, devices, modelloader
|
from modules import shared, devices, modelloader
|
||||||
from modules.paths import models_path
|
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
|
# 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.
|
# 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.
|
# I am making a choice to include some files from codeformer to work around this issue.
|
||||||
model_dir = "Codeformer"
|
model_dir = "Codeformer"
|
||||||
model_path = os.path.join(models_path, model_dir)
|
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'
|
model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
|
||||||
|
|
||||||
have_codeformer = False
|
have_codeformer = False
|
||||||
|
|
|
@ -7,10 +7,13 @@ import gfpgan
|
||||||
|
|
||||||
import modules.face_restoration
|
import modules.face_restoration
|
||||||
from modules import paths, shared, devices, modelloader
|
from modules import paths, shared, devices, modelloader
|
||||||
|
from modules.shared import cmd_opts
|
||||||
|
|
||||||
model_dir = "GFPGAN"
|
model_dir = "GFPGAN"
|
||||||
user_path = None
|
user_path = None
|
||||||
model_path = os.path.join(paths.models_path, model_dir)
|
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"
|
model_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
|
||||||
have_gfpgan = False
|
have_gfpgan = False
|
||||||
loaded_gfpgan_model = None
|
loaded_gfpgan_model = None
|
||||||
|
|
|
@ -39,6 +39,22 @@ class Upscaler:
|
||||||
|
|
||||||
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)
|
||||||
|
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:
|
if self.model_path and create_dirs:
|
||||||
os.makedirs(self.model_path, exist_ok=True)
|
os.makedirs(self.model_path, exist_ok=True)
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue
Block a user