diff --git a/.gitignore b/.gitignore index 9d78853af..57fa25062 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,7 @@ __pycache__ /venv /tmp /model.ckpt +/models/**/* /models/*.ckpt /GFPGANv1.3.pth /gfpgan/weights/*.pth diff --git a/modules/codeformer_model.py b/modules/codeformer_model.py index dc0a5eee0..efd881eb6 100644 --- a/modules/codeformer_model.py +++ b/modules/codeformer_model.py @@ -5,14 +5,13 @@ import traceback import cv2 import torch +import modules.face_restoration +import modules.shared from modules import shared, devices, modelloader from modules.paths import script_path, models_path -import modules.shared -import modules.face_restoration -from importlib import reload -# codeformer people made a choice to include modified basicsr library to their project, which makes -# it utterly impossible to use it alongside other libraries that also use basicsr, like GFPGAN. +# 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) @@ -31,11 +30,6 @@ def setup_model(dirname): if path is None: return - - # both GFPGAN and CodeFormer use bascisr, one has it installed from pip the other uses its own - #stored_sys_path = sys.path - #sys.path = [path] + sys.path - try: from torchvision.transforms.functional import normalize from modules.codeformer.codeformer_arch import CodeFormer @@ -67,7 +61,6 @@ def setup_model(dirname): print("Unable to load codeformer model.") return None, None net = net_class(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256']).to(devices.device_codeformer) - ckpt_path = load_file_from_url(url=pretrain_model_url, model_dir=os.path.join(path, 'weights/CodeFormer'), progress=True) checkpoint = torch.load(ckpt_path)['params_ema'] net.load_state_dict(checkpoint) net.eval() diff --git a/modules/esrgan_model.py b/modules/esrgan_model.py index dd0ee6291..5e10c49c3 100644 --- a/modules/esrgan_model.py +++ b/modules/esrgan_model.py @@ -18,7 +18,7 @@ from modules.shared import opts model_dir = "ESRGAN" model_path = os.path.join(models_path, model_dir) model_url = "https://drive.google.com/u/0/uc?id=1TPrz5QKd8DHHt1k8SRtm6tMiPjz_Qene&export=download" -model_name = "ESRGAN_x4.pth" +model_name = "ESRGAN_x4" def load_model(path: str, name: str): @@ -27,7 +27,7 @@ def load_model(path: str, name: str): global model_dir global model_name if "http" in path: - filename = load_file_from_url(url=model_url, model_dir=model_path, file_name=model_name, progress=True) + filename = load_file_from_url(url=model_url, model_dir=model_path, file_name="%s.pth" % model_name, progress=True) else: filename = path if not os.path.exists(filename) or filename is None: diff --git a/modules/ldsr_model.py b/modules/ldsr_model.py index e6e7ff74c..4f9b16579 100644 --- a/modules/ldsr_model.py +++ b/modules/ldsr_model.py @@ -19,7 +19,7 @@ have_ldsr = False LDSR_obj = None -class UpscalerLDSR(modules.images.Upscaler): +class UpscalerLDSR(images.Upscaler): def __init__(self, steps): self.steps = steps self.name = "LDSR" diff --git a/modules/modelloader.py b/modules/modelloader.py index d59fbe05c..9520a6819 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -1,8 +1,11 @@ import os +import shutil from urllib.parse import urlparse from basicsr.utils.download_util import load_file_from_url +from modules.paths import script_path, models_path + def load_models(model_path: str, model_url: str = None, command_path: str = None, dl_name: str = None, existing=None, ext_filter=None) -> list: @@ -63,3 +66,38 @@ def friendly_name(file: str): model_name, extension = os.path.splitext(file) model_name = model_name.replace("_", " ").title() return model_name + + +def cleanup_models(): + root_path = script_path + src_path = os.path.join(root_path, "ESRGAN") + dest_path = os.path.join(models_path, "ESRGAN") + move_files(src_path, dest_path) + src_path = os.path.join(root_path, "gfpgan") + dest_path = os.path.join(models_path, "GFPGAN") + move_files(src_path, dest_path) + src_path = os.path.join(root_path, "SwinIR") + dest_path = os.path.join(models_path, "SwinIR") + move_files(src_path, dest_path) + src_path = os.path.join(root_path, "repositories/latent-diffusion/experiments/pretrained_models/") + dest_path = os.path.join(models_path, "LDSR") + move_files(src_path, dest_path) + + +def move_files(src_path: str, dest_path: str): + try: + if not os.path.exists(dest_path): + os.makedirs(dest_path) + if os.path.exists(src_path): + for file in os.listdir(src_path): + if os.path.isfile(file): + fullpath = os.path.join(src_path, file) + print("Moving file: %s to %s" % (fullpath, dest_path)) + try: + shutil.move(fullpath, dest_path) + except: + pass + print("Removing folder: %s" % src_path) + shutil.rmtree(src_path, True) + except: + pass \ No newline at end of file diff --git a/webui.py b/webui.py index 7e0b32962..e71a217c6 100644 --- a/webui.py +++ b/webui.py @@ -18,9 +18,11 @@ import modules.shared as shared import modules.swinir_model as swinir import modules.txt2img import modules.ui +from modules import modelloader from modules.paths import script_path from modules.shared import cmd_opts +modelloader.cleanup_models() codeformer.setup_model(cmd_opts.codeformer_models_path) gfpgan.setup_model(cmd_opts.gfpgan_models_path) shared.face_restorers.append(modules.face_restoration.FaceRestoration())