diff --git a/modules/bsrgan_model.py b/modules/bsrgan_model.py index 47346f318..e62c66577 100644 --- a/modules/bsrgan_model.py +++ b/modules/bsrgan_model.py @@ -69,7 +69,7 @@ class UpscalerBSRGAN(modules.upscaler.Upscaler): if not os.path.exists(filename) or filename is None: print(f"BSRGAN: Unable to load model from {filename}", file=sys.stderr) return None - model = RRDBNet(in_nc=3, out_nc=3, nf=64, nb=23, gc=32, sf=2) # define network + model = RRDBNet(in_nc=3, out_nc=3, nf=64, nb=23, gc=32, sf=4) # define network model.load_state_dict(torch.load(filename), strict=True) model.eval() for k, v in model.named_parameters(): diff --git a/modules/ldsr_model.py b/modules/ldsr_model.py index 877e7e73e..1c1070fc6 100644 --- a/modules/ldsr_model.py +++ b/modules/ldsr_model.py @@ -22,8 +22,20 @@ class UpscalerLDSR(Upscaler): self.scalers = [scaler_data] def load_model(self, path: str): + # Remove incorrect project.yaml file if too big + yaml_path = os.path.join(self.model_path, "project.yaml") + old_model_path = os.path.join(self.model_path, "model.pth") + new_model_path = os.path.join(self.model_path, "model.ckpt") + if os.path.exists(yaml_path): + statinfo = os.stat(yaml_path) + if statinfo.st_size >= 10485760: + print("Removing invalid LDSR YAML file.") + os.remove(yaml_path) + if os.path.exists(old_model_path): + print("Renaming model from model.pth to model.ckpt") + os.rename(old_model_path, new_model_path) model = load_file_from_url(url=self.model_url, model_dir=self.model_path, - file_name="model.pth", progress=True) + file_name="model.ckpt", progress=True) yaml = load_file_from_url(url=self.yaml_url, model_dir=self.model_path, file_name="project.yaml", progress=True) @@ -41,5 +53,4 @@ class UpscalerLDSR(Upscaler): print("NO LDSR!") return img ddim_steps = shared.opts.ldsr_steps - pre_scale = shared.opts.ldsr_pre_down return ldsr.super_resolution(img, ddim_steps, self.scale) diff --git a/modules/ldsr_model_arch.py b/modules/ldsr_model_arch.py index 7faac6e18..14db50766 100644 --- a/modules/ldsr_model_arch.py +++ b/modules/ldsr_model_arch.py @@ -98,9 +98,7 @@ class LDSR: im_og = image width_og, height_og = im_og.size # If we can adjust the max upscale size, then the 4 below should be our variable - print("Foo") down_sample_rate = target_scale / 4 - print(f"Downsample rate is {down_sample_rate}") wd = width_og * down_sample_rate hd = height_og * down_sample_rate width_downsampled_pre = int(wd) @@ -111,7 +109,7 @@ class LDSR: f'Downsampling from [{width_og}, {height_og}] to [{width_downsampled_pre}, {height_downsampled_pre}]') im_og = im_og.resize((width_downsampled_pre, height_downsampled_pre), Image.LANCZOS) else: - print(f"Down sample rate is 1 from {target_scale} / 4") + print(f"Down sample rate is 1 from {target_scale} / 4 (Not downsampling)") logs = self.run(model["model"], im_og, diffusion_steps, eta) sample = logs["sample"] diff --git a/modules/modelloader.py b/modules/modelloader.py index 1106aeb7f..8c862b42f 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -1,3 +1,4 @@ +import glob import os import shutil import importlib @@ -41,7 +42,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None for place in places: if os.path.exists(place): - for file in os.listdir(place): + for file in glob.iglob(place + '**/**', recursive=True): full_path = os.path.join(place, file) if os.path.isdir(full_path): continue