add actual latent upscalers

2026-09-19 17:24:32 +02:00 · 2023-09-25 11:23:04 -04:00
parent 4ba2e23dd2
commit ff28ac35e0
11 changed files with 494 additions and 91 deletions
@@ -0,0 +1,64 @@
+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 = "StableDiffusion"
+        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,
+        ]
+
+    def load_model(self, path: str):
+        from modules.sd_models import set_diffuser_options
+        scaler = [x for x in self.scalers if x.data_path == path][0]
+        if scaler.model is None:
+            devices.set_cuda_params()
+            scaler.model = diffusers.DiffusionPipeline.from_pretrained(path, cache_dir=shared.opts.diffusers_dir, torch_dtype=devices.dtype)
+            if hasattr(scaler.model, "set_progress_bar_config"):
+                scaler.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')
+        return scaler.model
+
+    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]
+        return image