Diffusers subseeds

modules/processing_diffusers.py

@@ -170,6 +170,9 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         if shared.opts.diffusers_generator_device == "Unset":
             generator_device = None
             generator = None
+        elif getattr(p, "generator", None) is not None:
+            generator_device = devices.cpu if shared.opts.diffusers_generator_device == "CPU" else shared.device
+            generator = p.generator
         else:
             generator_device = devices.cpu if shared.opts.diffusers_generator_device == "CPU" else shared.device
             generator = [torch.Generator(generator_device).manual_seed(s) for s in p.seeds]
@@ -222,6 +225,8 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
             args['guidance_scale'] = p.cfg_scale
         if 'generator' in possible and generator is not None:
             args['generator'] = generator
+        if 'latents' in possible and getattr(p, "init_latent", None) is not None:
+            args['latents'] = p.init_latent
         if 'output_type' in possible:
             if hasattr(model, 'vae'):
                 args['output_type'] = 'np' # only set latent if model has vae

modules/ui_sections.py

@@ -85,7 +85,7 @@ def create_seed_inputs(tab, reuse_visible=True):
             seed = gr.Number(label='Initial seed', value=-1, elem_id=f"{tab}_seed", container=True)
             random_seed = ToolButton(ui_symbols.random, elem_id=f"{tab}_random_seed", label='Random seed')
             reuse_seed = ToolButton(ui_symbols.reuse, elem_id=f"{tab}_reuse_seed", label='Reuse seed', visible=reuse_visible)
-        with gr.Row(elem_id=f"{tab}_subseed_row", variant="compact", visible=shared.backend==shared.Backend.ORIGINAL):
+        with gr.Row(elem_id=f"{tab}_subseed_row", variant="compact", visible=True):
             subseed = gr.Number(label='Variation', value=-1, elem_id=f"{tab}_subseed", container=True)
             random_subseed = ToolButton(ui_symbols.random, elem_id=f"{tab}_random_subseed")
             reuse_subseed = ToolButton(ui_symbols.reuse, elem_id=f"{tab}_reuse_subseed", visible=reuse_visible)

scripts/init_latents.py

@@ -0,0 +1,52 @@
+# from PIL import Image
+# import gradio as gr
+from modules import scripts, processing, shared, devices
+from modules.processing_helpers import slerp
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+
+
+class Script(scripts.Script):
+    standalone = False
+
+    def title(self):
+        return 'Init Latents'
+
+    def show(self, is_img2img):
+        return scripts.AlwaysVisible if shared.backend == shared.Backend.DIFFUSERS else False
+
+    @staticmethod
+    def get_latents(p):
+        generator_device = devices.cpu if shared.opts.diffusers_generator_device == "CPU" else shared.device
+        generator = [torch.Generator(generator_device).manual_seed(s) for s in p.seeds]
+        shape = (len(generator), shared.sd_model.unet.config.in_channels, p.height // shared.sd_model.vae_scale_factor,
+                 p.width // shared.sd_model.vae_scale_factor)
+        latents = randn_tensor(shape, generator=generator, device=shared.sd_model._execution_device,
+                               dtype=shared.sd_model.unet.dtype)
+        var_generator = [torch.Generator(generator_device).manual_seed(ss) for ss in p.subseeds]
+        var_latents = randn_tensor(shape, generator=var_generator, device=shared.sd_model._execution_device,
+                                   dtype=shared.sd_model.unet.dtype)
+        return latents, var_latents, generator, var_generator
+
+    @staticmethod
+    def set_slerp(p, latents, var_latents, generator, var_generator):
+        if p.subseed_strength < 1:
+            p.init_latent = slerp(p.subseed_strength, latents, var_latents)
+        if p.subseed_strength == 1:
+            p.init_latent = var_latents
+        if 0 < p.subseed_strength <= 0.5:
+            p.generator = generator
+        if 0.5 < p.subseed_strength <= 1:
+            p.generator = var_generator
+
+
+    def process_batch(self, p: processing.StableDiffusionProcessing, *args, **kwargs): # pylint: disable=arguments-differ
+        if shared.backend != shared.Backend.DIFFUSERS:
+            return
+        args = list(args)
+        if p.subseed_strength != 0:
+            latents, var_latents, generator, var_generator = self.get_latents(p)
+            self.set_slerp(p, latents, var_latents, generator, var_generator)
+
+
+