stable-diffusion-webui
79 строк · 2.6 Кб
1import os2
3from modules import modelloader, errors4from modules.shared import cmd_opts, opts5from modules.upscaler import Upscaler, UpscalerData6from modules.upscaler_utils import upscale_with_model7
8
9class UpscalerDAT(Upscaler):10def __init__(self, user_path):11self.name = "DAT"12self.user_path = user_path13self.scalers = []14super().__init__()15
16for file in self.find_models(ext_filter=[".pt", ".pth"]):17name = modelloader.friendly_name(file)18scaler_data = UpscalerData(name, file, upscaler=self, scale=None)19self.scalers.append(scaler_data)20
21for model in get_dat_models(self):22if model.name in opts.dat_enabled_models:23self.scalers.append(model)24
25def do_upscale(self, img, path):26try:27info = self.load_model(path)28except Exception:29errors.report(f"Unable to load DAT model {path}", exc_info=True)30return img31
32model_descriptor = modelloader.load_spandrel_model(33info.local_data_path,34device=self.device,35prefer_half=(not cmd_opts.no_half and not cmd_opts.upcast_sampling),36expected_architecture="DAT",37)38return upscale_with_model(39model_descriptor,40img,41tile_size=opts.DAT_tile,42tile_overlap=opts.DAT_tile_overlap,43)44
45def load_model(self, path):46for scaler in self.scalers:47if scaler.data_path == path:48if scaler.local_data_path.startswith("http"):49scaler.local_data_path = modelloader.load_file_from_url(50scaler.data_path,51model_dir=self.model_download_path,52)53if not os.path.exists(scaler.local_data_path):54raise FileNotFoundError(f"DAT data missing: {scaler.local_data_path}")55return scaler56raise ValueError(f"Unable to find model info: {path}")57
58
59def get_dat_models(scaler):60return [61UpscalerData(62name="DAT x2",63path="https://github.com/n0kovo/dat_upscaler_models/raw/main/DAT/DAT_x2.pth",64scale=2,65upscaler=scaler,66),67UpscalerData(68name="DAT x3",69path="https://github.com/n0kovo/dat_upscaler_models/raw/main/DAT/DAT_x3.pth",70scale=3,71upscaler=scaler,72),73UpscalerData(74name="DAT x4",75path="https://github.com/n0kovo/dat_upscaler_models/raw/main/DAT/DAT_x4.pth",76scale=4,77upscaler=scaler,78),79]80