diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 7a56c66ce..7e005c358 100644 --- a/modules/processing_diffusers.py +++ b/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 diff --git a/modules/ui_sections.py b/modules/ui_sections.py index 3c026d3ca..7840789f5 100644 --- a/modules/ui_sections.py +++ b/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) diff --git a/scripts/init_latents.py b/scripts/init_latents.py new file mode 100644 index 000000000..08a2db205 --- /dev/null +++ b/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) + + +