stable-diffusion-webui
62 строки · 2.1 Кб
1from modules import modelloader, devices, errors2from modules.shared import opts3from modules.upscaler import Upscaler, UpscalerData4from modules.upscaler_utils import upscale_with_model5
6
7class UpscalerESRGAN(Upscaler):8def __init__(self, dirname):9self.name = "ESRGAN"10self.model_url = "https://github.com/cszn/KAIR/releases/download/v1.0/ESRGAN.pth"11self.model_name = "ESRGAN_4x"12self.scalers = []13self.user_path = dirname14super().__init__()15model_paths = self.find_models(ext_filter=[".pt", ".pth"])16scalers = []17if len(model_paths) == 0:18scaler_data = UpscalerData(self.model_name, self.model_url, self, 4)19scalers.append(scaler_data)20for file in model_paths:21if file.startswith("http"):22name = self.model_name23else:24name = modelloader.friendly_name(file)25
26scaler_data = UpscalerData(name, file, self, 4)27self.scalers.append(scaler_data)28
29def do_upscale(self, img, selected_model):30try:31model = self.load_model(selected_model)32except Exception:33errors.report(f"Unable to load ESRGAN model {selected_model}", exc_info=True)34return img35model.to(devices.device_esrgan)36return esrgan_upscale(model, img)37
38def load_model(self, path: str):39if path.startswith("http"):40# TODO: this doesn't use `path` at all?41filename = modelloader.load_file_from_url(42url=self.model_url,43model_dir=self.model_download_path,44file_name=f"{self.model_name}.pth",45)46else:47filename = path48
49return modelloader.load_spandrel_model(50filename,51device=('cpu' if devices.device_esrgan.type == 'mps' else None),52expected_architecture='ESRGAN',53)54
55
56def esrgan_upscale(model, img):57return upscale_with_model(58model,59img,60tile_size=opts.ESRGAN_tile,61tile_overlap=opts.ESRGAN_tile_overlap,62)63