Spaces:
Runtime error
Runtime error
| import torch | |
| import diffusers | |
| from PIL import Image | |
| from modules import shared, devices | |
| from modules.upscaler import Upscaler, UpscalerData | |
| class UpscalerSD(Upscaler): | |
| def __init__(self, dirname): # pylint: disable=super-init-not-called | |
| self.name = "SDUpscale" | |
| self.user_path = dirname | |
| if shared.backend != shared.Backend.DIFFUSERS: | |
| super().__init__() | |
| return | |
| self.scalers = [ | |
| UpscalerData(name="SD Latent 2x", path="stabilityai/sd-x2-latent-upscaler", upscaler=self, model=None, scale=4), | |
| UpscalerData(name="SD Latent 4x", path="stabilityai/stable-diffusion-x4-upscaler", upscaler=self, model=None, scale=4), | |
| ] | |
| self.pipelines = [ | |
| None, | |
| None, | |
| ] | |
| self.models = {} | |
| def load_model(self, path: str): | |
| from modules.sd_models import set_diffuser_options | |
| scaler: UpscalerData = [x for x in self.scalers if x.data_path == path][0] | |
| if self.models.get(path, None) is not None: | |
| shared.log.debug(f"Upscaler cached: type={scaler.name} model={path}") | |
| return self.models[path] | |
| else: | |
| devices.set_cuda_params() | |
| model = diffusers.DiffusionPipeline.from_pretrained(path, cache_dir=shared.opts.diffusers_dir, torch_dtype=devices.dtype) | |
| if hasattr(model, "set_progress_bar_config"): | |
| model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + 'Upscale', ncols=80, colour='#327fba') | |
| set_diffuser_options(scaler.model, vae=None, op='upscaler') | |
| self.models[path] = model | |
| return self.models[path] | |
| def callback(self, _step: int, _timestep: int, _latents: torch.FloatTensor): | |
| pass | |
| def do_upscale(self, img: Image.Image, selected_model): | |
| devices.torch_gc() | |
| model = self.load_model(selected_model) | |
| if model is None: | |
| return img | |
| seeds = [torch.randint(0, 2 ** 32, (1,)).item() for _ in range(1)] | |
| generator_device = devices.cpu if shared.opts.diffusers_generator_device == "CPU" else devices.device | |
| generator = [torch.Generator(generator_device).manual_seed(s) for s in seeds] | |
| args = { | |
| 'prompt': '', | |
| 'negative_prompt': '', | |
| 'image': img, | |
| 'num_inference_steps': 20, | |
| 'guidance_scale': 7.5, | |
| 'generator': generator, | |
| 'latents': None, | |
| 'return_dict': True, | |
| 'callback': self.callback, | |
| 'callback_steps': 1, | |
| # 'noise_level': 100, | |
| # 'num_images_per_prompt': 1, | |
| # 'eta': 0.0, | |
| # 'cross_attention_kwargs': None, | |
| } | |
| model = model.to(devices.device) | |
| output = model(**args) | |
| image = output.images[0] | |
| if shared.opts.upscaler_unload and selected_model in self.models: | |
| del self.models[selected_model] | |
| shared.log.debug(f"Upscaler unloaded: type={self.name} model={selected_model}") | |
| devices.torch_gc(force=True) | |
| return image | |