diff --git a/extensions-builtin/ScuNET/.priority b/extensions-builtin/ScuNET/.priority
deleted file mode 100644
index 3ad5abd03..000000000
--- a/extensions-builtin/ScuNET/.priority
+++ /dev/null
@@ -1 +0,0 @@
-99
diff --git a/extensions-builtin/ScuNET/preload.py b/extensions-builtin/ScuNET/preload.py
deleted file mode 100644
index 4ce82b1d4..000000000
--- a/extensions-builtin/ScuNET/preload.py
+++ /dev/null
@@ -1,6 +0,0 @@
-import os
-from modules import paths
-
-
-def preload(parser):
- parser.add_argument("--scunet-models-path", type=str, help="Path to directory with ScuNET model file(s).", default=os.path.join(paths.models_path, 'ScuNET'))
diff --git a/extensions-builtin/SwinIR/preload.py b/extensions-builtin/SwinIR/preload.py
deleted file mode 100644
index e912c6402..000000000
--- a/extensions-builtin/SwinIR/preload.py
+++ /dev/null
@@ -1,6 +0,0 @@
-import os
-from modules import paths
-
-
-def preload(parser):
- parser.add_argument("--swinir-models-path", type=str, help="Path to directory with SwinIR model file(s).", default=os.path.join(paths.models_path, 'SwinIR'))
diff --git a/html/upscalers.json b/html/upscalers.json
new file mode 100644
index 000000000..eb1abf7f1
--- /dev/null
+++ b/html/upscalers.json
@@ -0,0 +1,32 @@
+{
+ "SwinIR": [
+ ["4x GAN", "https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth"],
+ ["4x PSNR", "https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_PSNR.pth"]
+ ],
+ "ESRGAN": [
+ ["4x GAN", "https://github.com/cszn/KAIR/releases/download/v1.0/ESRGAN.pth"],
+ ["4x Ultrasharp", "https://huggingface.co/Aptronym/SDNext/resolve/main/Upscalers/4x-UltraSharp.pth"],
+ ["4x Valar", "https://huggingface.co/Aptronym/SDNext/resolve/main/Upscalers/4x_Valar_v1.pth"],
+ ["4x Box", "https://huggingface.co/Aptronym/SDNext/resolve/main/Upscalers/4xBox.pth"],
+ ["4x BigFace V3", "https://huggingface.co/Aptronym/SDNext/resolve/main/Upscalers/4x_BigFace_v3.pth"],
+ ["4x Remacri", "https://huggingface.co/Aptronym/SDNext/resolve/main/Upscalers/4x_foolhardy_Remacri.pth"],
+ ["4x NMKD Siax", "https://huggingface.co/Aptronym/SDNext/resolve/main/Upscalers/4x_NMKD-Siax_200k.pth"],
+ ["4x NMKD Superscale", "https://huggingface.co/Aptronym/SDNext/resolve/main/Upscalers/4x_NMKD-Superscale-SP_178000_G.pth"],
+ ["4x NMKD YandereNeo", "https://huggingface.co/Aptronym/SDNext/resolve/main/Upscalers/4x_NMKD-YandereNeoXL_200k.pth"],
+ ["8x NMKD Faces", "https://huggingface.co/Zabin/Resizers/resolve/main/8x_NMKD-Faces_160000_G.pth"],
+ ["8x Superscale", "https://huggingface.co/Aptronym/SDNext/resolve/main/Upscalers/8x_NMKD-Superscale_150000_G.pth"],
+ ["8x HugePaint", "https://huggingface.co/Aptronym/SDNext/resolve/main/Upscalers/8x_HugePaint_v1.pth"]
+ ],
+ "RealESRGAN": [
+ ["4x General V3", "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"],
+ ["4x General WDN V3", "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth"],
+ ["AnimeVideo V3", "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth"],
+ ["4x+ Anime6B", "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"],
+ ["4x+", "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"],
+ ["2x+", "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"]
+ ],
+ "SCUNet": [
+ ["GAN", "https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_gan.pth"],
+ ["PSNR", "https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_psnr.pth"]
+ ]
+}
diff --git a/modules/api/api.py b/modules/api/api.py
index e750fccb4..6a053cb03 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -22,7 +22,6 @@ from modules.textual_inversion.preprocess import preprocess
from modules.hypernetworks.hypernetwork import create_hypernetwork, train_hypernetwork
from modules.sd_models import checkpoints_list, unload_model_weights, reload_model_weights
from modules.sd_models_config import find_checkpoint_config_near_filename
-from modules.realesrgan_model import get_realesrgan_models
from modules import devices
errors.install()
@@ -478,6 +477,7 @@ class Api:
return [{"name":x.name(), "cmd_dir": getattr(x, "cmd_dir", None)} for x in shared.face_restorers]
def get_realesrgan_models(self):
+ from modules.postprocess.realesrgan_model import get_realesrgan_models
return [{"name":x.name,"path":x.data_path, "scale":x.scale} for x in get_realesrgan_models(None)]
def get_prompt_styles(self):
diff --git a/modules/modelloader.py b/modules/modelloader.py
index 7d99f1714..8963f9a19 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -325,7 +325,53 @@ def extension_filter(ext_filter=None, ext_blacklist=None):
return (not ext_filter or any(fp.upper().endswith(ew) for ew in ext_filter)) and (not ext_blacklist or not any(fp.upper().endswith(ew) for ew in ext_blacklist))
return filter
-def load_file_from_url(url: str, *, model_dir: str, progress: bool = True, file_name = None):
+
+def download_url_to_file(url: str, dst: str):
+ # based on torch.hub.download_url_to_file
+ import uuid
+ import tempfile
+ from urllib.request import urlopen, Request
+ from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn
+
+ file_size = None
+ req = Request(url, headers={"User-Agent": "sdnext"})
+ u = urlopen(req)
+ meta = u.info()
+ if hasattr(meta, 'getheaders'):
+ content_length = meta.getheaders("Content-Length")
+ else:
+ content_length = meta.get_all("Content-Length")
+ if content_length is not None and len(content_length) > 0:
+ file_size = int(content_length[0])
+ dst = os.path.expanduser(dst)
+ for _seq in range(tempfile.TMP_MAX):
+ tmp_dst = dst + '.' + uuid.uuid4().hex + '.partial'
+ try:
+ f = open(tmp_dst, 'w+b')
+ except FileExistsError:
+ continue
+ break
+ else:
+ shared.log.error('Error downloading: url={url} no usable temporary filename found')
+ return
+ try:
+ with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=shared.console) as progress:
+ task = progress.add_task(description="Downloading", total=file_size)
+ while True:
+ buffer = u.read(8192)
+ if len(buffer) == 0:
+ break
+ f.write(buffer)
+ progress.update(task, advance=len(buffer))
+ f.close()
+ shutil.move(f.name, dst)
+ finally:
+ f.close()
+ if os.path.exists(f.name):
+ os.remove(f.name)
+
+
+def load_file_from_url(url: str, *, model_dir: str, progress: bool = True, file_name = None): # pylint: disable=unused-argument
"""Download a file from url into model_dir, using the file present if possible. Returns the path to the downloaded file."""
os.makedirs(model_dir, exist_ok=True)
if not file_name:
@@ -334,8 +380,7 @@ def load_file_from_url(url: str, *, model_dir: str, progress: bool = True, file_
cached_file = os.path.abspath(os.path.join(model_dir, file_name))
if not os.path.exists(cached_file):
shared.log.info(f'Downloading: url="{url}" file={cached_file}')
- from torch.hub import download_url_to_file
- download_url_to_file(url, cached_file, progress=progress)
+ download_url_to_file(url, cached_file)
return cached_file
@@ -367,12 +412,18 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
def friendly_name(file: str):
if "http" in file:
file = urlparse(file).path
-
file = os.path.basename(file)
model_name, _extension = os.path.splitext(file)
return model_name
+def friendly_fullname(file: str):
+ if "http" in file:
+ file = urlparse(file).path
+ file = os.path.basename(file)
+ return file
+
+
def cleanup_models():
# This code could probably be more efficient if we used a tuple list or something to store the src/destinations
# and then enumerate that, but this works for now. In the future, it'd be nice to just have every "model" scaler
@@ -427,15 +478,15 @@ def move_files(src_path: str, dest_path: str, ext_filter: str = None):
def load_upscalers():
# We can only do this 'magic' method to dynamically load upscalers if they are referenced, so we'll try to import any _model.py files before looking in __subclasses__
- modules_dir = os.path.join(shared.script_path, "modules")
+ modules_dir = os.path.join(shared.script_path, "modules", "postprocess")
for file in os.listdir(modules_dir):
if "_model.py" in file:
model_name = file.replace("_model.py", "")
- full_model = f"modules.{model_name}_model"
+ full_model = f"modules.postprocess.{model_name}_model"
try:
importlib.import_module(full_model)
- except Exception:
- pass
+ except Exception as e:
+ shared.log.error(f'Error loading upscaler: {model_name} {e}')
datas = []
commandline_options = vars(shared.cmd_opts)
# some of upscaler classes will not go away after reloading their modules, and we'll end up with two copies of those classes. The newest copy will always be the last in the list, so we go from end to beginning and ignore duplicates
@@ -456,4 +507,4 @@ def load_upscalers():
datas,
key=lambda x: x.name.lower() if not isinstance(x.scaler, (UpscalerNone, UpscalerLanczos, UpscalerNearest)) else "" # Special case for UpscalerNone keeps it at the beginning of the list.
)
- shared.log.debug(f"Loaded upscalers: items={len(shared.sd_upscalers)}")
+ shared.log.debug(f"Loaded upscalers: total={len(shared.sd_upscalers)} downloaded={len([x for x in shared.sd_upscalers if x.data_path is not None and os.path.isfile(x.data_path)])} user={len([x for x in shared.sd_upscalers if x.custom])}")
diff --git a/modules/codeformer/codeformer_arch.py b/modules/postprocess/codeformer_arch.py
similarity index 97%
rename from modules/codeformer/codeformer_arch.py
rename to modules/postprocess/codeformer_arch.py
index 0ccd2a7ce..f1a6bf601 100644
--- a/modules/codeformer/codeformer_arch.py
+++ b/modules/postprocess/codeformer_arch.py
@@ -1,15 +1,11 @@
# this file is copied from CodeFormer repository. Please see comment in modules/codeformer_model.py
import math
-import numpy as np
+from typing import Optional
import torch
from torch import nn, Tensor
import torch.nn.functional as F
-from typing import Optional, List
-
-from modules.codeformer.vqgan_arch import VQAutoEncoder, ResBlock
-from basicsr.utils import get_root_logger
-from basicsr.utils.registry import ARCH_REGISTRY
+from modules.postprocess.vqgan_arch import VQAutoEncoder, ResBlock
def calc_mean_std(feat, eps=1e-5):
"""Calculate mean and std for adaptive_instance_normalization.
@@ -158,7 +154,6 @@ class Fuse_sft_block(nn.Module):
return out
-@ARCH_REGISTRY.register()
class CodeFormer(VQAutoEncoder):
def __init__(self, dim_embd=512, n_head=8, n_layers=9,
codebook_size=1024, latent_size=256,
diff --git a/modules/codeformer_model.py b/modules/postprocess/codeformer_model.py
similarity index 91%
rename from modules/codeformer_model.py
rename to modules/postprocess/codeformer_model.py
index 34e3a7a5a..5f274b089 100644
--- a/modules/codeformer_model.py
+++ b/modules/postprocess/codeformer_model.py
@@ -28,11 +28,10 @@ def setup_model(dirname):
try:
from torchvision.transforms.functional import normalize
- from modules.codeformer.codeformer_arch import CodeFormer
+ from modules.postprocess.codeformer_arch import CodeFormer
from basicsr.utils import img2tensor, tensor2img
from facelib.utils.face_restoration_helper import FaceRestoreHelper
from facelib.detection.retinaface import retinaface
- from modules.shared import cmd_opts
net_class = CodeFormer
@@ -54,12 +53,13 @@ def setup_model(dirname):
if len(model_paths) != 0:
ckpt_path = model_paths[0]
else:
- print("Unable to load codeformer model.")
+ shared.log.error(f"Model failed loading: type=CodeFormer model={model_path}")
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)
checkpoint = torch.load(ckpt_path)['params_ema']
net.load_state_dict(checkpoint)
net.eval()
+ shared.log.info(f"Model loaded: type=CodeFormer model={ckpt_path}")
if hasattr(retinaface, 'device'):
retinaface.device = devices.device_codeformer
@@ -102,8 +102,8 @@ def setup_model(dirname):
restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
del output
devices.torch_gc()
- except Exception as error:
- print(f'\tFailed inference for CodeFormer: {error}', file=sys.stderr)
+ except Exception as e:
+ shared.log.error(f'CodeForomer error: {e}')
restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))
restored_face = restored_face.astype('uint8')
diff --git a/modules/esrgan_model.py b/modules/postprocess/esrgan_model.py
similarity index 80%
rename from modules/esrgan_model.py
rename to modules/postprocess/esrgan_model.py
index a9682bceb..3532ed4fa 100644
--- a/modules/esrgan_model.py
+++ b/modules/postprocess/esrgan_model.py
@@ -1,17 +1,13 @@
-import os
-
import numpy as np
import torch
from PIL import Image
from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn
-
-import modules.esrgan_model_arch as arch
-from modules import modelloader, images, devices
-from modules.upscaler import Upscaler, UpscalerData
+import modules.postprocess.esrgan_model_arch as arch
+from modules import images, devices
+from modules.upscaler import Upscaler
from modules.shared import opts, log, console
-
def mod2normal(state_dict):
# this code is copied from https://github.com/victorca25/iNNfer
if 'conv_first.weight' in state_dict:
@@ -123,24 +119,10 @@ def infer_params(state_dict):
class UpscalerESRGAN(Upscaler):
def __init__(self, dirname):
self.name = "ESRGAN"
- self.model_url = "https://github.com/cszn/KAIR/releases/download/v1.0/ESRGAN.pth"
- self.model_name = "ESRGAN_4x"
- self.scalers = []
self.user_path = dirname
super().__init__()
- model_paths = self.find_models(ext_filter=[".pt", ".pth"])
- scalers = []
- if len(model_paths) == 0:
- scaler_data = UpscalerData(self.model_name, self.model_url, self, 4)
- scalers.append(scaler_data)
- for file in model_paths:
- if "http" in file:
- name = self.model_name
- else:
- name = modelloader.friendly_name(file)
+ self.scalers = self.find_scalers()
- scaler_data = UpscalerData(name, file, self, 4)
- self.scalers.append(scaler_data)
def do_upscale(self, img, selected_model):
model = self.load_model(selected_model)
@@ -151,35 +133,23 @@ class UpscalerESRGAN(Upscaler):
return img
def load_model(self, path: str):
- if "http" in path:
- from modules.modelloader import load_file_from_url
- filename = load_file_from_url(
- url=self.model_url,
- model_dir=self.model_download_path,
- file_name=f"{self.model_name}.pth",
- progress=True,
- )
- else:
- filename = path
- if not os.path.exists(filename) or filename is None:
- log.error(f"Model failed loading: type=ESRGAN model={filename}")
- return None
-
- state_dict = torch.load(filename, map_location='cpu' if devices.device_esrgan.type == 'mps' else None)
- log.info(f"Model loaded: type=ESRGAN model={filename}")
+ info = self.find_model(path)
+ if info is None:
+ return
+ state_dict = torch.load(info.local_data_path, map_location='cpu' if devices.device_esrgan.type == 'mps' else None)
+ log.info(f"Upscaler loaded: type={self.name} model={info.local_data_path}")
if "params_ema" in state_dict:
state_dict = state_dict["params_ema"]
elif "params" in state_dict:
state_dict = state_dict["params"]
- num_conv = 16 if "realesr-animevideov3" in filename else 32
+ num_conv = 16 if "realesr-animevideov3" in info.local_data_path else 32
model = arch.SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=num_conv, upscale=4, act_type='prelu')
model.load_state_dict(state_dict)
model.eval()
return model
-
if "body.0.rdb1.conv1.weight" in state_dict and "conv_first.weight" in state_dict:
- nb = 6 if "RealESRGAN_x4plus_anime_6B" in filename else 23
+ nb = 6 if "RealESRGAN_x4plus_anime_6B" in info.local_data_path else 23
state_dict = resrgan2normal(state_dict, nb)
elif "conv_first.weight" in state_dict:
state_dict = mod2normal(state_dict)
diff --git a/modules/esrgan_model_arch.py b/modules/postprocess/esrgan_model_arch.py
similarity index 100%
rename from modules/esrgan_model_arch.py
rename to modules/postprocess/esrgan_model_arch.py
diff --git a/modules/gfpgan_model.py b/modules/postprocess/gfpgan_model.py
similarity index 97%
rename from modules/gfpgan_model.py
rename to modules/postprocess/gfpgan_model.py
index 2988cd407..f8e2b2cc9 100644
--- a/modules/gfpgan_model.py
+++ b/modules/postprocess/gfpgan_model.py
@@ -18,10 +18,8 @@ def gfpgann():
if loaded_gfpgan_model is not None:
loaded_gfpgan_model.gfpgan.to(devices.device_gfpgan)
return loaded_gfpgan_model
-
if gfpgan_constructor is None:
return None
-
models = modelloader.load_models(model_path, model_url, user_path, ext_filter="GFPGAN")
if len(models) == 1 and "http" in models[0]:
model_file = models[0]
@@ -36,7 +34,6 @@ def gfpgann():
model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None, device=devices.device_gfpgan)
loaded_gfpgan_model = model
shared.log.info(f"Model loaded: type=GFPGAN model={model_file}")
-
return model
diff --git a/modules/postprocess/realesrgan_model.py b/modules/postprocess/realesrgan_model.py
new file mode 100644
index 000000000..0646ce339
--- /dev/null
+++ b/modules/postprocess/realesrgan_model.py
@@ -0,0 +1,62 @@
+import os
+import numpy as np
+from PIL import Image
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from modules.postprocess.realesrgan_model_arch import SRVGGNetCompact
+from modules.upscaler import Upscaler
+from modules.shared import opts, device, log
+
+
+class UpscalerRealESRGAN(Upscaler):
+ def __init__(self, dirname):
+ self.name = "RealESRGAN"
+ self.user_path = dirname
+ super().__init__()
+ self.scalers = self.find_scalers()
+ for scaler in self.scalers:
+ if scaler.name == 'RealESRGAN 2x+':
+ scaler.model = lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+ scaler.scale = 2
+ elif scaler.name == 'RealESRGAN 4x+ Anime6B':
+ scaler.model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+ elif scaler.name == 'RealESRGAN 4x General V3':
+ scaler.model = lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+ elif scaler.name == 'RealESRGAN 4x General WDN V3':
+ scaler.model = lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+ elif scaler.name == 'RealESRGAN AnimeVideo V3':
+ scaler.model=lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
+ elif scaler.name == 'RealESRGAN 4x+':
+ scaler.model = lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+ else:
+ log.error(f"Upscaler unrecognized model: type={self.name} model={scaler.name}")
+
+ def load_model(self, path): # pylint: disable=unused-argument
+ pass
+
+ def do_upscale(self, img, selected_model):
+ if not self.enable:
+ return img
+ try:
+ from modules.postprocess.realesrgan_model_arch import RealESRGANer
+ except Exception:
+ log.error("Error importing Real-ESRGAN:")
+ return img
+
+ info = self.find_model(selected_model)
+ if info is None or not os.path.exists(info.local_data_path):
+ return img
+
+ upsampler = RealESRGANer(
+ scale=info.scale,
+ model_path=info.local_data_path,
+ model=info.model(),
+ half=not opts.no_half and not opts.upcast_sampling,
+ tile=opts.ESRGAN_tile,
+ tile_pad=opts.ESRGAN_tile_overlap,
+ device=device,
+ )
+
+ upsampled = upsampler.enhance(np.array(img), outscale=info.scale)[0]
+
+ image = Image.fromarray(upsampled)
+ return image
diff --git a/modules/realesrgan_model_arch.py b/modules/postprocess/realesrgan_model_arch.py
similarity index 99%
rename from modules/realesrgan_model_arch.py
rename to modules/postprocess/realesrgan_model_arch.py
index 50f085255..b0da8d522 100644
--- a/modules/realesrgan_model_arch.py
+++ b/modules/postprocess/realesrgan_model_arch.py
@@ -5,6 +5,7 @@ import threading
import cv2
import numpy as np
import torch
+from torch import nn
from torch.nn import functional as F
from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn
from modules.shared import log, console
@@ -313,10 +314,6 @@ class IOConsumer(threading.Thread):
save_path = msg['save_path']
cv2.imwrite(save_path, output)
-from basicsr.utils.registry import ARCH_REGISTRY
-from torch import nn as nn
-from torch.nn import functional as F # noqa
-
class SRVGGNetCompact(nn.Module):
"""A compact VGG-style network structure for super-resolution.
diff --git a/extensions-builtin/ScuNET/scripts/scunet_model.py b/modules/postprocess/scunet_model.py
similarity index 58%
rename from extensions-builtin/ScuNET/scripts/scunet_model.py
rename to modules/postprocess/scunet_model.py
index 120ec0ffc..200fd3d65 100644
--- a/extensions-builtin/ScuNET/scripts/scunet_model.py
+++ b/modules/postprocess/scunet_model.py
@@ -1,47 +1,32 @@
-import os.path
-import sys
-import traceback
-
import PIL.Image
import numpy as np
import torch
-from tqdm import tqdm
-
-import modules.upscaler
-from modules import devices, modelloader, script_callbacks
-from scunet_model_arch import SCUNet as net
-from modules.shared import opts
+from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn
+from modules import devices, script_callbacks
+from modules.postprocess.scunet_model_arch import SCUNet as net
+from modules.shared import opts, log, console, device
+from modules.upscaler import Upscaler
-class UpscalerScuNET(modules.upscaler.Upscaler):
+class UpscalerScuNET(Upscaler):
def __init__(self, dirname):
- self.name = "ScuNET"
- self.model_name = "ScuNET GAN"
- self.model_name2 = "ScuNET PSNR"
- self.model_url = "https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_gan.pth"
- self.model_url2 = "https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_psnr.pth"
+ self.name = "SCUNet"
self.user_path = dirname
super().__init__()
- model_paths = self.find_models(ext_filter=[".pth"])
- scalers = []
- add_model2 = True
- for file in model_paths:
- if "http" in file:
- name = self.model_name
- else:
- name = modelloader.friendly_name(file)
- if name == self.model_name2 or file == self.model_url2:
- add_model2 = False
- try:
- scaler_data = modules.upscaler.UpscalerData(name, file, self, 4)
- scalers.append(scaler_data)
- except Exception:
- print(f"Error loading ScuNET model: {file}", file=sys.stderr)
- print(traceback.format_exc(), file=sys.stderr)
- if add_model2:
- scaler_data2 = modules.upscaler.UpscalerData(self.model_name2, self.model_url2, self)
- scalers.append(scaler_data2)
- self.scalers = scalers
+ self.scalers = self.find_scalers()
+
+ def load_model(self, path: str):
+ info = self.find_model(path)
+ if info is None:
+ return
+ model = net(in_nc=3, config=[4, 4, 4, 4, 4, 4, 4], dim=64)
+ model.load_state_dict(torch.load(info.local_data_path), strict=True)
+ model.eval()
+ log.info(f"Upscaler loaded: type={self.name} model={info.local_data_path}")
+ for _, v in model.named_parameters():
+ v.requires_grad = False
+ model = model.to(device)
+ return model
@staticmethod
@torch.no_grad()
@@ -52,95 +37,58 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
tile_overlap = opts.SCUNET_tile_overlap
if tile == 0:
return model(img)
-
- device = devices.get_device_for('scunet')
assert tile % 8 == 0, "tile size should be a multiple of window_size"
sf = 1
-
stride = tile - tile_overlap
h_idx_list = list(range(0, h - tile, stride)) + [h - tile]
w_idx_list = list(range(0, w - tile, stride)) + [w - tile]
E = torch.zeros(1, 3, h * sf, w * sf, dtype=img.dtype, device=device)
W = torch.zeros_like(E, dtype=devices.dtype, device=device)
-
- with tqdm(total=len(h_idx_list) * len(w_idx_list), desc="ScuNET tiles") as pbar:
+ with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=console) as progress:
+ task = progress.add_task(description="Upscaling", total=len(h_idx_list) * len(w_idx_list))
for h_idx in h_idx_list:
-
for w_idx in w_idx_list:
-
in_patch = img[..., h_idx: h_idx + tile, w_idx: w_idx + tile]
-
out_patch = model(in_patch)
out_patch_mask = torch.ones_like(out_patch)
-
E[
..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
].add_(out_patch)
W[
..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
].add_(out_patch_mask)
- pbar.update(1)
+ progress.update(task, advance=1, description="Upscaling")
output = E.div_(W)
-
return output
- def do_upscale(self, img: PIL.Image.Image, selected_file):
-
+ def do_upscale(self, img: PIL.Image.Image, selected_file): # pylint: disable=arguments-renamed
devices.torch_gc()
-
model = self.load_model(selected_file)
if model is None:
- print(f"ScuNET: Unable to load model from {selected_file}", file=sys.stderr)
return img
-
- device = devices.get_device_for('scunet')
tile = opts.SCUNET_tile
h, w = img.height, img.width
np_img = np.array(img)
np_img = np_img[:, :, ::-1] # RGB to BGR
np_img = np_img.transpose((2, 0, 1)) / 255 # HWC to CHW
torch_img = torch.from_numpy(np_img).float().unsqueeze(0).to(device) # type: ignore
-
if tile > h or tile > w:
_img = torch.zeros(1, 3, max(h, tile), max(w, tile), dtype=torch_img.dtype, device=torch_img.device)
_img[:, :, :h, :w] = torch_img # pad image
torch_img = _img
-
torch_output = self.tiled_inference(torch_img, model).squeeze(0)
torch_output = torch_output[:, :h * 1, :w * 1] # remove padding, if any
np_output: np.ndarray = torch_output.float().cpu().clamp_(0, 1).numpy()
del torch_img, torch_output
devices.torch_gc()
-
output = np_output.transpose((1, 2, 0)) # CHW to HWC
output = output[:, :, ::-1] # BGR to RGB
return PIL.Image.fromarray((output * 255).astype(np.uint8))
- def load_model(self, path: str):
- device = devices.get_device_for('scunet')
- if "http" in path:
- from modules.modelloader import load_file_from_url
- filename = load_file_from_url(url=self.model_url, model_dir=self.model_download_path, file_name="%s.pth" % self.name, progress=True)
- else:
- filename = path
- if not os.path.exists(os.path.join(self.model_path, filename)) or filename is None:
- print(f"ScuNET: Unable to load model from {filename}", file=sys.stderr)
- return None
-
- model = net(in_nc=3, config=[4, 4, 4, 4, 4, 4, 4], dim=64)
- model.load_state_dict(torch.load(filename), strict=True)
- model.eval()
- for _, v in model.named_parameters():
- v.requires_grad = False
- model = model.to(device)
-
- return model
-
def on_ui_settings():
import gradio as gr
from modules import shared
-
shared.opts.add_option("SCUNET_tile", shared.OptionInfo(256, "Tile size for SCUNET upscalers", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}, section=('postprocessing', "Postprocessing")).info("0 = no tiling"))
shared.opts.add_option("SCUNET_tile_overlap", shared.OptionInfo(8, "Tile overlap for SCUNET upscalers", gr.Slider, {"minimum": 0, "maximum": 64, "step": 1}, section=('postprocessing', "Postprocessing")).info("Low values = visible seam"))
diff --git a/extensions-builtin/ScuNET/scunet_model_arch.py b/modules/postprocess/scunet_model_arch.py
similarity index 100%
rename from extensions-builtin/ScuNET/scunet_model_arch.py
rename to modules/postprocess/scunet_model_arch.py
diff --git a/extensions-builtin/SwinIR/scripts/swinir_model.py b/modules/postprocess/swinir_model.py
similarity index 63%
rename from extensions-builtin/SwinIR/scripts/swinir_model.py
rename to modules/postprocess/swinir_model.py
index 4527b4c79..712e59dd7 100644
--- a/extensions-builtin/SwinIR/scripts/swinir_model.py
+++ b/modules/postprocess/swinir_model.py
@@ -1,57 +1,25 @@
-import os
import numpy as np
import torch
from PIL import Image
from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn
-from swinir_model_arch import SwinIR as net
-from swinir_model_arch_v2 import Swin2SR as net2
-from modules import modelloader, devices, script_callbacks, shared
-from modules.shared import opts, state
-from modules.upscaler import Upscaler, UpscalerData
-
-
-device_swinir = devices.get_device_for('swinir')
+from modules.postprocess.swinir_model_arch import SwinIR as net
+from modules.postprocess.swinir_model_arch_v2 import Swin2SR as net2
+from modules import devices, script_callbacks, shared
+from modules.upscaler import Upscaler
class UpscalerSwinIR(Upscaler):
def __init__(self, dirname):
self.name = "SwinIR"
- self.model_url = "https://github.com/JingyunLiang/SwinIR/releases/download/v0.0" \
- "/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR" \
- "-L_x4_GAN.pth "
- self.model_name = "SwinIR 4x"
self.user_path = dirname
super().__init__()
- scalers = []
- model_files = self.find_models(ext_filter=[".pt", ".pth"])
- for model in model_files:
- if "http" in model:
- name = self.model_name
- else:
- name = modelloader.friendly_name(model)
- model_data = UpscalerData(name, model, self)
- scalers.append(model_data)
- self.scalers = scalers
-
- def do_upscale(self, img, selected_model):
- model = self.load_model(selected_model)
- if model is None:
- return img
- model = model.to(device_swinir, dtype=devices.dtype)
- img = upscale(img, model)
- devices.torch_gc()
- return img
+ self.scalers = self.find_scalers()
def load_model(self, path, scale=4):
- if "http" in path:
- from modules.modelloader import load_file_from_url
- dl_name = "%s%s" % (self.model_name.replace(" ", "_"), ".pth") # pylint: disable=consider-using-f-string
- filename = load_file_from_url(url=path, model_dir=self.model_download_path, file_name=dl_name, progress=True)
- else:
- filename = path
- if filename is None or not os.path.exists(filename):
- shared.log.error(f"Model failed loading: type=SwinIR model={filename}")
- return None
+ info = self.find_model(path)
+ if info is None:
+ return
+ pretrained_model = torch.load(info.local_data_path)
model_v2 = net2(
upscale=scale,
in_chans=3,
@@ -78,9 +46,6 @@ class UpscalerSwinIR(Upscaler):
upsampler="nearest+conv",
resi_connection="3conv",
)
- pretrained_model = torch.load(filename)
- shared.log.info(f"Model loaded: type=SwinIR model={filename}")
-
for model in [model_v1, model_v2]:
for param in ["params_ema", "params", None]:
try:
@@ -88,13 +53,21 @@ class UpscalerSwinIR(Upscaler):
model.load_state_dict(pretrained_model[param], strict=True)
else:
model.load_state_dict(pretrained_model, strict=True)
- shared.log.info(f'Loaded SwinIR model: {filename} param={param}')
+ shared.log.info(f"Upscaler loaded: type={self.name} model={info.local_data_path} param={param}")
return model
- except Exception:
- pass
- shared.log.error(f'Could not determine SwinIR model parameters: {filename}')
+ except Exception as e:
+ shared.log.error(f'Upscaler invalid parameters: type={self.name} model={info.local_data_path} {e}')
return model
+ def do_upscale(self, img, selected_model):
+ model = self.load_model(selected_model)
+ if model is None:
+ return img
+ model = model.to(shared.device, dtype=devices.dtype)
+ img = upscale(img, model)
+ devices.torch_gc()
+ return img
+
def upscale(
img,
@@ -104,15 +77,13 @@ def upscale(
window_size=8,
scale=4,
):
- tile = tile or opts.SWIN_tile
- tile_overlap = tile_overlap or opts.SWIN_tile_overlap
-
-
+ tile = tile or shared.opts.SWIN_tile
+ tile_overlap = tile_overlap or shared.opts.SWIN_tile_overlap
img = np.array(img)
img = img[:, :, ::-1]
img = np.moveaxis(img, 2, 0) / 255
img = torch.from_numpy(img).float()
- img = img.unsqueeze(0).to(device_swinir, dtype=devices.dtype)
+ img = img.unsqueeze(0).to(shared.device, dtype=devices.dtype)
with torch.no_grad(), devices.autocast():
_, _, h_old, w_old = img.size()
h_pad = (h_old // window_size + 1) * window_size - h_old
@@ -136,23 +107,18 @@ def inference(img, model, tile, tile_overlap, window_size, scale):
tile = min(tile, h, w)
assert tile % window_size == 0, "tile size should be a multiple of window_size"
sf = scale
-
stride = tile - tile_overlap
h_idx_list = list(range(0, h - tile, stride)) + [h - tile]
w_idx_list = list(range(0, w - tile, stride)) + [w - tile]
- E = torch.zeros(b, c, h * sf, w * sf, dtype=devices.dtype, device=device_swinir).type_as(img)
- W = torch.zeros_like(E, dtype=devices.dtype, device=device_swinir)
+ E = torch.zeros(b, c, h * sf, w * sf, dtype=devices.dtype, device=shared.device).type_as(img)
+ W = torch.zeros_like(E, dtype=devices.dtype, device=shared.device)
with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=shared.console) as progress:
task = progress.add_task(description="Upscaling Initializing", total=len(h_idx_list) * len(w_idx_list))
for h_idx in h_idx_list:
- if state.interrupted or state.skipped:
- break
-
for w_idx in w_idx_list:
- if state.interrupted or state.skipped:
+ if shared.state.interrupted or shared.state.skipped:
break
-
in_patch = img[..., h_idx: h_idx + tile, w_idx: w_idx + tile]
out_patch = model(in_patch)
out_patch_mask = torch.ones_like(out_patch)
@@ -165,15 +131,13 @@ def inference(img, model, tile, tile_overlap, window_size, scale):
].add_(out_patch_mask)
progress.update(task, advance=1, description="Upscaling")
output = E.div_(W)
-
return output
def on_ui_settings():
import gradio as gr
-
- shared.opts.add_option("SWIN_tile", shared.OptionInfo(192, "Tile size for all SwinIR", gr.Slider, {"minimum": 16, "maximum": 512, "step": 16}, section=('postprocessing', "Postprocessing")))
- shared.opts.add_option("SWIN_tile_overlap", shared.OptionInfo(8, "Tile overlap, in pixels for SwinIR. Low values = visible seam", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}, section=('postprocessing', "Postprocessing")))
+ shared.opts.add_option("SWIN_tile", shared.OptionInfo(192, "Tile size for SwinIR upscaler", gr.Slider, {"minimum": 16, "maximum": 512, "step": 16}, section=('postprocessing', "Postprocessing")))
+ shared.opts.add_option("SWIN_tile_overlap", shared.OptionInfo(8, "Tile overlap for SwinIR upscaler", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}, section=('postprocessing', "Postprocessing")))
script_callbacks.on_ui_settings(on_ui_settings)
diff --git a/extensions-builtin/SwinIR/swinir_model_arch.py b/modules/postprocess/swinir_model_arch.py
similarity index 99%
rename from extensions-builtin/SwinIR/swinir_model_arch.py
rename to modules/postprocess/swinir_model_arch.py
index 93b932747..4f6696a46 100644
--- a/extensions-builtin/SwinIR/swinir_model_arch.py
+++ b/modules/postprocess/swinir_model_arch.py
@@ -5,7 +5,7 @@
import math
import torch
-import torch.nn as nn
+from torch import nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
diff --git a/extensions-builtin/SwinIR/swinir_model_arch_v2.py b/modules/postprocess/swinir_model_arch_v2.py
similarity index 97%
rename from extensions-builtin/SwinIR/swinir_model_arch_v2.py
rename to modules/postprocess/swinir_model_arch_v2.py
index dad22cca2..991e3212d 100644
--- a/extensions-builtin/SwinIR/swinir_model_arch_v2.py
+++ b/modules/postprocess/swinir_model_arch_v2.py
@@ -6,7 +6,7 @@
import math
import numpy as np
import torch
-import torch.nn as nn
+from torch import nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
diff --git a/modules/codeformer/vqgan_arch.py b/modules/postprocess/vqgan_arch.py
similarity index 99%
rename from modules/codeformer/vqgan_arch.py
rename to modules/postprocess/vqgan_arch.py
index 3acb121b6..87911f8cf 100644
--- a/modules/codeformer/vqgan_arch.py
+++ b/modules/postprocess/vqgan_arch.py
@@ -5,11 +5,9 @@ VQGAN code, adapted from the original created by the Unleashing Transformers aut
https://github.com/samb-t/unleashing-transformers/blob/master/models/vqgan.py
'''
-import numpy as np
import torch
-import torch.nn as nn
+from torch import nn
import torch.nn.functional as F
-import copy
from basicsr.utils import get_root_logger
from basicsr.utils.registry import ARCH_REGISTRY
diff --git a/modules/processing.py b/modules/processing.py
index f325e6000..f1fda1941 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -965,16 +965,13 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
if (self.hr_upscale_to_x == self.width and self.hr_upscale_to_y == self.height) or self.hr_upscaler is None or self.hr_upscaler == 'None':
self.is_hr_pass = False
return
- if self.denoising_strength == 0:
- self.is_hr_pass = False
- return
self.is_hr_pass = True
if not shared.state.processing_has_refined_job_count:
if shared.state.job_count == -1:
shared.state.job_count = self.n_iter
shared.state.job_count = shared.state.job_count * 2
shared.state.processing_has_refined_job_count = True
- shared.log.debug(f'Init hires: upscaler={self.hr_upscaler} sampler={self.latent_sampler} resize={self.hr_resize_x}x{self.hr_resize_y} upscale={self.hr_upscale_to_x}x{self.hr_upscale_to_y}')
+ shared.log.debug(f'Init hires: upscaler="{self.hr_upscaler}" sampler="{self.latent_sampler}" resize={self.hr_resize_x}x{self.hr_resize_y} upscale={self.hr_upscale_to_x}x{self.hr_upscale_to_y}')
def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):
@@ -1013,7 +1010,6 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
self.init_hr()
if self.is_hr_pass:
- self.ops.append('hires')
target_width = self.hr_upscale_to_x
target_height = self.hr_upscale_to_y
for i in range(samples.shape[0]):
@@ -1051,13 +1047,17 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
if self.latent_sampler == "PLMS":
self.latent_sampler = 'UniPC'
if self.hr_force or latent_scale_mode is not None:
- devices.torch_gc() # GC now before running the next img2img to prevent running out of memory
- self.sampler = modules.sd_samplers.create_sampler(self.latent_sampler or self.sampler_name, self.sd_model)
- samples = samples[:, :, self.truncate_y//2:samples.shape[2]-(self.truncate_y+1)//2, self.truncate_x//2:samples.shape[3]-(self.truncate_x+1)//2]
- noise = create_random_tensors(samples.shape[1:], seeds=seeds, subseeds=subseeds, subseed_strength=subseed_strength, p=self)
- modules.sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio(for_hr=True))
- samples = self.sampler.sample_img2img(self, samples, noise, conditioning, unconditional_conditioning, steps=self.hr_second_pass_steps or self.steps, image_conditioning=image_conditioning)
- modules.sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio())
+ if self.denoising_strength > 0:
+ self.ops.append('hires')
+ devices.torch_gc() # GC now before running the next img2img to prevent running out of memory
+ self.sampler = modules.sd_samplers.create_sampler(self.latent_sampler or self.sampler_name, self.sd_model)
+ samples = samples[:, :, self.truncate_y//2:samples.shape[2]-(self.truncate_y+1)//2, self.truncate_x//2:samples.shape[3]-(self.truncate_x+1)//2]
+ noise = create_random_tensors(samples.shape[1:], seeds=seeds, subseeds=subseeds, subseed_strength=subseed_strength, p=self)
+ modules.sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio(for_hr=True))
+ samples = self.sampler.sample_img2img(self, samples, noise, conditioning, unconditional_conditioning, steps=self.hr_second_pass_steps or self.steps, image_conditioning=image_conditioning)
+ modules.sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio())
+ else:
+ self.ops.append('upscale')
x = None
shared.state.nextjob()
self.is_hr_pass = False
diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py
deleted file mode 100644
index 5a88f2330..000000000
--- a/modules/realesrgan_model.py
+++ /dev/null
@@ -1,132 +0,0 @@
-import os
-import numpy as np
-from PIL import Image
-from modules.upscaler import Upscaler, UpscalerData
-from modules.shared import opts, device, log
-from modules import modelloader
-
-
-class UpscalerRealESRGAN(Upscaler):
- def __init__(self, path):
- self.name = "RealESRGAN"
- self.model_path = path
- super().__init__()
- try:
- from basicsr.archs.rrdbnet_arch import RRDBNet # pylint: disable=unused-import
- from modules.realesrgan_model_arch import RealESRGANer, SRVGGNetCompact # pylint: disable=unused-import
- self.enable = True
- self.scalers = []
- scalers = self.load_models(path)
- local_model_paths = self.find_models(ext_filter=[".pth"])
- for scaler in scalers:
- if scaler.local_data_path.startswith("http"):
- filename = modelloader.friendly_name(scaler.local_data_path)
- local_model_candidates = [local_model for local_model in local_model_paths if local_model.endswith(f"{filename}.pth")]
- if local_model_candidates:
- scaler.local_data_path = local_model_candidates[0]
- self.scalers.append(scaler)
- except Exception as e:
- log.error(f"Error loading Real-ESRGAN: model={path} {e}")
- self.enable = False
- self.scalers = []
-
- def do_upscale(self, img, selected_model):
- if not self.enable:
- return img
-
- try:
- from modules.realesrgan_model_arch import RealESRGANer
- except Exception:
- log.error("Error importing Real-ESRGAN:")
- return img
-
- info = self.load_model(selected_model)
- if info is None or not os.path.exists(info.local_data_path):
- return img
-
- upsampler = RealESRGANer(
- scale=info.scale,
- model_path=info.local_data_path,
- model=info.model(),
- half=not opts.no_half and not opts.upcast_sampling,
- tile=opts.ESRGAN_tile,
- tile_pad=opts.ESRGAN_tile_overlap,
- device=device,
- )
-
- upsampled = upsampler.enhance(np.array(img), outscale=info.scale)[0]
-
- image = Image.fromarray(upsampled)
- return image
-
- def load_model(self, path):
- try:
- info = next(iter([scaler for scaler in self.scalers if scaler.data_path == path]), None)
- if info is None:
- log.error(f"Model failed loading: type=R-ESRGAN model={info.name}")
- return None
- if info.local_data_path.startswith("http"):
- from modules.modelloader import load_file_from_url
- info.local_data_path = load_file_from_url(url=info.data_path, model_dir=self.model_download_path, progress=True)
- log.info(f"Model loaded: type=R-ESRGAN model={info.name}")
- return info
- except Exception as e:
- log.error(f"Model failed loading: type=R-ESRGAN model={info.name} {e}")
- return None
-
- def load_models(self, _):
- return get_realesrgan_models(self)
-
-
-def get_realesrgan_models(scaler):
- try:
- from basicsr.archs.rrdbnet_arch import RRDBNet
- from modules.realesrgan_model_arch import SRVGGNetCompact # pylint: disable=unused-import
- models = [
- UpscalerData(
- name="R-ESRGAN General 4xV3",
- path="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
- scale=4,
- upscaler=scaler,
- model=lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
- ),
- UpscalerData(
- name="R-ESRGAN General WDN 4xV3",
- path="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth",
- scale=4,
- upscaler=scaler,
- model=lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
- ),
- UpscalerData(
- name="R-ESRGAN AnimeVideo",
- path="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth",
- scale=4,
- upscaler=scaler,
- model=lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
- ),
- UpscalerData(
- name="R-ESRGAN 4x+",
- path="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
- scale=4,
- upscaler=scaler,
- model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
- ),
- UpscalerData(
- name="R-ESRGAN 4x+ Anime6B",
- path="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
- scale=4,
- upscaler=scaler,
- model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
- ),
- UpscalerData(
- name="R-ESRGAN 2x+",
- path="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
- scale=2,
- upscaler=scaler,
- model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
- ),
- ]
- return models
- except Exception as e:
- log.error(f'Error creating Real-ESRGAN models list: {e}')
- return []
diff --git a/modules/shared.py b/modules/shared.py
index d69f0ff3b..bd792ee5c 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -292,6 +292,7 @@ def list_builtin_themes():
def list_themes():
+ log.info('Browser session started')
fn = os.path.join('html', 'themes.json')
if not os.path.exists(fn):
refresh_themes()
@@ -303,7 +304,7 @@ def list_themes():
builtin = list_builtin_themes()
default = ["gradio/default", "gradio/base", "gradio/glass", "gradio/monochrome", "gradio/soft"]
external = {x['id'] for x in res if x['status'] == 'RUNNING' and 'test' not in x['id'].lower()}
- log.info(f'Themes: builtin={len(builtin)} default={len(default)} external={len(external)}')
+ log.debug(f'Themes: builtin={len(builtin)} default={len(default)} external={len(external)}')
themes = sorted(builtin) + sorted(default) + sorted(external, key=str.casefold)
return themes
@@ -506,7 +507,7 @@ options_templates.update(options_section(('saving-paths', "Image Naming & Paths"
"saving_sep_images": OptionInfo("Images
", "", gr.HTML),
"save_images_add_number": OptionInfo(True, "Add number to filename when saving", component_args=hide_dirs),
"use_original_name_batch": OptionInfo(True, "Use original name for output filename during batch process"),
- "use_upscaler_name_as_suffix": OptionInfo(True, "Use upscaler name as filename suffix in the extras tab"),
+ "use_upscaler_name_as_suffix": OptionInfo(True, "Use upscaler name as filename suffix in the extras tab", gr.Checkbox, {"visible": False}),
"samples_filename_pattern": OptionInfo("[seq]-[prompt_words]", "Images filename pattern", component_args=hide_dirs),
"outdir_sep_dirs": OptionInfo("Directories
", "", gr.HTML),
diff --git a/modules/ui.py b/modules/ui.py
index 898da57d1..292935024 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -18,9 +18,7 @@ from modules.dml import directml_override_opts
from modules import prompt_parser
from modules import timer
import modules.ui_symbols as symbols
-import modules.codeformer_model
import modules.generation_parameters_copypaste as parameters_copypaste
-import modules.gfpgan_model
import modules.hypernetworks.ui
import modules.scripts
import modules.shared
diff --git a/modules/upscaler.py b/modules/upscaler.py
index 34bb83230..c008dbca5 100644
--- a/modules/upscaler.py
+++ b/modules/upscaler.py
@@ -1,18 +1,18 @@
import os
from abc import abstractmethod
-
import PIL
from PIL import Image
-
import modules.shared
from modules import modelloader
+
LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
NEAREST = (Image.Resampling.NEAREST if hasattr(Image, 'Resampling') else Image.NEAREST)
-
+models = None
class Upscaler:
name = None
+ folder = None
model_path = None
model_name = None
model_url = None
@@ -23,7 +23,10 @@ class Upscaler:
scalers = []
tile = True
- def __init__(self, create_dirs=False):
+ def __init__(self, create_dirs=True):
+ global models # pylint: disable=global-statement
+ if models is None:
+ models = modules.shared.readfile('html/upscalers.json')
self.mod_pad_h = None
self.tile_size = modules.shared.opts.ESRGAN_tile
self.tile_pad = modules.shared.opts.ESRGAN_tile_overlap
@@ -35,18 +38,42 @@ class Upscaler:
self.pre_pad = 0
self.mod_scale = None
self.model_download_path = None
-
if self.model_path is None and self.name:
self.model_path = os.path.join(modules.shared.models_path, self.name)
if self.model_path and create_dirs:
os.makedirs(self.model_path, exist_ok=True)
-
try:
import cv2 # pylint: disable=unused-import
self.can_tile = True
except Exception:
pass
+ def find_scalers(self):
+ scalers = []
+ loaded = []
+ for k, v in models.items(): # from config
+ if k != self.name:
+ continue
+ for model in v:
+ local_name = os.path.join(self.user_path, modelloader.friendly_fullname(model[1]))
+ model_path = local_name if os.path.exists(local_name) else model[1]
+ scaler = UpscalerData(name=f'{k} {model[0]}', path=model_path, upscaler=self)
+ scalers.append(scaler)
+ loaded.append(model_path)
+ # modules.shared.log.debug(f'Upscaler type={self.name} folder="{self.user_path}" model="{model[0]}" path="{model_path}"')
+ for fn in os.listdir(self.user_path): # from folder
+ if not fn.endswith('.pth') and not fn.endswith('.pt'):
+ continue
+ file_name = os.path.join(self.user_path, fn)
+ if file_name not in loaded:
+ model_name = os.path.splitext(fn)[0]
+ scaler = UpscalerData(name=f'{self.name} {model_name}', path=file_name, upscaler=self)
+ scaler.custom = True
+ scalers.append(scaler)
+ loaded.append(file_name)
+ # modules.shared.log.debug(f'Upscaler type={self.name} folder="{self.user_path}" model="{model_name}" path="{file_name}"')
+ return scalers
+
@abstractmethod
def do_upscale(self, img: PIL.Image, selected_model: str):
return img
@@ -76,8 +103,26 @@ class Upscaler:
def update_status(self, prompt):
print(f"\nextras: {prompt}", file=modules.shared.progress_print_out)
+ def find_model(self, path):
+ info = None
+ for scaler in self.scalers:
+ if scaler.data_path == path:
+ info = scaler
+ break
+ if info is None:
+ modules.shared.log.error(f'Upscaler cannot match model: type={self.name} model="{path}"')
+ return None
+ if info.local_data_path.startswith("http"):
+ from modules.modelloader import load_file_from_url
+ info.local_data_path = load_file_from_url(url=info.data_path, model_dir=self.model_download_path, progress=True)
+ if not os.path.isfile(info.local_data_path):
+ modules.shared.log.error(f'Upscaler cannot find model: type={self.name} model="{info.local_data_path}"')
+ return None
+ return info
+
class UpscalerData:
+ custom: bool = False
name = None
data_path = None
scale: int = 4
diff --git a/requirements.txt b/requirements.txt
index 59d4a3f55..3fc7a4336 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -46,7 +46,7 @@ requests==2.31.0
tqdm==4.66.1
accelerate==0.20.3
opencv-python-headless==4.7.0.72
-diffusers==0.20.2
+diffusers==0.21.2
einops==0.4.1
gradio==3.43.2
huggingface_hub==0.17.1
diff --git a/scripts/postprocessing_codeformer.py b/scripts/postprocessing_codeformer.py
index 9b1644c05..09e004e16 100644
--- a/scripts/postprocessing_codeformer.py
+++ b/scripts/postprocessing_codeformer.py
@@ -1,7 +1,8 @@
from PIL import Image
import numpy as np
import gradio as gr
-from modules import scripts_postprocessing, codeformer_model
+from modules import scripts_postprocessing
+from modules.postprocess import codeformer_model
from modules.ui_components import FormRow
diff --git a/scripts/postprocessing_gfpgan.py b/scripts/postprocessing_gfpgan.py
index 7cea285d7..8214eb712 100644
--- a/scripts/postprocessing_gfpgan.py
+++ b/scripts/postprocessing_gfpgan.py
@@ -1,7 +1,8 @@
from PIL import Image
import numpy as np
import gradio as gr
-from modules import scripts_postprocessing, gfpgan_model
+from modules import scripts_postprocessing
+from modules.postprocess import gfpgan_model
from modules.ui_components import FormRow
diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index 28c46a999..a50f904ff 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -98,13 +98,6 @@ def apply_clip_skip(p, x, xs):
shared.opts.data["clip_skip"] = x
-def apply_upscale_latent_space(p, x, xs):
- if x.lower().strip() != '0':
- shared.opts.data["use_scale_latent_for_hires_fix"] = True
- else:
- shared.opts.data["use_scale_latent_for_hires_fix"] = False
-
-
def find_vae(name: str):
if name.lower() in ['auto', 'automatic']:
return sd_vae.unspecified
@@ -131,6 +124,13 @@ def apply_schedulers_solver_order(p, x, xs):
shared.opts.data["schedulers_solver_order"] = min(x, p.steps - 1)
+def apply_upscaler(p: StableDiffusionProcessingTxt2Img, opt, x):
+ p.enable_hr = True
+ p.hr_force = True
+ p.denoising_strength = 0.0
+ p.hr_upscaler = opt
+
+
def apply_face_restore(p, opt, x):
opt = opt.lower()
if opt == 'codeformer':
@@ -203,35 +203,33 @@ class AxisOptionTxt2Img(AxisOption):
axis_options = [
AxisOption("Nothing", str, do_nothing, fmt=format_nothing),
- AxisOption("Checkpoint name", str, apply_checkpoint, fmt=format_value, cost=1.0, choices=lambda: sorted(sd_models.checkpoints_list)),
- AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: ['None'] + list(sd_vae.vae_dict)),
- AxisOption("Dict name", str, apply_dict, fmt=format_value, cost=1.0, choices=lambda: ['None'] + list(sd_models.checkpoints_list)),
AxisOption("Prompt S/R", str, apply_prompt, fmt=format_value),
+ AxisOption("Model", str, apply_checkpoint, fmt=format_value, cost=1.0, choices=lambda: sorted(sd_models.checkpoints_list)),
+ AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: ['None'] + list(sd_vae.vae_dict)),
AxisOption("Styles", str, apply_styles, choices=lambda: list(shared.prompt_styles.styles)),
AxisOptionTxt2Img("Sampler", str, apply_sampler, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]),
AxisOptionImg2Img("Sampler", str, apply_sampler, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers_for_img2img]),
AxisOption("Seed", int, apply_field("seed")),
AxisOption("Steps", int, apply_field("steps")),
AxisOption("CFG Scale", float, apply_field("cfg_scale")),
- AxisOption("Var. seed", int, apply_field("subseed")),
- AxisOption("Var. strength", float, apply_field("subseed_strength")),
+ AxisOption("Variation seed", int, apply_field("subseed")),
+ AxisOption("Variation strength", float, apply_field("subseed_strength")),
AxisOption("Clip skip", int, apply_clip_skip),
- AxisOption("Denoising", float, apply_field("denoising_strength")),
- AxisOptionTxt2Img("Hires steps", int, apply_field("hr_second_pass_steps")),
- AxisOptionImg2Img("Image CFG scale", float, apply_field("image_cfg_scale")),
+ AxisOption("Denoising strength", float, apply_field("denoising_strength")),
AxisOption("Prompt order", str_permutations, apply_order, fmt=format_value_join_list),
- AxisOption("Sampler Sigma Churn", float, apply_field("s_churn")),
- AxisOption("Sampler Sigma min", float, apply_field("s_tmin")),
- AxisOption("Sampler Sigma max", float, apply_field("s_tmax")),
- AxisOption("Sampler Sigma noise", float, apply_field("s_noise")),
- AxisOption("Sampler Eta", float, apply_field("eta")),
- AxisOptionTxt2Img("Hires upscaler", str, apply_field("hr_upscaler"), choices=lambda: [*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]]),
- AxisOptionImg2Img("Image Mask Weight", float, apply_field("inpainting_mask_weight")),
- AxisOption("Sampler Solver Order", int, apply_schedulers_solver_order, cost=0.5),
+ AxisOption("Upscaler", str, apply_upscaler, choices=lambda: [x.name for x in shared.sd_upscalers][1:]),
AxisOption("Face restore", str, apply_face_restore, fmt=format_value),
AxisOption("Token merging ratio", float, apply_override('token_merging_ratio')),
- AxisOption("Token merging ratio high-res", float, apply_override('token_merging_ratio_hr')),
- #Second PASS
+ # AxisOption("Sampler Sigma Churn", float, apply_field("s_churn")),
+ # AxisOption("Sampler Sigma min", float, apply_field("s_tmin")),
+ # AxisOption("Sampler Sigma max", float, apply_field("s_tmax")),
+ # AxisOption("Sampler Sigma noise", float, apply_field("s_noise")),
+ # AxisOption("Sampler Eta", float, apply_field("eta")),
+ # AxisOption("Sampler Solver Order", int, apply_schedulers_solver_order, cost=0.5),
+ # AxisOption("Token merging ratio high-res", float, apply_override('token_merging_ratio_hr')),
+ AxisOptionImg2Img("Image mask weight", float, apply_field("inpainting_mask_weight")),
+ AxisOption("Model dictionary", str, apply_dict, fmt=format_value, cost=1.0, choices=lambda: ['None'] + list(sd_models.checkpoints_list)),
+ AxisOption("SecondPass Upscaler", str, apply_field("hr_upscaler"), choices=lambda: [*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]]),
AxisOption("SecondPass Sampler", str, apply_latent_sampler, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]),
AxisOption("SecondPass Denoising Strength", float, apply_field("denoising_strength")),
AxisOption("SecondPass Steps", int, apply_field("hr_second_pass_steps")),
diff --git a/webui.py b/webui.py
index 5a1699abe..5ffd0673c 100644
--- a/webui.py
+++ b/webui.py
@@ -20,9 +20,6 @@ from modules.call_queue import queue_lock, wrap_queued_call, wrap_gradio_gpu_cal
import modules.devices
import modules.sd_samplers
import modules.upscaler
-import modules.codeformer_model as codeformer
-import modules.face_restoration
-import modules.gfpgan_model as gfpgan
import modules.img2img
import modules.lowvram
import modules.scripts
@@ -96,11 +93,11 @@ def initialize():
modules.sd_models.setup_model()
timer.startup.record("models")
+ import modules.postprocess.codeformer_model as codeformer
codeformer.setup_model(opts.codeformer_models_path)
- timer.startup.record("codeformer")
-
+ import modules.postprocess.gfpgan_model as gfpgan
gfpgan.setup_model(opts.gfpgan_models_path)
- timer.startup.record("gfpgan")
+ timer.startup.record("face-restore")
log.debug('Loading extensions')
t_timer, t_total = modules.scripts.load_scripts()