diff --git a/LICENSE-HiDiffusion.txt b/LICENSE-HiDiffusion.txt new file mode 100644 index 00000000000..73095184ff3 --- /dev/null +++ b/LICENSE-HiDiffusion.txt @@ -0,0 +1,65 @@ +HiDiffusion - License notice + +Original project: https://github.com/megvii-research/HiDiffusion +Vendored from: https://github.com/monofy-org/HiDiffusion +Vendored under: invokeai/backend/hidiffusion/ + +================================================================================ +CODE (Apache License 2.0) +================================================================================ + +The vendored HiDiffusion source code is licensed under the Apache License, +Version 2.0. The Apache 2.0 license text is available in the root LICENSE file +and at: + + https://www.apache.org/licenses/LICENSE-2.0 + +================================================================================ +MODULE KEYS (BSD 3-Clause License) +================================================================================ + +The SD 1.5 and SDXL module-key lists under +invokeai/backend/hidiffusion/sd_module_key/ were sourced from dgenerate at: + + https://github.com/Teriks/dgenerate/tree/d83b839033cc22c5101fb0f987bd4eb2de3d5d12/dgenerate/extras/hidiffusion/sd_module_key + +BSD 3-Clause License + +Copyright (c) 2023, Teriks + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +================================================================================ +LOCAL MODIFICATIONS +================================================================================ + +The following changes were applied when integrating HiDiffusion into InvokeAI: + +* The code was vendored under the invokeai.backend.hidiffusion package. +* apply_hidiffusion() was extended to accept a torch.Generator, enabling + deterministic image generation from a seed. +* HiDiffusion patching was integrated with InvokeAI's model loading and + generation lifecycle. diff --git a/docs/src/content/docs/features/hidiffusion.mdx b/docs/src/content/docs/features/hidiffusion.mdx new file mode 100644 index 00000000000..8158f4980a1 --- /dev/null +++ b/docs/src/content/docs/features/hidiffusion.mdx @@ -0,0 +1,39 @@ +--- +title: HiDiffusion +--- + +# HiDiffusion + +HiDiffusion is an optional denoising enhancement that can improve detail and structure at higher resolutions for SD 1.5, SD 2.1, and SDXL. It modifies the UNet during denoising and is most noticeable at 1536px and above. + +The vendored implementation includes module-key definitions for SD 1.5 and SDXL. SD 2.1 intentionally reuses the SD 1.5 module-key mapping. + +Learn more: https://github.com/megvii-research/HiDiffusion + +## Where to find the switches + +1. Open the **Canvas** tab. +2. Expand **Advanced Settings**. +3. In the **Advanced** grid, enable **HiDiffusion** and optionally adjust the two sub‑toggles and ratios: + - **HiDiffusion: RAU‑Net** + - **HiDiffusion: Window Attention** + - **HiDiffusion: T1 Ratio** + - **HiDiffusion: T2 Ratio** + +## What the switches do + +- **HiDiffusion**: Enables the HiDiffusion patch for denoising. Use this for high‑resolution generations; the effect is subtle at lower sizes. + +- **HiDiffusion: RAU‑Net**: Enables RAU‑Net blocks. This typically improves structure and mid‑frequency detail, especially at larger resolutions. + +- **HiDiffusion: Window Attention**: Enables windowed attention blocks. This can boost local texture/detail, but may slightly affect global coherence in some prompts. + +- **HiDiffusion: T1 Ratio**: Controls when HiDiffusion switches into its mid‑stage behavior. Lower values switch earlier; higher values preserve global structure longer. + +- **HiDiffusion: T2 Ratio**: Controls when HiDiffusion switches into its late‑stage behavior. Higher values keep window attention active longer and can sharpen local detail. + +## Tips + +- Try **1536–2048 px** for the clearest benefits (SDXL). +- If results look worse, disable **Window Attention** first, then RAU‑Net. +- Effects vary by scheduler and model; compare with the same seed for a fair test. diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index 89508dfca68..68e8dfc0dcd 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -1,7 +1,7 @@ # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654) import inspect import os -from contextlib import ExitStack +from contextlib import ExitStack, nullcontext from typing import Any, Dict, Iterator, List, Optional, Tuple, Union import torch @@ -65,6 +65,7 @@ from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType from invokeai.backend.stable_diffusion.extensions.controlnet import ControlNetExt from invokeai.backend.stable_diffusion.extensions.freeu import FreeUExt +from invokeai.backend.stable_diffusion.extensions.hidiffusion import HiDiffusionExt from invokeai.backend.stable_diffusion.extensions.inpaint import InpaintExt from invokeai.backend.stable_diffusion.extensions.inpaint_model import InpaintModelExt from invokeai.backend.stable_diffusion.extensions.lora import LoRAExt @@ -73,6 +74,7 @@ from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt from invokeai.backend.stable_diffusion.extensions.t2i_adapter import T2IAdapterExt from invokeai.backend.stable_diffusion.extensions_manager import ExtensionsManager +from invokeai.backend.stable_diffusion.hidiffusion_utils import hidiffusion_patch from invokeai.backend.stable_diffusion.schedulers import SCHEDULER_MAP from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES from invokeai.backend.util.devices import TorchDevice @@ -130,7 +132,7 @@ def get_scheduler( title="Denoise - SD1.5, SDXL", tags=["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"], category="latents", - version="1.5.4", + version="1.6.0", ) class DenoiseLatentsInvocation(BaseInvocation): """Denoises noisy latents to decodable images""" @@ -191,6 +193,35 @@ class DenoiseLatentsInvocation(BaseInvocation): cfg_rescale_multiplier: float = InputField( title="CFG Rescale Multiplier", default=0, ge=0, lt=1, description=FieldDescriptions.cfg_rescale_multiplier ) + hidiffusion: bool = InputField( + default=False, + description=FieldDescriptions.hidiffusion, + title="HiDiffusion", + ) + hidiffusion_raunet: bool = InputField( + default=True, + description=FieldDescriptions.hidiffusion_raunet, + title="HiDiffusion: RAU-Net", + ) + hidiffusion_window_attn: bool = InputField( + default=True, + description=FieldDescriptions.hidiffusion_window_attn, + title="HiDiffusion: Window Attention", + ) + hidiffusion_t1_ratio: float = InputField( + default=0.4, + ge=0, + le=1, + description=FieldDescriptions.hidiffusion_t1_ratio, + title="HiDiffusion: T1 Ratio", + ) + hidiffusion_t2_ratio: float = InputField( + default=0.0, + ge=0, + le=1, + description=FieldDescriptions.hidiffusion_t2_ratio, + title="HiDiffusion: T2 Ratio", + ) latents: Optional[LatentsField] = InputField( default=None, description=FieldDescriptions.latents, @@ -486,6 +517,14 @@ def prep_control_data( return controlnet_data + @staticmethod + def _get_hidiffusion_name_or_path(unet_config: AnyModelConfig) -> Optional[str]: + return ( + getattr(unet_config, "source", None) + or getattr(unet_config, "path", None) + or getattr(unet_config, "name", None) + ) + @staticmethod def parse_controlnet_field( exit_stack: ExitStack, @@ -837,6 +876,7 @@ def _new_invoke(self, context: InvocationContext) -> LatentsOutput: # get the unet's config so that we can pass the base to sd_step_callback() unet_config = context.models.get_config(self.unet.unet.key) + hidiffusion_name_or_path = self._get_hidiffusion_name_or_path(unet_config) conditioning_data = self.get_conditioning_data( context=context, @@ -874,6 +914,17 @@ def step_callback(state: PipelineIntermediateState) -> None: context.util.sd_step_callback(state, unet_config.base) ext_manager.add_extension(PreviewExt(step_callback)) + if self.hidiffusion: + ext_manager.add_extension( + HiDiffusionExt( + name_or_path=hidiffusion_name_or_path, + apply_raunet=self.hidiffusion_raunet, + apply_window_attn=self.hidiffusion_window_attn, + t1_ratio=self.hidiffusion_t1_ratio, + t2_ratio=self.hidiffusion_t2_ratio, + generator=torch.Generator(device="cpu").manual_seed(seed), + ) + ) ### cfg rescale if self.cfg_rescale_multiplier > 0: @@ -940,14 +991,17 @@ def step_callback(state: PipelineIntermediateState) -> None: # ext: t2i/ip adapter ext_manager.run_callback(ExtensionCallbackType.SETUP, denoise_ctx) - with ( - context.models.load(self.unet.unet).model_on_device() as (cached_weights, unet), - ModelPatcher.patch_unet_attention_processor(unet, denoise_ctx.inputs.attention_processor_cls), + with ExitStack() as unet_stack: + cached_weights, unet = unet_stack.enter_context(context.models.load(self.unet.unet).model_on_device()) + unet._num_timesteps = timesteps.shape[0] + unet_stack.enter_context( + ModelPatcher.patch_unet_attention_processor(unet, denoise_ctx.inputs.attention_processor_cls) + ) # ext: controlnet - ext_manager.patch_extensions(denoise_ctx), - # ext: freeu, seamless, ip adapter, lora - ext_manager.patch_unet(unet, cached_weights), - ): + unet_stack.enter_context(ext_manager.patch_extensions(denoise_ctx)) + # ext: freeu, seamless, ip adapter, lora, hidiffusion + unet_stack.enter_context(ext_manager.patch_unet(unet, cached_weights)) + sd_backend = StableDiffusionBackend(unet, scheduler) denoise_ctx.unet = unet result_latents = sd_backend.latents_from_embeddings(denoise_ctx, ext_manager) @@ -997,6 +1051,7 @@ def _old_invoke(self, context: InvocationContext) -> LatentsOutput: # get the unet's config so that we can pass the base to sd_step_callback() unet_config = context.models.get_config(self.unet.unet.key) + hidiffusion_name_or_path = self._get_hidiffusion_name_or_path(unet_config) def step_callback(state: PipelineIntermediateState) -> None: context.util.sd_step_callback(state, unet_config.base) @@ -1083,23 +1138,37 @@ def _lora_loader() -> Iterator[PatchSpec]: denoising_end=self.denoising_end, seed=seed, ) + pipeline._num_timesteps = timesteps.shape[0] - result_latents = pipeline.latents_from_embeddings( - latents=latents, - timesteps=timesteps, - init_timestep=init_timestep, - noise=noise, - seed=seed, - mask=mask, - masked_latents=masked_latents, - is_gradient_mask=gradient_mask, - scheduler_step_kwargs=scheduler_step_kwargs, - conditioning_data=conditioning_data, - control_data=controlnet_data, - ip_adapter_data=ip_adapter_data, - t2i_adapter_data=t2i_adapter_data, - callback=step_callback, - ) + with ( + hidiffusion_patch( + pipeline, + name_or_path=hidiffusion_name_or_path, + apply_raunet=self.hidiffusion_raunet, + apply_window_attn=self.hidiffusion_window_attn, + t1_ratio=self.hidiffusion_t1_ratio, + t2_ratio=self.hidiffusion_t2_ratio, + generator=torch.Generator(device="cpu").manual_seed(seed), + ) + if self.hidiffusion + else nullcontext() + ): + result_latents = pipeline.latents_from_embeddings( + latents=latents, + timesteps=timesteps, + init_timestep=init_timestep, + noise=noise, + seed=seed, + mask=mask, + masked_latents=masked_latents, + is_gradient_mask=gradient_mask, + scheduler_step_kwargs=scheduler_step_kwargs, + conditioning_data=conditioning_data, + control_data=controlnet_data, + ip_adapter_data=ip_adapter_data, + t2i_adapter_data=t2i_adapter_data, + callback=step_callback, + ) # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699 result_latents = result_latents.to("cpu") diff --git a/invokeai/app/invocations/fields.py b/invokeai/app/invocations/fields.py index 8c866188184..256b0f20376 100644 --- a/invokeai/app/invocations/fields.py +++ b/invokeai/app/invocations/fields.py @@ -149,6 +149,11 @@ class FieldDescriptions: denoising_end = "When to stop denoising, expressed a percentage of total steps" cfg_scale = "Classifier-Free Guidance scale" cfg_rescale_multiplier = "Rescale multiplier for CFG guidance, used for models trained with zero-terminal SNR" + hidiffusion = "Apply HiDiffusion (RAU-Net + MSW-MSA) for higher-resolution denoising" + hidiffusion_raunet = "Apply HiDiffusion RAU-Net blocks" + hidiffusion_window_attn = "Apply HiDiffusion window attention blocks" + hidiffusion_t1_ratio = "Override HiDiffusion early switch threshold (T1 ratio)" + hidiffusion_t2_ratio = "Override HiDiffusion late switch threshold (T2 ratio)" scheduler = "Scheduler to use during inference" positive_cond = "Positive conditioning tensor" negative_cond = "Negative conditioning tensor" diff --git a/invokeai/app/invocations/metadata_linked.py b/invokeai/app/invocations/metadata_linked.py index cd733fab648..3ee70440436 100644 --- a/invokeai/app/invocations/metadata_linked.py +++ b/invokeai/app/invocations/metadata_linked.py @@ -622,7 +622,7 @@ class LatentsMetaOutput(LatentsOutput, MetadataOutput): title=f"{DenoiseLatentsInvocation.UIConfig.title} + Metadata", tags=["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"], category="metadata", - version="1.1.1", + version="1.2.0", ) class DenoiseLatentsMetaInvocation(DenoiseLatentsInvocation, WithMetadata): def invoke(self, context: InvocationContext) -> LatentsMetaOutput: @@ -661,6 +661,11 @@ def _loras_to_json(obj: Union[Any, list[Any]]): md.update({"denoising_end": self.denoising_end}) md.update({"scheduler": self.scheduler}) md.update({"model": self.unet.unet}) + md.update({"hidiffusion": self.hidiffusion}) + md.update({"hidiffusion_raunet": self.hidiffusion_raunet}) + md.update({"hidiffusion_window_attn": self.hidiffusion_window_attn}) + md.update({"hidiffusion_t1_ratio": self.hidiffusion_t1_ratio}) + md.update({"hidiffusion_t2_ratio": self.hidiffusion_t2_ratio}) if isinstance(self.control, ControlField) or (isinstance(self.control, list) and len(self.control) > 0): md.update({"controlnets": _to_json(self.control)}) if isinstance(self.ip_adapter, IPAdapterField) or ( diff --git a/invokeai/backend/hidiffusion/__init__.py b/invokeai/backend/hidiffusion/__init__.py new file mode 100644 index 00000000000..c4b4d23d94c --- /dev/null +++ b/invokeai/backend/hidiffusion/__init__.py @@ -0,0 +1,3 @@ +from invokeai.backend.hidiffusion.hidiffusion import apply_hidiffusion, remove_hidiffusion + +__all__ = ["apply_hidiffusion", "remove_hidiffusion"] diff --git a/invokeai/backend/hidiffusion/hidiffusion.py b/invokeai/backend/hidiffusion/hidiffusion.py new file mode 100644 index 00000000000..5d67f2554ec --- /dev/null +++ b/invokeai/backend/hidiffusion/hidiffusion.py @@ -0,0 +1,2228 @@ +import importlib.resources +import math +import warnings +from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union + +import diffusers +import torch +import torch.nn.functional as F +from diffusers.image_processor import PipelineImageInput +from diffusers.models import ControlNetModel +from diffusers.models.attention import _chunked_feed_forward +from diffusers.pipelines import auto_pipeline +from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel +from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput +from diffusers.utils import USE_PEFT_BACKEND, deprecate, scale_lora_layers, unscale_lora_layers +from diffusers.utils.torch_utils import apply_freeu, is_compiled_module, is_torch_version + +from invokeai.backend.hidiffusion.utils import isinstance_str + +diffusers_version = diffusers.__version__ +if diffusers_version < "0.27.0": + from diffusers.models.unet_2d_condition import UNet2DConditionOutput + + old_diffusers = True +else: + from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput + + old_diffusers = False + + +def sd15_hidiffusion_key(): + modified_key = {} + modified_key["down_module_key"] = ["down_blocks.0.downsamplers.0.conv"] + modified_key["down_module_key_extra"] = ["down_blocks.1"] + modified_key["up_module_key"] = ["up_blocks.2.upsamplers.0.conv"] + modified_key["up_module_key_extra"] = ["up_blocks.2"] + modified_key["windown_attn_module_key"] = [ + "down_blocks.0.attentions.0.transformer_blocks.0", + "down_blocks.0.attentions.1.transformer_blocks.0", + "up_blocks.3.attentions.0.transformer_blocks.0", + "up_blocks.3.attentions.1.transformer_blocks.0", + "up_blocks.3.attentions.2.transformer_blocks.0", + ] + return modified_key + + +def sdxl_hidiffusion_key(): + modified_key = {} + modified_key["down_module_key"] = ["down_blocks.1"] + modified_key["down_module_key_extra"] = ["down_blocks.1.downsamplers.0.conv"] + modified_key["up_module_key"] = ["up_blocks.1"] + modified_key["up_module_key_extra"] = ["up_blocks.0.upsamplers.0.conv"] + modified_key["windown_attn_module_key"] = [ + "down_blocks.1.attentions.0.transformer_blocks.0", + "down_blocks.1.attentions.0.transformer_blocks.1", + "down_blocks.1.attentions.1.transformer_blocks.0", + "down_blocks.1.attentions.1.transformer_blocks.1", + "up_blocks.1.attentions.0.transformer_blocks.0", + "up_blocks.1.attentions.0.transformer_blocks.1", + "up_blocks.1.attentions.1.transformer_blocks.0", + "up_blocks.1.attentions.1.transformer_blocks.1", + "up_blocks.1.attentions.2.transformer_blocks.0", + "up_blocks.1.attentions.2.transformer_blocks.1", + ] + + return modified_key + + +def sdxl_turbo_hidiffusion_key(): + modified_key = {} + modified_key["down_module_key"] = ["down_blocks.1"] + modified_key["up_module_key"] = ["up_blocks.1"] + modified_key["windown_attn_module_key"] = [ + "down_blocks.1.attentions.0.transformer_blocks.0", + "down_blocks.1.attentions.0.transformer_blocks.1", + "down_blocks.1.attentions.1.transformer_blocks.0", + "down_blocks.1.attentions.1.transformer_blocks.1", + "up_blocks.1.attentions.0.transformer_blocks.0", + "up_blocks.1.attentions.0.transformer_blocks.1", + "up_blocks.1.attentions.1.transformer_blocks.0", + "up_blocks.1.attentions.1.transformer_blocks.1", + "up_blocks.1.attentions.2.transformer_blocks.0", + "up_blocks.1.attentions.2.transformer_blocks.1", + ] + + return modified_key + + +# supported official model. If you use non-official model based on the following models/pipelines, hidiffusion will automatically select the best strategy to fit it. +supported_official_model = [ + "runwayml/stable-diffusion-v1-5", + "stabilityai/stable-diffusion-2-1-base", + "stabilityai/stable-diffusion-xl-base-1.0", + "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", + "stabilityai/sdxl-turbo", +] + + +# T1_ratio: see T1 introduced in the main paper. T1 = number_inference_step * T1_ratio. A higher T1_ratio can better mitigate object duplication. We set T1_ratio=0.4 by default. You'd better adjust it to fit your prompt. Only active when apply_raunet=True. +# T2_ratio: see T2 introduced in the appendix, used in extreme resolution image generation. T2 = number_inference_step * T2_ratio. A higher T2_ratio can better mitigate object duplication. Only active when apply_raunet=True +switching_threshold_ratio_dict = { + "sd15_1024": {"T1_ratio": 0.4, "T2_ratio": 0.0}, + "sd15_2048": {"T1_ratio": 0.7, "T2_ratio": 0.3}, + "sdxl_2048": {"T1_ratio": 0.4, "T2_ratio": 0.0}, + "sdxl_4096": {"T1_ratio": 0.7, "T2_ratio": 0.3}, + "sdxl_turbo_1024": {"T1_ratio": 0.5, "T2_ratio": 0.0}, +} + +text_to_img_controlnet_switching_threshold_ratio_dict = { + "sdxl_2048": {"T1_ratio": 0.5, "T2_ratio": 0.0}, +} +controlnet_apply_steps_rate = 0.6 + +is_aggressive_raunet = True +aggressive_step = 8 + +inpainting_is_aggressive_raunet = False +playground_is_aggressive_raunet = False + + +with importlib.resources.open_text(f"{__package__}.sd_module_key", "sd15_module_key.txt", encoding="utf-8") as f: + sd15_module_key = f.read().splitlines() + +with importlib.resources.open_text(f"{__package__}.sd_module_key", "sdxl_module_key.txt", encoding="utf-8") as f: + sdxl_module_key = f.read().splitlines() + + +def _get_max_timesteps(info_dict: dict) -> int: + """ + Helper function to get the maximum number of timesteps from a pipeline. + """ + pipeline = info_dict["pipeline"] + if hasattr(pipeline, "_num_timesteps"): + return pipeline._num_timesteps + else: + return len(pipeline.scheduler.timesteps) + + +def make_diffusers_sdxl_controlnet_ppl(block_class): + class sdxl_controlnet_ppl(block_class): + # Save for unpatching later + _parent = block_class + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]] = None, + prompt_2: Optional[Union[str, List[str]]] = None, + image: PipelineImageInput = None, + control_image: PipelineImageInput = None, + height: Optional[int] = None, + width: Optional[int] = None, + strength: float = 0.8, + num_inference_steps: int = 50, + guidance_scale: float = 5.0, + negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt_2: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + controlnet_conditioning_scale: Union[float, List[float]] = 0.8, + guess_mode: bool = False, + control_guidance_start: Union[float, List[float]] = 0.0, + control_guidance_end: Union[float, List[float]] = 1.0, + original_size: Tuple[int, int] = None, + crops_coords_top_left: Tuple[int, int] = (0, 0), + target_size: Tuple[int, int] = None, + negative_original_size: Optional[Tuple[int, int]] = None, + negative_crops_coords_top_left: Tuple[int, int] = (0, 0), + negative_target_size: Optional[Tuple[int, int]] = None, + aesthetic_score: float = 6.0, + negative_aesthetic_score: float = 2.5, + clip_skip: Optional[int] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: Optional[List[str]] = None, + **kwargs, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is + used in both text-encoders + image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: + `List[List[torch.FloatTensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + The initial image will be used as the starting point for the image generation process. Can also accept + image latents as `image`, if passing latents directly, it will not be encoded again. + control_image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: + `List[List[torch.FloatTensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If + the type is specified as `Torch.FloatTensor`, it is passed to ControlNet as is. `PIL.Image.Image` can + also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If + height and/or width are passed, `image` is resized according to them. If multiple ControlNets are + specified in init, images must be passed as a list such that each element of the list can be correctly + batched for input to a single controlnet. + height (`int`, *optional*, defaults to the size of control_image): + The height in pixels of the generated image. Anything below 512 pixels won't work well for + [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) + and checkpoints that are not specifically fine-tuned on low resolutions. + width (`int`, *optional*, defaults to the size of control_image): + The width in pixels of the generated image. Anything below 512 pixels won't work well for + [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) + and checkpoints that are not specifically fine-tuned on low resolutions. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + strength (`float`, *optional*, defaults to 0.3): + Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image` + will be used as a starting point, adding more noise to it the larger the `strength`. The number of + denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will + be maximum and the denoising process will run for the full number of iterations specified in + `num_inference_steps`. + guidance_scale (`float`, *optional*, defaults to 7.5): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + negative_prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and + `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to + [`schedulers.DDIMScheduler`], will be ignored for others. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, pooled text embeddings will be generated from `prompt` input argument. + negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` + input argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a + plain tuple. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added + to the residual in the original unet. If multiple ControlNets are specified in init, you can set the + corresponding scale as a list. + guess_mode (`bool`, *optional*, defaults to `False`): + In this mode, the ControlNet encoder will try best to recognize the content of the input image even if + you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended. + control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + The percentage of total steps at which the controlnet starts applying. + control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + The percentage of total steps at which the controlnet stops applying. + original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. + `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as + explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). + crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position + `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting + `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). + target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + For most cases, `target_size` should be set to the desired height and width of the generated image. If + not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in + section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). + negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + To negatively condition the generation process based on a specific image resolution. Part of SDXL's + micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more + information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. + negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's + micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more + information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. + negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + To negatively condition the generation process based on a target image resolution. It should be as same + as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more + information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. + aesthetic_score (`float`, *optional*, defaults to 6.0): + Used to simulate an aesthetic score of the generated image by influencing the positive text condition. + Part of SDXL's micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). + negative_aesthetic_score (`float`, *optional*, defaults to 2.5): + Part of SDXL's micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). Can be used to + simulate an aesthetic score of the generated image by influencing the negative text condition. + clip_skip (`int`, *optional*): + Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that + the output of the pre-final layer will be used for computing the prompt embeddings. + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeine class. + + Examples: + + Returns: + [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: + [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple` + containing the output images. + """ + if callback_on_step_end_tensor_inputs is None: + callback_on_step_end_tensor_inputs = ["latents"] + + # convert image to control_image to fit sdxl_controlnet ppl. + if control_image is None: + control_image = image + image = None + self.info["text_to_img_controlnet"] = True + else: + self.info["text_to_img_controlnet"] = False + + callback = kwargs.pop("callback", None) + callback_steps = kwargs.pop("callback_steps", None) + + if callback is not None: + deprecate( + "callback", + "1.0.0", + "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", + ) + if callback_steps is not None: + deprecate( + "callback_steps", + "1.0.0", + "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", + ) + + controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet + + # align format for control guidance + if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): + control_guidance_start = len(control_guidance_end) * [control_guidance_start] + elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): + control_guidance_end = len(control_guidance_start) * [control_guidance_end] + elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): + mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 + control_guidance_start, control_guidance_end = ( + mult * [control_guidance_start], + mult * [control_guidance_end], + ) + + # 1. Check inputs. Raise error if not correct + if image is not None: + # image-to-image controlnet + if old_diffusers: + self.check_inputs( + prompt, + prompt_2, + control_image, + strength, + num_inference_steps, + callback_steps, + negative_prompt, + negative_prompt_2, + prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + controlnet_conditioning_scale, + control_guidance_start, + control_guidance_end, + callback_on_step_end_tensor_inputs, + ) + else: + self.check_inputs( + prompt, + prompt_2, + control_image, + strength, + num_inference_steps, + callback_steps, + negative_prompt, + negative_prompt_2, + prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + None, + None, + controlnet_conditioning_scale, + control_guidance_start, + control_guidance_end, + callback_on_step_end_tensor_inputs, + ) + else: + # text-to-image controlnet + if old_diffusers: + self.check_inputs( + prompt, + prompt_2, + control_image, + callback_steps, + negative_prompt, + negative_prompt_2, + prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + controlnet_conditioning_scale, + control_guidance_start, + control_guidance_end, + callback_on_step_end_tensor_inputs, + ) + else: + self.check_inputs( + prompt, + prompt_2, + control_image, + callback_steps, + negative_prompt, + negative_prompt_2, + prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + None, + None, + negative_pooled_prompt_embeds, + controlnet_conditioning_scale, + control_guidance_start, + control_guidance_end, + callback_on_step_end_tensor_inputs, + ) + + self._guidance_scale = guidance_scale + self._clip_skip = clip_skip + self._cross_attention_kwargs = cross_attention_kwargs + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): + controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) + + global_pool_conditions = ( + controlnet.config.global_pool_conditions + if isinstance(controlnet, ControlNetModel) + else controlnet.nets[0].config.global_pool_conditions + ) + guess_mode = guess_mode or global_pool_conditions + + # 3. Encode input prompt + text_encoder_lora_scale = ( + self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None + ) + ( + prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ) = self.encode_prompt( + prompt, + prompt_2, + device, + num_images_per_prompt, + self.do_classifier_free_guidance, + negative_prompt, + negative_prompt_2, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + lora_scale=text_encoder_lora_scale, + clip_skip=self.clip_skip, + ) + + # 4. Prepare image and controlnet_conditioning_image + if image is not None: + image = self.image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) + if isinstance(controlnet, ControlNetModel): + control_image = self.prepare_control_image( + image=control_image, + width=width, + height=height, + batch_size=batch_size * num_images_per_prompt, + num_images_per_prompt=num_images_per_prompt, + device=device, + dtype=controlnet.dtype, + do_classifier_free_guidance=self.do_classifier_free_guidance, + guess_mode=guess_mode, + ) + height, width = control_image.shape[-2:] + elif isinstance(controlnet, MultiControlNetModel): + control_images = [] + + for control_image_ in control_image: + control_image_ = self.prepare_control_image( + image=control_image_, + width=width, + height=height, + batch_size=batch_size * num_images_per_prompt, + num_images_per_prompt=num_images_per_prompt, + device=device, + dtype=controlnet.dtype, + do_classifier_free_guidance=self.do_classifier_free_guidance, + guess_mode=guess_mode, + ) + + control_images.append(control_image_) + + control_image = control_images + height, width = control_image[0].shape[-2:] + else: + raise AssertionError("Unsupported controlnet type for control image preprocessing.") + else: + if isinstance(controlnet, ControlNetModel): + control_image = self.prepare_image( + image=control_image, + width=width, + height=height, + batch_size=batch_size * num_images_per_prompt, + num_images_per_prompt=num_images_per_prompt, + device=device, + dtype=controlnet.dtype, + do_classifier_free_guidance=self.do_classifier_free_guidance, + guess_mode=guess_mode, + ) + height, width = control_image.shape[-2:] + elif isinstance(controlnet, MultiControlNetModel): + images = [] + + for image_ in control_image: + image_ = self.prepare_image( + image=image_, + width=width, + height=height, + batch_size=batch_size * num_images_per_prompt, + num_images_per_prompt=num_images_per_prompt, + device=device, + dtype=controlnet.dtype, + do_classifier_free_guidance=self.do_classifier_free_guidance, + guess_mode=guess_mode, + ) + + images.append(image_) + + control_image = images + height, width = image[0].shape[-2:] + else: + raise AssertionError("Unsupported controlnet type for image preprocessing.") + # 5. Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + if image is not None: + timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) + latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) + else: + timesteps = self.scheduler.timesteps + self._num_timesteps = len(timesteps) + + # 6. Prepare latent variables + if image is not None: + # image-to-image controlnet + latents = self.prepare_latents( + image, + latent_timestep, + batch_size, + num_images_per_prompt, + prompt_embeds.dtype, + device, + generator, + True, + ) + else: + # text-to-image controlnet + num_channels_latents = self.unet.config.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + # num_channels_latents = self.unet.config.in_channels + # shape = (batch_size * num_images_per_prompt, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor) + # if isinstance(generator, list) and len(generator) != batch_size: + # raise ValueError( + # f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + # f" size of {batch_size}. Make sure the batch size matches the length of the generators." + # ) + + # if latents is None: + # latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + # else: + # latents = latents.to(device) + + # # scale the initial noise by the standard deviation required by the scheduler + # latents = latents * self.scheduler.init_noise_sigma + + # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 7.1 Create tensor stating which controlnets to keep + controlnet_keep = [] + for i in range(len(timesteps)): + keeps = [ + 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) + for s, e in zip(control_guidance_start, control_guidance_end, strict=False) + ] + controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) + + # 7.2 Prepare added time ids & embeddings + if image is not None: + if isinstance(control_image, list): + original_size = original_size or control_image[0].shape[-2:] + else: + original_size = original_size or control_image.shape[-2:] + target_size = target_size or (height, width) + + if negative_original_size is None: + negative_original_size = original_size + if negative_target_size is None: + negative_target_size = target_size + add_text_embeds = pooled_prompt_embeds + + if self.text_encoder_2 is None: + text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1]) + else: + text_encoder_projection_dim = self.text_encoder_2.config.projection_dim + + add_time_ids, add_neg_time_ids = self._get_add_time_ids( + original_size, + crops_coords_top_left, + target_size, + aesthetic_score, + negative_aesthetic_score, + negative_original_size, + negative_crops_coords_top_left, + negative_target_size, + dtype=prompt_embeds.dtype, + text_encoder_projection_dim=text_encoder_projection_dim, + ) + add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1) + + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0) + add_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1) + add_time_ids = torch.cat([add_neg_time_ids, add_time_ids], dim=0) + + prompt_embeds = prompt_embeds.to(device) + add_text_embeds = add_text_embeds.to(device) + add_time_ids = add_time_ids.to(device) + else: + if isinstance(control_image, list): + original_size = original_size or control_image[0].shape[-2:] + else: + original_size = original_size or control_image.shape[-2:] + target_size = target_size or (height, width) + + add_text_embeds = pooled_prompt_embeds + if self.text_encoder_2 is None: + text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1]) + else: + text_encoder_projection_dim = self.text_encoder_2.config.projection_dim + + add_time_ids = self._get_add_time_ids( + original_size, + crops_coords_top_left, + target_size, + dtype=prompt_embeds.dtype, + text_encoder_projection_dim=text_encoder_projection_dim, + ) + + if negative_original_size is not None and negative_target_size is not None: + negative_add_time_ids = self._get_add_time_ids( + negative_original_size, + negative_crops_coords_top_left, + negative_target_size, + dtype=prompt_embeds.dtype, + text_encoder_projection_dim=text_encoder_projection_dim, + ) + else: + negative_add_time_ids = add_time_ids + + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0) + add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0) + + prompt_embeds = prompt_embeds.to(device) + add_text_embeds = add_text_embeds.to(device) + add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1) + + # 8. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} + + # controlnet(s) inference + if guess_mode and self.do_classifier_free_guidance: + # Infer ControlNet only for the conditional batch. + control_model_input = latents + control_model_input = self.scheduler.scale_model_input(control_model_input, t) + controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] + controlnet_added_cond_kwargs = { + "text_embeds": add_text_embeds.chunk(2)[1], + "time_ids": add_time_ids.chunk(2)[1], + } + else: + control_model_input = latent_model_input + controlnet_prompt_embeds = prompt_embeds + controlnet_added_cond_kwargs = added_cond_kwargs + + if isinstance(controlnet_keep[i], list): + cond_scale = [ + c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i], strict=False) + ] + else: + controlnet_cond_scale = controlnet_conditioning_scale + if isinstance(controlnet_cond_scale, list): + controlnet_cond_scale = controlnet_cond_scale[0] + cond_scale = controlnet_cond_scale * controlnet_keep[i] + + if i < controlnet_apply_steps_rate * num_inference_steps: + original_h, original_w = (128, 128) + _, _, model_input_h, model_input_w = control_model_input.shape + downsample_factor = max(model_input_h / original_h, model_input_w / original_w) + downsample_size = ( + int(model_input_h // downsample_factor) // 8 * 8, + int(model_input_w // downsample_factor) // 8 * 8, + ) + + # original_pixel_h, original_pixel_w = (1024,1024) + # _, _, pixel_h, pixel_w = control_image.shape + # downsample_pixel_factor = max(pixel_h/original_pixel_h, pixel_w/original_pixel_w) + # downsample_pixel_size = (int(pixel_h//downsample_pixel_factor)//8*8, int(pixel_w//downsample_pixel_factor)//8*8) + downsample_pixel_size = [downsample_size[0] * 8, downsample_size[1] * 8] + + down_block_res_samples, mid_block_res_sample = self.controlnet( + F.interpolate(control_model_input, downsample_size), + # control_model_input, + t, + encoder_hidden_states=controlnet_prompt_embeds, + controlnet_cond=F.interpolate(control_image, downsample_pixel_size), + # controlnet_cond=control_image, + conditioning_scale=cond_scale, + guess_mode=guess_mode, + added_cond_kwargs=controlnet_added_cond_kwargs, + return_dict=False, + ) + + if guess_mode and self.do_classifier_free_guidance: + # Infered ControlNet only for the conditional batch. + # To apply the output of ControlNet to both the unconditional and conditional batches, + # add 0 to the unconditional batch to keep it unchanged. + down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] + mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) + + # predict the noise residual + if i < controlnet_apply_steps_rate * num_inference_steps: + noise_pred = self.unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + cross_attention_kwargs=self.cross_attention_kwargs, + down_block_additional_residuals=down_block_res_samples, + mid_block_additional_residual=mid_block_res_sample, + added_cond_kwargs=added_cond_kwargs, + return_dict=False, + )[0] + else: + noise_pred = self.unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + cross_attention_kwargs=self.cross_attention_kwargs, + down_block_additional_residuals=None, + mid_block_additional_residual=None, + added_cond_kwargs=added_cond_kwargs, + return_dict=False, + )[0] + + # perform guidance + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + if callback is not None and i % callback_steps == 0: + step_idx = i // getattr(self.scheduler, "order", 1) + callback(step_idx, t, latents) + + # If we do sequential model offloading, let's offload unet and controlnet + # manually for max memory savings + if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: + self.unet.to("cpu") + self.controlnet.to("cpu") + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + if not output_type == "latent": + # make sure the VAE is in float32 mode, as it overflows in float16 + needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast + + if needs_upcasting: + self.upcast_vae() + latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype) + + image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] + + # cast back to fp16 if needed + if needs_upcasting: + self.vae.to(dtype=torch.float16) + else: + image = latents + return StableDiffusionXLPipelineOutput(images=image) + + # apply watermark if available + if self.watermark is not None: + image = self.watermark.apply_watermark(image) + + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return StableDiffusionXLPipelineOutput(images=image) + + # let's be nice and not change the __name__ of the pipeline class + # this messes up some important pipeline detection code in dgenerate. + + sdxl_controlnet_ppl.__name__ = block_class.__name__ + + return sdxl_controlnet_ppl + + +def make_diffusers_unet_2d_condition(block_class): + class unet_2d_condition(block_class): + # Save for unpatching later + _parent = block_class + + def forward( + self, + sample: torch.FloatTensor, + timestep: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + class_labels: Optional[torch.Tensor] = None, + timestep_cond: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, + down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, + mid_block_additional_residual: Optional[torch.Tensor] = None, + down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + return_dict: bool = True, + ) -> Union[UNet2DConditionOutput, Tuple]: + r""" + The [`UNet2DConditionModel`] forward method. + + Args: + sample (`torch.FloatTensor`): + The noisy input tensor with the following shape `(batch, channel, height, width)`. + timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input. + encoder_hidden_states (`torch.FloatTensor`): + The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. + class_labels (`torch.Tensor`, *optional*, defaults to `None`): + Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. + timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`): + Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed + through the `self.time_embedding` layer to obtain the timestep embeddings. + attention_mask (`torch.Tensor`, *optional*, defaults to `None`): + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask + is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large + negative values to the attention scores corresponding to "discard" tokens. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + added_cond_kwargs: (`dict`, *optional*): + A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that + are passed along to the UNet blocks. + down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*): + A tuple of tensors that if specified are added to the residuals of down unet blocks. + mid_block_additional_residual: (`torch.Tensor`, *optional*): + A tensor that if specified is added to the residual of the middle unet block. + encoder_attention_mask (`torch.Tensor`): + A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If + `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias, + which adds large negative values to the attention scores corresponding to "discard" tokens. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain + tuple. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the [`AttnProcessor`]. + added_cond_kwargs: (`dict`, *optional*): + A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that + are passed along to the UNet blocks. + down_block_additional_residuals (`tuple` of `torch.Tensor`, *optional*): + additional residuals to be added to UNet long skip connections from down blocks to up blocks for + example from ControlNet side model(s) + mid_block_additional_residual (`torch.Tensor`, *optional*): + additional residual to be added to UNet mid block output, for example from ControlNet side model + down_intrablock_additional_residuals (`tuple` of `torch.Tensor`, *optional*): + additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s) + + Returns: + [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: + If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise + a `tuple` is returned where the first element is the sample tensor. + """ + # By default samples have to be AT least a multiple of the overall upsampling factor. + # The overall upsampling factor is equal to 2 ** (# num of upsampling layers). + # However, the upsampling interpolation output size can be forced to fit any upsampling size + # on the fly if necessary. + default_overall_up_factor = 2**self.num_upsamplers + + # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` + forward_upsample_size = False + upsample_size = None + + for dim in sample.shape[-2:]: + if dim % default_overall_up_factor != 0: + # Forward upsample size to force interpolation output size. + forward_upsample_size = True + break + + # ensure attention_mask is a bias, and give it a singleton query_tokens dimension + # expects mask of shape: + # [batch, key_tokens] + # adds singleton query_tokens dimension: + # [batch, 1, key_tokens] + # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: + # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) + # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) + if attention_mask is not None: + # assume that mask is expressed as: + # (1 = keep, 0 = discard) + # convert mask into a bias that can be added to attention scores: + # (keep = +0, discard = -10000.0) + attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + # convert encoder_attention_mask to a bias the same way we do for attention_mask + if encoder_attention_mask is not None: + encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) + + # 0. center input if necessary + if self.config.center_input_sample: + sample = 2 * sample - 1.0 + + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = sample.device.type == "mps" + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timesteps = timesteps.expand(sample.shape[0]) + + t_emb = self.time_proj(timesteps) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # but time_embedding might actually be running in fp16. so we need to cast here. + # there might be better ways to encapsulate this. + t_emb = t_emb.to(dtype=sample.dtype) + + emb = self.time_embedding(t_emb, timestep_cond) + aug_emb = None + + if self.class_embedding is not None: + if class_labels is None: + raise ValueError("class_labels should be provided when num_class_embeds > 0") + + if self.config.class_embed_type == "timestep": + class_labels = self.time_proj(class_labels) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # there might be better ways to encapsulate this. + class_labels = class_labels.to(dtype=sample.dtype) + + class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype) + + if self.config.class_embeddings_concat: + emb = torch.cat([emb, class_emb], dim=-1) + else: + emb = emb + class_emb + + if self.config.addition_embed_type == "text": + aug_emb = self.add_embedding(encoder_hidden_states) + elif self.config.addition_embed_type == "text_image": + # Kandinsky 2.1 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" + ) + + image_embs = added_cond_kwargs.get("image_embeds") + text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states) + aug_emb = self.add_embedding(text_embs, image_embs) + elif self.config.addition_embed_type == "text_time": + # SDXL - style + if "text_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" + ) + text_embeds = added_cond_kwargs.get("text_embeds") + if "time_ids" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" + ) + time_ids = added_cond_kwargs.get("time_ids") + time_embeds = self.add_time_proj(time_ids.flatten()) + time_embeds = time_embeds.reshape((text_embeds.shape[0], -1)) + add_embeds = torch.concat([text_embeds, time_embeds], dim=-1) + add_embeds = add_embeds.to(emb.dtype) + aug_emb = self.add_embedding(add_embeds) + elif self.config.addition_embed_type == "image": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" + ) + image_embs = added_cond_kwargs.get("image_embeds") + aug_emb = self.add_embedding(image_embs) + elif self.config.addition_embed_type == "image_hint": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`" + ) + image_embs = added_cond_kwargs.get("image_embeds") + hint = added_cond_kwargs.get("hint") + aug_emb, hint = self.add_embedding(image_embs, hint) + sample = torch.cat([sample, hint], dim=1) + + emb = emb + aug_emb if aug_emb is not None else emb + + if self.time_embed_act is not None: + emb = self.time_embed_act(emb) + + if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj": + encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states) + elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj": + # Kadinsky 2.1 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" + ) + + image_embeds = added_cond_kwargs.get("image_embeds") + encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds) + elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" + ) + image_embeds = added_cond_kwargs.get("image_embeds") + encoder_hidden_states = self.encoder_hid_proj(image_embeds) + elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "ip_image_proj": + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'ip_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" + ) + image_embeds = added_cond_kwargs.get("image_embeds") + image_embeds = self.encoder_hid_proj(image_embeds).to(encoder_hidden_states.dtype) + encoder_hidden_states = torch.cat([encoder_hidden_states, image_embeds], dim=1) + + # 2. pre-process + sample = self.conv_in(sample) + + # 2.5 GLIGEN position net + if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None: + cross_attention_kwargs = cross_attention_kwargs.copy() + gligen_args = cross_attention_kwargs.pop("gligen") + cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)} + + # 3. down + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + if USE_PEFT_BACKEND: + # weight the lora layers by setting `lora_scale` for each PEFT layer + scale_lora_layers(self, lora_scale) + + is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None + # using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets + is_adapter = down_intrablock_additional_residuals is not None + # maintain backward compatibility for legacy usage, where + # T2I-Adapter and ControlNet both use down_block_additional_residuals arg + # but can only use one or the other + if not is_adapter and mid_block_additional_residual is None and down_block_additional_residuals is not None: + deprecate( + "T2I should not use down_block_additional_residuals", + "1.3.0", + "Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \ + and will be removed in diffusers 1.3.0. `down_block_additional_residuals` should only be used \ + for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ", + standard_warn=False, + ) + down_intrablock_additional_residuals = down_block_additional_residuals + is_adapter = True + + down_block_res_samples = (sample,) + for downsample_block in self.down_blocks: + if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: + # For t2i-adapter CrossAttnDownBlock2D + additional_residuals = {} + if is_adapter and len(down_intrablock_additional_residuals) > 0: + additional_residuals["additional_residuals"] = down_intrablock_additional_residuals.pop(0) + + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + **additional_residuals, + ) + else: + # sample, res_samples = downsample_block(hidden_states=sample, temb=emb, scale=lora_scale) + sample, res_samples = downsample_block(hidden_states=sample, temb=emb) + if is_adapter and len(down_intrablock_additional_residuals) > 0: + sample += down_intrablock_additional_residuals.pop(0) + + down_block_res_samples += res_samples + + if is_controlnet: + new_down_block_res_samples = () + + for down_block_res_sample, down_block_additional_residual in zip( + down_block_res_samples, down_block_additional_residuals, strict=False + ): + _, _, ori_H, ori_W = down_block_res_sample.shape + down_block_additional_residual = F.interpolate( + down_block_additional_residual, (ori_H, ori_W), mode="bicubic" + ) + down_block_res_sample = down_block_res_sample + down_block_additional_residual + new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,) + + down_block_res_samples = new_down_block_res_samples + + # 4. mid + if self.mid_block is not None: + if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention: + sample = self.mid_block( + sample, + emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + else: + sample = self.mid_block(sample, emb) + + # To support T2I-Adapter-XL + if ( + is_adapter + and len(down_intrablock_additional_residuals) > 0 + and sample.shape == down_intrablock_additional_residuals[0].shape + ): + sample += down_intrablock_additional_residuals.pop(0) + + if is_controlnet: + _, _, ori_H, ori_W = sample.shape + mid_block_additional_residual = F.interpolate( + mid_block_additional_residual, (ori_H, ori_W), mode="bicubic" + ) + sample = sample + mid_block_additional_residual + + # 5. up + for i, upsample_block in enumerate(self.up_blocks): + is_final_block = i == len(self.up_blocks) - 1 + + res_samples = down_block_res_samples[-len(upsample_block.resnets) :] + down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] + + # if we have not reached the final block and need to forward the + # upsample size, we do it here + if not is_final_block and forward_upsample_size: + upsample_size = down_block_res_samples[-1].shape[2:] + + if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + upsample_size=upsample_size, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + ) + else: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + upsample_size=upsample_size, + # scale=lora_scale, + ) + # sample = upsample_block( + # hidden_states=sample, + # temb=emb, + # res_hidden_states_tuple=res_samples, + # upsample_size=upsample_size, + # scale=lora_scale, + # ) + + # 6. post-process + if self.conv_norm_out: + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + if USE_PEFT_BACKEND: + # remove `lora_scale` from each PEFT layer + unscale_lora_layers(self, lora_scale) + + if not return_dict: + return (sample,) + + return UNet2DConditionOutput(sample=sample) + + return unet_2d_condition + + +def make_diffusers_transformer_block( + block_class: Type[torch.nn.Module], generator: torch.Generator +) -> Type[torch.nn.Module]: + # replace global self-attention with MSW-MSA + class transformer_block(block_class): + # Save for unpatching later + _parent = block_class + + def forward( + self, + hidden_states: torch.FloatTensor, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + timestep: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + class_labels: Optional[torch.LongTensor] = None, + added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, + ) -> torch.FloatTensor: + # reference: https://github.com/microsoft/Swin-Transformer + def window_partition(x, window_size, shift_size, H, W): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, N, C = x.shape + x = x.view(B, H, W, C) + if H % 2 != 0 or W % 2 != 0: + warnings.warn( + f"HiDiffusion Warning: The feature size is {(H, W)} and cannot be directly partitioned into windows. We interpolate the size to {(window_size[0] * 2, window_size[1] * 2)} " + f"to enable the window partition. Even though the generation is OK, the image quality would be largely decreased. " + f"We suggest removing window attention by setting apply_hidiffusion(pipe, apply_window_attn=False) for better image quality.", + stacklevel=2, + ) + x = ( + F.interpolate( + x.permute(0, 3, 1, 2).contiguous(), + size=(window_size[0] * 2, window_size[1] * 2), + mode="bicubic", + ) + .permute(0, 2, 3, 1) + .contiguous() + ) + if type(shift_size) is list or type(shift_size) is tuple: + if shift_size[0] > 0: + x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1]), dims=(1, 2)) + else: + if shift_size > 0: + x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2)) + x = x.view(B, 2, window_size[0], 2, window_size[1], C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size[0], window_size[1], C) + windows = windows.view(-1, window_size[0] * window_size[1], C) + return windows + + def window_reverse(windows, window_size, H, W, shift_size): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + + Returns: + x: (B, H, W, C) + """ + B, N, C = windows.shape + windows = windows.view(-1, window_size[0], window_size[1], C) + B = int(windows.shape[0] / 4) # 2x2 + x = windows.view(B, 2, 2, window_size[0], window_size[1], -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, window_size[0] * 2, window_size[1] * 2, -1) + if type(shift_size) is list or type(shift_size) is tuple: + if shift_size[0] > 0: + x = torch.roll(x, shifts=(shift_size[0], shift_size[1]), dims=(1, 2)) + else: + if shift_size > 0: + x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2)) + if H % 2 != 0 or W % 2 != 0: + x = ( + F.interpolate(x.permute(0, 3, 1, 2).contiguous(), size=(H, W), mode="bicubic") + .permute(0, 2, 3, 1) + .contiguous() + ) + x = x.view(B, H * W, C) + return x + + # Notice that normalization is always applied before the real computation in the following blocks. + # 0. Self-Attention + batch_size = hidden_states.shape[0] + if self.use_ada_layer_norm: + norm_hidden_states = self.norm1(hidden_states, timestep) + elif self.use_ada_layer_norm_zero: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + elif self.use_layer_norm: + norm_hidden_states = self.norm1(hidden_states) + elif self.use_ada_layer_norm_continuous: + norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs["pooled_text_emb"]) + elif self.use_ada_layer_norm_single: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) + ).chunk(6, dim=1) + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + norm_hidden_states = norm_hidden_states.squeeze(1) + else: + raise ValueError("Incorrect norm used") + + if self.pos_embed is not None: + norm_hidden_states = self.pos_embed(norm_hidden_states) + + # MSW-MSA + if generator is not None: + rand_num = torch.rand(1, generator=generator, device=generator.device) + else: + rand_num = torch.rand(1) + + B, N, C = hidden_states.shape + ori_H, ori_W = self.info["size"] + downsample_ratio = round(((ori_H * ori_W) / N) ** 0.5) + H, W = (math.ceil(ori_H / downsample_ratio), math.ceil(ori_W / downsample_ratio)) + widow_size = (math.ceil(H / 2), math.ceil(W / 2)) + if rand_num <= 0.25: + shift_size = (0, 0) + if rand_num > 0.25 and rand_num <= 0.5: + shift_size = (widow_size[0] // 4, widow_size[1] // 4) + if rand_num > 0.5 and rand_num <= 0.75: + shift_size = (widow_size[0] // 4 * 2, widow_size[1] // 4 * 2) + if rand_num > 0.75 and rand_num <= 1: + shift_size = (widow_size[0] // 4 * 3, widow_size[1] // 4 * 3) + norm_hidden_states = window_partition(norm_hidden_states, widow_size, shift_size, H, W) + # 2. Prepare GLIGEN inputs + cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} + gligen_kwargs = cross_attention_kwargs.pop("gligen", None) + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + if self.use_ada_layer_norm_zero: + attn_output = gate_msa.unsqueeze(1) * attn_output + elif self.use_ada_layer_norm_single: + attn_output = gate_msa * attn_output + + attn_output = window_reverse(attn_output, widow_size, H, W, shift_size) + + hidden_states = attn_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + # 2.5 GLIGEN Control + if gligen_kwargs is not None: + hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"]) + + # 3. Cross-Attention + if self.attn2 is not None: + if self.use_ada_layer_norm: + norm_hidden_states = self.norm2(hidden_states, timestep) + elif self.use_ada_layer_norm_zero or self.use_layer_norm: + norm_hidden_states = self.norm2(hidden_states) + elif self.use_ada_layer_norm_single: + # For PixArt norm2 isn't applied here: + # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103 + norm_hidden_states = hidden_states + elif self.use_ada_layer_norm_continuous: + norm_hidden_states = self.norm2(hidden_states, added_cond_kwargs["pooled_text_emb"]) + else: + raise ValueError("Incorrect norm") + + if self.pos_embed is not None and self.use_ada_layer_norm_single is False: + norm_hidden_states = self.pos_embed(norm_hidden_states) + + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + **cross_attention_kwargs, + ) + hidden_states = attn_output + hidden_states + + # 4. Feed-forward + if self.use_ada_layer_norm_continuous: + norm_hidden_states = self.norm3(hidden_states, added_cond_kwargs["pooled_text_emb"]) + elif not self.use_ada_layer_norm_single: + norm_hidden_states = self.norm3(hidden_states) + + if self.use_ada_layer_norm_zero: + norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + + if self.use_ada_layer_norm_single: + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + + if self._chunk_size is not None: + # "feed_forward_chunk_size" can be used to save memory + ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size) + # ff_output = _chunked_feed_forward( + # self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size, lora_scale=lora_scale + # ) + else: + ff_output = self.ff(norm_hidden_states) + # ff_output = self.ff(norm_hidden_states, scale=lora_scale) + + if self.use_ada_layer_norm_zero: + ff_output = gate_mlp.unsqueeze(1) * ff_output + elif self.use_ada_layer_norm_single: + ff_output = gate_mlp * ff_output + + hidden_states = ff_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + return hidden_states + + return transformer_block + + +def make_diffusers_cross_attn_down_block(block_class: Type[torch.nn.Module]) -> Type[torch.nn.Module]: + # replace conventional downsampler with resolution-aware downsampler + class cross_attn_down_block(block_class): + # Save for unpatching later + _parent = block_class + timestep = 0 + aggressive_raunet = False + T1_ratio = 0 + T1_start = 0 + T1_end = 0 + aggressive_raunet = False + T1 = 0 # to avoid confict with sdxl-turbo + max_timestep = 50 + info: dict = None + model: str = None + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + additional_residuals: Optional[torch.FloatTensor] = None, + ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]: + self.max_timestep = _get_max_timesteps(self.info) + ori_H, ori_W = self.info["size"] + if self.model == "sd15": + if ori_H < 256 or ori_W < 256: + self.T1_ratio = switching_threshold_ratio_dict["sd15_1024"][self.switching_threshold_ratio] + else: + self.T1_ratio = switching_threshold_ratio_dict["sd15_2048"][self.switching_threshold_ratio] + elif self.model == "sdxl": + if ori_H < 512 or ori_W < 512: + if self.info["text_to_img_controlnet"]: + self.T1_ratio = text_to_img_controlnet_switching_threshold_ratio_dict["sdxl_2048"][ + self.switching_threshold_ratio + ] + else: + self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] + + if self.info["is_inpainting_task"]: + self.aggressive_raunet = inpainting_is_aggressive_raunet + elif self.info["is_playground"]: + self.aggressive_raunet = playground_is_aggressive_raunet + else: + self.aggressive_raunet = is_aggressive_raunet + else: + self.T1_ratio = switching_threshold_ratio_dict["sdxl_4096"][self.switching_threshold_ratio] + elif self.model == "sdxl_turbo": + self.T1_ratio = switching_threshold_ratio_dict["sdxl_turbo_1024"][self.switching_threshold_ratio] + else: + raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") + + if self.aggressive_raunet: + # self.T1_start = min(int(self.max_timestep * self.T1_ratio * 0.4), int(8/50 * self.max_timestep)) + self.T1_start = int(aggressive_step / 50 * self.max_timestep) + self.T1_end = int(self.max_timestep * self.T1_ratio) + self.T1 = 0 # to avoid confict with sdxl-turbo + else: + self.T1 = int(self.max_timestep * self.T1_ratio) + + output_states = () + + blocks = list(zip(self.resnets, self.attentions, strict=False)) + + for i, (resnet, attn) in enumerate(blocks): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + else: + # hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = resnet(hidden_states, temb) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + + # apply additional residuals to the output of the last pair of resnet and attention blocks + if i == len(blocks) - 1 and additional_residuals is not None: + hidden_states = hidden_states + additional_residuals + + if i == 0: + if self.aggressive_raunet and self.timestep >= self.T1_start and self.timestep < self.T1_end: + self.info["upsample_size"] = (hidden_states.shape[2], hidden_states.shape[3]) + hidden_states = F.avg_pool2d(hidden_states, kernel_size=(2, 2), ceil_mode=True) + elif self.timestep < self.T1: + self.info["upsample_size"] = (hidden_states.shape[2], hidden_states.shape[3]) + hidden_states = F.avg_pool2d(hidden_states, kernel_size=(2, 2), ceil_mode=True) + output_states = output_states + (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + # hidden_states = downsampler(hidden_states, scale=lora_scale) + + output_states = output_states + (hidden_states,) + + self.timestep += 1 + if self.timestep == self.max_timestep: + self.timestep = 0 + + return hidden_states, output_states + + return cross_attn_down_block + + +def make_diffusers_cross_attn_up_block(block_class: Type[torch.nn.Module]) -> Type[torch.nn.Module]: + # replace conventional downsampler with resolution-aware downsampler + class cross_attn_up_block(block_class): + # Save for unpatching later + _parent = block_class + timestep = 0 + aggressive_raunet = False + T1_ratio = 0 + T1_start = 0 + T1_end = 0 + aggressive_raunet = False + T1 = 0 # to avoid confict with sdxl-turbo + max_timestep = 50 + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + upsample_size: Optional[int] = None, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + self.max_timestep = _get_max_timesteps(self.info) + ori_H, ori_W = self.info["size"] + if self.model == "sd15": + if ori_H < 256 or ori_W < 256: + self.T1_ratio = switching_threshold_ratio_dict["sd15_1024"][self.switching_threshold_ratio] + else: + self.T1_ratio = switching_threshold_ratio_dict["sd15_2048"][self.switching_threshold_ratio] + elif self.model == "sdxl": + if ori_H < 512 or ori_W < 512: + if self.info["text_to_img_controlnet"]: + self.T1_ratio = text_to_img_controlnet_switching_threshold_ratio_dict["sdxl_2048"][ + self.switching_threshold_ratio + ] + else: + self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] + + if self.info["is_inpainting_task"]: + self.aggressive_raunet = inpainting_is_aggressive_raunet + elif self.info["is_playground"]: + self.aggressive_raunet = playground_is_aggressive_raunet + else: + self.aggressive_raunet = is_aggressive_raunet + + else: + self.T1_ratio = switching_threshold_ratio_dict["sdxl_4096"][self.switching_threshold_ratio] + elif self.model == "sdxl_turbo": + self.T1_ratio = switching_threshold_ratio_dict["sdxl_turbo_1024"][self.switching_threshold_ratio] + else: + raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") + + if self.aggressive_raunet: + # self.T1_start = min(int(self.max_timestep * self.T1_ratio * 0.4), int(8/50 * self.max_timestep)) + self.T1_start = int(aggressive_step / 50 * self.max_timestep) + self.T1_end = int(self.max_timestep * self.T1_ratio) + self.T1 = 0 # to avoid confict with sdxl-turbo + else: + self.T1 = int(self.max_timestep * self.T1_ratio) + + is_freeu_enabled = ( + getattr(self, "s1", None) + and getattr(self, "s2", None) + and getattr(self, "b1", None) + and getattr(self, "b2", None) + ) + + for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions, strict=False)): + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + # FreeU: Only operate on the first two stages + if is_freeu_enabled: + hidden_states, res_hidden_states = apply_freeu( + self.resolution_idx, + hidden_states, + res_hidden_states, + s1=self.s1, + s2=self.s2, + b1=self.b1, + b2=self.b2, + ) + + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + else: + hidden_states = resnet(hidden_states, temb) + # hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + + if i == 1: + if self.aggressive_raunet and self.timestep >= self.T1_start and self.timestep < self.T1_end: + hidden_states = F.interpolate( + hidden_states, size=self.info["upsample_size"], mode="bicubic" + ) + elif self.timestep < self.T1: + hidden_states = F.interpolate( + hidden_states, size=self.info["upsample_size"], mode="bicubic" + ) + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size) + # hidden_states = upsampler(hidden_states, upsample_size, scale=lora_scale) + + self.timestep += 1 + if self.timestep == self.max_timestep: + self.timestep = 0 + + return hidden_states + + return cross_attn_up_block + + +def make_diffusers_downsampler_block(block_class: Type[torch.nn.Module]) -> Type[torch.nn.Module]: + # replace conventional downsampler with resolution-aware downsampler + class downsampler_block(block_class): + # Save for unpatching later + _parent = block_class + T1_ratio = 0 + T1 = 0 + timestep = 0 + aggressive_raunet = False + max_timestep = 50 + + def forward(self, hidden_states: torch.Tensor, scale=1.0) -> torch.Tensor: + self.max_timestep = _get_max_timesteps(self.info) + ori_H, ori_W = self.info["size"] + if self.model == "sd15": + if ori_H < 256 or ori_W < 256: + self.T1_ratio = switching_threshold_ratio_dict["sd15_1024"][self.switching_threshold_ratio] + else: + self.T1_ratio = switching_threshold_ratio_dict["sd15_2048"][self.switching_threshold_ratio] + elif self.model == "sdxl": + if ori_H < 512 or ori_W < 512: + if self.info["text_to_img_controlnet"]: + self.T1_ratio = text_to_img_controlnet_switching_threshold_ratio_dict["sdxl_2048"][ + self.switching_threshold_ratio + ] + else: + self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] + + if self.info["is_inpainting_task"]: + self.aggressive_raunet = inpainting_is_aggressive_raunet + elif self.info["is_playground"]: + self.aggressive_raunet = playground_is_aggressive_raunet + else: + self.aggressive_raunet = is_aggressive_raunet + else: + self.T1_ratio = switching_threshold_ratio_dict["sdxl_4096"][self.switching_threshold_ratio] + elif self.model == "sdxl_turbo": + self.T1_ratio = switching_threshold_ratio_dict["sdxl_turbo_1024"][self.switching_threshold_ratio] + else: + raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") + + if self.aggressive_raunet: + # self.T1 = min(int(self.max_timestep * self.T1_ratio), int(8/50 * self.max_timestep)) + self.T1 = int(aggressive_step / 50 * self.max_timestep) + else: + self.T1 = int(self.max_timestep * self.T1_ratio) + stride = self.stride + padding = self.padding + dilation = self.dilation + if self.timestep < self.T1: + stride = (4, 4) + padding = (2, 2) + dilation = (2, 2) + + if old_diffusers: + if self.lora_layer is None: + # make sure to the functional Conv2D function as otherwise torch.compile's graph will break + # see: https://github.com/huggingface/diffusers/pull/4315 + hidden_states = F.conv2d( + hidden_states, self.weight, self.bias, stride, padding, dilation, self.groups + ) + self.timestep += 1 + if self.timestep == self.max_timestep: + self.timestep = 0 + return hidden_states + else: + original_outputs = F.conv2d( + hidden_states, self.weight, self.bias, stride, padding, dilation, self.groups + ) + return original_outputs + (scale * self.lora_layer(hidden_states)) + else: + hidden_states = F.conv2d(hidden_states, self.weight, self.bias, stride, padding, dilation, self.groups) + self.timestep += 1 + if self.timestep == self.max_timestep: + self.timestep = 0 + return hidden_states + + return downsampler_block + + +def make_diffusers_upsampler_block(block_class: Type[torch.nn.Module]) -> Type[torch.nn.Module]: + # replace conventional upsampler with resolution-aware downsampler + class upsampler_block(block_class): + # Save for unpatching later + _parent = block_class + T1_ratio = 0 + T1 = 0 + timestep = 0 + aggressive_raunet = False + max_timestep = 50 + info: dict = None + + def forward(self, hidden_states: torch.Tensor, scale=1.0) -> torch.Tensor: + self.max_timestep = _get_max_timesteps(self.info) + ori_H, ori_W = self.info["size"] + if self.model == "sd15": + if ori_H < 256 or ori_W < 256: + self.T1_ratio = switching_threshold_ratio_dict["sd15_1024"][self.switching_threshold_ratio] + else: + self.T1_ratio = switching_threshold_ratio_dict["sd15_2048"][self.switching_threshold_ratio] + elif self.model == "sdxl": + if ori_H < 512 or ori_W < 512: + if self.info["text_to_img_controlnet"]: + self.T1_ratio = text_to_img_controlnet_switching_threshold_ratio_dict["sdxl_2048"][ + self.switching_threshold_ratio + ] + else: + self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] + + if self.info["is_inpainting_task"]: + self.aggressive_raunet = inpainting_is_aggressive_raunet + elif self.info["is_playground"]: + self.aggressive_raunet = playground_is_aggressive_raunet + else: + self.aggressive_raunet = is_aggressive_raunet + else: + self.T1_ratio = switching_threshold_ratio_dict["sdxl_4096"][self.switching_threshold_ratio] + elif self.model == "sdxl_turbo": + self.T1_ratio = switching_threshold_ratio_dict["sdxl_turbo_1024"][self.switching_threshold_ratio] + else: + raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") + + if self.aggressive_raunet: + # self.T1 = min(int(self.max_timestep * self.T1_ratio), int(8/50 * self.max_timestep)) + self.T1 = int(aggressive_step / 50 * self.max_timestep) + else: + self.T1 = int(self.max_timestep * self.T1_ratio) + self.timestep += 1 + if self.timestep == self.max_timestep: + self.timestep = 0 + + if old_diffusers: + if self.lora_layer is None: + # make sure to the functional Conv2D function as otherwise torch.compile's graph will break + # see: https://github.com/huggingface/diffusers/pull/4315 + return F.conv2d( + hidden_states, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups + ) + else: + original_outputs = F.conv2d( + hidden_states, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups + ) + return original_outputs + (scale * self.lora_layer(hidden_states)) + else: + return F.conv2d( + hidden_states, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups + ) + + return upsampler_block + + +def hook_diffusion_model(model: torch.nn.Module): + """Adds a forward pre hook to get the image size. This hook can be removed with remove_hidiffusion.""" + + def hook(module, args): + module.info["size"] = (args[0].shape[2], args[0].shape[3]) + return None + + model.info["hooks"].append(model.register_forward_pre_hook(hook)) + + +_HIDIFFUSION_RUNTIME_ATTRIBUTES = ( + "timestep", + "aggressive_raunet", + "T1_ratio", + "T1", + "T1_start", + "T1_end", + "max_timestep", + "ori_stride", + "ori_padding", + "ori_dilation", +) +_HIDIFFUSION_STATE_ATTRIBUTES = ( + "stride", + "padding", + "dilation", + *_HIDIFFUSION_RUNTIME_ATTRIBUTES, + "switching_threshold_ratio", + "model", + "info", +) +_HIDIFFUSION_STATE_SNAPSHOT = "_hidiffusion_original_state" + + +def _snapshot_hidiffusion_state(module: torch.nn.Module) -> None: + """Snapshot module state that HiDiffusion may mutate.""" + if _HIDIFFUSION_STATE_SNAPSHOT in module.__dict__: + return + module.__dict__[_HIDIFFUSION_STATE_SNAPSHOT] = { + attribute: (attribute in module.__dict__, module.__dict__.get(attribute)) + for attribute in _HIDIFFUSION_STATE_ATTRIBUTES + } + + +def _reset_hidiffusion_runtime_state(module: torch.nn.Module) -> None: + """Clear runtime state left on a cached module by a previous HiDiffusion patch.""" + for attribute, original_attribute in ( + ("stride", "ori_stride"), + ("padding", "ori_padding"), + ("dilation", "ori_dilation"), + ): + if original_attribute in module.__dict__: + setattr(module, attribute, module.__dict__[original_attribute]) + + for attribute in _HIDIFFUSION_RUNTIME_ATTRIBUTES: + module.__dict__.pop(attribute, None) + + +def _restore_hidiffusion_state(module: torch.nn.Module) -> None: + """Restore the exact per-module state captured before HiDiffusion was applied.""" + snapshot = module.__dict__.pop(_HIDIFFUSION_STATE_SNAPSHOT, None) + if snapshot is None: + _reset_hidiffusion_runtime_state(module) + return + + for attribute, (was_present, value) in snapshot.items(): + if was_present: + setattr(module, attribute, value) + else: + module.__dict__.pop(attribute, None) + + +def apply_hidiffusion( + model: torch.nn.Module, + apply_raunet: bool = True, + apply_window_attn: bool = True, + is_playground=False, + generator: torch.Generator | None = None, +): + """ + model: diffusers model. We support SD 1.5, 2.1, XL, XL Turbo. + + apply_raunet: whether to apply RAU-Net + + apply_window_attn: whether to apply MSW-MSA. + """ + + # Make sure the module is not currently patched + remove_hidiffusion(model) + + is_diffusers = isinstance_str(model, "DiffusionPipeline") or isinstance_str(model, "ModelMixin") + + if not is_diffusers: + raise RuntimeError("Provided model was not a diffusers model/pipeline, as expected.") + else: + # Check if the pipeline is a ControlNet pipeline + is_sdxl_controlnet = hasattr(model, "controlnet") and isinstance_str( + model, "StableDiffusionXLControlNet", prefix=True + ) + is_sd_controlnet = hasattr(model, "controlnet") and isinstance_str( + model, "StableDiffusionControlNet", prefix=True + ) + + # Check for ControlNet Inpaint pipelines + is_sdxl_controlnet_inpaint = is_sdxl_controlnet and isinstance_str(model, "Inpaint", contains=True) + is_sd_controlnet_inpaint = is_sd_controlnet and isinstance_str(model, "Inpaint", contains=True) + + if is_sdxl_controlnet_inpaint or is_sd_controlnet_inpaint: + # For ControlNet Inpaint pipelines, we don't patch the pipeline class + # because they already have all the necessary inpainting logic + # We only patch the UNet for HiDiffusion optimizations + make_block_fn = make_diffusers_unet_2d_condition + model.unet.__class__ = make_block_fn(model.unet.__class__) + elif is_sdxl_controlnet: + make_ppl_fn = make_diffusers_sdxl_controlnet_ppl + model.__class__ = make_ppl_fn(model.__class__) + + make_block_fn = make_diffusers_unet_2d_condition + model.unet.__class__ = make_block_fn(model.unet.__class__) + elif is_sd_controlnet: + # For SD 1.5 ControlNet, we don't need to patch the pipeline class + # Just patch the UNet for consistency + make_block_fn = make_diffusers_unet_2d_condition + model.unet.__class__ = make_block_fn(model.unet.__class__) + + diffusion_model = model.unet if hasattr(model, "unet") else model + + for _, module in diffusion_model.named_modules(): + _snapshot_hidiffusion_state(module) + + # Hack, avoid non-square problem. See unet_2d_condition.py in diffusers + diffusion_model.num_upsamplers += 12 + + name_or_path = getattr(model, "name_or_path", None) or getattr(model, "_name_or_path", "") + diffusion_model_module_key = [] + if name_or_path not in supported_official_model: + for key, _module in diffusion_model.named_modules(): + diffusion_model_module_key.append(key) + if set(sd15_module_key) < set(diffusion_model_module_key): + name_or_path = "runwayml/stable-diffusion-v1-5" + elif set(sdxl_module_key) < set(diffusion_model_module_key): + name_or_path = "stabilityai/stable-diffusion-xl-base-1.0" + + diffusion_model.info = { + "size": None, + "upsample_size": None, + "hooks": [], + "text_to_img_controlnet": hasattr(model, "controlnet"), + "is_inpainting_task": model.__class__ in auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING.values(), + "is_playground": is_playground, + "pipeline": model, + } + model.info = diffusion_model.info + hook_diffusion_model(diffusion_model) + + if name_or_path in ["runwayml/stable-diffusion-v1-5", "stabilityai/stable-diffusion-2-1-base"]: + modified_key = sd15_hidiffusion_key() + for key, module in diffusion_model.named_modules(): + if apply_raunet and key in modified_key["down_module_key"]: + _reset_hidiffusion_runtime_state(module) + make_block_fn = make_diffusers_downsampler_block + module.__class__ = make_block_fn(module.__class__) + module.switching_threshold_ratio = "T1_ratio" + if apply_raunet and key in modified_key["down_module_key_extra"]: + _reset_hidiffusion_runtime_state(module) + make_block_fn = make_diffusers_cross_attn_down_block + module.__class__ = make_block_fn(module.__class__) + module.switching_threshold_ratio = "T2_ratio" + if apply_raunet and key in modified_key["up_module_key"]: + _reset_hidiffusion_runtime_state(module) + make_block_fn = make_diffusers_upsampler_block + module.__class__ = make_block_fn(module.__class__) + module.switching_threshold_ratio = "T1_ratio" + if apply_raunet and key in modified_key["up_module_key_extra"]: + _reset_hidiffusion_runtime_state(module) + make_block_fn = make_diffusers_cross_attn_up_block + module.__class__ = make_block_fn(module.__class__) + module.switching_threshold_ratio = "T2_ratio" + if apply_window_attn and key in modified_key["windown_attn_module_key"]: + make_block_fn = make_diffusers_transformer_block + module.__class__ = make_block_fn(module.__class__, generator) + module.model = "sd15" + module.info = diffusion_model.info + + elif name_or_path in [ + "stabilityai/stable-diffusion-xl-base-1.0", + "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", + ]: + modified_key = sdxl_hidiffusion_key() + for key, module in diffusion_model.named_modules(): + if apply_raunet and key in modified_key["down_module_key"]: + _reset_hidiffusion_runtime_state(module) + make_block_fn = make_diffusers_cross_attn_down_block + module.__class__ = make_block_fn(module.__class__) + module.switching_threshold_ratio = "T1_ratio" + + if apply_raunet and key in modified_key["down_module_key_extra"]: + _reset_hidiffusion_runtime_state(module) + make_block_fn = make_diffusers_downsampler_block + module.__class__ = make_block_fn(module.__class__) + module.switching_threshold_ratio = "T2_ratio" + + if apply_raunet and key in modified_key["up_module_key"]: + _reset_hidiffusion_runtime_state(module) + make_block_fn = make_diffusers_cross_attn_up_block + module.__class__ = make_block_fn(module.__class__) + module.switching_threshold_ratio = "T1_ratio" + + if apply_raunet and key in modified_key["up_module_key_extra"]: + _reset_hidiffusion_runtime_state(module) + make_block_fn = make_diffusers_upsampler_block + module.__class__ = make_block_fn(module.__class__) + module.switching_threshold_ratio = "T2_ratio" + + if apply_window_attn and key in modified_key["windown_attn_module_key"]: + make_block_fn = make_diffusers_transformer_block + module.__class__ = make_block_fn(module.__class__, generator) + module.model = "sdxl" + module.info = diffusion_model.info + + elif name_or_path == "stabilityai/sdxl-turbo": + modified_key = sdxl_turbo_hidiffusion_key() + for key, module in diffusion_model.named_modules(): + if apply_raunet and key in modified_key["down_module_key"]: + _reset_hidiffusion_runtime_state(module) + make_block_fn = make_diffusers_cross_attn_down_block + module.__class__ = make_block_fn(module.__class__) + module.switching_threshold_ratio = "T1_ratio" + + if apply_raunet and key in modified_key["up_module_key"]: + _reset_hidiffusion_runtime_state(module) + make_block_fn = make_diffusers_cross_attn_up_block + module.__class__ = make_block_fn(module.__class__) + module.switching_threshold_ratio = "T1_ratio" + + if apply_window_attn and key in modified_key["windown_attn_module_key"]: + make_block_fn = make_diffusers_transformer_block + module.__class__ = make_block_fn(module.__class__, generator) + + module.model = "sdxl_turbo" + module.info = diffusion_model.info + else: + raise Exception( + f"{name_or_path} is not a supported model. HiDiffusion now only supports runwayml/stable-diffusion-v1-5, stabilityai/stable-diffusion-2-1-base, stabilityai/stable-diffusion-xl-base-1.0, stabilityai/sdxl-turbo, diffusers/stable-diffusion-xl-1.0-inpainting-0.1 and their derivative models/pipelines." + ) + return model + + +def remove_hidiffusion(model: torch.nn.Module): + """Removes hidiffusion from a Diffusion module if it was already patched.""" + # For diffusers + model = model.unet if hasattr(model, "unet") else model + + for _, module in model.named_modules(): + if hasattr(module, "info"): + for hook in module.info["hooks"]: + hook.remove() + module.info["hooks"].clear() + + is_patched = hasattr(module, "_parent") + if _HIDIFFUSION_STATE_SNAPSHOT in module.__dict__: + _restore_hidiffusion_state(module) + elif is_patched: + _reset_hidiffusion_runtime_state(module) + + if is_patched: + module.__class__ = module._parent + + return model diff --git a/invokeai/backend/hidiffusion/sd_module_key/sd15_module_key.txt b/invokeai/backend/hidiffusion/sd_module_key/sd15_module_key.txt new file mode 100644 index 00000000000..3c805da82a0 --- /dev/null +++ b/invokeai/backend/hidiffusion/sd_module_key/sd15_module_key.txt @@ -0,0 +1,685 @@ +conv_in +time_proj +time_embedding +time_embedding.linear_1 +time_embedding.act +time_embedding.linear_2 +down_blocks +down_blocks.0 +down_blocks.0.attentions +down_blocks.0.attentions.0 +down_blocks.0.attentions.0.norm +down_blocks.0.attentions.0.proj_in +down_blocks.0.attentions.0.transformer_blocks +down_blocks.0.attentions.0.transformer_blocks.0 +down_blocks.0.attentions.0.transformer_blocks.0.norm1 +down_blocks.0.attentions.0.transformer_blocks.0.attn1 +down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q +down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_k +down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_v +down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_out +down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_out.0 +down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_out.1 +down_blocks.0.attentions.0.transformer_blocks.0.norm2 +down_blocks.0.attentions.0.transformer_blocks.0.attn2 +down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_q +down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_k +down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_v +down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_out +down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_out.0 +down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_out.1 +down_blocks.0.attentions.0.transformer_blocks.0.norm3 +down_blocks.0.attentions.0.transformer_blocks.0.ff +down_blocks.0.attentions.0.transformer_blocks.0.ff.net +down_blocks.0.attentions.0.transformer_blocks.0.ff.net.0 +down_blocks.0.attentions.0.transformer_blocks.0.ff.net.0.proj +down_blocks.0.attentions.0.transformer_blocks.0.ff.net.1 +down_blocks.0.attentions.0.transformer_blocks.0.ff.net.2 +down_blocks.0.attentions.0.proj_out +down_blocks.0.attentions.1 +down_blocks.0.attentions.1.norm +down_blocks.0.attentions.1.proj_in +down_blocks.0.attentions.1.transformer_blocks +down_blocks.0.attentions.1.transformer_blocks.0 +down_blocks.0.attentions.1.transformer_blocks.0.norm1 +down_blocks.0.attentions.1.transformer_blocks.0.attn1 +down_blocks.0.attentions.1.transformer_blocks.0.attn1.to_q +down_blocks.0.attentions.1.transformer_blocks.0.attn1.to_k +down_blocks.0.attentions.1.transformer_blocks.0.attn1.to_v +down_blocks.0.attentions.1.transformer_blocks.0.attn1.to_out +down_blocks.0.attentions.1.transformer_blocks.0.attn1.to_out.0 +down_blocks.0.attentions.1.transformer_blocks.0.attn1.to_out.1 +down_blocks.0.attentions.1.transformer_blocks.0.norm2 +down_blocks.0.attentions.1.transformer_blocks.0.attn2 +down_blocks.0.attentions.1.transformer_blocks.0.attn2.to_q +down_blocks.0.attentions.1.transformer_blocks.0.attn2.to_k +down_blocks.0.attentions.1.transformer_blocks.0.attn2.to_v +down_blocks.0.attentions.1.transformer_blocks.0.attn2.to_out +down_blocks.0.attentions.1.transformer_blocks.0.attn2.to_out.0 +down_blocks.0.attentions.1.transformer_blocks.0.attn2.to_out.1 +down_blocks.0.attentions.1.transformer_blocks.0.norm3 +down_blocks.0.attentions.1.transformer_blocks.0.ff +down_blocks.0.attentions.1.transformer_blocks.0.ff.net +down_blocks.0.attentions.1.transformer_blocks.0.ff.net.0 +down_blocks.0.attentions.1.transformer_blocks.0.ff.net.0.proj +down_blocks.0.attentions.1.transformer_blocks.0.ff.net.1 +down_blocks.0.attentions.1.transformer_blocks.0.ff.net.2 +down_blocks.0.attentions.1.proj_out +down_blocks.0.resnets +down_blocks.0.resnets.0 +down_blocks.0.resnets.0.norm1 +down_blocks.0.resnets.0.conv1 +down_blocks.0.resnets.0.time_emb_proj +down_blocks.0.resnets.0.norm2 +down_blocks.0.resnets.0.dropout +down_blocks.0.resnets.0.conv2 +down_blocks.0.resnets.1 +down_blocks.0.resnets.1.norm1 +down_blocks.0.resnets.1.conv1 +down_blocks.0.resnets.1.time_emb_proj +down_blocks.0.resnets.1.norm2 +down_blocks.0.resnets.1.dropout +down_blocks.0.resnets.1.conv2 +down_blocks.0.downsamplers +down_blocks.0.downsamplers.0 +down_blocks.0.downsamplers.0.conv +down_blocks.1 +down_blocks.1.attentions +down_blocks.1.attentions.0 +down_blocks.1.attentions.0.norm +down_blocks.1.attentions.0.proj_in +down_blocks.1.attentions.0.transformer_blocks +down_blocks.1.attentions.0.transformer_blocks.0 +down_blocks.1.attentions.0.transformer_blocks.0.norm1 +down_blocks.1.attentions.0.transformer_blocks.0.attn1 +down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_q +down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_k +down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_v +down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out +down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out.0 +down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out.1 +down_blocks.1.attentions.0.transformer_blocks.0.norm2 +down_blocks.1.attentions.0.transformer_blocks.0.attn2 +down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_q +down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k +down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_v +down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out +down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out.0 +down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out.1 +down_blocks.1.attentions.0.transformer_blocks.0.norm3 +down_blocks.1.attentions.0.transformer_blocks.0.ff +down_blocks.1.attentions.0.transformer_blocks.0.ff.net +down_blocks.1.attentions.0.transformer_blocks.0.ff.net.0 +down_blocks.1.attentions.0.transformer_blocks.0.ff.net.0.proj +down_blocks.1.attentions.0.transformer_blocks.0.ff.net.1 +down_blocks.1.attentions.0.transformer_blocks.0.ff.net.2 +down_blocks.1.attentions.0.proj_out +down_blocks.1.attentions.1 +down_blocks.1.attentions.1.norm +down_blocks.1.attentions.1.proj_in +down_blocks.1.attentions.1.transformer_blocks +down_blocks.1.attentions.1.transformer_blocks.0 +down_blocks.1.attentions.1.transformer_blocks.0.norm1 +down_blocks.1.attentions.1.transformer_blocks.0.attn1 +down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_q +down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_k +down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_v +down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out +down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out.0 +down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out.1 +down_blocks.1.attentions.1.transformer_blocks.0.norm2 +down_blocks.1.attentions.1.transformer_blocks.0.attn2 +down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_q +down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_k +down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_v +down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out +down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out.0 +down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out.1 +down_blocks.1.attentions.1.transformer_blocks.0.norm3 +down_blocks.1.attentions.1.transformer_blocks.0.ff +down_blocks.1.attentions.1.transformer_blocks.0.ff.net +down_blocks.1.attentions.1.transformer_blocks.0.ff.net.0 +down_blocks.1.attentions.1.transformer_blocks.0.ff.net.0.proj +down_blocks.1.attentions.1.transformer_blocks.0.ff.net.1 +down_blocks.1.attentions.1.transformer_blocks.0.ff.net.2 +down_blocks.1.attentions.1.proj_out +down_blocks.1.resnets +down_blocks.1.resnets.0 +down_blocks.1.resnets.0.norm1 +down_blocks.1.resnets.0.conv1 +down_blocks.1.resnets.0.time_emb_proj +down_blocks.1.resnets.0.norm2 +down_blocks.1.resnets.0.dropout +down_blocks.1.resnets.0.conv2 +down_blocks.1.resnets.0.conv_shortcut +down_blocks.1.resnets.1 +down_blocks.1.resnets.1.norm1 +down_blocks.1.resnets.1.conv1 +down_blocks.1.resnets.1.time_emb_proj +down_blocks.1.resnets.1.norm2 +down_blocks.1.resnets.1.dropout +down_blocks.1.resnets.1.conv2 +down_blocks.1.downsamplers +down_blocks.1.downsamplers.0 +down_blocks.1.downsamplers.0.conv +down_blocks.2 +down_blocks.2.attentions +down_blocks.2.attentions.0 +down_blocks.2.attentions.0.norm +down_blocks.2.attentions.0.proj_in +down_blocks.2.attentions.0.transformer_blocks +down_blocks.2.attentions.0.transformer_blocks.0 +down_blocks.2.attentions.0.transformer_blocks.0.norm1 +down_blocks.2.attentions.0.transformer_blocks.0.attn1 +down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_q +down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_k +down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_v +down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_out +down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.0.norm2 +down_blocks.2.attentions.0.transformer_blocks.0.attn2 +down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_q +down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_k +down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_v +down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_out +down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.0.norm3 +down_blocks.2.attentions.0.transformer_blocks.0.ff +down_blocks.2.attentions.0.transformer_blocks.0.ff.net +down_blocks.2.attentions.0.transformer_blocks.0.ff.net.0 +down_blocks.2.attentions.0.transformer_blocks.0.ff.net.0.proj +down_blocks.2.attentions.0.transformer_blocks.0.ff.net.1 +down_blocks.2.attentions.0.transformer_blocks.0.ff.net.2 +down_blocks.2.attentions.0.proj_out +down_blocks.2.attentions.1 +down_blocks.2.attentions.1.norm +down_blocks.2.attentions.1.proj_in +down_blocks.2.attentions.1.transformer_blocks +down_blocks.2.attentions.1.transformer_blocks.0 +down_blocks.2.attentions.1.transformer_blocks.0.norm1 +down_blocks.2.attentions.1.transformer_blocks.0.attn1 +down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_q +down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_k +down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_v +down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_out +down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.0.norm2 +down_blocks.2.attentions.1.transformer_blocks.0.attn2 +down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_q +down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_k +down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_v +down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_out +down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.0.norm3 +down_blocks.2.attentions.1.transformer_blocks.0.ff +down_blocks.2.attentions.1.transformer_blocks.0.ff.net +down_blocks.2.attentions.1.transformer_blocks.0.ff.net.0 +down_blocks.2.attentions.1.transformer_blocks.0.ff.net.0.proj +down_blocks.2.attentions.1.transformer_blocks.0.ff.net.1 +down_blocks.2.attentions.1.transformer_blocks.0.ff.net.2 +down_blocks.2.attentions.1.proj_out +down_blocks.2.resnets +down_blocks.2.resnets.0 +down_blocks.2.resnets.0.norm1 +down_blocks.2.resnets.0.conv1 +down_blocks.2.resnets.0.time_emb_proj +down_blocks.2.resnets.0.norm2 +down_blocks.2.resnets.0.dropout +down_blocks.2.resnets.0.conv2 +down_blocks.2.resnets.0.conv_shortcut +down_blocks.2.resnets.1 +down_blocks.2.resnets.1.norm1 +down_blocks.2.resnets.1.conv1 +down_blocks.2.resnets.1.time_emb_proj +down_blocks.2.resnets.1.norm2 +down_blocks.2.resnets.1.dropout +down_blocks.2.resnets.1.conv2 +down_blocks.2.downsamplers +down_blocks.2.downsamplers.0 +down_blocks.2.downsamplers.0.conv +down_blocks.3 +down_blocks.3.resnets +down_blocks.3.resnets.0 +down_blocks.3.resnets.0.norm1 +down_blocks.3.resnets.0.conv1 +down_blocks.3.resnets.0.time_emb_proj +down_blocks.3.resnets.0.norm2 +down_blocks.3.resnets.0.dropout +down_blocks.3.resnets.0.conv2 +down_blocks.3.resnets.1 +down_blocks.3.resnets.1.norm1 +down_blocks.3.resnets.1.conv1 +down_blocks.3.resnets.1.time_emb_proj +down_blocks.3.resnets.1.norm2 +down_blocks.3.resnets.1.dropout +down_blocks.3.resnets.1.conv2 +up_blocks +up_blocks.0 +up_blocks.0.resnets +up_blocks.0.resnets.0 +up_blocks.0.resnets.0.norm1 +up_blocks.0.resnets.0.conv1 +up_blocks.0.resnets.0.time_emb_proj +up_blocks.0.resnets.0.norm2 +up_blocks.0.resnets.0.dropout +up_blocks.0.resnets.0.conv2 +up_blocks.0.resnets.0.conv_shortcut +up_blocks.0.resnets.1 +up_blocks.0.resnets.1.norm1 +up_blocks.0.resnets.1.conv1 +up_blocks.0.resnets.1.time_emb_proj +up_blocks.0.resnets.1.norm2 +up_blocks.0.resnets.1.dropout +up_blocks.0.resnets.1.conv2 +up_blocks.0.resnets.1.conv_shortcut +up_blocks.0.resnets.2 +up_blocks.0.resnets.2.norm1 +up_blocks.0.resnets.2.conv1 +up_blocks.0.resnets.2.time_emb_proj +up_blocks.0.resnets.2.norm2 +up_blocks.0.resnets.2.dropout +up_blocks.0.resnets.2.conv2 +up_blocks.0.resnets.2.conv_shortcut +up_blocks.0.upsamplers +up_blocks.0.upsamplers.0 +up_blocks.0.upsamplers.0.conv +up_blocks.1 +up_blocks.1.attentions +up_blocks.1.attentions.0 +up_blocks.1.attentions.0.norm +up_blocks.1.attentions.0.proj_in +up_blocks.1.attentions.0.transformer_blocks +up_blocks.1.attentions.0.transformer_blocks.0 +up_blocks.1.attentions.0.transformer_blocks.0.norm1 +up_blocks.1.attentions.0.transformer_blocks.0.attn1 +up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_q +up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_k +up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_v +up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out +up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out.0 +up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out.1 +up_blocks.1.attentions.0.transformer_blocks.0.norm2 +up_blocks.1.attentions.0.transformer_blocks.0.attn2 +up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_q +up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k +up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_v +up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out +up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out.0 +up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out.1 +up_blocks.1.attentions.0.transformer_blocks.0.norm3 +up_blocks.1.attentions.0.transformer_blocks.0.ff +up_blocks.1.attentions.0.transformer_blocks.0.ff.net +up_blocks.1.attentions.0.transformer_blocks.0.ff.net.0 +up_blocks.1.attentions.0.transformer_blocks.0.ff.net.0.proj +up_blocks.1.attentions.0.transformer_blocks.0.ff.net.1 +up_blocks.1.attentions.0.transformer_blocks.0.ff.net.2 +up_blocks.1.attentions.0.proj_out +up_blocks.1.attentions.1 +up_blocks.1.attentions.1.norm +up_blocks.1.attentions.1.proj_in +up_blocks.1.attentions.1.transformer_blocks +up_blocks.1.attentions.1.transformer_blocks.0 +up_blocks.1.attentions.1.transformer_blocks.0.norm1 +up_blocks.1.attentions.1.transformer_blocks.0.attn1 +up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_q +up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_k +up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_v +up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out +up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out.0 +up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out.1 +up_blocks.1.attentions.1.transformer_blocks.0.norm2 +up_blocks.1.attentions.1.transformer_blocks.0.attn2 +up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_q +up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_k +up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_v +up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out +up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out.0 +up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out.1 +up_blocks.1.attentions.1.transformer_blocks.0.norm3 +up_blocks.1.attentions.1.transformer_blocks.0.ff +up_blocks.1.attentions.1.transformer_blocks.0.ff.net +up_blocks.1.attentions.1.transformer_blocks.0.ff.net.0 +up_blocks.1.attentions.1.transformer_blocks.0.ff.net.0.proj +up_blocks.1.attentions.1.transformer_blocks.0.ff.net.1 +up_blocks.1.attentions.1.transformer_blocks.0.ff.net.2 +up_blocks.1.attentions.1.proj_out +up_blocks.1.attentions.2 +up_blocks.1.attentions.2.norm +up_blocks.1.attentions.2.proj_in +up_blocks.1.attentions.2.transformer_blocks +up_blocks.1.attentions.2.transformer_blocks.0 +up_blocks.1.attentions.2.transformer_blocks.0.norm1 +up_blocks.1.attentions.2.transformer_blocks.0.attn1 +up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_q +up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_k +up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_v +up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_out +up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_out.0 +up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_out.1 +up_blocks.1.attentions.2.transformer_blocks.0.norm2 +up_blocks.1.attentions.2.transformer_blocks.0.attn2 +up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_q +up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_k +up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_v +up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_out +up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_out.0 +up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_out.1 +up_blocks.1.attentions.2.transformer_blocks.0.norm3 +up_blocks.1.attentions.2.transformer_blocks.0.ff +up_blocks.1.attentions.2.transformer_blocks.0.ff.net +up_blocks.1.attentions.2.transformer_blocks.0.ff.net.0 +up_blocks.1.attentions.2.transformer_blocks.0.ff.net.0.proj +up_blocks.1.attentions.2.transformer_blocks.0.ff.net.1 +up_blocks.1.attentions.2.transformer_blocks.0.ff.net.2 +up_blocks.1.attentions.2.proj_out +up_blocks.1.resnets +up_blocks.1.resnets.0 +up_blocks.1.resnets.0.norm1 +up_blocks.1.resnets.0.conv1 +up_blocks.1.resnets.0.time_emb_proj +up_blocks.1.resnets.0.norm2 +up_blocks.1.resnets.0.dropout +up_blocks.1.resnets.0.conv2 +up_blocks.1.resnets.0.conv_shortcut +up_blocks.1.resnets.1 +up_blocks.1.resnets.1.norm1 +up_blocks.1.resnets.1.conv1 +up_blocks.1.resnets.1.time_emb_proj +up_blocks.1.resnets.1.norm2 +up_blocks.1.resnets.1.dropout +up_blocks.1.resnets.1.conv2 +up_blocks.1.resnets.1.conv_shortcut +up_blocks.1.resnets.2 +up_blocks.1.resnets.2.norm1 +up_blocks.1.resnets.2.conv1 +up_blocks.1.resnets.2.time_emb_proj +up_blocks.1.resnets.2.norm2 +up_blocks.1.resnets.2.dropout +up_blocks.1.resnets.2.conv2 +up_blocks.1.resnets.2.conv_shortcut +up_blocks.1.upsamplers +up_blocks.1.upsamplers.0 +up_blocks.1.upsamplers.0.conv +up_blocks.2 +up_blocks.2.attentions +up_blocks.2.attentions.0 +up_blocks.2.attentions.0.norm +up_blocks.2.attentions.0.proj_in +up_blocks.2.attentions.0.transformer_blocks +up_blocks.2.attentions.0.transformer_blocks.0 +up_blocks.2.attentions.0.transformer_blocks.0.norm1 +up_blocks.2.attentions.0.transformer_blocks.0.attn1 +up_blocks.2.attentions.0.transformer_blocks.0.attn1.to_q +up_blocks.2.attentions.0.transformer_blocks.0.attn1.to_k +up_blocks.2.attentions.0.transformer_blocks.0.attn1.to_v +up_blocks.2.attentions.0.transformer_blocks.0.attn1.to_out +up_blocks.2.attentions.0.transformer_blocks.0.attn1.to_out.0 +up_blocks.2.attentions.0.transformer_blocks.0.attn1.to_out.1 +up_blocks.2.attentions.0.transformer_blocks.0.norm2 +up_blocks.2.attentions.0.transformer_blocks.0.attn2 +up_blocks.2.attentions.0.transformer_blocks.0.attn2.to_q +up_blocks.2.attentions.0.transformer_blocks.0.attn2.to_k +up_blocks.2.attentions.0.transformer_blocks.0.attn2.to_v +up_blocks.2.attentions.0.transformer_blocks.0.attn2.to_out +up_blocks.2.attentions.0.transformer_blocks.0.attn2.to_out.0 +up_blocks.2.attentions.0.transformer_blocks.0.attn2.to_out.1 +up_blocks.2.attentions.0.transformer_blocks.0.norm3 +up_blocks.2.attentions.0.transformer_blocks.0.ff +up_blocks.2.attentions.0.transformer_blocks.0.ff.net +up_blocks.2.attentions.0.transformer_blocks.0.ff.net.0 +up_blocks.2.attentions.0.transformer_blocks.0.ff.net.0.proj +up_blocks.2.attentions.0.transformer_blocks.0.ff.net.1 +up_blocks.2.attentions.0.transformer_blocks.0.ff.net.2 +up_blocks.2.attentions.0.proj_out +up_blocks.2.attentions.1 +up_blocks.2.attentions.1.norm +up_blocks.2.attentions.1.proj_in +up_blocks.2.attentions.1.transformer_blocks +up_blocks.2.attentions.1.transformer_blocks.0 +up_blocks.2.attentions.1.transformer_blocks.0.norm1 +up_blocks.2.attentions.1.transformer_blocks.0.attn1 +up_blocks.2.attentions.1.transformer_blocks.0.attn1.to_q +up_blocks.2.attentions.1.transformer_blocks.0.attn1.to_k +up_blocks.2.attentions.1.transformer_blocks.0.attn1.to_v +up_blocks.2.attentions.1.transformer_blocks.0.attn1.to_out +up_blocks.2.attentions.1.transformer_blocks.0.attn1.to_out.0 +up_blocks.2.attentions.1.transformer_blocks.0.attn1.to_out.1 +up_blocks.2.attentions.1.transformer_blocks.0.norm2 +up_blocks.2.attentions.1.transformer_blocks.0.attn2 +up_blocks.2.attentions.1.transformer_blocks.0.attn2.to_q +up_blocks.2.attentions.1.transformer_blocks.0.attn2.to_k +up_blocks.2.attentions.1.transformer_blocks.0.attn2.to_v +up_blocks.2.attentions.1.transformer_blocks.0.attn2.to_out +up_blocks.2.attentions.1.transformer_blocks.0.attn2.to_out.0 +up_blocks.2.attentions.1.transformer_blocks.0.attn2.to_out.1 +up_blocks.2.attentions.1.transformer_blocks.0.norm3 +up_blocks.2.attentions.1.transformer_blocks.0.ff +up_blocks.2.attentions.1.transformer_blocks.0.ff.net +up_blocks.2.attentions.1.transformer_blocks.0.ff.net.0 +up_blocks.2.attentions.1.transformer_blocks.0.ff.net.0.proj +up_blocks.2.attentions.1.transformer_blocks.0.ff.net.1 +up_blocks.2.attentions.1.transformer_blocks.0.ff.net.2 +up_blocks.2.attentions.1.proj_out +up_blocks.2.attentions.2 +up_blocks.2.attentions.2.norm +up_blocks.2.attentions.2.proj_in +up_blocks.2.attentions.2.transformer_blocks +up_blocks.2.attentions.2.transformer_blocks.0 +up_blocks.2.attentions.2.transformer_blocks.0.norm1 +up_blocks.2.attentions.2.transformer_blocks.0.attn1 +up_blocks.2.attentions.2.transformer_blocks.0.attn1.to_q +up_blocks.2.attentions.2.transformer_blocks.0.attn1.to_k +up_blocks.2.attentions.2.transformer_blocks.0.attn1.to_v +up_blocks.2.attentions.2.transformer_blocks.0.attn1.to_out +up_blocks.2.attentions.2.transformer_blocks.0.attn1.to_out.0 +up_blocks.2.attentions.2.transformer_blocks.0.attn1.to_out.1 +up_blocks.2.attentions.2.transformer_blocks.0.norm2 +up_blocks.2.attentions.2.transformer_blocks.0.attn2 +up_blocks.2.attentions.2.transformer_blocks.0.attn2.to_q +up_blocks.2.attentions.2.transformer_blocks.0.attn2.to_k +up_blocks.2.attentions.2.transformer_blocks.0.attn2.to_v +up_blocks.2.attentions.2.transformer_blocks.0.attn2.to_out +up_blocks.2.attentions.2.transformer_blocks.0.attn2.to_out.0 +up_blocks.2.attentions.2.transformer_blocks.0.attn2.to_out.1 +up_blocks.2.attentions.2.transformer_blocks.0.norm3 +up_blocks.2.attentions.2.transformer_blocks.0.ff +up_blocks.2.attentions.2.transformer_blocks.0.ff.net +up_blocks.2.attentions.2.transformer_blocks.0.ff.net.0 +up_blocks.2.attentions.2.transformer_blocks.0.ff.net.0.proj +up_blocks.2.attentions.2.transformer_blocks.0.ff.net.1 +up_blocks.2.attentions.2.transformer_blocks.0.ff.net.2 +up_blocks.2.attentions.2.proj_out +up_blocks.2.resnets +up_blocks.2.resnets.0 +up_blocks.2.resnets.0.norm1 +up_blocks.2.resnets.0.conv1 +up_blocks.2.resnets.0.time_emb_proj +up_blocks.2.resnets.0.norm2 +up_blocks.2.resnets.0.dropout +up_blocks.2.resnets.0.conv2 +up_blocks.2.resnets.0.conv_shortcut +up_blocks.2.resnets.1 +up_blocks.2.resnets.1.norm1 +up_blocks.2.resnets.1.conv1 +up_blocks.2.resnets.1.time_emb_proj +up_blocks.2.resnets.1.norm2 +up_blocks.2.resnets.1.dropout +up_blocks.2.resnets.1.conv2 +up_blocks.2.resnets.1.conv_shortcut +up_blocks.2.resnets.2 +up_blocks.2.resnets.2.norm1 +up_blocks.2.resnets.2.conv1 +up_blocks.2.resnets.2.time_emb_proj +up_blocks.2.resnets.2.norm2 +up_blocks.2.resnets.2.dropout +up_blocks.2.resnets.2.conv2 +up_blocks.2.resnets.2.conv_shortcut +up_blocks.2.upsamplers +up_blocks.2.upsamplers.0 +up_blocks.2.upsamplers.0.conv +up_blocks.3 +up_blocks.3.attentions +up_blocks.3.attentions.0 +up_blocks.3.attentions.0.norm +up_blocks.3.attentions.0.proj_in +up_blocks.3.attentions.0.transformer_blocks +up_blocks.3.attentions.0.transformer_blocks.0 +up_blocks.3.attentions.0.transformer_blocks.0.norm1 +up_blocks.3.attentions.0.transformer_blocks.0.attn1 +up_blocks.3.attentions.0.transformer_blocks.0.attn1.to_q +up_blocks.3.attentions.0.transformer_blocks.0.attn1.to_k +up_blocks.3.attentions.0.transformer_blocks.0.attn1.to_v +up_blocks.3.attentions.0.transformer_blocks.0.attn1.to_out +up_blocks.3.attentions.0.transformer_blocks.0.attn1.to_out.0 +up_blocks.3.attentions.0.transformer_blocks.0.attn1.to_out.1 +up_blocks.3.attentions.0.transformer_blocks.0.norm2 +up_blocks.3.attentions.0.transformer_blocks.0.attn2 +up_blocks.3.attentions.0.transformer_blocks.0.attn2.to_q +up_blocks.3.attentions.0.transformer_blocks.0.attn2.to_k +up_blocks.3.attentions.0.transformer_blocks.0.attn2.to_v +up_blocks.3.attentions.0.transformer_blocks.0.attn2.to_out +up_blocks.3.attentions.0.transformer_blocks.0.attn2.to_out.0 +up_blocks.3.attentions.0.transformer_blocks.0.attn2.to_out.1 +up_blocks.3.attentions.0.transformer_blocks.0.norm3 +up_blocks.3.attentions.0.transformer_blocks.0.ff +up_blocks.3.attentions.0.transformer_blocks.0.ff.net +up_blocks.3.attentions.0.transformer_blocks.0.ff.net.0 +up_blocks.3.attentions.0.transformer_blocks.0.ff.net.0.proj +up_blocks.3.attentions.0.transformer_blocks.0.ff.net.1 +up_blocks.3.attentions.0.transformer_blocks.0.ff.net.2 +up_blocks.3.attentions.0.proj_out +up_blocks.3.attentions.1 +up_blocks.3.attentions.1.norm +up_blocks.3.attentions.1.proj_in +up_blocks.3.attentions.1.transformer_blocks +up_blocks.3.attentions.1.transformer_blocks.0 +up_blocks.3.attentions.1.transformer_blocks.0.norm1 +up_blocks.3.attentions.1.transformer_blocks.0.attn1 +up_blocks.3.attentions.1.transformer_blocks.0.attn1.to_q +up_blocks.3.attentions.1.transformer_blocks.0.attn1.to_k +up_blocks.3.attentions.1.transformer_blocks.0.attn1.to_v +up_blocks.3.attentions.1.transformer_blocks.0.attn1.to_out +up_blocks.3.attentions.1.transformer_blocks.0.attn1.to_out.0 +up_blocks.3.attentions.1.transformer_blocks.0.attn1.to_out.1 +up_blocks.3.attentions.1.transformer_blocks.0.norm2 +up_blocks.3.attentions.1.transformer_blocks.0.attn2 +up_blocks.3.attentions.1.transformer_blocks.0.attn2.to_q +up_blocks.3.attentions.1.transformer_blocks.0.attn2.to_k +up_blocks.3.attentions.1.transformer_blocks.0.attn2.to_v +up_blocks.3.attentions.1.transformer_blocks.0.attn2.to_out +up_blocks.3.attentions.1.transformer_blocks.0.attn2.to_out.0 +up_blocks.3.attentions.1.transformer_blocks.0.attn2.to_out.1 +up_blocks.3.attentions.1.transformer_blocks.0.norm3 +up_blocks.3.attentions.1.transformer_blocks.0.ff +up_blocks.3.attentions.1.transformer_blocks.0.ff.net +up_blocks.3.attentions.1.transformer_blocks.0.ff.net.0 +up_blocks.3.attentions.1.transformer_blocks.0.ff.net.0.proj +up_blocks.3.attentions.1.transformer_blocks.0.ff.net.1 +up_blocks.3.attentions.1.transformer_blocks.0.ff.net.2 +up_blocks.3.attentions.1.proj_out +up_blocks.3.attentions.2 +up_blocks.3.attentions.2.norm +up_blocks.3.attentions.2.proj_in +up_blocks.3.attentions.2.transformer_blocks +up_blocks.3.attentions.2.transformer_blocks.0 +up_blocks.3.attentions.2.transformer_blocks.0.norm1 +up_blocks.3.attentions.2.transformer_blocks.0.attn1 +up_blocks.3.attentions.2.transformer_blocks.0.attn1.to_q +up_blocks.3.attentions.2.transformer_blocks.0.attn1.to_k +up_blocks.3.attentions.2.transformer_blocks.0.attn1.to_v +up_blocks.3.attentions.2.transformer_blocks.0.attn1.to_out +up_blocks.3.attentions.2.transformer_blocks.0.attn1.to_out.0 +up_blocks.3.attentions.2.transformer_blocks.0.attn1.to_out.1 +up_blocks.3.attentions.2.transformer_blocks.0.norm2 +up_blocks.3.attentions.2.transformer_blocks.0.attn2 +up_blocks.3.attentions.2.transformer_blocks.0.attn2.to_q +up_blocks.3.attentions.2.transformer_blocks.0.attn2.to_k +up_blocks.3.attentions.2.transformer_blocks.0.attn2.to_v +up_blocks.3.attentions.2.transformer_blocks.0.attn2.to_out +up_blocks.3.attentions.2.transformer_blocks.0.attn2.to_out.0 +up_blocks.3.attentions.2.transformer_blocks.0.attn2.to_out.1 +up_blocks.3.attentions.2.transformer_blocks.0.norm3 +up_blocks.3.attentions.2.transformer_blocks.0.ff +up_blocks.3.attentions.2.transformer_blocks.0.ff.net +up_blocks.3.attentions.2.transformer_blocks.0.ff.net.0 +up_blocks.3.attentions.2.transformer_blocks.0.ff.net.0.proj +up_blocks.3.attentions.2.transformer_blocks.0.ff.net.1 +up_blocks.3.attentions.2.transformer_blocks.0.ff.net.2 +up_blocks.3.attentions.2.proj_out +up_blocks.3.resnets +up_blocks.3.resnets.0 +up_blocks.3.resnets.0.norm1 +up_blocks.3.resnets.0.conv1 +up_blocks.3.resnets.0.time_emb_proj +up_blocks.3.resnets.0.norm2 +up_blocks.3.resnets.0.dropout +up_blocks.3.resnets.0.conv2 +up_blocks.3.resnets.0.conv_shortcut +up_blocks.3.resnets.1 +up_blocks.3.resnets.1.norm1 +up_blocks.3.resnets.1.conv1 +up_blocks.3.resnets.1.time_emb_proj +up_blocks.3.resnets.1.norm2 +up_blocks.3.resnets.1.dropout +up_blocks.3.resnets.1.conv2 +up_blocks.3.resnets.1.conv_shortcut +up_blocks.3.resnets.2 +up_blocks.3.resnets.2.norm1 +up_blocks.3.resnets.2.conv1 +up_blocks.3.resnets.2.time_emb_proj +up_blocks.3.resnets.2.norm2 +up_blocks.3.resnets.2.dropout +up_blocks.3.resnets.2.conv2 +up_blocks.3.resnets.2.conv_shortcut +mid_block +mid_block.attentions +mid_block.attentions.0 +mid_block.attentions.0.norm +mid_block.attentions.0.proj_in +mid_block.attentions.0.transformer_blocks +mid_block.attentions.0.transformer_blocks.0 +mid_block.attentions.0.transformer_blocks.0.norm1 +mid_block.attentions.0.transformer_blocks.0.attn1 +mid_block.attentions.0.transformer_blocks.0.attn1.to_q +mid_block.attentions.0.transformer_blocks.0.attn1.to_k +mid_block.attentions.0.transformer_blocks.0.attn1.to_v +mid_block.attentions.0.transformer_blocks.0.attn1.to_out +mid_block.attentions.0.transformer_blocks.0.attn1.to_out.0 +mid_block.attentions.0.transformer_blocks.0.attn1.to_out.1 +mid_block.attentions.0.transformer_blocks.0.norm2 +mid_block.attentions.0.transformer_blocks.0.attn2 +mid_block.attentions.0.transformer_blocks.0.attn2.to_q +mid_block.attentions.0.transformer_blocks.0.attn2.to_k +mid_block.attentions.0.transformer_blocks.0.attn2.to_v +mid_block.attentions.0.transformer_blocks.0.attn2.to_out +mid_block.attentions.0.transformer_blocks.0.attn2.to_out.0 +mid_block.attentions.0.transformer_blocks.0.attn2.to_out.1 +mid_block.attentions.0.transformer_blocks.0.norm3 +mid_block.attentions.0.transformer_blocks.0.ff +mid_block.attentions.0.transformer_blocks.0.ff.net +mid_block.attentions.0.transformer_blocks.0.ff.net.0 +mid_block.attentions.0.transformer_blocks.0.ff.net.0.proj +mid_block.attentions.0.transformer_blocks.0.ff.net.1 +mid_block.attentions.0.transformer_blocks.0.ff.net.2 +mid_block.attentions.0.proj_out +mid_block.resnets +mid_block.resnets.0 +mid_block.resnets.0.norm1 +mid_block.resnets.0.conv1 +mid_block.resnets.0.time_emb_proj +mid_block.resnets.0.norm2 +mid_block.resnets.0.dropout +mid_block.resnets.0.conv2 +mid_block.resnets.1 +mid_block.resnets.1.norm1 +mid_block.resnets.1.conv1 +mid_block.resnets.1.time_emb_proj +mid_block.resnets.1.norm2 +mid_block.resnets.1.dropout +mid_block.resnets.1.conv2 +conv_norm_out +conv_out diff --git a/invokeai/backend/hidiffusion/sd_module_key/sdxl_module_key.txt b/invokeai/backend/hidiffusion/sd_module_key/sdxl_module_key.txt new file mode 100644 index 00000000000..075d480dee8 --- /dev/null +++ b/invokeai/backend/hidiffusion/sd_module_key/sdxl_module_key.txt @@ -0,0 +1,1910 @@ +conv_in +time_proj +time_embedding +time_embedding.linear_1 +time_embedding.act +time_embedding.linear_2 +add_time_proj +add_embedding +add_embedding.linear_1 +add_embedding.linear_2 +down_blocks +down_blocks.0 +down_blocks.0.resnets +down_blocks.0.resnets.0 +down_blocks.0.resnets.0.norm1 +down_blocks.0.resnets.0.conv1 +down_blocks.0.resnets.0.time_emb_proj +down_blocks.0.resnets.0.norm2 +down_blocks.0.resnets.0.dropout +down_blocks.0.resnets.0.conv2 +down_blocks.0.resnets.1 +down_blocks.0.resnets.1.norm1 +down_blocks.0.resnets.1.conv1 +down_blocks.0.resnets.1.time_emb_proj +down_blocks.0.resnets.1.norm2 +down_blocks.0.resnets.1.dropout +down_blocks.0.resnets.1.conv2 +down_blocks.0.downsamplers +down_blocks.0.downsamplers.0 +down_blocks.0.downsamplers.0.conv +down_blocks.1 +down_blocks.1.attentions +down_blocks.1.attentions.0 +down_blocks.1.attentions.0.norm +down_blocks.1.attentions.0.proj_in +down_blocks.1.attentions.0.transformer_blocks +down_blocks.1.attentions.0.transformer_blocks.0 +down_blocks.1.attentions.0.transformer_blocks.0.norm1 +down_blocks.1.attentions.0.transformer_blocks.0.attn1 +down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_q +down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_k +down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_v +down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out +down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out.0 +down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out.1 +down_blocks.1.attentions.0.transformer_blocks.0.norm2 +down_blocks.1.attentions.0.transformer_blocks.0.attn2 +down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_q +down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k +down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_v +down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out +down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out.0 +down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out.1 +down_blocks.1.attentions.0.transformer_blocks.0.norm3 +down_blocks.1.attentions.0.transformer_blocks.0.ff +down_blocks.1.attentions.0.transformer_blocks.0.ff.net +down_blocks.1.attentions.0.transformer_blocks.0.ff.net.0 +down_blocks.1.attentions.0.transformer_blocks.0.ff.net.0.proj +down_blocks.1.attentions.0.transformer_blocks.0.ff.net.1 +down_blocks.1.attentions.0.transformer_blocks.0.ff.net.2 +down_blocks.1.attentions.0.transformer_blocks.1 +down_blocks.1.attentions.0.transformer_blocks.1.norm1 +down_blocks.1.attentions.0.transformer_blocks.1.attn1 +down_blocks.1.attentions.0.transformer_blocks.1.attn1.to_q +down_blocks.1.attentions.0.transformer_blocks.1.attn1.to_k +down_blocks.1.attentions.0.transformer_blocks.1.attn1.to_v +down_blocks.1.attentions.0.transformer_blocks.1.attn1.to_out +down_blocks.1.attentions.0.transformer_blocks.1.attn1.to_out.0 +down_blocks.1.attentions.0.transformer_blocks.1.attn1.to_out.1 +down_blocks.1.attentions.0.transformer_blocks.1.norm2 +down_blocks.1.attentions.0.transformer_blocks.1.attn2 +down_blocks.1.attentions.0.transformer_blocks.1.attn2.to_q +down_blocks.1.attentions.0.transformer_blocks.1.attn2.to_k +down_blocks.1.attentions.0.transformer_blocks.1.attn2.to_v +down_blocks.1.attentions.0.transformer_blocks.1.attn2.to_out +down_blocks.1.attentions.0.transformer_blocks.1.attn2.to_out.0 +down_blocks.1.attentions.0.transformer_blocks.1.attn2.to_out.1 +down_blocks.1.attentions.0.transformer_blocks.1.norm3 +down_blocks.1.attentions.0.transformer_blocks.1.ff +down_blocks.1.attentions.0.transformer_blocks.1.ff.net +down_blocks.1.attentions.0.transformer_blocks.1.ff.net.0 +down_blocks.1.attentions.0.transformer_blocks.1.ff.net.0.proj +down_blocks.1.attentions.0.transformer_blocks.1.ff.net.1 +down_blocks.1.attentions.0.transformer_blocks.1.ff.net.2 +down_blocks.1.attentions.0.proj_out +down_blocks.1.attentions.1 +down_blocks.1.attentions.1.norm +down_blocks.1.attentions.1.proj_in +down_blocks.1.attentions.1.transformer_blocks +down_blocks.1.attentions.1.transformer_blocks.0 +down_blocks.1.attentions.1.transformer_blocks.0.norm1 +down_blocks.1.attentions.1.transformer_blocks.0.attn1 +down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_q +down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_k +down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_v +down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out +down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out.0 +down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out.1 +down_blocks.1.attentions.1.transformer_blocks.0.norm2 +down_blocks.1.attentions.1.transformer_blocks.0.attn2 +down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_q +down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_k +down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_v +down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out +down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out.0 +down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out.1 +down_blocks.1.attentions.1.transformer_blocks.0.norm3 +down_blocks.1.attentions.1.transformer_blocks.0.ff +down_blocks.1.attentions.1.transformer_blocks.0.ff.net +down_blocks.1.attentions.1.transformer_blocks.0.ff.net.0 +down_blocks.1.attentions.1.transformer_blocks.0.ff.net.0.proj +down_blocks.1.attentions.1.transformer_blocks.0.ff.net.1 +down_blocks.1.attentions.1.transformer_blocks.0.ff.net.2 +down_blocks.1.attentions.1.transformer_blocks.1 +down_blocks.1.attentions.1.transformer_blocks.1.norm1 +down_blocks.1.attentions.1.transformer_blocks.1.attn1 +down_blocks.1.attentions.1.transformer_blocks.1.attn1.to_q +down_blocks.1.attentions.1.transformer_blocks.1.attn1.to_k +down_blocks.1.attentions.1.transformer_blocks.1.attn1.to_v +down_blocks.1.attentions.1.transformer_blocks.1.attn1.to_out +down_blocks.1.attentions.1.transformer_blocks.1.attn1.to_out.0 +down_blocks.1.attentions.1.transformer_blocks.1.attn1.to_out.1 +down_blocks.1.attentions.1.transformer_blocks.1.norm2 +down_blocks.1.attentions.1.transformer_blocks.1.attn2 +down_blocks.1.attentions.1.transformer_blocks.1.attn2.to_q +down_blocks.1.attentions.1.transformer_blocks.1.attn2.to_k +down_blocks.1.attentions.1.transformer_blocks.1.attn2.to_v +down_blocks.1.attentions.1.transformer_blocks.1.attn2.to_out +down_blocks.1.attentions.1.transformer_blocks.1.attn2.to_out.0 +down_blocks.1.attentions.1.transformer_blocks.1.attn2.to_out.1 +down_blocks.1.attentions.1.transformer_blocks.1.norm3 +down_blocks.1.attentions.1.transformer_blocks.1.ff +down_blocks.1.attentions.1.transformer_blocks.1.ff.net +down_blocks.1.attentions.1.transformer_blocks.1.ff.net.0 +down_blocks.1.attentions.1.transformer_blocks.1.ff.net.0.proj +down_blocks.1.attentions.1.transformer_blocks.1.ff.net.1 +down_blocks.1.attentions.1.transformer_blocks.1.ff.net.2 +down_blocks.1.attentions.1.proj_out +down_blocks.1.resnets +down_blocks.1.resnets.0 +down_blocks.1.resnets.0.norm1 +down_blocks.1.resnets.0.conv1 +down_blocks.1.resnets.0.time_emb_proj +down_blocks.1.resnets.0.norm2 +down_blocks.1.resnets.0.dropout +down_blocks.1.resnets.0.conv2 +down_blocks.1.resnets.0.conv_shortcut +down_blocks.1.resnets.1 +down_blocks.1.resnets.1.norm1 +down_blocks.1.resnets.1.conv1 +down_blocks.1.resnets.1.time_emb_proj +down_blocks.1.resnets.1.norm2 +down_blocks.1.resnets.1.dropout +down_blocks.1.resnets.1.conv2 +down_blocks.1.downsamplers +down_blocks.1.downsamplers.0 +down_blocks.1.downsamplers.0.conv +down_blocks.2 +down_blocks.2.attentions +down_blocks.2.attentions.0 +down_blocks.2.attentions.0.norm +down_blocks.2.attentions.0.proj_in +down_blocks.2.attentions.0.transformer_blocks +down_blocks.2.attentions.0.transformer_blocks.0 +down_blocks.2.attentions.0.transformer_blocks.0.norm1 +down_blocks.2.attentions.0.transformer_blocks.0.attn1 +down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_q +down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_k +down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_v +down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_out +down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.0.norm2 +down_blocks.2.attentions.0.transformer_blocks.0.attn2 +down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_q +down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_k +down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_v +down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_out +down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.0.norm3 +down_blocks.2.attentions.0.transformer_blocks.0.ff +down_blocks.2.attentions.0.transformer_blocks.0.ff.net +down_blocks.2.attentions.0.transformer_blocks.0.ff.net.0 +down_blocks.2.attentions.0.transformer_blocks.0.ff.net.0.proj +down_blocks.2.attentions.0.transformer_blocks.0.ff.net.1 +down_blocks.2.attentions.0.transformer_blocks.0.ff.net.2 +down_blocks.2.attentions.0.transformer_blocks.1 +down_blocks.2.attentions.0.transformer_blocks.1.norm1 +down_blocks.2.attentions.0.transformer_blocks.1.attn1 +down_blocks.2.attentions.0.transformer_blocks.1.attn1.to_q +down_blocks.2.attentions.0.transformer_blocks.1.attn1.to_k +down_blocks.2.attentions.0.transformer_blocks.1.attn1.to_v +down_blocks.2.attentions.0.transformer_blocks.1.attn1.to_out +down_blocks.2.attentions.0.transformer_blocks.1.attn1.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.1.attn1.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.1.norm2 +down_blocks.2.attentions.0.transformer_blocks.1.attn2 +down_blocks.2.attentions.0.transformer_blocks.1.attn2.to_q +down_blocks.2.attentions.0.transformer_blocks.1.attn2.to_k +down_blocks.2.attentions.0.transformer_blocks.1.attn2.to_v +down_blocks.2.attentions.0.transformer_blocks.1.attn2.to_out +down_blocks.2.attentions.0.transformer_blocks.1.attn2.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.1.attn2.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.1.norm3 +down_blocks.2.attentions.0.transformer_blocks.1.ff +down_blocks.2.attentions.0.transformer_blocks.1.ff.net +down_blocks.2.attentions.0.transformer_blocks.1.ff.net.0 +down_blocks.2.attentions.0.transformer_blocks.1.ff.net.0.proj +down_blocks.2.attentions.0.transformer_blocks.1.ff.net.1 +down_blocks.2.attentions.0.transformer_blocks.1.ff.net.2 +down_blocks.2.attentions.0.transformer_blocks.2 +down_blocks.2.attentions.0.transformer_blocks.2.norm1 +down_blocks.2.attentions.0.transformer_blocks.2.attn1 +down_blocks.2.attentions.0.transformer_blocks.2.attn1.to_q +down_blocks.2.attentions.0.transformer_blocks.2.attn1.to_k +down_blocks.2.attentions.0.transformer_blocks.2.attn1.to_v +down_blocks.2.attentions.0.transformer_blocks.2.attn1.to_out +down_blocks.2.attentions.0.transformer_blocks.2.attn1.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.2.attn1.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.2.norm2 +down_blocks.2.attentions.0.transformer_blocks.2.attn2 +down_blocks.2.attentions.0.transformer_blocks.2.attn2.to_q +down_blocks.2.attentions.0.transformer_blocks.2.attn2.to_k +down_blocks.2.attentions.0.transformer_blocks.2.attn2.to_v +down_blocks.2.attentions.0.transformer_blocks.2.attn2.to_out +down_blocks.2.attentions.0.transformer_blocks.2.attn2.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.2.attn2.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.2.norm3 +down_blocks.2.attentions.0.transformer_blocks.2.ff +down_blocks.2.attentions.0.transformer_blocks.2.ff.net +down_blocks.2.attentions.0.transformer_blocks.2.ff.net.0 +down_blocks.2.attentions.0.transformer_blocks.2.ff.net.0.proj +down_blocks.2.attentions.0.transformer_blocks.2.ff.net.1 +down_blocks.2.attentions.0.transformer_blocks.2.ff.net.2 +down_blocks.2.attentions.0.transformer_blocks.3 +down_blocks.2.attentions.0.transformer_blocks.3.norm1 +down_blocks.2.attentions.0.transformer_blocks.3.attn1 +down_blocks.2.attentions.0.transformer_blocks.3.attn1.to_q +down_blocks.2.attentions.0.transformer_blocks.3.attn1.to_k +down_blocks.2.attentions.0.transformer_blocks.3.attn1.to_v +down_blocks.2.attentions.0.transformer_blocks.3.attn1.to_out +down_blocks.2.attentions.0.transformer_blocks.3.attn1.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.3.attn1.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.3.norm2 +down_blocks.2.attentions.0.transformer_blocks.3.attn2 +down_blocks.2.attentions.0.transformer_blocks.3.attn2.to_q +down_blocks.2.attentions.0.transformer_blocks.3.attn2.to_k +down_blocks.2.attentions.0.transformer_blocks.3.attn2.to_v +down_blocks.2.attentions.0.transformer_blocks.3.attn2.to_out +down_blocks.2.attentions.0.transformer_blocks.3.attn2.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.3.attn2.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.3.norm3 +down_blocks.2.attentions.0.transformer_blocks.3.ff +down_blocks.2.attentions.0.transformer_blocks.3.ff.net +down_blocks.2.attentions.0.transformer_blocks.3.ff.net.0 +down_blocks.2.attentions.0.transformer_blocks.3.ff.net.0.proj +down_blocks.2.attentions.0.transformer_blocks.3.ff.net.1 +down_blocks.2.attentions.0.transformer_blocks.3.ff.net.2 +down_blocks.2.attentions.0.transformer_blocks.4 +down_blocks.2.attentions.0.transformer_blocks.4.norm1 +down_blocks.2.attentions.0.transformer_blocks.4.attn1 +down_blocks.2.attentions.0.transformer_blocks.4.attn1.to_q +down_blocks.2.attentions.0.transformer_blocks.4.attn1.to_k +down_blocks.2.attentions.0.transformer_blocks.4.attn1.to_v +down_blocks.2.attentions.0.transformer_blocks.4.attn1.to_out +down_blocks.2.attentions.0.transformer_blocks.4.attn1.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.4.attn1.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.4.norm2 +down_blocks.2.attentions.0.transformer_blocks.4.attn2 +down_blocks.2.attentions.0.transformer_blocks.4.attn2.to_q +down_blocks.2.attentions.0.transformer_blocks.4.attn2.to_k +down_blocks.2.attentions.0.transformer_blocks.4.attn2.to_v +down_blocks.2.attentions.0.transformer_blocks.4.attn2.to_out +down_blocks.2.attentions.0.transformer_blocks.4.attn2.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.4.attn2.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.4.norm3 +down_blocks.2.attentions.0.transformer_blocks.4.ff +down_blocks.2.attentions.0.transformer_blocks.4.ff.net +down_blocks.2.attentions.0.transformer_blocks.4.ff.net.0 +down_blocks.2.attentions.0.transformer_blocks.4.ff.net.0.proj +down_blocks.2.attentions.0.transformer_blocks.4.ff.net.1 +down_blocks.2.attentions.0.transformer_blocks.4.ff.net.2 +down_blocks.2.attentions.0.transformer_blocks.5 +down_blocks.2.attentions.0.transformer_blocks.5.norm1 +down_blocks.2.attentions.0.transformer_blocks.5.attn1 +down_blocks.2.attentions.0.transformer_blocks.5.attn1.to_q +down_blocks.2.attentions.0.transformer_blocks.5.attn1.to_k +down_blocks.2.attentions.0.transformer_blocks.5.attn1.to_v +down_blocks.2.attentions.0.transformer_blocks.5.attn1.to_out +down_blocks.2.attentions.0.transformer_blocks.5.attn1.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.5.attn1.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.5.norm2 +down_blocks.2.attentions.0.transformer_blocks.5.attn2 +down_blocks.2.attentions.0.transformer_blocks.5.attn2.to_q +down_blocks.2.attentions.0.transformer_blocks.5.attn2.to_k +down_blocks.2.attentions.0.transformer_blocks.5.attn2.to_v +down_blocks.2.attentions.0.transformer_blocks.5.attn2.to_out +down_blocks.2.attentions.0.transformer_blocks.5.attn2.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.5.attn2.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.5.norm3 +down_blocks.2.attentions.0.transformer_blocks.5.ff +down_blocks.2.attentions.0.transformer_blocks.5.ff.net +down_blocks.2.attentions.0.transformer_blocks.5.ff.net.0 +down_blocks.2.attentions.0.transformer_blocks.5.ff.net.0.proj +down_blocks.2.attentions.0.transformer_blocks.5.ff.net.1 +down_blocks.2.attentions.0.transformer_blocks.5.ff.net.2 +down_blocks.2.attentions.0.transformer_blocks.6 +down_blocks.2.attentions.0.transformer_blocks.6.norm1 +down_blocks.2.attentions.0.transformer_blocks.6.attn1 +down_blocks.2.attentions.0.transformer_blocks.6.attn1.to_q +down_blocks.2.attentions.0.transformer_blocks.6.attn1.to_k +down_blocks.2.attentions.0.transformer_blocks.6.attn1.to_v +down_blocks.2.attentions.0.transformer_blocks.6.attn1.to_out +down_blocks.2.attentions.0.transformer_blocks.6.attn1.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.6.attn1.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.6.norm2 +down_blocks.2.attentions.0.transformer_blocks.6.attn2 +down_blocks.2.attentions.0.transformer_blocks.6.attn2.to_q +down_blocks.2.attentions.0.transformer_blocks.6.attn2.to_k +down_blocks.2.attentions.0.transformer_blocks.6.attn2.to_v +down_blocks.2.attentions.0.transformer_blocks.6.attn2.to_out +down_blocks.2.attentions.0.transformer_blocks.6.attn2.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.6.attn2.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.6.norm3 +down_blocks.2.attentions.0.transformer_blocks.6.ff +down_blocks.2.attentions.0.transformer_blocks.6.ff.net +down_blocks.2.attentions.0.transformer_blocks.6.ff.net.0 +down_blocks.2.attentions.0.transformer_blocks.6.ff.net.0.proj +down_blocks.2.attentions.0.transformer_blocks.6.ff.net.1 +down_blocks.2.attentions.0.transformer_blocks.6.ff.net.2 +down_blocks.2.attentions.0.transformer_blocks.7 +down_blocks.2.attentions.0.transformer_blocks.7.norm1 +down_blocks.2.attentions.0.transformer_blocks.7.attn1 +down_blocks.2.attentions.0.transformer_blocks.7.attn1.to_q +down_blocks.2.attentions.0.transformer_blocks.7.attn1.to_k +down_blocks.2.attentions.0.transformer_blocks.7.attn1.to_v +down_blocks.2.attentions.0.transformer_blocks.7.attn1.to_out +down_blocks.2.attentions.0.transformer_blocks.7.attn1.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.7.attn1.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.7.norm2 +down_blocks.2.attentions.0.transformer_blocks.7.attn2 +down_blocks.2.attentions.0.transformer_blocks.7.attn2.to_q +down_blocks.2.attentions.0.transformer_blocks.7.attn2.to_k +down_blocks.2.attentions.0.transformer_blocks.7.attn2.to_v +down_blocks.2.attentions.0.transformer_blocks.7.attn2.to_out +down_blocks.2.attentions.0.transformer_blocks.7.attn2.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.7.attn2.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.7.norm3 +down_blocks.2.attentions.0.transformer_blocks.7.ff +down_blocks.2.attentions.0.transformer_blocks.7.ff.net +down_blocks.2.attentions.0.transformer_blocks.7.ff.net.0 +down_blocks.2.attentions.0.transformer_blocks.7.ff.net.0.proj +down_blocks.2.attentions.0.transformer_blocks.7.ff.net.1 +down_blocks.2.attentions.0.transformer_blocks.7.ff.net.2 +down_blocks.2.attentions.0.transformer_blocks.8 +down_blocks.2.attentions.0.transformer_blocks.8.norm1 +down_blocks.2.attentions.0.transformer_blocks.8.attn1 +down_blocks.2.attentions.0.transformer_blocks.8.attn1.to_q +down_blocks.2.attentions.0.transformer_blocks.8.attn1.to_k +down_blocks.2.attentions.0.transformer_blocks.8.attn1.to_v +down_blocks.2.attentions.0.transformer_blocks.8.attn1.to_out +down_blocks.2.attentions.0.transformer_blocks.8.attn1.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.8.attn1.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.8.norm2 +down_blocks.2.attentions.0.transformer_blocks.8.attn2 +down_blocks.2.attentions.0.transformer_blocks.8.attn2.to_q +down_blocks.2.attentions.0.transformer_blocks.8.attn2.to_k +down_blocks.2.attentions.0.transformer_blocks.8.attn2.to_v +down_blocks.2.attentions.0.transformer_blocks.8.attn2.to_out +down_blocks.2.attentions.0.transformer_blocks.8.attn2.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.8.attn2.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.8.norm3 +down_blocks.2.attentions.0.transformer_blocks.8.ff +down_blocks.2.attentions.0.transformer_blocks.8.ff.net +down_blocks.2.attentions.0.transformer_blocks.8.ff.net.0 +down_blocks.2.attentions.0.transformer_blocks.8.ff.net.0.proj +down_blocks.2.attentions.0.transformer_blocks.8.ff.net.1 +down_blocks.2.attentions.0.transformer_blocks.8.ff.net.2 +down_blocks.2.attentions.0.transformer_blocks.9 +down_blocks.2.attentions.0.transformer_blocks.9.norm1 +down_blocks.2.attentions.0.transformer_blocks.9.attn1 +down_blocks.2.attentions.0.transformer_blocks.9.attn1.to_q +down_blocks.2.attentions.0.transformer_blocks.9.attn1.to_k +down_blocks.2.attentions.0.transformer_blocks.9.attn1.to_v +down_blocks.2.attentions.0.transformer_blocks.9.attn1.to_out +down_blocks.2.attentions.0.transformer_blocks.9.attn1.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.9.attn1.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.9.norm2 +down_blocks.2.attentions.0.transformer_blocks.9.attn2 +down_blocks.2.attentions.0.transformer_blocks.9.attn2.to_q +down_blocks.2.attentions.0.transformer_blocks.9.attn2.to_k +down_blocks.2.attentions.0.transformer_blocks.9.attn2.to_v +down_blocks.2.attentions.0.transformer_blocks.9.attn2.to_out +down_blocks.2.attentions.0.transformer_blocks.9.attn2.to_out.0 +down_blocks.2.attentions.0.transformer_blocks.9.attn2.to_out.1 +down_blocks.2.attentions.0.transformer_blocks.9.norm3 +down_blocks.2.attentions.0.transformer_blocks.9.ff +down_blocks.2.attentions.0.transformer_blocks.9.ff.net +down_blocks.2.attentions.0.transformer_blocks.9.ff.net.0 +down_blocks.2.attentions.0.transformer_blocks.9.ff.net.0.proj +down_blocks.2.attentions.0.transformer_blocks.9.ff.net.1 +down_blocks.2.attentions.0.transformer_blocks.9.ff.net.2 +down_blocks.2.attentions.0.proj_out +down_blocks.2.attentions.1 +down_blocks.2.attentions.1.norm +down_blocks.2.attentions.1.proj_in +down_blocks.2.attentions.1.transformer_blocks +down_blocks.2.attentions.1.transformer_blocks.0 +down_blocks.2.attentions.1.transformer_blocks.0.norm1 +down_blocks.2.attentions.1.transformer_blocks.0.attn1 +down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_q +down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_k +down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_v +down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_out +down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.0.norm2 +down_blocks.2.attentions.1.transformer_blocks.0.attn2 +down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_q +down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_k +down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_v +down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_out +down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.0.norm3 +down_blocks.2.attentions.1.transformer_blocks.0.ff +down_blocks.2.attentions.1.transformer_blocks.0.ff.net +down_blocks.2.attentions.1.transformer_blocks.0.ff.net.0 +down_blocks.2.attentions.1.transformer_blocks.0.ff.net.0.proj +down_blocks.2.attentions.1.transformer_blocks.0.ff.net.1 +down_blocks.2.attentions.1.transformer_blocks.0.ff.net.2 +down_blocks.2.attentions.1.transformer_blocks.1 +down_blocks.2.attentions.1.transformer_blocks.1.norm1 +down_blocks.2.attentions.1.transformer_blocks.1.attn1 +down_blocks.2.attentions.1.transformer_blocks.1.attn1.to_q +down_blocks.2.attentions.1.transformer_blocks.1.attn1.to_k +down_blocks.2.attentions.1.transformer_blocks.1.attn1.to_v +down_blocks.2.attentions.1.transformer_blocks.1.attn1.to_out +down_blocks.2.attentions.1.transformer_blocks.1.attn1.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.1.attn1.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.1.norm2 +down_blocks.2.attentions.1.transformer_blocks.1.attn2 +down_blocks.2.attentions.1.transformer_blocks.1.attn2.to_q +down_blocks.2.attentions.1.transformer_blocks.1.attn2.to_k +down_blocks.2.attentions.1.transformer_blocks.1.attn2.to_v +down_blocks.2.attentions.1.transformer_blocks.1.attn2.to_out +down_blocks.2.attentions.1.transformer_blocks.1.attn2.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.1.attn2.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.1.norm3 +down_blocks.2.attentions.1.transformer_blocks.1.ff +down_blocks.2.attentions.1.transformer_blocks.1.ff.net +down_blocks.2.attentions.1.transformer_blocks.1.ff.net.0 +down_blocks.2.attentions.1.transformer_blocks.1.ff.net.0.proj +down_blocks.2.attentions.1.transformer_blocks.1.ff.net.1 +down_blocks.2.attentions.1.transformer_blocks.1.ff.net.2 +down_blocks.2.attentions.1.transformer_blocks.2 +down_blocks.2.attentions.1.transformer_blocks.2.norm1 +down_blocks.2.attentions.1.transformer_blocks.2.attn1 +down_blocks.2.attentions.1.transformer_blocks.2.attn1.to_q +down_blocks.2.attentions.1.transformer_blocks.2.attn1.to_k +down_blocks.2.attentions.1.transformer_blocks.2.attn1.to_v +down_blocks.2.attentions.1.transformer_blocks.2.attn1.to_out +down_blocks.2.attentions.1.transformer_blocks.2.attn1.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.2.attn1.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.2.norm2 +down_blocks.2.attentions.1.transformer_blocks.2.attn2 +down_blocks.2.attentions.1.transformer_blocks.2.attn2.to_q +down_blocks.2.attentions.1.transformer_blocks.2.attn2.to_k +down_blocks.2.attentions.1.transformer_blocks.2.attn2.to_v +down_blocks.2.attentions.1.transformer_blocks.2.attn2.to_out +down_blocks.2.attentions.1.transformer_blocks.2.attn2.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.2.attn2.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.2.norm3 +down_blocks.2.attentions.1.transformer_blocks.2.ff +down_blocks.2.attentions.1.transformer_blocks.2.ff.net +down_blocks.2.attentions.1.transformer_blocks.2.ff.net.0 +down_blocks.2.attentions.1.transformer_blocks.2.ff.net.0.proj +down_blocks.2.attentions.1.transformer_blocks.2.ff.net.1 +down_blocks.2.attentions.1.transformer_blocks.2.ff.net.2 +down_blocks.2.attentions.1.transformer_blocks.3 +down_blocks.2.attentions.1.transformer_blocks.3.norm1 +down_blocks.2.attentions.1.transformer_blocks.3.attn1 +down_blocks.2.attentions.1.transformer_blocks.3.attn1.to_q +down_blocks.2.attentions.1.transformer_blocks.3.attn1.to_k +down_blocks.2.attentions.1.transformer_blocks.3.attn1.to_v +down_blocks.2.attentions.1.transformer_blocks.3.attn1.to_out +down_blocks.2.attentions.1.transformer_blocks.3.attn1.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.3.attn1.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.3.norm2 +down_blocks.2.attentions.1.transformer_blocks.3.attn2 +down_blocks.2.attentions.1.transformer_blocks.3.attn2.to_q +down_blocks.2.attentions.1.transformer_blocks.3.attn2.to_k +down_blocks.2.attentions.1.transformer_blocks.3.attn2.to_v +down_blocks.2.attentions.1.transformer_blocks.3.attn2.to_out +down_blocks.2.attentions.1.transformer_blocks.3.attn2.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.3.attn2.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.3.norm3 +down_blocks.2.attentions.1.transformer_blocks.3.ff +down_blocks.2.attentions.1.transformer_blocks.3.ff.net +down_blocks.2.attentions.1.transformer_blocks.3.ff.net.0 +down_blocks.2.attentions.1.transformer_blocks.3.ff.net.0.proj +down_blocks.2.attentions.1.transformer_blocks.3.ff.net.1 +down_blocks.2.attentions.1.transformer_blocks.3.ff.net.2 +down_blocks.2.attentions.1.transformer_blocks.4 +down_blocks.2.attentions.1.transformer_blocks.4.norm1 +down_blocks.2.attentions.1.transformer_blocks.4.attn1 +down_blocks.2.attentions.1.transformer_blocks.4.attn1.to_q +down_blocks.2.attentions.1.transformer_blocks.4.attn1.to_k +down_blocks.2.attentions.1.transformer_blocks.4.attn1.to_v +down_blocks.2.attentions.1.transformer_blocks.4.attn1.to_out +down_blocks.2.attentions.1.transformer_blocks.4.attn1.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.4.attn1.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.4.norm2 +down_blocks.2.attentions.1.transformer_blocks.4.attn2 +down_blocks.2.attentions.1.transformer_blocks.4.attn2.to_q +down_blocks.2.attentions.1.transformer_blocks.4.attn2.to_k +down_blocks.2.attentions.1.transformer_blocks.4.attn2.to_v +down_blocks.2.attentions.1.transformer_blocks.4.attn2.to_out +down_blocks.2.attentions.1.transformer_blocks.4.attn2.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.4.attn2.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.4.norm3 +down_blocks.2.attentions.1.transformer_blocks.4.ff +down_blocks.2.attentions.1.transformer_blocks.4.ff.net +down_blocks.2.attentions.1.transformer_blocks.4.ff.net.0 +down_blocks.2.attentions.1.transformer_blocks.4.ff.net.0.proj +down_blocks.2.attentions.1.transformer_blocks.4.ff.net.1 +down_blocks.2.attentions.1.transformer_blocks.4.ff.net.2 +down_blocks.2.attentions.1.transformer_blocks.5 +down_blocks.2.attentions.1.transformer_blocks.5.norm1 +down_blocks.2.attentions.1.transformer_blocks.5.attn1 +down_blocks.2.attentions.1.transformer_blocks.5.attn1.to_q +down_blocks.2.attentions.1.transformer_blocks.5.attn1.to_k +down_blocks.2.attentions.1.transformer_blocks.5.attn1.to_v +down_blocks.2.attentions.1.transformer_blocks.5.attn1.to_out +down_blocks.2.attentions.1.transformer_blocks.5.attn1.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.5.attn1.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.5.norm2 +down_blocks.2.attentions.1.transformer_blocks.5.attn2 +down_blocks.2.attentions.1.transformer_blocks.5.attn2.to_q +down_blocks.2.attentions.1.transformer_blocks.5.attn2.to_k +down_blocks.2.attentions.1.transformer_blocks.5.attn2.to_v +down_blocks.2.attentions.1.transformer_blocks.5.attn2.to_out +down_blocks.2.attentions.1.transformer_blocks.5.attn2.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.5.attn2.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.5.norm3 +down_blocks.2.attentions.1.transformer_blocks.5.ff +down_blocks.2.attentions.1.transformer_blocks.5.ff.net +down_blocks.2.attentions.1.transformer_blocks.5.ff.net.0 +down_blocks.2.attentions.1.transformer_blocks.5.ff.net.0.proj +down_blocks.2.attentions.1.transformer_blocks.5.ff.net.1 +down_blocks.2.attentions.1.transformer_blocks.5.ff.net.2 +down_blocks.2.attentions.1.transformer_blocks.6 +down_blocks.2.attentions.1.transformer_blocks.6.norm1 +down_blocks.2.attentions.1.transformer_blocks.6.attn1 +down_blocks.2.attentions.1.transformer_blocks.6.attn1.to_q +down_blocks.2.attentions.1.transformer_blocks.6.attn1.to_k +down_blocks.2.attentions.1.transformer_blocks.6.attn1.to_v +down_blocks.2.attentions.1.transformer_blocks.6.attn1.to_out +down_blocks.2.attentions.1.transformer_blocks.6.attn1.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.6.attn1.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.6.norm2 +down_blocks.2.attentions.1.transformer_blocks.6.attn2 +down_blocks.2.attentions.1.transformer_blocks.6.attn2.to_q +down_blocks.2.attentions.1.transformer_blocks.6.attn2.to_k +down_blocks.2.attentions.1.transformer_blocks.6.attn2.to_v +down_blocks.2.attentions.1.transformer_blocks.6.attn2.to_out +down_blocks.2.attentions.1.transformer_blocks.6.attn2.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.6.attn2.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.6.norm3 +down_blocks.2.attentions.1.transformer_blocks.6.ff +down_blocks.2.attentions.1.transformer_blocks.6.ff.net +down_blocks.2.attentions.1.transformer_blocks.6.ff.net.0 +down_blocks.2.attentions.1.transformer_blocks.6.ff.net.0.proj +down_blocks.2.attentions.1.transformer_blocks.6.ff.net.1 +down_blocks.2.attentions.1.transformer_blocks.6.ff.net.2 +down_blocks.2.attentions.1.transformer_blocks.7 +down_blocks.2.attentions.1.transformer_blocks.7.norm1 +down_blocks.2.attentions.1.transformer_blocks.7.attn1 +down_blocks.2.attentions.1.transformer_blocks.7.attn1.to_q +down_blocks.2.attentions.1.transformer_blocks.7.attn1.to_k +down_blocks.2.attentions.1.transformer_blocks.7.attn1.to_v +down_blocks.2.attentions.1.transformer_blocks.7.attn1.to_out +down_blocks.2.attentions.1.transformer_blocks.7.attn1.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.7.attn1.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.7.norm2 +down_blocks.2.attentions.1.transformer_blocks.7.attn2 +down_blocks.2.attentions.1.transformer_blocks.7.attn2.to_q +down_blocks.2.attentions.1.transformer_blocks.7.attn2.to_k +down_blocks.2.attentions.1.transformer_blocks.7.attn2.to_v +down_blocks.2.attentions.1.transformer_blocks.7.attn2.to_out +down_blocks.2.attentions.1.transformer_blocks.7.attn2.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.7.attn2.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.7.norm3 +down_blocks.2.attentions.1.transformer_blocks.7.ff +down_blocks.2.attentions.1.transformer_blocks.7.ff.net +down_blocks.2.attentions.1.transformer_blocks.7.ff.net.0 +down_blocks.2.attentions.1.transformer_blocks.7.ff.net.0.proj +down_blocks.2.attentions.1.transformer_blocks.7.ff.net.1 +down_blocks.2.attentions.1.transformer_blocks.7.ff.net.2 +down_blocks.2.attentions.1.transformer_blocks.8 +down_blocks.2.attentions.1.transformer_blocks.8.norm1 +down_blocks.2.attentions.1.transformer_blocks.8.attn1 +down_blocks.2.attentions.1.transformer_blocks.8.attn1.to_q +down_blocks.2.attentions.1.transformer_blocks.8.attn1.to_k +down_blocks.2.attentions.1.transformer_blocks.8.attn1.to_v +down_blocks.2.attentions.1.transformer_blocks.8.attn1.to_out +down_blocks.2.attentions.1.transformer_blocks.8.attn1.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.8.attn1.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.8.norm2 +down_blocks.2.attentions.1.transformer_blocks.8.attn2 +down_blocks.2.attentions.1.transformer_blocks.8.attn2.to_q +down_blocks.2.attentions.1.transformer_blocks.8.attn2.to_k +down_blocks.2.attentions.1.transformer_blocks.8.attn2.to_v +down_blocks.2.attentions.1.transformer_blocks.8.attn2.to_out +down_blocks.2.attentions.1.transformer_blocks.8.attn2.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.8.attn2.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.8.norm3 +down_blocks.2.attentions.1.transformer_blocks.8.ff +down_blocks.2.attentions.1.transformer_blocks.8.ff.net +down_blocks.2.attentions.1.transformer_blocks.8.ff.net.0 +down_blocks.2.attentions.1.transformer_blocks.8.ff.net.0.proj +down_blocks.2.attentions.1.transformer_blocks.8.ff.net.1 +down_blocks.2.attentions.1.transformer_blocks.8.ff.net.2 +down_blocks.2.attentions.1.transformer_blocks.9 +down_blocks.2.attentions.1.transformer_blocks.9.norm1 +down_blocks.2.attentions.1.transformer_blocks.9.attn1 +down_blocks.2.attentions.1.transformer_blocks.9.attn1.to_q +down_blocks.2.attentions.1.transformer_blocks.9.attn1.to_k +down_blocks.2.attentions.1.transformer_blocks.9.attn1.to_v +down_blocks.2.attentions.1.transformer_blocks.9.attn1.to_out +down_blocks.2.attentions.1.transformer_blocks.9.attn1.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.9.attn1.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.9.norm2 +down_blocks.2.attentions.1.transformer_blocks.9.attn2 +down_blocks.2.attentions.1.transformer_blocks.9.attn2.to_q +down_blocks.2.attentions.1.transformer_blocks.9.attn2.to_k +down_blocks.2.attentions.1.transformer_blocks.9.attn2.to_v +down_blocks.2.attentions.1.transformer_blocks.9.attn2.to_out +down_blocks.2.attentions.1.transformer_blocks.9.attn2.to_out.0 +down_blocks.2.attentions.1.transformer_blocks.9.attn2.to_out.1 +down_blocks.2.attentions.1.transformer_blocks.9.norm3 +down_blocks.2.attentions.1.transformer_blocks.9.ff +down_blocks.2.attentions.1.transformer_blocks.9.ff.net +down_blocks.2.attentions.1.transformer_blocks.9.ff.net.0 +down_blocks.2.attentions.1.transformer_blocks.9.ff.net.0.proj +down_blocks.2.attentions.1.transformer_blocks.9.ff.net.1 +down_blocks.2.attentions.1.transformer_blocks.9.ff.net.2 +down_blocks.2.attentions.1.proj_out +down_blocks.2.resnets +down_blocks.2.resnets.0 +down_blocks.2.resnets.0.norm1 +down_blocks.2.resnets.0.conv1 +down_blocks.2.resnets.0.time_emb_proj +down_blocks.2.resnets.0.norm2 +down_blocks.2.resnets.0.dropout +down_blocks.2.resnets.0.conv2 +down_blocks.2.resnets.0.conv_shortcut +down_blocks.2.resnets.1 +down_blocks.2.resnets.1.norm1 +down_blocks.2.resnets.1.conv1 +down_blocks.2.resnets.1.time_emb_proj +down_blocks.2.resnets.1.norm2 +down_blocks.2.resnets.1.dropout +down_blocks.2.resnets.1.conv2 +up_blocks +up_blocks.0 +up_blocks.0.attentions +up_blocks.0.attentions.0 +up_blocks.0.attentions.0.norm +up_blocks.0.attentions.0.proj_in +up_blocks.0.attentions.0.transformer_blocks +up_blocks.0.attentions.0.transformer_blocks.0 +up_blocks.0.attentions.0.transformer_blocks.0.norm1 +up_blocks.0.attentions.0.transformer_blocks.0.attn1 +up_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q +up_blocks.0.attentions.0.transformer_blocks.0.attn1.to_k +up_blocks.0.attentions.0.transformer_blocks.0.attn1.to_v +up_blocks.0.attentions.0.transformer_blocks.0.attn1.to_out +up_blocks.0.attentions.0.transformer_blocks.0.attn1.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.0.attn1.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.0.norm2 +up_blocks.0.attentions.0.transformer_blocks.0.attn2 +up_blocks.0.attentions.0.transformer_blocks.0.attn2.to_q +up_blocks.0.attentions.0.transformer_blocks.0.attn2.to_k +up_blocks.0.attentions.0.transformer_blocks.0.attn2.to_v +up_blocks.0.attentions.0.transformer_blocks.0.attn2.to_out +up_blocks.0.attentions.0.transformer_blocks.0.attn2.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.0.attn2.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.0.norm3 +up_blocks.0.attentions.0.transformer_blocks.0.ff +up_blocks.0.attentions.0.transformer_blocks.0.ff.net +up_blocks.0.attentions.0.transformer_blocks.0.ff.net.0 +up_blocks.0.attentions.0.transformer_blocks.0.ff.net.0.proj +up_blocks.0.attentions.0.transformer_blocks.0.ff.net.1 +up_blocks.0.attentions.0.transformer_blocks.0.ff.net.2 +up_blocks.0.attentions.0.transformer_blocks.1 +up_blocks.0.attentions.0.transformer_blocks.1.norm1 +up_blocks.0.attentions.0.transformer_blocks.1.attn1 +up_blocks.0.attentions.0.transformer_blocks.1.attn1.to_q +up_blocks.0.attentions.0.transformer_blocks.1.attn1.to_k +up_blocks.0.attentions.0.transformer_blocks.1.attn1.to_v +up_blocks.0.attentions.0.transformer_blocks.1.attn1.to_out +up_blocks.0.attentions.0.transformer_blocks.1.attn1.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.1.attn1.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.1.norm2 +up_blocks.0.attentions.0.transformer_blocks.1.attn2 +up_blocks.0.attentions.0.transformer_blocks.1.attn2.to_q +up_blocks.0.attentions.0.transformer_blocks.1.attn2.to_k +up_blocks.0.attentions.0.transformer_blocks.1.attn2.to_v +up_blocks.0.attentions.0.transformer_blocks.1.attn2.to_out +up_blocks.0.attentions.0.transformer_blocks.1.attn2.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.1.attn2.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.1.norm3 +up_blocks.0.attentions.0.transformer_blocks.1.ff +up_blocks.0.attentions.0.transformer_blocks.1.ff.net +up_blocks.0.attentions.0.transformer_blocks.1.ff.net.0 +up_blocks.0.attentions.0.transformer_blocks.1.ff.net.0.proj +up_blocks.0.attentions.0.transformer_blocks.1.ff.net.1 +up_blocks.0.attentions.0.transformer_blocks.1.ff.net.2 +up_blocks.0.attentions.0.transformer_blocks.2 +up_blocks.0.attentions.0.transformer_blocks.2.norm1 +up_blocks.0.attentions.0.transformer_blocks.2.attn1 +up_blocks.0.attentions.0.transformer_blocks.2.attn1.to_q +up_blocks.0.attentions.0.transformer_blocks.2.attn1.to_k +up_blocks.0.attentions.0.transformer_blocks.2.attn1.to_v +up_blocks.0.attentions.0.transformer_blocks.2.attn1.to_out +up_blocks.0.attentions.0.transformer_blocks.2.attn1.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.2.attn1.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.2.norm2 +up_blocks.0.attentions.0.transformer_blocks.2.attn2 +up_blocks.0.attentions.0.transformer_blocks.2.attn2.to_q +up_blocks.0.attentions.0.transformer_blocks.2.attn2.to_k +up_blocks.0.attentions.0.transformer_blocks.2.attn2.to_v +up_blocks.0.attentions.0.transformer_blocks.2.attn2.to_out +up_blocks.0.attentions.0.transformer_blocks.2.attn2.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.2.attn2.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.2.norm3 +up_blocks.0.attentions.0.transformer_blocks.2.ff +up_blocks.0.attentions.0.transformer_blocks.2.ff.net +up_blocks.0.attentions.0.transformer_blocks.2.ff.net.0 +up_blocks.0.attentions.0.transformer_blocks.2.ff.net.0.proj +up_blocks.0.attentions.0.transformer_blocks.2.ff.net.1 +up_blocks.0.attentions.0.transformer_blocks.2.ff.net.2 +up_blocks.0.attentions.0.transformer_blocks.3 +up_blocks.0.attentions.0.transformer_blocks.3.norm1 +up_blocks.0.attentions.0.transformer_blocks.3.attn1 +up_blocks.0.attentions.0.transformer_blocks.3.attn1.to_q +up_blocks.0.attentions.0.transformer_blocks.3.attn1.to_k +up_blocks.0.attentions.0.transformer_blocks.3.attn1.to_v +up_blocks.0.attentions.0.transformer_blocks.3.attn1.to_out +up_blocks.0.attentions.0.transformer_blocks.3.attn1.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.3.attn1.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.3.norm2 +up_blocks.0.attentions.0.transformer_blocks.3.attn2 +up_blocks.0.attentions.0.transformer_blocks.3.attn2.to_q +up_blocks.0.attentions.0.transformer_blocks.3.attn2.to_k +up_blocks.0.attentions.0.transformer_blocks.3.attn2.to_v +up_blocks.0.attentions.0.transformer_blocks.3.attn2.to_out +up_blocks.0.attentions.0.transformer_blocks.3.attn2.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.3.attn2.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.3.norm3 +up_blocks.0.attentions.0.transformer_blocks.3.ff +up_blocks.0.attentions.0.transformer_blocks.3.ff.net +up_blocks.0.attentions.0.transformer_blocks.3.ff.net.0 +up_blocks.0.attentions.0.transformer_blocks.3.ff.net.0.proj +up_blocks.0.attentions.0.transformer_blocks.3.ff.net.1 +up_blocks.0.attentions.0.transformer_blocks.3.ff.net.2 +up_blocks.0.attentions.0.transformer_blocks.4 +up_blocks.0.attentions.0.transformer_blocks.4.norm1 +up_blocks.0.attentions.0.transformer_blocks.4.attn1 +up_blocks.0.attentions.0.transformer_blocks.4.attn1.to_q +up_blocks.0.attentions.0.transformer_blocks.4.attn1.to_k +up_blocks.0.attentions.0.transformer_blocks.4.attn1.to_v +up_blocks.0.attentions.0.transformer_blocks.4.attn1.to_out +up_blocks.0.attentions.0.transformer_blocks.4.attn1.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.4.attn1.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.4.norm2 +up_blocks.0.attentions.0.transformer_blocks.4.attn2 +up_blocks.0.attentions.0.transformer_blocks.4.attn2.to_q +up_blocks.0.attentions.0.transformer_blocks.4.attn2.to_k +up_blocks.0.attentions.0.transformer_blocks.4.attn2.to_v +up_blocks.0.attentions.0.transformer_blocks.4.attn2.to_out +up_blocks.0.attentions.0.transformer_blocks.4.attn2.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.4.attn2.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.4.norm3 +up_blocks.0.attentions.0.transformer_blocks.4.ff +up_blocks.0.attentions.0.transformer_blocks.4.ff.net +up_blocks.0.attentions.0.transformer_blocks.4.ff.net.0 +up_blocks.0.attentions.0.transformer_blocks.4.ff.net.0.proj +up_blocks.0.attentions.0.transformer_blocks.4.ff.net.1 +up_blocks.0.attentions.0.transformer_blocks.4.ff.net.2 +up_blocks.0.attentions.0.transformer_blocks.5 +up_blocks.0.attentions.0.transformer_blocks.5.norm1 +up_blocks.0.attentions.0.transformer_blocks.5.attn1 +up_blocks.0.attentions.0.transformer_blocks.5.attn1.to_q +up_blocks.0.attentions.0.transformer_blocks.5.attn1.to_k +up_blocks.0.attentions.0.transformer_blocks.5.attn1.to_v +up_blocks.0.attentions.0.transformer_blocks.5.attn1.to_out +up_blocks.0.attentions.0.transformer_blocks.5.attn1.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.5.attn1.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.5.norm2 +up_blocks.0.attentions.0.transformer_blocks.5.attn2 +up_blocks.0.attentions.0.transformer_blocks.5.attn2.to_q +up_blocks.0.attentions.0.transformer_blocks.5.attn2.to_k +up_blocks.0.attentions.0.transformer_blocks.5.attn2.to_v +up_blocks.0.attentions.0.transformer_blocks.5.attn2.to_out +up_blocks.0.attentions.0.transformer_blocks.5.attn2.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.5.attn2.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.5.norm3 +up_blocks.0.attentions.0.transformer_blocks.5.ff +up_blocks.0.attentions.0.transformer_blocks.5.ff.net +up_blocks.0.attentions.0.transformer_blocks.5.ff.net.0 +up_blocks.0.attentions.0.transformer_blocks.5.ff.net.0.proj +up_blocks.0.attentions.0.transformer_blocks.5.ff.net.1 +up_blocks.0.attentions.0.transformer_blocks.5.ff.net.2 +up_blocks.0.attentions.0.transformer_blocks.6 +up_blocks.0.attentions.0.transformer_blocks.6.norm1 +up_blocks.0.attentions.0.transformer_blocks.6.attn1 +up_blocks.0.attentions.0.transformer_blocks.6.attn1.to_q +up_blocks.0.attentions.0.transformer_blocks.6.attn1.to_k +up_blocks.0.attentions.0.transformer_blocks.6.attn1.to_v +up_blocks.0.attentions.0.transformer_blocks.6.attn1.to_out +up_blocks.0.attentions.0.transformer_blocks.6.attn1.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.6.attn1.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.6.norm2 +up_blocks.0.attentions.0.transformer_blocks.6.attn2 +up_blocks.0.attentions.0.transformer_blocks.6.attn2.to_q +up_blocks.0.attentions.0.transformer_blocks.6.attn2.to_k +up_blocks.0.attentions.0.transformer_blocks.6.attn2.to_v +up_blocks.0.attentions.0.transformer_blocks.6.attn2.to_out +up_blocks.0.attentions.0.transformer_blocks.6.attn2.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.6.attn2.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.6.norm3 +up_blocks.0.attentions.0.transformer_blocks.6.ff +up_blocks.0.attentions.0.transformer_blocks.6.ff.net +up_blocks.0.attentions.0.transformer_blocks.6.ff.net.0 +up_blocks.0.attentions.0.transformer_blocks.6.ff.net.0.proj +up_blocks.0.attentions.0.transformer_blocks.6.ff.net.1 +up_blocks.0.attentions.0.transformer_blocks.6.ff.net.2 +up_blocks.0.attentions.0.transformer_blocks.7 +up_blocks.0.attentions.0.transformer_blocks.7.norm1 +up_blocks.0.attentions.0.transformer_blocks.7.attn1 +up_blocks.0.attentions.0.transformer_blocks.7.attn1.to_q +up_blocks.0.attentions.0.transformer_blocks.7.attn1.to_k +up_blocks.0.attentions.0.transformer_blocks.7.attn1.to_v +up_blocks.0.attentions.0.transformer_blocks.7.attn1.to_out +up_blocks.0.attentions.0.transformer_blocks.7.attn1.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.7.attn1.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.7.norm2 +up_blocks.0.attentions.0.transformer_blocks.7.attn2 +up_blocks.0.attentions.0.transformer_blocks.7.attn2.to_q +up_blocks.0.attentions.0.transformer_blocks.7.attn2.to_k +up_blocks.0.attentions.0.transformer_blocks.7.attn2.to_v +up_blocks.0.attentions.0.transformer_blocks.7.attn2.to_out +up_blocks.0.attentions.0.transformer_blocks.7.attn2.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.7.attn2.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.7.norm3 +up_blocks.0.attentions.0.transformer_blocks.7.ff +up_blocks.0.attentions.0.transformer_blocks.7.ff.net +up_blocks.0.attentions.0.transformer_blocks.7.ff.net.0 +up_blocks.0.attentions.0.transformer_blocks.7.ff.net.0.proj +up_blocks.0.attentions.0.transformer_blocks.7.ff.net.1 +up_blocks.0.attentions.0.transformer_blocks.7.ff.net.2 +up_blocks.0.attentions.0.transformer_blocks.8 +up_blocks.0.attentions.0.transformer_blocks.8.norm1 +up_blocks.0.attentions.0.transformer_blocks.8.attn1 +up_blocks.0.attentions.0.transformer_blocks.8.attn1.to_q +up_blocks.0.attentions.0.transformer_blocks.8.attn1.to_k +up_blocks.0.attentions.0.transformer_blocks.8.attn1.to_v +up_blocks.0.attentions.0.transformer_blocks.8.attn1.to_out +up_blocks.0.attentions.0.transformer_blocks.8.attn1.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.8.attn1.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.8.norm2 +up_blocks.0.attentions.0.transformer_blocks.8.attn2 +up_blocks.0.attentions.0.transformer_blocks.8.attn2.to_q +up_blocks.0.attentions.0.transformer_blocks.8.attn2.to_k +up_blocks.0.attentions.0.transformer_blocks.8.attn2.to_v +up_blocks.0.attentions.0.transformer_blocks.8.attn2.to_out +up_blocks.0.attentions.0.transformer_blocks.8.attn2.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.8.attn2.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.8.norm3 +up_blocks.0.attentions.0.transformer_blocks.8.ff +up_blocks.0.attentions.0.transformer_blocks.8.ff.net +up_blocks.0.attentions.0.transformer_blocks.8.ff.net.0 +up_blocks.0.attentions.0.transformer_blocks.8.ff.net.0.proj +up_blocks.0.attentions.0.transformer_blocks.8.ff.net.1 +up_blocks.0.attentions.0.transformer_blocks.8.ff.net.2 +up_blocks.0.attentions.0.transformer_blocks.9 +up_blocks.0.attentions.0.transformer_blocks.9.norm1 +up_blocks.0.attentions.0.transformer_blocks.9.attn1 +up_blocks.0.attentions.0.transformer_blocks.9.attn1.to_q +up_blocks.0.attentions.0.transformer_blocks.9.attn1.to_k +up_blocks.0.attentions.0.transformer_blocks.9.attn1.to_v +up_blocks.0.attentions.0.transformer_blocks.9.attn1.to_out +up_blocks.0.attentions.0.transformer_blocks.9.attn1.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.9.attn1.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.9.norm2 +up_blocks.0.attentions.0.transformer_blocks.9.attn2 +up_blocks.0.attentions.0.transformer_blocks.9.attn2.to_q +up_blocks.0.attentions.0.transformer_blocks.9.attn2.to_k +up_blocks.0.attentions.0.transformer_blocks.9.attn2.to_v +up_blocks.0.attentions.0.transformer_blocks.9.attn2.to_out +up_blocks.0.attentions.0.transformer_blocks.9.attn2.to_out.0 +up_blocks.0.attentions.0.transformer_blocks.9.attn2.to_out.1 +up_blocks.0.attentions.0.transformer_blocks.9.norm3 +up_blocks.0.attentions.0.transformer_blocks.9.ff +up_blocks.0.attentions.0.transformer_blocks.9.ff.net +up_blocks.0.attentions.0.transformer_blocks.9.ff.net.0 +up_blocks.0.attentions.0.transformer_blocks.9.ff.net.0.proj +up_blocks.0.attentions.0.transformer_blocks.9.ff.net.1 +up_blocks.0.attentions.0.transformer_blocks.9.ff.net.2 +up_blocks.0.attentions.0.proj_out +up_blocks.0.attentions.1 +up_blocks.0.attentions.1.norm +up_blocks.0.attentions.1.proj_in +up_blocks.0.attentions.1.transformer_blocks +up_blocks.0.attentions.1.transformer_blocks.0 +up_blocks.0.attentions.1.transformer_blocks.0.norm1 +up_blocks.0.attentions.1.transformer_blocks.0.attn1 +up_blocks.0.attentions.1.transformer_blocks.0.attn1.to_q +up_blocks.0.attentions.1.transformer_blocks.0.attn1.to_k +up_blocks.0.attentions.1.transformer_blocks.0.attn1.to_v +up_blocks.0.attentions.1.transformer_blocks.0.attn1.to_out +up_blocks.0.attentions.1.transformer_blocks.0.attn1.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.0.attn1.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.0.norm2 +up_blocks.0.attentions.1.transformer_blocks.0.attn2 +up_blocks.0.attentions.1.transformer_blocks.0.attn2.to_q +up_blocks.0.attentions.1.transformer_blocks.0.attn2.to_k +up_blocks.0.attentions.1.transformer_blocks.0.attn2.to_v +up_blocks.0.attentions.1.transformer_blocks.0.attn2.to_out +up_blocks.0.attentions.1.transformer_blocks.0.attn2.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.0.attn2.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.0.norm3 +up_blocks.0.attentions.1.transformer_blocks.0.ff +up_blocks.0.attentions.1.transformer_blocks.0.ff.net +up_blocks.0.attentions.1.transformer_blocks.0.ff.net.0 +up_blocks.0.attentions.1.transformer_blocks.0.ff.net.0.proj +up_blocks.0.attentions.1.transformer_blocks.0.ff.net.1 +up_blocks.0.attentions.1.transformer_blocks.0.ff.net.2 +up_blocks.0.attentions.1.transformer_blocks.1 +up_blocks.0.attentions.1.transformer_blocks.1.norm1 +up_blocks.0.attentions.1.transformer_blocks.1.attn1 +up_blocks.0.attentions.1.transformer_blocks.1.attn1.to_q +up_blocks.0.attentions.1.transformer_blocks.1.attn1.to_k +up_blocks.0.attentions.1.transformer_blocks.1.attn1.to_v +up_blocks.0.attentions.1.transformer_blocks.1.attn1.to_out +up_blocks.0.attentions.1.transformer_blocks.1.attn1.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.1.attn1.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.1.norm2 +up_blocks.0.attentions.1.transformer_blocks.1.attn2 +up_blocks.0.attentions.1.transformer_blocks.1.attn2.to_q +up_blocks.0.attentions.1.transformer_blocks.1.attn2.to_k +up_blocks.0.attentions.1.transformer_blocks.1.attn2.to_v +up_blocks.0.attentions.1.transformer_blocks.1.attn2.to_out +up_blocks.0.attentions.1.transformer_blocks.1.attn2.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.1.attn2.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.1.norm3 +up_blocks.0.attentions.1.transformer_blocks.1.ff +up_blocks.0.attentions.1.transformer_blocks.1.ff.net +up_blocks.0.attentions.1.transformer_blocks.1.ff.net.0 +up_blocks.0.attentions.1.transformer_blocks.1.ff.net.0.proj +up_blocks.0.attentions.1.transformer_blocks.1.ff.net.1 +up_blocks.0.attentions.1.transformer_blocks.1.ff.net.2 +up_blocks.0.attentions.1.transformer_blocks.2 +up_blocks.0.attentions.1.transformer_blocks.2.norm1 +up_blocks.0.attentions.1.transformer_blocks.2.attn1 +up_blocks.0.attentions.1.transformer_blocks.2.attn1.to_q +up_blocks.0.attentions.1.transformer_blocks.2.attn1.to_k +up_blocks.0.attentions.1.transformer_blocks.2.attn1.to_v +up_blocks.0.attentions.1.transformer_blocks.2.attn1.to_out +up_blocks.0.attentions.1.transformer_blocks.2.attn1.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.2.attn1.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.2.norm2 +up_blocks.0.attentions.1.transformer_blocks.2.attn2 +up_blocks.0.attentions.1.transformer_blocks.2.attn2.to_q +up_blocks.0.attentions.1.transformer_blocks.2.attn2.to_k +up_blocks.0.attentions.1.transformer_blocks.2.attn2.to_v +up_blocks.0.attentions.1.transformer_blocks.2.attn2.to_out +up_blocks.0.attentions.1.transformer_blocks.2.attn2.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.2.attn2.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.2.norm3 +up_blocks.0.attentions.1.transformer_blocks.2.ff +up_blocks.0.attentions.1.transformer_blocks.2.ff.net +up_blocks.0.attentions.1.transformer_blocks.2.ff.net.0 +up_blocks.0.attentions.1.transformer_blocks.2.ff.net.0.proj +up_blocks.0.attentions.1.transformer_blocks.2.ff.net.1 +up_blocks.0.attentions.1.transformer_blocks.2.ff.net.2 +up_blocks.0.attentions.1.transformer_blocks.3 +up_blocks.0.attentions.1.transformer_blocks.3.norm1 +up_blocks.0.attentions.1.transformer_blocks.3.attn1 +up_blocks.0.attentions.1.transformer_blocks.3.attn1.to_q +up_blocks.0.attentions.1.transformer_blocks.3.attn1.to_k +up_blocks.0.attentions.1.transformer_blocks.3.attn1.to_v +up_blocks.0.attentions.1.transformer_blocks.3.attn1.to_out +up_blocks.0.attentions.1.transformer_blocks.3.attn1.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.3.attn1.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.3.norm2 +up_blocks.0.attentions.1.transformer_blocks.3.attn2 +up_blocks.0.attentions.1.transformer_blocks.3.attn2.to_q +up_blocks.0.attentions.1.transformer_blocks.3.attn2.to_k +up_blocks.0.attentions.1.transformer_blocks.3.attn2.to_v +up_blocks.0.attentions.1.transformer_blocks.3.attn2.to_out +up_blocks.0.attentions.1.transformer_blocks.3.attn2.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.3.attn2.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.3.norm3 +up_blocks.0.attentions.1.transformer_blocks.3.ff +up_blocks.0.attentions.1.transformer_blocks.3.ff.net +up_blocks.0.attentions.1.transformer_blocks.3.ff.net.0 +up_blocks.0.attentions.1.transformer_blocks.3.ff.net.0.proj +up_blocks.0.attentions.1.transformer_blocks.3.ff.net.1 +up_blocks.0.attentions.1.transformer_blocks.3.ff.net.2 +up_blocks.0.attentions.1.transformer_blocks.4 +up_blocks.0.attentions.1.transformer_blocks.4.norm1 +up_blocks.0.attentions.1.transformer_blocks.4.attn1 +up_blocks.0.attentions.1.transformer_blocks.4.attn1.to_q +up_blocks.0.attentions.1.transformer_blocks.4.attn1.to_k +up_blocks.0.attentions.1.transformer_blocks.4.attn1.to_v +up_blocks.0.attentions.1.transformer_blocks.4.attn1.to_out +up_blocks.0.attentions.1.transformer_blocks.4.attn1.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.4.attn1.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.4.norm2 +up_blocks.0.attentions.1.transformer_blocks.4.attn2 +up_blocks.0.attentions.1.transformer_blocks.4.attn2.to_q +up_blocks.0.attentions.1.transformer_blocks.4.attn2.to_k +up_blocks.0.attentions.1.transformer_blocks.4.attn2.to_v +up_blocks.0.attentions.1.transformer_blocks.4.attn2.to_out +up_blocks.0.attentions.1.transformer_blocks.4.attn2.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.4.attn2.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.4.norm3 +up_blocks.0.attentions.1.transformer_blocks.4.ff +up_blocks.0.attentions.1.transformer_blocks.4.ff.net +up_blocks.0.attentions.1.transformer_blocks.4.ff.net.0 +up_blocks.0.attentions.1.transformer_blocks.4.ff.net.0.proj +up_blocks.0.attentions.1.transformer_blocks.4.ff.net.1 +up_blocks.0.attentions.1.transformer_blocks.4.ff.net.2 +up_blocks.0.attentions.1.transformer_blocks.5 +up_blocks.0.attentions.1.transformer_blocks.5.norm1 +up_blocks.0.attentions.1.transformer_blocks.5.attn1 +up_blocks.0.attentions.1.transformer_blocks.5.attn1.to_q +up_blocks.0.attentions.1.transformer_blocks.5.attn1.to_k +up_blocks.0.attentions.1.transformer_blocks.5.attn1.to_v +up_blocks.0.attentions.1.transformer_blocks.5.attn1.to_out +up_blocks.0.attentions.1.transformer_blocks.5.attn1.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.5.attn1.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.5.norm2 +up_blocks.0.attentions.1.transformer_blocks.5.attn2 +up_blocks.0.attentions.1.transformer_blocks.5.attn2.to_q +up_blocks.0.attentions.1.transformer_blocks.5.attn2.to_k +up_blocks.0.attentions.1.transformer_blocks.5.attn2.to_v +up_blocks.0.attentions.1.transformer_blocks.5.attn2.to_out +up_blocks.0.attentions.1.transformer_blocks.5.attn2.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.5.attn2.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.5.norm3 +up_blocks.0.attentions.1.transformer_blocks.5.ff +up_blocks.0.attentions.1.transformer_blocks.5.ff.net +up_blocks.0.attentions.1.transformer_blocks.5.ff.net.0 +up_blocks.0.attentions.1.transformer_blocks.5.ff.net.0.proj +up_blocks.0.attentions.1.transformer_blocks.5.ff.net.1 +up_blocks.0.attentions.1.transformer_blocks.5.ff.net.2 +up_blocks.0.attentions.1.transformer_blocks.6 +up_blocks.0.attentions.1.transformer_blocks.6.norm1 +up_blocks.0.attentions.1.transformer_blocks.6.attn1 +up_blocks.0.attentions.1.transformer_blocks.6.attn1.to_q +up_blocks.0.attentions.1.transformer_blocks.6.attn1.to_k +up_blocks.0.attentions.1.transformer_blocks.6.attn1.to_v +up_blocks.0.attentions.1.transformer_blocks.6.attn1.to_out +up_blocks.0.attentions.1.transformer_blocks.6.attn1.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.6.attn1.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.6.norm2 +up_blocks.0.attentions.1.transformer_blocks.6.attn2 +up_blocks.0.attentions.1.transformer_blocks.6.attn2.to_q +up_blocks.0.attentions.1.transformer_blocks.6.attn2.to_k +up_blocks.0.attentions.1.transformer_blocks.6.attn2.to_v +up_blocks.0.attentions.1.transformer_blocks.6.attn2.to_out +up_blocks.0.attentions.1.transformer_blocks.6.attn2.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.6.attn2.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.6.norm3 +up_blocks.0.attentions.1.transformer_blocks.6.ff +up_blocks.0.attentions.1.transformer_blocks.6.ff.net +up_blocks.0.attentions.1.transformer_blocks.6.ff.net.0 +up_blocks.0.attentions.1.transformer_blocks.6.ff.net.0.proj +up_blocks.0.attentions.1.transformer_blocks.6.ff.net.1 +up_blocks.0.attentions.1.transformer_blocks.6.ff.net.2 +up_blocks.0.attentions.1.transformer_blocks.7 +up_blocks.0.attentions.1.transformer_blocks.7.norm1 +up_blocks.0.attentions.1.transformer_blocks.7.attn1 +up_blocks.0.attentions.1.transformer_blocks.7.attn1.to_q +up_blocks.0.attentions.1.transformer_blocks.7.attn1.to_k +up_blocks.0.attentions.1.transformer_blocks.7.attn1.to_v +up_blocks.0.attentions.1.transformer_blocks.7.attn1.to_out +up_blocks.0.attentions.1.transformer_blocks.7.attn1.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.7.attn1.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.7.norm2 +up_blocks.0.attentions.1.transformer_blocks.7.attn2 +up_blocks.0.attentions.1.transformer_blocks.7.attn2.to_q +up_blocks.0.attentions.1.transformer_blocks.7.attn2.to_k +up_blocks.0.attentions.1.transformer_blocks.7.attn2.to_v +up_blocks.0.attentions.1.transformer_blocks.7.attn2.to_out +up_blocks.0.attentions.1.transformer_blocks.7.attn2.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.7.attn2.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.7.norm3 +up_blocks.0.attentions.1.transformer_blocks.7.ff +up_blocks.0.attentions.1.transformer_blocks.7.ff.net +up_blocks.0.attentions.1.transformer_blocks.7.ff.net.0 +up_blocks.0.attentions.1.transformer_blocks.7.ff.net.0.proj +up_blocks.0.attentions.1.transformer_blocks.7.ff.net.1 +up_blocks.0.attentions.1.transformer_blocks.7.ff.net.2 +up_blocks.0.attentions.1.transformer_blocks.8 +up_blocks.0.attentions.1.transformer_blocks.8.norm1 +up_blocks.0.attentions.1.transformer_blocks.8.attn1 +up_blocks.0.attentions.1.transformer_blocks.8.attn1.to_q +up_blocks.0.attentions.1.transformer_blocks.8.attn1.to_k +up_blocks.0.attentions.1.transformer_blocks.8.attn1.to_v +up_blocks.0.attentions.1.transformer_blocks.8.attn1.to_out +up_blocks.0.attentions.1.transformer_blocks.8.attn1.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.8.attn1.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.8.norm2 +up_blocks.0.attentions.1.transformer_blocks.8.attn2 +up_blocks.0.attentions.1.transformer_blocks.8.attn2.to_q +up_blocks.0.attentions.1.transformer_blocks.8.attn2.to_k +up_blocks.0.attentions.1.transformer_blocks.8.attn2.to_v +up_blocks.0.attentions.1.transformer_blocks.8.attn2.to_out +up_blocks.0.attentions.1.transformer_blocks.8.attn2.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.8.attn2.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.8.norm3 +up_blocks.0.attentions.1.transformer_blocks.8.ff +up_blocks.0.attentions.1.transformer_blocks.8.ff.net +up_blocks.0.attentions.1.transformer_blocks.8.ff.net.0 +up_blocks.0.attentions.1.transformer_blocks.8.ff.net.0.proj +up_blocks.0.attentions.1.transformer_blocks.8.ff.net.1 +up_blocks.0.attentions.1.transformer_blocks.8.ff.net.2 +up_blocks.0.attentions.1.transformer_blocks.9 +up_blocks.0.attentions.1.transformer_blocks.9.norm1 +up_blocks.0.attentions.1.transformer_blocks.9.attn1 +up_blocks.0.attentions.1.transformer_blocks.9.attn1.to_q +up_blocks.0.attentions.1.transformer_blocks.9.attn1.to_k +up_blocks.0.attentions.1.transformer_blocks.9.attn1.to_v +up_blocks.0.attentions.1.transformer_blocks.9.attn1.to_out +up_blocks.0.attentions.1.transformer_blocks.9.attn1.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.9.attn1.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.9.norm2 +up_blocks.0.attentions.1.transformer_blocks.9.attn2 +up_blocks.0.attentions.1.transformer_blocks.9.attn2.to_q +up_blocks.0.attentions.1.transformer_blocks.9.attn2.to_k +up_blocks.0.attentions.1.transformer_blocks.9.attn2.to_v +up_blocks.0.attentions.1.transformer_blocks.9.attn2.to_out +up_blocks.0.attentions.1.transformer_blocks.9.attn2.to_out.0 +up_blocks.0.attentions.1.transformer_blocks.9.attn2.to_out.1 +up_blocks.0.attentions.1.transformer_blocks.9.norm3 +up_blocks.0.attentions.1.transformer_blocks.9.ff +up_blocks.0.attentions.1.transformer_blocks.9.ff.net +up_blocks.0.attentions.1.transformer_blocks.9.ff.net.0 +up_blocks.0.attentions.1.transformer_blocks.9.ff.net.0.proj +up_blocks.0.attentions.1.transformer_blocks.9.ff.net.1 +up_blocks.0.attentions.1.transformer_blocks.9.ff.net.2 +up_blocks.0.attentions.1.proj_out +up_blocks.0.attentions.2 +up_blocks.0.attentions.2.norm +up_blocks.0.attentions.2.proj_in +up_blocks.0.attentions.2.transformer_blocks +up_blocks.0.attentions.2.transformer_blocks.0 +up_blocks.0.attentions.2.transformer_blocks.0.norm1 +up_blocks.0.attentions.2.transformer_blocks.0.attn1 +up_blocks.0.attentions.2.transformer_blocks.0.attn1.to_q +up_blocks.0.attentions.2.transformer_blocks.0.attn1.to_k +up_blocks.0.attentions.2.transformer_blocks.0.attn1.to_v +up_blocks.0.attentions.2.transformer_blocks.0.attn1.to_out +up_blocks.0.attentions.2.transformer_blocks.0.attn1.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.0.attn1.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.0.norm2 +up_blocks.0.attentions.2.transformer_blocks.0.attn2 +up_blocks.0.attentions.2.transformer_blocks.0.attn2.to_q +up_blocks.0.attentions.2.transformer_blocks.0.attn2.to_k +up_blocks.0.attentions.2.transformer_blocks.0.attn2.to_v +up_blocks.0.attentions.2.transformer_blocks.0.attn2.to_out +up_blocks.0.attentions.2.transformer_blocks.0.attn2.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.0.attn2.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.0.norm3 +up_blocks.0.attentions.2.transformer_blocks.0.ff +up_blocks.0.attentions.2.transformer_blocks.0.ff.net +up_blocks.0.attentions.2.transformer_blocks.0.ff.net.0 +up_blocks.0.attentions.2.transformer_blocks.0.ff.net.0.proj +up_blocks.0.attentions.2.transformer_blocks.0.ff.net.1 +up_blocks.0.attentions.2.transformer_blocks.0.ff.net.2 +up_blocks.0.attentions.2.transformer_blocks.1 +up_blocks.0.attentions.2.transformer_blocks.1.norm1 +up_blocks.0.attentions.2.transformer_blocks.1.attn1 +up_blocks.0.attentions.2.transformer_blocks.1.attn1.to_q +up_blocks.0.attentions.2.transformer_blocks.1.attn1.to_k +up_blocks.0.attentions.2.transformer_blocks.1.attn1.to_v +up_blocks.0.attentions.2.transformer_blocks.1.attn1.to_out +up_blocks.0.attentions.2.transformer_blocks.1.attn1.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.1.attn1.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.1.norm2 +up_blocks.0.attentions.2.transformer_blocks.1.attn2 +up_blocks.0.attentions.2.transformer_blocks.1.attn2.to_q +up_blocks.0.attentions.2.transformer_blocks.1.attn2.to_k +up_blocks.0.attentions.2.transformer_blocks.1.attn2.to_v +up_blocks.0.attentions.2.transformer_blocks.1.attn2.to_out +up_blocks.0.attentions.2.transformer_blocks.1.attn2.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.1.attn2.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.1.norm3 +up_blocks.0.attentions.2.transformer_blocks.1.ff +up_blocks.0.attentions.2.transformer_blocks.1.ff.net +up_blocks.0.attentions.2.transformer_blocks.1.ff.net.0 +up_blocks.0.attentions.2.transformer_blocks.1.ff.net.0.proj +up_blocks.0.attentions.2.transformer_blocks.1.ff.net.1 +up_blocks.0.attentions.2.transformer_blocks.1.ff.net.2 +up_blocks.0.attentions.2.transformer_blocks.2 +up_blocks.0.attentions.2.transformer_blocks.2.norm1 +up_blocks.0.attentions.2.transformer_blocks.2.attn1 +up_blocks.0.attentions.2.transformer_blocks.2.attn1.to_q +up_blocks.0.attentions.2.transformer_blocks.2.attn1.to_k +up_blocks.0.attentions.2.transformer_blocks.2.attn1.to_v +up_blocks.0.attentions.2.transformer_blocks.2.attn1.to_out +up_blocks.0.attentions.2.transformer_blocks.2.attn1.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.2.attn1.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.2.norm2 +up_blocks.0.attentions.2.transformer_blocks.2.attn2 +up_blocks.0.attentions.2.transformer_blocks.2.attn2.to_q +up_blocks.0.attentions.2.transformer_blocks.2.attn2.to_k +up_blocks.0.attentions.2.transformer_blocks.2.attn2.to_v +up_blocks.0.attentions.2.transformer_blocks.2.attn2.to_out +up_blocks.0.attentions.2.transformer_blocks.2.attn2.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.2.attn2.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.2.norm3 +up_blocks.0.attentions.2.transformer_blocks.2.ff +up_blocks.0.attentions.2.transformer_blocks.2.ff.net +up_blocks.0.attentions.2.transformer_blocks.2.ff.net.0 +up_blocks.0.attentions.2.transformer_blocks.2.ff.net.0.proj +up_blocks.0.attentions.2.transformer_blocks.2.ff.net.1 +up_blocks.0.attentions.2.transformer_blocks.2.ff.net.2 +up_blocks.0.attentions.2.transformer_blocks.3 +up_blocks.0.attentions.2.transformer_blocks.3.norm1 +up_blocks.0.attentions.2.transformer_blocks.3.attn1 +up_blocks.0.attentions.2.transformer_blocks.3.attn1.to_q +up_blocks.0.attentions.2.transformer_blocks.3.attn1.to_k +up_blocks.0.attentions.2.transformer_blocks.3.attn1.to_v +up_blocks.0.attentions.2.transformer_blocks.3.attn1.to_out +up_blocks.0.attentions.2.transformer_blocks.3.attn1.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.3.attn1.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.3.norm2 +up_blocks.0.attentions.2.transformer_blocks.3.attn2 +up_blocks.0.attentions.2.transformer_blocks.3.attn2.to_q +up_blocks.0.attentions.2.transformer_blocks.3.attn2.to_k +up_blocks.0.attentions.2.transformer_blocks.3.attn2.to_v +up_blocks.0.attentions.2.transformer_blocks.3.attn2.to_out +up_blocks.0.attentions.2.transformer_blocks.3.attn2.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.3.attn2.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.3.norm3 +up_blocks.0.attentions.2.transformer_blocks.3.ff +up_blocks.0.attentions.2.transformer_blocks.3.ff.net +up_blocks.0.attentions.2.transformer_blocks.3.ff.net.0 +up_blocks.0.attentions.2.transformer_blocks.3.ff.net.0.proj +up_blocks.0.attentions.2.transformer_blocks.3.ff.net.1 +up_blocks.0.attentions.2.transformer_blocks.3.ff.net.2 +up_blocks.0.attentions.2.transformer_blocks.4 +up_blocks.0.attentions.2.transformer_blocks.4.norm1 +up_blocks.0.attentions.2.transformer_blocks.4.attn1 +up_blocks.0.attentions.2.transformer_blocks.4.attn1.to_q +up_blocks.0.attentions.2.transformer_blocks.4.attn1.to_k +up_blocks.0.attentions.2.transformer_blocks.4.attn1.to_v +up_blocks.0.attentions.2.transformer_blocks.4.attn1.to_out +up_blocks.0.attentions.2.transformer_blocks.4.attn1.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.4.attn1.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.4.norm2 +up_blocks.0.attentions.2.transformer_blocks.4.attn2 +up_blocks.0.attentions.2.transformer_blocks.4.attn2.to_q +up_blocks.0.attentions.2.transformer_blocks.4.attn2.to_k +up_blocks.0.attentions.2.transformer_blocks.4.attn2.to_v +up_blocks.0.attentions.2.transformer_blocks.4.attn2.to_out +up_blocks.0.attentions.2.transformer_blocks.4.attn2.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.4.attn2.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.4.norm3 +up_blocks.0.attentions.2.transformer_blocks.4.ff +up_blocks.0.attentions.2.transformer_blocks.4.ff.net +up_blocks.0.attentions.2.transformer_blocks.4.ff.net.0 +up_blocks.0.attentions.2.transformer_blocks.4.ff.net.0.proj +up_blocks.0.attentions.2.transformer_blocks.4.ff.net.1 +up_blocks.0.attentions.2.transformer_blocks.4.ff.net.2 +up_blocks.0.attentions.2.transformer_blocks.5 +up_blocks.0.attentions.2.transformer_blocks.5.norm1 +up_blocks.0.attentions.2.transformer_blocks.5.attn1 +up_blocks.0.attentions.2.transformer_blocks.5.attn1.to_q +up_blocks.0.attentions.2.transformer_blocks.5.attn1.to_k +up_blocks.0.attentions.2.transformer_blocks.5.attn1.to_v +up_blocks.0.attentions.2.transformer_blocks.5.attn1.to_out +up_blocks.0.attentions.2.transformer_blocks.5.attn1.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.5.attn1.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.5.norm2 +up_blocks.0.attentions.2.transformer_blocks.5.attn2 +up_blocks.0.attentions.2.transformer_blocks.5.attn2.to_q +up_blocks.0.attentions.2.transformer_blocks.5.attn2.to_k +up_blocks.0.attentions.2.transformer_blocks.5.attn2.to_v +up_blocks.0.attentions.2.transformer_blocks.5.attn2.to_out +up_blocks.0.attentions.2.transformer_blocks.5.attn2.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.5.attn2.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.5.norm3 +up_blocks.0.attentions.2.transformer_blocks.5.ff +up_blocks.0.attentions.2.transformer_blocks.5.ff.net +up_blocks.0.attentions.2.transformer_blocks.5.ff.net.0 +up_blocks.0.attentions.2.transformer_blocks.5.ff.net.0.proj +up_blocks.0.attentions.2.transformer_blocks.5.ff.net.1 +up_blocks.0.attentions.2.transformer_blocks.5.ff.net.2 +up_blocks.0.attentions.2.transformer_blocks.6 +up_blocks.0.attentions.2.transformer_blocks.6.norm1 +up_blocks.0.attentions.2.transformer_blocks.6.attn1 +up_blocks.0.attentions.2.transformer_blocks.6.attn1.to_q +up_blocks.0.attentions.2.transformer_blocks.6.attn1.to_k +up_blocks.0.attentions.2.transformer_blocks.6.attn1.to_v +up_blocks.0.attentions.2.transformer_blocks.6.attn1.to_out +up_blocks.0.attentions.2.transformer_blocks.6.attn1.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.6.attn1.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.6.norm2 +up_blocks.0.attentions.2.transformer_blocks.6.attn2 +up_blocks.0.attentions.2.transformer_blocks.6.attn2.to_q +up_blocks.0.attentions.2.transformer_blocks.6.attn2.to_k +up_blocks.0.attentions.2.transformer_blocks.6.attn2.to_v +up_blocks.0.attentions.2.transformer_blocks.6.attn2.to_out +up_blocks.0.attentions.2.transformer_blocks.6.attn2.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.6.attn2.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.6.norm3 +up_blocks.0.attentions.2.transformer_blocks.6.ff +up_blocks.0.attentions.2.transformer_blocks.6.ff.net +up_blocks.0.attentions.2.transformer_blocks.6.ff.net.0 +up_blocks.0.attentions.2.transformer_blocks.6.ff.net.0.proj +up_blocks.0.attentions.2.transformer_blocks.6.ff.net.1 +up_blocks.0.attentions.2.transformer_blocks.6.ff.net.2 +up_blocks.0.attentions.2.transformer_blocks.7 +up_blocks.0.attentions.2.transformer_blocks.7.norm1 +up_blocks.0.attentions.2.transformer_blocks.7.attn1 +up_blocks.0.attentions.2.transformer_blocks.7.attn1.to_q +up_blocks.0.attentions.2.transformer_blocks.7.attn1.to_k +up_blocks.0.attentions.2.transformer_blocks.7.attn1.to_v +up_blocks.0.attentions.2.transformer_blocks.7.attn1.to_out +up_blocks.0.attentions.2.transformer_blocks.7.attn1.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.7.attn1.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.7.norm2 +up_blocks.0.attentions.2.transformer_blocks.7.attn2 +up_blocks.0.attentions.2.transformer_blocks.7.attn2.to_q +up_blocks.0.attentions.2.transformer_blocks.7.attn2.to_k +up_blocks.0.attentions.2.transformer_blocks.7.attn2.to_v +up_blocks.0.attentions.2.transformer_blocks.7.attn2.to_out +up_blocks.0.attentions.2.transformer_blocks.7.attn2.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.7.attn2.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.7.norm3 +up_blocks.0.attentions.2.transformer_blocks.7.ff +up_blocks.0.attentions.2.transformer_blocks.7.ff.net +up_blocks.0.attentions.2.transformer_blocks.7.ff.net.0 +up_blocks.0.attentions.2.transformer_blocks.7.ff.net.0.proj +up_blocks.0.attentions.2.transformer_blocks.7.ff.net.1 +up_blocks.0.attentions.2.transformer_blocks.7.ff.net.2 +up_blocks.0.attentions.2.transformer_blocks.8 +up_blocks.0.attentions.2.transformer_blocks.8.norm1 +up_blocks.0.attentions.2.transformer_blocks.8.attn1 +up_blocks.0.attentions.2.transformer_blocks.8.attn1.to_q +up_blocks.0.attentions.2.transformer_blocks.8.attn1.to_k +up_blocks.0.attentions.2.transformer_blocks.8.attn1.to_v +up_blocks.0.attentions.2.transformer_blocks.8.attn1.to_out +up_blocks.0.attentions.2.transformer_blocks.8.attn1.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.8.attn1.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.8.norm2 +up_blocks.0.attentions.2.transformer_blocks.8.attn2 +up_blocks.0.attentions.2.transformer_blocks.8.attn2.to_q +up_blocks.0.attentions.2.transformer_blocks.8.attn2.to_k +up_blocks.0.attentions.2.transformer_blocks.8.attn2.to_v +up_blocks.0.attentions.2.transformer_blocks.8.attn2.to_out +up_blocks.0.attentions.2.transformer_blocks.8.attn2.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.8.attn2.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.8.norm3 +up_blocks.0.attentions.2.transformer_blocks.8.ff +up_blocks.0.attentions.2.transformer_blocks.8.ff.net +up_blocks.0.attentions.2.transformer_blocks.8.ff.net.0 +up_blocks.0.attentions.2.transformer_blocks.8.ff.net.0.proj +up_blocks.0.attentions.2.transformer_blocks.8.ff.net.1 +up_blocks.0.attentions.2.transformer_blocks.8.ff.net.2 +up_blocks.0.attentions.2.transformer_blocks.9 +up_blocks.0.attentions.2.transformer_blocks.9.norm1 +up_blocks.0.attentions.2.transformer_blocks.9.attn1 +up_blocks.0.attentions.2.transformer_blocks.9.attn1.to_q +up_blocks.0.attentions.2.transformer_blocks.9.attn1.to_k +up_blocks.0.attentions.2.transformer_blocks.9.attn1.to_v +up_blocks.0.attentions.2.transformer_blocks.9.attn1.to_out +up_blocks.0.attentions.2.transformer_blocks.9.attn1.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.9.attn1.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.9.norm2 +up_blocks.0.attentions.2.transformer_blocks.9.attn2 +up_blocks.0.attentions.2.transformer_blocks.9.attn2.to_q +up_blocks.0.attentions.2.transformer_blocks.9.attn2.to_k +up_blocks.0.attentions.2.transformer_blocks.9.attn2.to_v +up_blocks.0.attentions.2.transformer_blocks.9.attn2.to_out +up_blocks.0.attentions.2.transformer_blocks.9.attn2.to_out.0 +up_blocks.0.attentions.2.transformer_blocks.9.attn2.to_out.1 +up_blocks.0.attentions.2.transformer_blocks.9.norm3 +up_blocks.0.attentions.2.transformer_blocks.9.ff +up_blocks.0.attentions.2.transformer_blocks.9.ff.net +up_blocks.0.attentions.2.transformer_blocks.9.ff.net.0 +up_blocks.0.attentions.2.transformer_blocks.9.ff.net.0.proj +up_blocks.0.attentions.2.transformer_blocks.9.ff.net.1 +up_blocks.0.attentions.2.transformer_blocks.9.ff.net.2 +up_blocks.0.attentions.2.proj_out +up_blocks.0.resnets +up_blocks.0.resnets.0 +up_blocks.0.resnets.0.norm1 +up_blocks.0.resnets.0.conv1 +up_blocks.0.resnets.0.time_emb_proj +up_blocks.0.resnets.0.norm2 +up_blocks.0.resnets.0.dropout +up_blocks.0.resnets.0.conv2 +up_blocks.0.resnets.0.conv_shortcut +up_blocks.0.resnets.1 +up_blocks.0.resnets.1.norm1 +up_blocks.0.resnets.1.conv1 +up_blocks.0.resnets.1.time_emb_proj +up_blocks.0.resnets.1.norm2 +up_blocks.0.resnets.1.dropout +up_blocks.0.resnets.1.conv2 +up_blocks.0.resnets.1.conv_shortcut +up_blocks.0.resnets.2 +up_blocks.0.resnets.2.norm1 +up_blocks.0.resnets.2.conv1 +up_blocks.0.resnets.2.time_emb_proj +up_blocks.0.resnets.2.norm2 +up_blocks.0.resnets.2.dropout +up_blocks.0.resnets.2.conv2 +up_blocks.0.resnets.2.conv_shortcut +up_blocks.0.upsamplers +up_blocks.0.upsamplers.0 +up_blocks.0.upsamplers.0.conv +up_blocks.1 +up_blocks.1.attentions +up_blocks.1.attentions.0 +up_blocks.1.attentions.0.norm +up_blocks.1.attentions.0.proj_in +up_blocks.1.attentions.0.transformer_blocks +up_blocks.1.attentions.0.transformer_blocks.0 +up_blocks.1.attentions.0.transformer_blocks.0.norm1 +up_blocks.1.attentions.0.transformer_blocks.0.attn1 +up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_q +up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_k +up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_v +up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out +up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out.0 +up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out.1 +up_blocks.1.attentions.0.transformer_blocks.0.norm2 +up_blocks.1.attentions.0.transformer_blocks.0.attn2 +up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_q +up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k +up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_v +up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out +up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out.0 +up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out.1 +up_blocks.1.attentions.0.transformer_blocks.0.norm3 +up_blocks.1.attentions.0.transformer_blocks.0.ff +up_blocks.1.attentions.0.transformer_blocks.0.ff.net +up_blocks.1.attentions.0.transformer_blocks.0.ff.net.0 +up_blocks.1.attentions.0.transformer_blocks.0.ff.net.0.proj +up_blocks.1.attentions.0.transformer_blocks.0.ff.net.1 +up_blocks.1.attentions.0.transformer_blocks.0.ff.net.2 +up_blocks.1.attentions.0.transformer_blocks.1 +up_blocks.1.attentions.0.transformer_blocks.1.norm1 +up_blocks.1.attentions.0.transformer_blocks.1.attn1 +up_blocks.1.attentions.0.transformer_blocks.1.attn1.to_q +up_blocks.1.attentions.0.transformer_blocks.1.attn1.to_k +up_blocks.1.attentions.0.transformer_blocks.1.attn1.to_v +up_blocks.1.attentions.0.transformer_blocks.1.attn1.to_out +up_blocks.1.attentions.0.transformer_blocks.1.attn1.to_out.0 +up_blocks.1.attentions.0.transformer_blocks.1.attn1.to_out.1 +up_blocks.1.attentions.0.transformer_blocks.1.norm2 +up_blocks.1.attentions.0.transformer_blocks.1.attn2 +up_blocks.1.attentions.0.transformer_blocks.1.attn2.to_q +up_blocks.1.attentions.0.transformer_blocks.1.attn2.to_k +up_blocks.1.attentions.0.transformer_blocks.1.attn2.to_v +up_blocks.1.attentions.0.transformer_blocks.1.attn2.to_out +up_blocks.1.attentions.0.transformer_blocks.1.attn2.to_out.0 +up_blocks.1.attentions.0.transformer_blocks.1.attn2.to_out.1 +up_blocks.1.attentions.0.transformer_blocks.1.norm3 +up_blocks.1.attentions.0.transformer_blocks.1.ff +up_blocks.1.attentions.0.transformer_blocks.1.ff.net +up_blocks.1.attentions.0.transformer_blocks.1.ff.net.0 +up_blocks.1.attentions.0.transformer_blocks.1.ff.net.0.proj +up_blocks.1.attentions.0.transformer_blocks.1.ff.net.1 +up_blocks.1.attentions.0.transformer_blocks.1.ff.net.2 +up_blocks.1.attentions.0.proj_out +up_blocks.1.attentions.1 +up_blocks.1.attentions.1.norm +up_blocks.1.attentions.1.proj_in +up_blocks.1.attentions.1.transformer_blocks +up_blocks.1.attentions.1.transformer_blocks.0 +up_blocks.1.attentions.1.transformer_blocks.0.norm1 +up_blocks.1.attentions.1.transformer_blocks.0.attn1 +up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_q +up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_k +up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_v +up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out +up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out.0 +up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out.1 +up_blocks.1.attentions.1.transformer_blocks.0.norm2 +up_blocks.1.attentions.1.transformer_blocks.0.attn2 +up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_q +up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_k +up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_v +up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out +up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out.0 +up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out.1 +up_blocks.1.attentions.1.transformer_blocks.0.norm3 +up_blocks.1.attentions.1.transformer_blocks.0.ff +up_blocks.1.attentions.1.transformer_blocks.0.ff.net +up_blocks.1.attentions.1.transformer_blocks.0.ff.net.0 +up_blocks.1.attentions.1.transformer_blocks.0.ff.net.0.proj +up_blocks.1.attentions.1.transformer_blocks.0.ff.net.1 +up_blocks.1.attentions.1.transformer_blocks.0.ff.net.2 +up_blocks.1.attentions.1.transformer_blocks.1 +up_blocks.1.attentions.1.transformer_blocks.1.norm1 +up_blocks.1.attentions.1.transformer_blocks.1.attn1 +up_blocks.1.attentions.1.transformer_blocks.1.attn1.to_q +up_blocks.1.attentions.1.transformer_blocks.1.attn1.to_k +up_blocks.1.attentions.1.transformer_blocks.1.attn1.to_v +up_blocks.1.attentions.1.transformer_blocks.1.attn1.to_out +up_blocks.1.attentions.1.transformer_blocks.1.attn1.to_out.0 +up_blocks.1.attentions.1.transformer_blocks.1.attn1.to_out.1 +up_blocks.1.attentions.1.transformer_blocks.1.norm2 +up_blocks.1.attentions.1.transformer_blocks.1.attn2 +up_blocks.1.attentions.1.transformer_blocks.1.attn2.to_q +up_blocks.1.attentions.1.transformer_blocks.1.attn2.to_k +up_blocks.1.attentions.1.transformer_blocks.1.attn2.to_v +up_blocks.1.attentions.1.transformer_blocks.1.attn2.to_out +up_blocks.1.attentions.1.transformer_blocks.1.attn2.to_out.0 +up_blocks.1.attentions.1.transformer_blocks.1.attn2.to_out.1 +up_blocks.1.attentions.1.transformer_blocks.1.norm3 +up_blocks.1.attentions.1.transformer_blocks.1.ff +up_blocks.1.attentions.1.transformer_blocks.1.ff.net +up_blocks.1.attentions.1.transformer_blocks.1.ff.net.0 +up_blocks.1.attentions.1.transformer_blocks.1.ff.net.0.proj +up_blocks.1.attentions.1.transformer_blocks.1.ff.net.1 +up_blocks.1.attentions.1.transformer_blocks.1.ff.net.2 +up_blocks.1.attentions.1.proj_out +up_blocks.1.attentions.2 +up_blocks.1.attentions.2.norm +up_blocks.1.attentions.2.proj_in +up_blocks.1.attentions.2.transformer_blocks +up_blocks.1.attentions.2.transformer_blocks.0 +up_blocks.1.attentions.2.transformer_blocks.0.norm1 +up_blocks.1.attentions.2.transformer_blocks.0.attn1 +up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_q +up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_k +up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_v +up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_out +up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_out.0 +up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_out.1 +up_blocks.1.attentions.2.transformer_blocks.0.norm2 +up_blocks.1.attentions.2.transformer_blocks.0.attn2 +up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_q +up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_k +up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_v +up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_out +up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_out.0 +up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_out.1 +up_blocks.1.attentions.2.transformer_blocks.0.norm3 +up_blocks.1.attentions.2.transformer_blocks.0.ff +up_blocks.1.attentions.2.transformer_blocks.0.ff.net +up_blocks.1.attentions.2.transformer_blocks.0.ff.net.0 +up_blocks.1.attentions.2.transformer_blocks.0.ff.net.0.proj +up_blocks.1.attentions.2.transformer_blocks.0.ff.net.1 +up_blocks.1.attentions.2.transformer_blocks.0.ff.net.2 +up_blocks.1.attentions.2.transformer_blocks.1 +up_blocks.1.attentions.2.transformer_blocks.1.norm1 +up_blocks.1.attentions.2.transformer_blocks.1.attn1 +up_blocks.1.attentions.2.transformer_blocks.1.attn1.to_q +up_blocks.1.attentions.2.transformer_blocks.1.attn1.to_k +up_blocks.1.attentions.2.transformer_blocks.1.attn1.to_v +up_blocks.1.attentions.2.transformer_blocks.1.attn1.to_out +up_blocks.1.attentions.2.transformer_blocks.1.attn1.to_out.0 +up_blocks.1.attentions.2.transformer_blocks.1.attn1.to_out.1 +up_blocks.1.attentions.2.transformer_blocks.1.norm2 +up_blocks.1.attentions.2.transformer_blocks.1.attn2 +up_blocks.1.attentions.2.transformer_blocks.1.attn2.to_q +up_blocks.1.attentions.2.transformer_blocks.1.attn2.to_k +up_blocks.1.attentions.2.transformer_blocks.1.attn2.to_v +up_blocks.1.attentions.2.transformer_blocks.1.attn2.to_out +up_blocks.1.attentions.2.transformer_blocks.1.attn2.to_out.0 +up_blocks.1.attentions.2.transformer_blocks.1.attn2.to_out.1 +up_blocks.1.attentions.2.transformer_blocks.1.norm3 +up_blocks.1.attentions.2.transformer_blocks.1.ff +up_blocks.1.attentions.2.transformer_blocks.1.ff.net +up_blocks.1.attentions.2.transformer_blocks.1.ff.net.0 +up_blocks.1.attentions.2.transformer_blocks.1.ff.net.0.proj +up_blocks.1.attentions.2.transformer_blocks.1.ff.net.1 +up_blocks.1.attentions.2.transformer_blocks.1.ff.net.2 +up_blocks.1.attentions.2.proj_out +up_blocks.1.resnets +up_blocks.1.resnets.0 +up_blocks.1.resnets.0.norm1 +up_blocks.1.resnets.0.conv1 +up_blocks.1.resnets.0.time_emb_proj +up_blocks.1.resnets.0.norm2 +up_blocks.1.resnets.0.dropout +up_blocks.1.resnets.0.conv2 +up_blocks.1.resnets.0.conv_shortcut +up_blocks.1.resnets.1 +up_blocks.1.resnets.1.norm1 +up_blocks.1.resnets.1.conv1 +up_blocks.1.resnets.1.time_emb_proj +up_blocks.1.resnets.1.norm2 +up_blocks.1.resnets.1.dropout +up_blocks.1.resnets.1.conv2 +up_blocks.1.resnets.1.conv_shortcut +up_blocks.1.resnets.2 +up_blocks.1.resnets.2.norm1 +up_blocks.1.resnets.2.conv1 +up_blocks.1.resnets.2.time_emb_proj +up_blocks.1.resnets.2.norm2 +up_blocks.1.resnets.2.dropout +up_blocks.1.resnets.2.conv2 +up_blocks.1.resnets.2.conv_shortcut +up_blocks.1.upsamplers +up_blocks.1.upsamplers.0 +up_blocks.1.upsamplers.0.conv +up_blocks.2 +up_blocks.2.resnets +up_blocks.2.resnets.0 +up_blocks.2.resnets.0.norm1 +up_blocks.2.resnets.0.conv1 +up_blocks.2.resnets.0.time_emb_proj +up_blocks.2.resnets.0.norm2 +up_blocks.2.resnets.0.dropout +up_blocks.2.resnets.0.conv2 +up_blocks.2.resnets.0.conv_shortcut +up_blocks.2.resnets.1 +up_blocks.2.resnets.1.norm1 +up_blocks.2.resnets.1.conv1 +up_blocks.2.resnets.1.time_emb_proj +up_blocks.2.resnets.1.norm2 +up_blocks.2.resnets.1.dropout +up_blocks.2.resnets.1.conv2 +up_blocks.2.resnets.1.conv_shortcut +up_blocks.2.resnets.2 +up_blocks.2.resnets.2.norm1 +up_blocks.2.resnets.2.conv1 +up_blocks.2.resnets.2.time_emb_proj +up_blocks.2.resnets.2.norm2 +up_blocks.2.resnets.2.dropout +up_blocks.2.resnets.2.conv2 +up_blocks.2.resnets.2.conv_shortcut +mid_block +mid_block.attentions +mid_block.attentions.0 +mid_block.attentions.0.norm +mid_block.attentions.0.proj_in +mid_block.attentions.0.transformer_blocks +mid_block.attentions.0.transformer_blocks.0 +mid_block.attentions.0.transformer_blocks.0.norm1 +mid_block.attentions.0.transformer_blocks.0.attn1 +mid_block.attentions.0.transformer_blocks.0.attn1.to_q +mid_block.attentions.0.transformer_blocks.0.attn1.to_k +mid_block.attentions.0.transformer_blocks.0.attn1.to_v +mid_block.attentions.0.transformer_blocks.0.attn1.to_out +mid_block.attentions.0.transformer_blocks.0.attn1.to_out.0 +mid_block.attentions.0.transformer_blocks.0.attn1.to_out.1 +mid_block.attentions.0.transformer_blocks.0.norm2 +mid_block.attentions.0.transformer_blocks.0.attn2 +mid_block.attentions.0.transformer_blocks.0.attn2.to_q +mid_block.attentions.0.transformer_blocks.0.attn2.to_k +mid_block.attentions.0.transformer_blocks.0.attn2.to_v +mid_block.attentions.0.transformer_blocks.0.attn2.to_out +mid_block.attentions.0.transformer_blocks.0.attn2.to_out.0 +mid_block.attentions.0.transformer_blocks.0.attn2.to_out.1 +mid_block.attentions.0.transformer_blocks.0.norm3 +mid_block.attentions.0.transformer_blocks.0.ff +mid_block.attentions.0.transformer_blocks.0.ff.net +mid_block.attentions.0.transformer_blocks.0.ff.net.0 +mid_block.attentions.0.transformer_blocks.0.ff.net.0.proj +mid_block.attentions.0.transformer_blocks.0.ff.net.1 +mid_block.attentions.0.transformer_blocks.0.ff.net.2 +mid_block.attentions.0.transformer_blocks.1 +mid_block.attentions.0.transformer_blocks.1.norm1 +mid_block.attentions.0.transformer_blocks.1.attn1 +mid_block.attentions.0.transformer_blocks.1.attn1.to_q +mid_block.attentions.0.transformer_blocks.1.attn1.to_k +mid_block.attentions.0.transformer_blocks.1.attn1.to_v +mid_block.attentions.0.transformer_blocks.1.attn1.to_out +mid_block.attentions.0.transformer_blocks.1.attn1.to_out.0 +mid_block.attentions.0.transformer_blocks.1.attn1.to_out.1 +mid_block.attentions.0.transformer_blocks.1.norm2 +mid_block.attentions.0.transformer_blocks.1.attn2 +mid_block.attentions.0.transformer_blocks.1.attn2.to_q +mid_block.attentions.0.transformer_blocks.1.attn2.to_k +mid_block.attentions.0.transformer_blocks.1.attn2.to_v +mid_block.attentions.0.transformer_blocks.1.attn2.to_out +mid_block.attentions.0.transformer_blocks.1.attn2.to_out.0 +mid_block.attentions.0.transformer_blocks.1.attn2.to_out.1 +mid_block.attentions.0.transformer_blocks.1.norm3 +mid_block.attentions.0.transformer_blocks.1.ff +mid_block.attentions.0.transformer_blocks.1.ff.net +mid_block.attentions.0.transformer_blocks.1.ff.net.0 +mid_block.attentions.0.transformer_blocks.1.ff.net.0.proj +mid_block.attentions.0.transformer_blocks.1.ff.net.1 +mid_block.attentions.0.transformer_blocks.1.ff.net.2 +mid_block.attentions.0.transformer_blocks.2 +mid_block.attentions.0.transformer_blocks.2.norm1 +mid_block.attentions.0.transformer_blocks.2.attn1 +mid_block.attentions.0.transformer_blocks.2.attn1.to_q +mid_block.attentions.0.transformer_blocks.2.attn1.to_k +mid_block.attentions.0.transformer_blocks.2.attn1.to_v +mid_block.attentions.0.transformer_blocks.2.attn1.to_out +mid_block.attentions.0.transformer_blocks.2.attn1.to_out.0 +mid_block.attentions.0.transformer_blocks.2.attn1.to_out.1 +mid_block.attentions.0.transformer_blocks.2.norm2 +mid_block.attentions.0.transformer_blocks.2.attn2 +mid_block.attentions.0.transformer_blocks.2.attn2.to_q +mid_block.attentions.0.transformer_blocks.2.attn2.to_k +mid_block.attentions.0.transformer_blocks.2.attn2.to_v +mid_block.attentions.0.transformer_blocks.2.attn2.to_out +mid_block.attentions.0.transformer_blocks.2.attn2.to_out.0 +mid_block.attentions.0.transformer_blocks.2.attn2.to_out.1 +mid_block.attentions.0.transformer_blocks.2.norm3 +mid_block.attentions.0.transformer_blocks.2.ff +mid_block.attentions.0.transformer_blocks.2.ff.net +mid_block.attentions.0.transformer_blocks.2.ff.net.0 +mid_block.attentions.0.transformer_blocks.2.ff.net.0.proj +mid_block.attentions.0.transformer_blocks.2.ff.net.1 +mid_block.attentions.0.transformer_blocks.2.ff.net.2 +mid_block.attentions.0.transformer_blocks.3 +mid_block.attentions.0.transformer_blocks.3.norm1 +mid_block.attentions.0.transformer_blocks.3.attn1 +mid_block.attentions.0.transformer_blocks.3.attn1.to_q +mid_block.attentions.0.transformer_blocks.3.attn1.to_k +mid_block.attentions.0.transformer_blocks.3.attn1.to_v +mid_block.attentions.0.transformer_blocks.3.attn1.to_out +mid_block.attentions.0.transformer_blocks.3.attn1.to_out.0 +mid_block.attentions.0.transformer_blocks.3.attn1.to_out.1 +mid_block.attentions.0.transformer_blocks.3.norm2 +mid_block.attentions.0.transformer_blocks.3.attn2 +mid_block.attentions.0.transformer_blocks.3.attn2.to_q +mid_block.attentions.0.transformer_blocks.3.attn2.to_k +mid_block.attentions.0.transformer_blocks.3.attn2.to_v +mid_block.attentions.0.transformer_blocks.3.attn2.to_out +mid_block.attentions.0.transformer_blocks.3.attn2.to_out.0 +mid_block.attentions.0.transformer_blocks.3.attn2.to_out.1 +mid_block.attentions.0.transformer_blocks.3.norm3 +mid_block.attentions.0.transformer_blocks.3.ff +mid_block.attentions.0.transformer_blocks.3.ff.net +mid_block.attentions.0.transformer_blocks.3.ff.net.0 +mid_block.attentions.0.transformer_blocks.3.ff.net.0.proj +mid_block.attentions.0.transformer_blocks.3.ff.net.1 +mid_block.attentions.0.transformer_blocks.3.ff.net.2 +mid_block.attentions.0.transformer_blocks.4 +mid_block.attentions.0.transformer_blocks.4.norm1 +mid_block.attentions.0.transformer_blocks.4.attn1 +mid_block.attentions.0.transformer_blocks.4.attn1.to_q +mid_block.attentions.0.transformer_blocks.4.attn1.to_k +mid_block.attentions.0.transformer_blocks.4.attn1.to_v +mid_block.attentions.0.transformer_blocks.4.attn1.to_out +mid_block.attentions.0.transformer_blocks.4.attn1.to_out.0 +mid_block.attentions.0.transformer_blocks.4.attn1.to_out.1 +mid_block.attentions.0.transformer_blocks.4.norm2 +mid_block.attentions.0.transformer_blocks.4.attn2 +mid_block.attentions.0.transformer_blocks.4.attn2.to_q +mid_block.attentions.0.transformer_blocks.4.attn2.to_k +mid_block.attentions.0.transformer_blocks.4.attn2.to_v +mid_block.attentions.0.transformer_blocks.4.attn2.to_out +mid_block.attentions.0.transformer_blocks.4.attn2.to_out.0 +mid_block.attentions.0.transformer_blocks.4.attn2.to_out.1 +mid_block.attentions.0.transformer_blocks.4.norm3 +mid_block.attentions.0.transformer_blocks.4.ff +mid_block.attentions.0.transformer_blocks.4.ff.net +mid_block.attentions.0.transformer_blocks.4.ff.net.0 +mid_block.attentions.0.transformer_blocks.4.ff.net.0.proj +mid_block.attentions.0.transformer_blocks.4.ff.net.1 +mid_block.attentions.0.transformer_blocks.4.ff.net.2 +mid_block.attentions.0.transformer_blocks.5 +mid_block.attentions.0.transformer_blocks.5.norm1 +mid_block.attentions.0.transformer_blocks.5.attn1 +mid_block.attentions.0.transformer_blocks.5.attn1.to_q +mid_block.attentions.0.transformer_blocks.5.attn1.to_k +mid_block.attentions.0.transformer_blocks.5.attn1.to_v +mid_block.attentions.0.transformer_blocks.5.attn1.to_out +mid_block.attentions.0.transformer_blocks.5.attn1.to_out.0 +mid_block.attentions.0.transformer_blocks.5.attn1.to_out.1 +mid_block.attentions.0.transformer_blocks.5.norm2 +mid_block.attentions.0.transformer_blocks.5.attn2 +mid_block.attentions.0.transformer_blocks.5.attn2.to_q +mid_block.attentions.0.transformer_blocks.5.attn2.to_k +mid_block.attentions.0.transformer_blocks.5.attn2.to_v +mid_block.attentions.0.transformer_blocks.5.attn2.to_out +mid_block.attentions.0.transformer_blocks.5.attn2.to_out.0 +mid_block.attentions.0.transformer_blocks.5.attn2.to_out.1 +mid_block.attentions.0.transformer_blocks.5.norm3 +mid_block.attentions.0.transformer_blocks.5.ff +mid_block.attentions.0.transformer_blocks.5.ff.net +mid_block.attentions.0.transformer_blocks.5.ff.net.0 +mid_block.attentions.0.transformer_blocks.5.ff.net.0.proj +mid_block.attentions.0.transformer_blocks.5.ff.net.1 +mid_block.attentions.0.transformer_blocks.5.ff.net.2 +mid_block.attentions.0.transformer_blocks.6 +mid_block.attentions.0.transformer_blocks.6.norm1 +mid_block.attentions.0.transformer_blocks.6.attn1 +mid_block.attentions.0.transformer_blocks.6.attn1.to_q +mid_block.attentions.0.transformer_blocks.6.attn1.to_k +mid_block.attentions.0.transformer_blocks.6.attn1.to_v +mid_block.attentions.0.transformer_blocks.6.attn1.to_out +mid_block.attentions.0.transformer_blocks.6.attn1.to_out.0 +mid_block.attentions.0.transformer_blocks.6.attn1.to_out.1 +mid_block.attentions.0.transformer_blocks.6.norm2 +mid_block.attentions.0.transformer_blocks.6.attn2 +mid_block.attentions.0.transformer_blocks.6.attn2.to_q +mid_block.attentions.0.transformer_blocks.6.attn2.to_k +mid_block.attentions.0.transformer_blocks.6.attn2.to_v +mid_block.attentions.0.transformer_blocks.6.attn2.to_out +mid_block.attentions.0.transformer_blocks.6.attn2.to_out.0 +mid_block.attentions.0.transformer_blocks.6.attn2.to_out.1 +mid_block.attentions.0.transformer_blocks.6.norm3 +mid_block.attentions.0.transformer_blocks.6.ff +mid_block.attentions.0.transformer_blocks.6.ff.net +mid_block.attentions.0.transformer_blocks.6.ff.net.0 +mid_block.attentions.0.transformer_blocks.6.ff.net.0.proj +mid_block.attentions.0.transformer_blocks.6.ff.net.1 +mid_block.attentions.0.transformer_blocks.6.ff.net.2 +mid_block.attentions.0.transformer_blocks.7 +mid_block.attentions.0.transformer_blocks.7.norm1 +mid_block.attentions.0.transformer_blocks.7.attn1 +mid_block.attentions.0.transformer_blocks.7.attn1.to_q +mid_block.attentions.0.transformer_blocks.7.attn1.to_k +mid_block.attentions.0.transformer_blocks.7.attn1.to_v +mid_block.attentions.0.transformer_blocks.7.attn1.to_out +mid_block.attentions.0.transformer_blocks.7.attn1.to_out.0 +mid_block.attentions.0.transformer_blocks.7.attn1.to_out.1 +mid_block.attentions.0.transformer_blocks.7.norm2 +mid_block.attentions.0.transformer_blocks.7.attn2 +mid_block.attentions.0.transformer_blocks.7.attn2.to_q +mid_block.attentions.0.transformer_blocks.7.attn2.to_k +mid_block.attentions.0.transformer_blocks.7.attn2.to_v +mid_block.attentions.0.transformer_blocks.7.attn2.to_out +mid_block.attentions.0.transformer_blocks.7.attn2.to_out.0 +mid_block.attentions.0.transformer_blocks.7.attn2.to_out.1 +mid_block.attentions.0.transformer_blocks.7.norm3 +mid_block.attentions.0.transformer_blocks.7.ff +mid_block.attentions.0.transformer_blocks.7.ff.net +mid_block.attentions.0.transformer_blocks.7.ff.net.0 +mid_block.attentions.0.transformer_blocks.7.ff.net.0.proj +mid_block.attentions.0.transformer_blocks.7.ff.net.1 +mid_block.attentions.0.transformer_blocks.7.ff.net.2 +mid_block.attentions.0.transformer_blocks.8 +mid_block.attentions.0.transformer_blocks.8.norm1 +mid_block.attentions.0.transformer_blocks.8.attn1 +mid_block.attentions.0.transformer_blocks.8.attn1.to_q +mid_block.attentions.0.transformer_blocks.8.attn1.to_k +mid_block.attentions.0.transformer_blocks.8.attn1.to_v +mid_block.attentions.0.transformer_blocks.8.attn1.to_out +mid_block.attentions.0.transformer_blocks.8.attn1.to_out.0 +mid_block.attentions.0.transformer_blocks.8.attn1.to_out.1 +mid_block.attentions.0.transformer_blocks.8.norm2 +mid_block.attentions.0.transformer_blocks.8.attn2 +mid_block.attentions.0.transformer_blocks.8.attn2.to_q +mid_block.attentions.0.transformer_blocks.8.attn2.to_k +mid_block.attentions.0.transformer_blocks.8.attn2.to_v +mid_block.attentions.0.transformer_blocks.8.attn2.to_out +mid_block.attentions.0.transformer_blocks.8.attn2.to_out.0 +mid_block.attentions.0.transformer_blocks.8.attn2.to_out.1 +mid_block.attentions.0.transformer_blocks.8.norm3 +mid_block.attentions.0.transformer_blocks.8.ff +mid_block.attentions.0.transformer_blocks.8.ff.net +mid_block.attentions.0.transformer_blocks.8.ff.net.0 +mid_block.attentions.0.transformer_blocks.8.ff.net.0.proj +mid_block.attentions.0.transformer_blocks.8.ff.net.1 +mid_block.attentions.0.transformer_blocks.8.ff.net.2 +mid_block.attentions.0.transformer_blocks.9 +mid_block.attentions.0.transformer_blocks.9.norm1 +mid_block.attentions.0.transformer_blocks.9.attn1 +mid_block.attentions.0.transformer_blocks.9.attn1.to_q +mid_block.attentions.0.transformer_blocks.9.attn1.to_k +mid_block.attentions.0.transformer_blocks.9.attn1.to_v +mid_block.attentions.0.transformer_blocks.9.attn1.to_out +mid_block.attentions.0.transformer_blocks.9.attn1.to_out.0 +mid_block.attentions.0.transformer_blocks.9.attn1.to_out.1 +mid_block.attentions.0.transformer_blocks.9.norm2 +mid_block.attentions.0.transformer_blocks.9.attn2 +mid_block.attentions.0.transformer_blocks.9.attn2.to_q +mid_block.attentions.0.transformer_blocks.9.attn2.to_k +mid_block.attentions.0.transformer_blocks.9.attn2.to_v +mid_block.attentions.0.transformer_blocks.9.attn2.to_out +mid_block.attentions.0.transformer_blocks.9.attn2.to_out.0 +mid_block.attentions.0.transformer_blocks.9.attn2.to_out.1 +mid_block.attentions.0.transformer_blocks.9.norm3 +mid_block.attentions.0.transformer_blocks.9.ff +mid_block.attentions.0.transformer_blocks.9.ff.net +mid_block.attentions.0.transformer_blocks.9.ff.net.0 +mid_block.attentions.0.transformer_blocks.9.ff.net.0.proj +mid_block.attentions.0.transformer_blocks.9.ff.net.1 +mid_block.attentions.0.transformer_blocks.9.ff.net.2 +mid_block.attentions.0.proj_out +mid_block.resnets +mid_block.resnets.0 +mid_block.resnets.0.norm1 +mid_block.resnets.0.conv1 +mid_block.resnets.0.time_emb_proj +mid_block.resnets.0.norm2 +mid_block.resnets.0.dropout +mid_block.resnets.0.conv2 +mid_block.resnets.1 +mid_block.resnets.1.norm1 +mid_block.resnets.1.conv1 +mid_block.resnets.1.time_emb_proj +mid_block.resnets.1.norm2 +mid_block.resnets.1.dropout +mid_block.resnets.1.conv2 +conv_norm_out +conv_out diff --git a/invokeai/backend/hidiffusion/utils.py b/invokeai/backend/hidiffusion/utils.py new file mode 100644 index 00000000000..b8db24238b6 --- /dev/null +++ b/invokeai/backend/hidiffusion/utils.py @@ -0,0 +1,35 @@ +import torch + + +def isinstance_str(x: object, cls_name: str, prefix: bool = False, contains: bool = False): + """ + Checks whether x has any class equal to, prefixed with, or contains (cls_name) in its ancestry. + Doesn't require access to the class's implementation. + + Useful for patching! + """ + + for _cls in x.__class__.__mro__: + if _cls.__name__.startswith(cls_name) and prefix: + return True + if contains and cls_name in _cls.__name__: + return True + if _cls.__name__ == cls_name: + return True + + return False + + +def init_generator(device: torch.device, fallback: torch.Generator | None = None): + """ + Forks the current default random generator given device. + """ + if device.type == "cpu": + return torch.Generator(device="cpu").set_state(torch.get_rng_state()) + elif device.type == "cuda": + return torch.Generator(device=device).set_state(torch.cuda.get_rng_state()) + else: + if fallback is None: + return init_generator(torch.device("cpu")) + else: + return fallback diff --git a/invokeai/backend/stable_diffusion/extensions/hidiffusion.py b/invokeai/backend/stable_diffusion/extensions/hidiffusion.py new file mode 100644 index 00000000000..444c90f5480 --- /dev/null +++ b/invokeai/backend/stable_diffusion/extensions/hidiffusion.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from contextlib import contextmanager +from typing import Optional + +import torch +from diffusers import UNet2DConditionModel + +from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase +from invokeai.backend.stable_diffusion.hidiffusion_utils import hidiffusion_patch +from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage + + +class HiDiffusionExt(ExtensionBase): + def __init__( + self, + name_or_path: Optional[str], + apply_raunet: bool = True, + apply_window_attn: bool = True, + t1_ratio: Optional[float] = None, + t2_ratio: Optional[float] = None, + generator: torch.Generator | None = None, + ): + super().__init__() + self._name_or_path = name_or_path + self._apply_raunet = apply_raunet + self._apply_window_attn = apply_window_attn + self._t1_ratio = t1_ratio + self._t2_ratio = t2_ratio + self._generator = generator + + @contextmanager + def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeightsStorage): + with hidiffusion_patch( + unet, + name_or_path=self._name_or_path, + apply_raunet=self._apply_raunet, + apply_window_attn=self._apply_window_attn, + t1_ratio=self._t1_ratio, + t2_ratio=self._t2_ratio, + generator=self._generator, + ): + yield None diff --git a/invokeai/backend/stable_diffusion/hidiffusion_utils.py b/invokeai/backend/stable_diffusion/hidiffusion_utils.py new file mode 100644 index 00000000000..327d7d083b1 --- /dev/null +++ b/invokeai/backend/stable_diffusion/hidiffusion_utils.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import copy +import sys +from contextlib import contextmanager +from typing import Any, Optional + +import torch + +from invokeai.backend.hidiffusion.hidiffusion import ( + switching_threshold_ratio_dict as _switching_threshold_ratio_dict, +) +from invokeai.backend.hidiffusion.hidiffusion import ( + text_to_img_controlnet_switching_threshold_ratio_dict as _text_to_img_controlnet_switching_threshold_ratio_dict, +) + + +@contextmanager +def hidiffusion_patch( + model: Any, + name_or_path: Optional[str], + apply_raunet: bool = True, + apply_window_attn: bool = True, + t1_ratio: Optional[float] = None, + t2_ratio: Optional[float] = None, + generator: torch.Generator | None = None, +): + """Context manager that applies HiDiffusion and restores the model on exit.""" + from invokeai.backend.hidiffusion.hidiffusion import apply_hidiffusion, remove_hidiffusion + + target = model.unet if hasattr(model, "unet") else model + + had_model_name_or_path = hasattr(model, "_name_or_path") + had_config = hasattr(model, "config") + config = model.config if had_config else None + had_config_name_or_path = bool(config) and hasattr(config, "_name_or_path") + + original_model_name_or_path = model._name_or_path if had_model_name_or_path else None + original_config_name_or_path = config._name_or_path if had_config_name_or_path else None + + effective_name_or_path = ( + name_or_path + or getattr(model, "name_or_path", None) + or original_model_name_or_path + or original_config_name_or_path + or "" + ) + + def _set_name_or_path_on_config(cfg, value: str) -> bool: + if cfg is None: + return False + if hasattr(cfg, "_internal_dict"): + try: + cfg._internal_dict["_name_or_path"] = value + return True + except Exception: + pass + try: + object.__setattr__(cfg, "_name_or_path", value) + return True + except Exception: + pass + try: + cfg.__dict__["_name_or_path"] = value + return True + except Exception: + return False + + original_num_upsamplers = getattr(target, "num_upsamplers", None) + + ratio_overrides = None + ratio_dicts = None + if t1_ratio is not None or t2_ratio is not None: + ratio_dicts = ( + _switching_threshold_ratio_dict, + _text_to_img_controlnet_switching_threshold_ratio_dict, + ) + + set_model_name_or_path = False + set_config_name_or_path = False + try: + try: + object.__setattr__(model, "_name_or_path", effective_name_or_path) + set_model_name_or_path = True + except Exception: + set_model_name_or_path = False + + set_config_name_or_path = _set_name_or_path_on_config(config, effective_name_or_path) + + # Ensure the property resolves to a non-None value before calling HiDiffusion. + try: + if getattr(model, "name_or_path", None) is None: + if not set_model_name_or_path: + try: + object.__setattr__(model, "_name_or_path", effective_name_or_path) + set_model_name_or_path = True + except Exception: + pass + if not set_config_name_or_path: + set_config_name_or_path = _set_name_or_path_on_config(config, effective_name_or_path) + except Exception: + pass + + if ratio_dicts is not None: + ratio_overrides = (copy.deepcopy(ratio_dicts[0]), copy.deepcopy(ratio_dicts[1])) + + def _apply_ratio_overrides(ratio_dict: dict) -> None: + for _, entry in ratio_dict.items(): + if t1_ratio is not None: + entry["T1_ratio"] = t1_ratio + if t2_ratio is not None and "T2_ratio" in entry: + entry["T2_ratio"] = t2_ratio + + _apply_ratio_overrides(ratio_dicts[0]) + _apply_ratio_overrides(ratio_dicts[1]) + + apply_hidiffusion( + model, + apply_raunet=apply_raunet, + apply_window_attn=apply_window_attn, + generator=generator, + ) + yield + finally: + had_active_exception = sys.exc_info()[0] is not None + teardown_error: Exception | None = None + try: + remove_hidiffusion(model) + except Exception as error: + if not had_active_exception: + teardown_error = error + if ratio_overrides is not None and ratio_dicts is not None: + ratio_dicts[0].clear() + ratio_dicts[0].update(ratio_overrides[0]) + ratio_dicts[1].clear() + ratio_dicts[1].update(ratio_overrides[1]) + if original_num_upsamplers is not None: + target.num_upsamplers = original_num_upsamplers + if set_model_name_or_path: + if had_model_name_or_path: + try: + object.__setattr__(model, "_name_or_path", original_model_name_or_path) + except Exception: + pass + else: + try: + delattr(model, "_name_or_path") + except AttributeError: + pass + if set_config_name_or_path and had_config: + if had_config_name_or_path: + _set_name_or_path_on_config(config, original_config_name_or_path) + elif config is not None: + internal_dict = getattr(config, "_internal_dict", None) + if internal_dict is not None: + try: + internal_dict.pop("_name_or_path", None) + except Exception: + pass + try: + delattr(config, "_name_or_path") + except AttributeError: + pass + if teardown_error is not None: + raise teardown_error diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index a509f79080c..e19d0163e31 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -24036,6 +24036,60 @@ "title": "CFG Rescale Multiplier", "type": "number" }, + "hidiffusion": { + "default": false, + "description": "Apply HiDiffusion (RAU-Net + MSW-MSA) for higher-resolution denoising", + "field_kind": "input", + "input": "any", + "orig_default": false, + "orig_required": false, + "title": "HiDiffusion", + "type": "boolean" + }, + "hidiffusion_raunet": { + "default": true, + "description": "Apply HiDiffusion RAU-Net blocks", + "field_kind": "input", + "input": "any", + "orig_default": true, + "orig_required": false, + "title": "HiDiffusion: RAU-Net", + "type": "boolean" + }, + "hidiffusion_window_attn": { + "default": true, + "description": "Apply HiDiffusion window attention blocks", + "field_kind": "input", + "input": "any", + "orig_default": true, + "orig_required": false, + "title": "HiDiffusion: Window Attention", + "type": "boolean" + }, + "hidiffusion_t1_ratio": { + "default": 0.4, + "description": "Override HiDiffusion early switch threshold (T1 ratio)", + "field_kind": "input", + "input": "any", + "maximum": 1, + "minimum": 0, + "orig_default": 0.4, + "orig_required": false, + "title": "HiDiffusion: T1 Ratio", + "type": "number" + }, + "hidiffusion_t2_ratio": { + "default": 0.0, + "description": "Override HiDiffusion late switch threshold (T2 ratio)", + "field_kind": "input", + "input": "any", + "maximum": 1, + "minimum": 0, + "orig_default": 0.0, + "orig_required": false, + "title": "HiDiffusion: T2 Ratio", + "type": "number" + }, "latents": { "anyOf": [ { @@ -24082,7 +24136,7 @@ "tags": ["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"], "title": "Denoise - SD1.5, SDXL", "type": "object", - "version": "1.5.4", + "version": "1.6.0", "output": { "$ref": "#/components/schemas/LatentsOutput" } @@ -24395,6 +24449,60 @@ "title": "CFG Rescale Multiplier", "type": "number" }, + "hidiffusion": { + "default": false, + "description": "Apply HiDiffusion (RAU-Net + MSW-MSA) for higher-resolution denoising", + "field_kind": "input", + "input": "any", + "orig_default": false, + "orig_required": false, + "title": "HiDiffusion", + "type": "boolean" + }, + "hidiffusion_raunet": { + "default": true, + "description": "Apply HiDiffusion RAU-Net blocks", + "field_kind": "input", + "input": "any", + "orig_default": true, + "orig_required": false, + "title": "HiDiffusion: RAU-Net", + "type": "boolean" + }, + "hidiffusion_window_attn": { + "default": true, + "description": "Apply HiDiffusion window attention blocks", + "field_kind": "input", + "input": "any", + "orig_default": true, + "orig_required": false, + "title": "HiDiffusion: Window Attention", + "type": "boolean" + }, + "hidiffusion_t1_ratio": { + "default": 0.4, + "description": "Override HiDiffusion early switch threshold (T1 ratio)", + "field_kind": "input", + "input": "any", + "maximum": 1, + "minimum": 0, + "orig_default": 0.4, + "orig_required": false, + "title": "HiDiffusion: T1 Ratio", + "type": "number" + }, + "hidiffusion_t2_ratio": { + "default": 0.0, + "description": "Override HiDiffusion late switch threshold (T2 ratio)", + "field_kind": "input", + "input": "any", + "maximum": 1, + "minimum": 0, + "orig_default": 0.0, + "orig_required": false, + "title": "HiDiffusion: T2 Ratio", + "type": "number" + }, "latents": { "anyOf": [ { @@ -24441,7 +24549,7 @@ "tags": ["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"], "title": "Denoise - SD1.5, SDXL + Metadata", "type": "object", - "version": "1.1.1", + "version": "1.2.0", "output": { "$ref": "#/components/schemas/LatentsMetaOutput" } diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json index 3d80e90ecfd..a4c7986b10c 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -1068,6 +1068,11 @@ "seedreamOptimizePrompt": "Seedream Optimize Prompt", "guidance": "Guidance", "height": "Height", + "hiDiffusion": "HiDiffusion", + "hiDiffusionRauNet": "HiDiffusion: RAU-Net", + "hiDiffusionWindowAttn": "HiDiffusion: Window Attention", + "hiDiffusionT1Ratio": "HiDiffusion: T1 Ratio", + "hiDiffusionT2Ratio": "HiDiffusion: T2 Ratio", "imageDetails": "Image Details", "imageDimensions": "Image Dimensions", "imageSize": "Image Size", @@ -1717,6 +1722,11 @@ "cfgScale": "CFG Scale", "cfgRescaleMultiplier": "CFG Rescale Multiplier", "clipSkip": "CLIP Skip", + "hiDiffusion": "HiDiffusion", + "hiDiffusionRauNet": "HiDiffusion: RAU-Net", + "hiDiffusionWindowAttn": "HiDiffusion: Window Attention", + "hiDiffusionT1Ratio": "HiDiffusion: T1 Ratio", + "hiDiffusionT2Ratio": "HiDiffusion: T2 Ratio", "coherenceMode": "Mode", "coherenceEdgeSize": "Edge Size", "coherenceMinDenoise": "Min Denoise", @@ -2131,6 +2141,41 @@ "invalidBboxDesc": "The bounding box has no valid dimensions" }, "popovers": { + "hidiffusion": { + "heading": "HiDiffusion", + "paragraphs": [ + "Enables HiDiffusion denoising for higher-resolution generation.", + "Try at 1536px or higher for the most visible effect." + ] + }, + "hidiffusionRauNet": { + "heading": "HiDiffusion: RAU-Net", + "paragraphs": [ + "Enables RAU-Net blocks used by HiDiffusion.", + "Can improve structure and detail at higher resolutions." + ] + }, + "hidiffusionWindowAttn": { + "heading": "HiDiffusion: Window Attention", + "paragraphs": [ + "Enables windowed attention used by HiDiffusion.", + "Can boost local detail, but may affect global coherence." + ] + }, + "hidiffusionT1Ratio": { + "heading": "HiDiffusion: T1 Ratio", + "paragraphs": [ + "Controls the early switch point for HiDiffusion (T1).", + "Lower values switch earlier; higher values preserve global structure longer." + ] + }, + "hidiffusionT2Ratio": { + "heading": "HiDiffusion: T2 Ratio", + "paragraphs": [ + "Controls the late switch point for HiDiffusion (T2).", + "Higher values keep window attention active longer." + ] + }, "clipSkip": { "heading": "CLIP Skip", "paragraphs": [ diff --git a/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts b/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts index 5f9fe8ffb74..d91077f6607 100644 --- a/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts +++ b/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts @@ -75,6 +75,11 @@ export type Feature = | 'scaleBeforeProcessing' | 'seamlessTilingXAxis' | 'seamlessTilingYAxis' + | 'hidiffusion' + | 'hidiffusionRauNet' + | 'hidiffusionWindowAttn' + | 'hidiffusionT1Ratio' + | 'hidiffusionT2Ratio' | 'colorCompensation' | 'upscaleModel' | 'scale' @@ -246,6 +251,13 @@ export const POPOVER_DATA: { [key in Feature]?: PopoverData } = { seamlessTilingYAxis: { href: 'https://support.invoke.ai/support/solutions/articles/151000178161-advanced-settings', }, + hidiffusion: { + href: 'https://github.com/megvii-research/HiDiffusion', + }, + hidiffusionRauNet: {}, + hidiffusionWindowAttn: {}, + hidiffusionT1Ratio: {}, + hidiffusionT2Ratio: {}, } as const; export const OPEN_DELAY = 1000; // in milliseconds diff --git a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.test.ts b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.test.ts index 8e46e7ba075..ff143b562b9 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.test.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.test.ts @@ -144,10 +144,10 @@ describe('paramsSlice selectors for external models', () => { describe('paramsSliceConfig persisted state migration', () => { const migrate = paramsSliceConfig.persistConfig?.migrate; - it('backfills new Qwen Image fields when migrating from v2 and preserves existing params', () => { + it('backfills Qwen Image and HiDiffusion fields when migrating from v2 and preserves existing params', () => { expect(migrate).toBeDefined(); - // Build a valid pre-PR v2 persisted state by removing the fields that were added in v3 + // Build a valid pre-PR v2 persisted state by removing the fields that were added later. const initial = getInitialParamsState(); const v2State: Record = { ...initial, @@ -159,6 +159,11 @@ describe('paramsSliceConfig persisted state migration', () => { }; delete v2State.qwenImageVaeModel; delete v2State.qwenImageQwenVLEncoderModel; + delete v2State.hiDiffusionEnabled; + delete v2State.hiDiffusionRauNetEnabled; + delete v2State.hiDiffusionWindowAttnEnabled; + delete v2State.hiDiffusionT1Ratio; + delete v2State.hiDiffusionT2Ratio; const result = migrate?.(v2State) as ReturnType; @@ -167,6 +172,11 @@ describe('paramsSliceConfig persisted state migration', () => { expect(result._version).toBe(4); expect(result.qwenImageVaeModel).toBeNull(); expect(result.qwenImageQwenVLEncoderModel).toBeNull(); + expect(result.hiDiffusionEnabled).toBe(false); + expect(result.hiDiffusionRauNetEnabled).toBe(true); + expect(result.hiDiffusionWindowAttnEnabled).toBe(true); + expect(result.hiDiffusionT1Ratio).toBe(0.4); + expect(result.hiDiffusionT2Ratio).toBe(0.0); // Existing params should be preserved expect(result.positivePrompt).toBe('a fluffy cat'); expect(result.seed).toBe(42); diff --git a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts index 21d2cc971b0..7884020abc6 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts @@ -170,6 +170,21 @@ const slice = createSlice({ setOptimizedDenoisingEnabled: (state, action: PayloadAction) => { state.optimizedDenoisingEnabled = action.payload; }, + setHiDiffusionEnabled: (state, action: PayloadAction) => { + state.hiDiffusionEnabled = action.payload; + }, + setHiDiffusionRauNetEnabled: (state, action: PayloadAction) => { + state.hiDiffusionRauNetEnabled = action.payload; + }, + setHiDiffusionWindowAttnEnabled: (state, action: PayloadAction) => { + state.hiDiffusionWindowAttnEnabled = action.payload; + }, + setHiDiffusionT1Ratio: (state, action: PayloadAction) => { + state.hiDiffusionT1Ratio = action.payload; + }, + setHiDiffusionT2Ratio: (state, action: PayloadAction) => { + state.hiDiffusionT2Ratio = action.payload; + }, setSeamlessXAxis: (state, action: PayloadAction) => { state.seamlessXAxis = action.payload; }, @@ -863,6 +878,11 @@ export const { setSeed, setImg2imgStrength, setOptimizedDenoisingEnabled, + setHiDiffusionEnabled, + setHiDiffusionRauNetEnabled, + setHiDiffusionWindowAttnEnabled, + setHiDiffusionT1Ratio, + setHiDiffusionT2Ratio, setSeamlessXAxis, setSeamlessYAxis, setShouldRandomizeSeed, @@ -985,6 +1005,22 @@ export const paramsSliceConfig: SliceConfig = { state.pidSteps = 4; } + if (!('hiDiffusionEnabled' in state)) { + state.hiDiffusionEnabled = false; + } + if (!('hiDiffusionRauNetEnabled' in state)) { + state.hiDiffusionRauNetEnabled = true; + } + if (!('hiDiffusionWindowAttnEnabled' in state)) { + state.hiDiffusionWindowAttnEnabled = true; + } + if (!('hiDiffusionT1Ratio' in state)) { + state.hiDiffusionT1Ratio = 0.4; + } + if (!('hiDiffusionT2Ratio' in state)) { + state.hiDiffusionT2Ratio = 0.0; + } + return zParamsState.parse(state); }, }, @@ -1069,6 +1105,11 @@ export const selectInfillPatchmatchDownscaleSize = createParamsSelector( export const selectInfillColorValue = createParamsSelector((params) => params.infillColorValue); export const selectImg2imgStrength = createParamsSelector((params) => params.img2imgStrength); export const selectOptimizedDenoisingEnabled = createParamsSelector((params) => params.optimizedDenoisingEnabled); +export const selectHiDiffusionEnabled = createParamsSelector((params) => params.hiDiffusionEnabled); +export const selectHiDiffusionRauNetEnabled = createParamsSelector((params) => params.hiDiffusionRauNetEnabled); +export const selectHiDiffusionWindowAttnEnabled = createParamsSelector((params) => params.hiDiffusionWindowAttnEnabled); +export const selectHiDiffusionT1Ratio = createParamsSelector((params) => params.hiDiffusionT1Ratio); +export const selectHiDiffusionT2Ratio = createParamsSelector((params) => params.hiDiffusionT2Ratio); export const selectPositivePrompt = createParamsSelector((params) => params.positivePrompt); export const selectNegativePrompt = createParamsSelector((params) => params.negativePrompt); export const selectNegativePromptWithFallback = createParamsSelector((params) => params.negativePrompt ?? ''); diff --git a/invokeai/frontend/web/src/features/controlLayers/store/types.ts b/invokeai/frontend/web/src/features/controlLayers/store/types.ts index 47e6543ad13..485b66fb25d 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/types.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/types.ts @@ -832,6 +832,11 @@ export const zParamsState = z.object({ guidance: zParameterGuidance, img2imgStrength: zParameterStrength, optimizedDenoisingEnabled: z.boolean(), + hiDiffusionEnabled: z.boolean(), + hiDiffusionRauNetEnabled: z.boolean(), + hiDiffusionWindowAttnEnabled: z.boolean(), + hiDiffusionT1Ratio: z.number(), + hiDiffusionT2Ratio: z.number(), iterations: z.number(), scheduler: zParameterScheduler, fluxScheduler: zParameterFluxScheduler, @@ -961,6 +966,11 @@ export const getInitialParamsState = (): ParamsState => ({ guidance: 4, img2imgStrength: 0.75, optimizedDenoisingEnabled: true, + hiDiffusionEnabled: false, + hiDiffusionRauNetEnabled: true, + hiDiffusionWindowAttnEnabled: true, + hiDiffusionT1Ratio: 0.4, + hiDiffusionT2Ratio: 0.0, iterations: 1, scheduler: 'dpmpp_3m_k', fluxScheduler: 'euler', diff --git a/invokeai/frontend/web/src/features/metadata/parsing.test.ts b/invokeai/frontend/web/src/features/metadata/parsing.test.ts index f4a542f658c..6072aa1c10f 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.test.ts +++ b/invokeai/frontend/web/src/features/metadata/parsing.test.ts @@ -1,3 +1,5 @@ +import type { AppStore } from 'app/store/store'; +import { setHiDiffusionEnabled } from 'features/controlLayers/store/paramsSlice'; import { describe, expect, it, vi } from 'vitest'; import { ImageMetadataHandlers, MetadataUtils, parseMetadataHandler } from './parsing'; @@ -159,3 +161,37 @@ describe('Qwen metadata parsing', () => { expect(mockStore.dispatch).toHaveBeenCalledTimes(1); }); }); + +describe('HiDiffusion metadata parsing', () => { + it('disables HiDiffusion when recalling all metadata from an older image', async () => { + let hiDiffusionEnabled = true; + const store = { + dispatch: vi.fn((action) => { + if (action.type === setHiDiffusionEnabled.type) { + hiDiffusionEnabled = action.payload; + } + return action; + }), + getState: vi.fn(() => ({ + params: { model: null }, + })), + } as unknown as AppStore; + + await MetadataUtils.recallAllImageMetadata( + { + generation_mode: 'txt2img', + width: 512, + height: 512, + steps: 20, + cfg_scale: 7.5, + scheduler: 'euler', + positive_prompt: 'an older image', + negative_prompt: '', + }, + store + ); + + expect(store.dispatch).toHaveBeenCalledWith(setHiDiffusionEnabled(false)); + expect(hiDiffusionEnabled).toBe(false); + }); +}); diff --git a/invokeai/frontend/web/src/features/metadata/parsing.tsx b/invokeai/frontend/web/src/features/metadata/parsing.tsx index 6a54ea7c031..c3be73c453f 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.tsx +++ b/invokeai/frontend/web/src/features/metadata/parsing.tsx @@ -42,6 +42,11 @@ import { setFluxDypeScale, setFluxScheduler, setGuidance, + setHiDiffusionEnabled, + setHiDiffusionRauNetEnabled, + setHiDiffusionT1Ratio, + setHiDiffusionT2Ratio, + setHiDiffusionWindowAttnEnabled, setIdeogram4ColorPalette, setIdeogram4GuidanceScale, setIdeogram4Mu, @@ -698,6 +703,96 @@ const SeamlessY: SingleMetadataHandler = { }; //#endregion SeamlessY +//#region HiDiffusion +const HiDiffusion: SingleMetadataHandler = { + [SingleMetadataKey]: true, + type: 'HiDiffusion', + parse: (metadata, _store) => { + const raw = getProperty(metadata, 'hidiffusion'); + const parsed = raw === undefined ? false : z.boolean().parse(raw); + return Promise.resolve(parsed); + }, + recall: (value, store) => { + store.dispatch(setHiDiffusionEnabled(value)); + }, + i18nKey: 'metadata.hiDiffusion', + LabelComponent: MetadataLabel, + ValueComponent: ({ value }: SingleMetadataValueProps) => , +}; +//#endregion HiDiffusion + +//#region HiDiffusionRAUNet +const HiDiffusionRauNet: SingleMetadataHandler = { + [SingleMetadataKey]: true, + type: 'HiDiffusionRauNet', + parse: (metadata, _store) => { + const raw = getProperty(metadata, 'hidiffusion_raunet'); + const parsed = z.boolean().parse(raw); + return Promise.resolve(parsed); + }, + recall: (value, store) => { + store.dispatch(setHiDiffusionRauNetEnabled(value)); + }, + i18nKey: 'metadata.hiDiffusionRauNet', + LabelComponent: MetadataLabel, + ValueComponent: ({ value }: SingleMetadataValueProps) => , +}; +//#endregion HiDiffusionRAUNet + +//#region HiDiffusionWindowAttn +const HiDiffusionWindowAttn: SingleMetadataHandler = { + [SingleMetadataKey]: true, + type: 'HiDiffusionWindowAttn', + parse: (metadata, _store) => { + const raw = getProperty(metadata, 'hidiffusion_window_attn'); + const parsed = z.boolean().parse(raw); + return Promise.resolve(parsed); + }, + recall: (value, store) => { + store.dispatch(setHiDiffusionWindowAttnEnabled(value)); + }, + i18nKey: 'metadata.hiDiffusionWindowAttn', + LabelComponent: MetadataLabel, + ValueComponent: ({ value }: SingleMetadataValueProps) => , +}; +//#endregion HiDiffusionWindowAttn + +//#region HiDiffusionT1Ratio +const HiDiffusionT1Ratio: SingleMetadataHandler = { + [SingleMetadataKey]: true, + type: 'HiDiffusionT1Ratio', + parse: (metadata, _store) => { + const raw = getProperty(metadata, 'hidiffusion_t1_ratio'); + const parsed = z.number().parse(raw); + return Promise.resolve(parsed); + }, + recall: (value, store) => { + store.dispatch(setHiDiffusionT1Ratio(value)); + }, + i18nKey: 'metadata.hiDiffusionT1Ratio', + LabelComponent: MetadataLabel, + ValueComponent: ({ value }: SingleMetadataValueProps) => , +}; +//#endregion HiDiffusionT1Ratio + +//#region HiDiffusionT2Ratio +const HiDiffusionT2Ratio: SingleMetadataHandler = { + [SingleMetadataKey]: true, + type: 'HiDiffusionT2Ratio', + parse: (metadata, _store) => { + const raw = getProperty(metadata, 'hidiffusion_t2_ratio'); + const parsed = z.number().parse(raw); + return Promise.resolve(parsed); + }, + recall: (value, store) => { + store.dispatch(setHiDiffusionT2Ratio(value)); + }, + i18nKey: 'metadata.hiDiffusionT2Ratio', + LabelComponent: MetadataLabel, + ValueComponent: ({ value }: SingleMetadataValueProps) => , +}; +//#endregion HiDiffusionT2Ratio + //#region ZImageSeedVarianceEnabled const ZImageSeedVarianceEnabled: SingleMetadataHandler = { [SingleMetadataKey]: true, @@ -2153,6 +2248,11 @@ export const ImageMetadataHandlers = { DenoisingStrength, SeamlessX, SeamlessY, + HiDiffusion, + HiDiffusionRauNet, + HiDiffusionWindowAttn, + HiDiffusionT1Ratio, + HiDiffusionT2Ratio, RefinerModel, RefinerSteps, RefinerCFGScale, diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.ts index 5485834db13..11a5333c7d2 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.ts @@ -23,6 +23,11 @@ export const addSDXLRefiner = async ( refinerScheduler, refinerCFGScale, refinerStart, + hiDiffusionEnabled, + hiDiffusionRauNetEnabled, + hiDiffusionT1Ratio, + hiDiffusionT2Ratio, + hiDiffusionWindowAttnEnabled, } = state.params; assert(refinerModel, 'No refiner model found in state'); @@ -57,6 +62,11 @@ export const addSDXLRefiner = async ( cfg_scale: refinerCFGScale, steps: refinerSteps, scheduler: refinerScheduler, + hidiffusion: hiDiffusionEnabled, + hidiffusion_raunet: hiDiffusionRauNetEnabled, + hidiffusion_window_attn: hiDiffusionWindowAttnEnabled, + hidiffusion_t1_ratio: hiDiffusionEnabled ? hiDiffusionT1Ratio : undefined, + hidiffusion_t2_ratio: hiDiffusionEnabled ? hiDiffusionT2Ratio : undefined, denoising_start: refinerStart, denoising_end: 1, }); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildHiDiffusionGraph.test.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildHiDiffusionGraph.test.ts new file mode 100644 index 00000000000..8853af97baa --- /dev/null +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildHiDiffusionGraph.test.ts @@ -0,0 +1,196 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +vi.mock('app/logging/logger', () => ({ + logger: () => ({ + debug: vi.fn(), + }), +})); + +let nextId = 0; +vi.mock('features/controlLayers/konva/util', () => ({ + getPrefixedId: (prefix: string) => `${prefix}:${nextId++}`, +})); + +const sd1Model = { + key: 'sd1-model', + hash: 'sd1-hash', + name: 'SD 1.5', + base: 'sd-1', + type: 'main', +}; + +const sdxlModel = { + key: 'sdxl-model', + hash: 'sdxl-hash', + name: 'SDXL', + base: 'sdxl', + type: 'main', +}; + +const defaultParams = { + cfgScale: 7.5, + cfgRescaleMultiplier: 0, + hiDiffusionEnabled: false, + hiDiffusionRauNetEnabled: false, + hiDiffusionT1Ratio: 0.25, + hiDiffusionT2Ratio: 0.1, + hiDiffusionWindowAttnEnabled: false, + scheduler: 'euler', + steps: 20, + clipSkip: 0, + shouldUseCpuNoise: false, + vaePrecision: 'fp16', + vae: null, + colorCompensation: false, + refinerModel: null, + pidMode: 'off', +}; + +let currentModel: typeof sd1Model | typeof sdxlModel = sd1Model; +let params = { ...defaultParams }; + +vi.mock('features/controlLayers/store/paramsSlice', () => ({ + selectMainModelConfig: vi.fn(() => currentModel), + selectParamsSlice: vi.fn(() => params), +})); + +vi.mock('features/controlLayers/store/refImagesSlice', () => ({ + selectRefImagesSlice: vi.fn(() => ({ entities: [] })), +})); + +vi.mock('features/controlLayers/store/selectors', () => ({ + selectCanvasSlice: vi.fn(() => ({ + bbox: { rect: { x: 0, y: 0, width: 1024, height: 1024 } }, + controlLayers: { entities: [] }, + regionalGuidance: { entities: [] }, + })), + selectCanvasMetadata: vi.fn(() => ({})), +})); + +vi.mock('features/nodes/util/graph/generation/addControlAdapters', () => ({ + addControlNets: vi.fn(() => Promise.resolve({ addedControlNets: 0 })), + addT2IAdapters: vi.fn(() => Promise.resolve({ addedT2IAdapters: 0 })), +})); + +vi.mock('features/nodes/util/graph/generation/addImageToImage', () => ({ + addImageToImage: vi.fn(), +})); + +vi.mock('features/nodes/util/graph/generation/addInpaint', () => ({ + addInpaint: vi.fn(), +})); + +vi.mock('features/nodes/util/graph/generation/addIPAdapters', () => ({ + addIPAdapters: vi.fn(() => ({ addedIPAdapters: 0 })), +})); + +vi.mock('features/nodes/util/graph/generation/addLoRAs', () => ({ + addLoRAs: vi.fn(), +})); + +vi.mock('features/nodes/util/graph/generation/addNSFWChecker', () => ({ + addNSFWChecker: vi.fn((_g, node) => node), +})); + +vi.mock('features/nodes/util/graph/generation/addOutpaint', () => ({ + addOutpaint: vi.fn(), +})); + +vi.mock('features/nodes/util/graph/generation/addSDXLLoRAs', () => ({ + addSDXLLoRAs: vi.fn(), +})); + +vi.mock('features/nodes/util/graph/generation/addSDXLRefiner', () => ({ + addSDXLRefiner: vi.fn(), +})); + +vi.mock('features/nodes/util/graph/generation/addSeamless', () => ({ + addSeamless: vi.fn(() => null), +})); + +vi.mock('features/nodes/util/graph/generation/addTextToImage', () => ({ + addTextToImage: vi.fn(({ l2i }) => l2i), +})); + +vi.mock('features/nodes/util/graph/generation/addWatermarker', () => ({ + addWatermarker: vi.fn((_g, node) => node), +})); + +vi.mock('features/nodes/util/graph/graphBuilderUtils', () => ({ + getOriginalAndScaledSizesForTextToImage: vi.fn(() => ({ + originalSize: { width: 1024, height: 1024 }, + scaledSize: { width: 1024, height: 1024 }, + aspectRatio: { id: 'Free', value: 1 }, + })), + selectCanvasOutputFields: vi.fn(() => ({})), + selectPresetModifiedPrompts: vi.fn(() => ({ + positive: 'a prompt', + negative: 'a negative prompt', + })), +})); + +vi.mock('features/ui/store/uiSelectors', () => ({ + selectActiveTab: vi.fn(() => 'generation'), +})); + +vi.mock('./addRegions', () => ({ + addRegions: vi.fn(() => Promise.resolve({ addedRegions: 0 })), +})); + +import type { GraphBuilderArg } from 'features/nodes/util/graph/types'; + +import { buildSD1Graph } from './buildSD1Graph'; +import { buildSDXLGraph } from './buildSDXLGraph'; +import type { Graph } from './Graph'; + +const buildGraphArg = (): GraphBuilderArg => + ({ + generationMode: 'txt2img', + manager: null, + state: { + system: { + shouldUseNSFWChecker: false, + shouldUseWatermarker: false, + }, + }, + }) as unknown as GraphBuilderArg; + +const getMetadata = (g: Graph): Record => + (g as unknown as { getMetadataNode: () => Record }).getMetadataNode(); + +const resetState = () => { + nextId = 0; + currentModel = sd1Model; + params = { ...defaultParams }; +}; + +beforeEach(resetState); +afterEach(resetState); + +describe('HiDiffusion graph metadata', () => { + it('persists disabled HiDiffusion settings in the SD1 metadata node', async () => { + currentModel = sd1Model; + + const { g } = await buildSD1Graph(buildGraphArg()); + const metadata = getMetadata(g); + + expect(metadata.hidiffusion).toBe(false); + expect(metadata.hidiffusion_raunet).toBe(false); + expect(metadata.hidiffusion_window_attn).toBe(false); + expect(metadata.hidiffusion_t1_ratio).toBe(0.25); + expect(metadata.hidiffusion_t2_ratio).toBe(0.1); + }); + + it('persists disabled HiDiffusion settings in the SDXL metadata node', async () => { + currentModel = sdxlModel; + + const { g } = await buildSDXLGraph(buildGraphArg()); + const metadata = getMetadata(g); + + expect(metadata.hidiffusion).toBe(false); + expect(metadata.hidiffusion_raunet).toBe(false); + expect(metadata.hidiffusion_window_attn).toBe(false); + expect(metadata.hidiffusion_t1_ratio).toBe(0.25); + expect(metadata.hidiffusion_t2_ratio).toBe(0.1); + }); +}); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts index 1dead116ddb..e3bb7ff2417 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts @@ -42,6 +42,11 @@ export const buildSD1Graph = async (arg: GraphBuilderArg): Promise { + const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled); + const dispatch = useAppDispatch(); + const { t } = useTranslation(); + + const onChange = useCallback( + (event: ChangeEvent) => { + dispatch(setHiDiffusionEnabled(event.target.checked)); + }, + [dispatch] + ); + + return ( + + + + {t('parameters.hiDiffusion')} + + + + + ); +}); + +ParamHiDiffusionToggle.displayName = 'ParamHiDiffusionToggle'; + +export const ParamHiDiffusionRauNetToggle = memo(() => { + const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled); + const hiDiffusionRauNetEnabled = useAppSelector(selectHiDiffusionRauNetEnabled); + const dispatch = useAppDispatch(); + const { t } = useTranslation(); + + const onChange = useCallback( + (event: ChangeEvent) => { + dispatch(setHiDiffusionRauNetEnabled(event.target.checked)); + }, + [dispatch] + ); + + return ( + + + + {t('parameters.hiDiffusionRauNet')} + + + + + ); +}); + +ParamHiDiffusionRauNetToggle.displayName = 'ParamHiDiffusionRauNetToggle'; + +export const ParamHiDiffusionWindowAttnToggle = memo(() => { + const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled); + const hiDiffusionWindowAttnEnabled = useAppSelector(selectHiDiffusionWindowAttnEnabled); + const dispatch = useAppDispatch(); + const { t } = useTranslation(); + + const onChange = useCallback( + (event: ChangeEvent) => { + dispatch(setHiDiffusionWindowAttnEnabled(event.target.checked)); + }, + [dispatch] + ); + + return ( + + + + {t('parameters.hiDiffusionWindowAttn')} + + + + + ); +}); + +ParamHiDiffusionWindowAttnToggle.displayName = 'ParamHiDiffusionWindowAttnToggle'; + +const RATIO_CONSTRAINTS = { + t1: { + initial: 0.4, + sliderMin: 0.1, + sliderMax: 1, + numberInputMin: 0.1, + numberInputMax: 1, + coarseStep: 0.05, + fineStep: 0.01, + }, + t2: { + initial: 0.0, + sliderMin: 0.0, + sliderMax: 1, + numberInputMin: 0.0, + numberInputMax: 1, + coarseStep: 0.05, + fineStep: 0.01, + }, +} as const; + +export const ParamHiDiffusionT1Ratio = memo(() => { + const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled); + const hiDiffusionT1Ratio = useAppSelector(selectHiDiffusionT1Ratio); + const dispatch = useAppDispatch(); + const { t } = useTranslation(); + + const onChange = useCallback((value: number) => dispatch(setHiDiffusionT1Ratio(value)), [dispatch]); + + return ( + + + + {t('parameters.hiDiffusionT1Ratio')} + + + + + + ); +}); + +ParamHiDiffusionT1Ratio.displayName = 'ParamHiDiffusionT1Ratio'; + +export const ParamHiDiffusionT2Ratio = memo(() => { + const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled); + const hiDiffusionT2Ratio = useAppSelector(selectHiDiffusionT2Ratio); + const dispatch = useAppDispatch(); + const { t } = useTranslation(); + + const onChange = useCallback((value: number) => dispatch(setHiDiffusionT2Ratio(value)), [dispatch]); + + return ( + + + + {t('parameters.hiDiffusionT2Ratio')} + + + + + + ); +}); + +ParamHiDiffusionT2Ratio.displayName = 'ParamHiDiffusionT2Ratio'; diff --git a/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx b/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx index 8aea0582d9e..979a206236a 100644 --- a/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx +++ b/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx @@ -1,5 +1,5 @@ import type { FormLabelProps } from '@invoke-ai/ui-library'; -import { Box, Flex, FormControlGroup, SimpleGrid, StandaloneAccordion } from '@invoke-ai/ui-library'; +import { Flex, FormControlGroup, SimpleGrid, StandaloneAccordion } from '@invoke-ai/ui-library'; import { skipToken } from '@reduxjs/toolkit/query'; import { createMemoizedSelector } from 'app/store/createMemoizedSelector'; import { useAppSelector } from 'app/store/storeHooks'; @@ -25,6 +25,13 @@ import ParamCLIPGEmbedModelSelect from 'features/parameters/components/Advanced/ import ParamCLIPLEmbedModelSelect from 'features/parameters/components/Advanced/ParamCLIPLEmbedModelSelect'; import ParamClipSkip from 'features/parameters/components/Advanced/ParamClipSkip'; import ParamFlux2KleinModelSelect from 'features/parameters/components/Advanced/ParamFlux2KleinModelSelect'; +import { + ParamHiDiffusionRauNetToggle, + ParamHiDiffusionT1Ratio, + ParamHiDiffusionT2Ratio, + ParamHiDiffusionToggle, + ParamHiDiffusionWindowAttnToggle, +} from 'features/parameters/components/Advanced/ParamHiDiffusionToggle'; import ParamKrea2ModelSelects from 'features/parameters/components/Advanced/ParamKrea2ModelSelects'; import ParamQwenImageComponentSourceSelect from 'features/parameters/components/Advanced/ParamQwenImageComponentSourceSelect'; import ParamQwenImageQuantization from 'features/parameters/components/Advanced/ParamQwenImageQuantization'; @@ -106,6 +113,9 @@ export const AdvancedSettingsAccordion = memo(() => { if (params.seamlessXAxis || params.seamlessYAxis) { badges.push('seamless'); } + if (params.hiDiffusionEnabled) { + badges.push('HiDiffusion'); + } } return badges; @@ -153,9 +163,12 @@ export const AdvancedSettingsAccordion = memo(() => { + - {/* Empty box for visual alignment. Replace with new option when needed. */} - + + + + diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 5e327936269..596472a0d43 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -9204,6 +9204,36 @@ export type components = { * @default 0 */ cfg_rescale_multiplier?: number; + /** + * HiDiffusion + * @description Apply HiDiffusion (RAU-Net + MSW-MSA) for higher-resolution denoising + * @default false + */ + hidiffusion?: boolean; + /** + * HiDiffusion: RAU-Net + * @description Apply HiDiffusion RAU-Net blocks + * @default true + */ + hidiffusion_raunet?: boolean; + /** + * HiDiffusion: Window Attention + * @description Apply HiDiffusion window attention blocks + * @default true + */ + hidiffusion_window_attn?: boolean; + /** + * HiDiffusion: T1 Ratio + * @description Override HiDiffusion early switch threshold (T1 ratio) + * @default 0.4 + */ + hidiffusion_t1_ratio?: number; + /** + * HiDiffusion: T2 Ratio + * @description Override HiDiffusion late switch threshold (T2 ratio) + * @default 0 + */ + hidiffusion_t2_ratio?: number; /** * @description Latents tensor * @default null @@ -9322,6 +9352,36 @@ export type components = { * @default 0 */ cfg_rescale_multiplier?: number; + /** + * HiDiffusion + * @description Apply HiDiffusion (RAU-Net + MSW-MSA) for higher-resolution denoising + * @default false + */ + hidiffusion?: boolean; + /** + * HiDiffusion: RAU-Net + * @description Apply HiDiffusion RAU-Net blocks + * @default true + */ + hidiffusion_raunet?: boolean; + /** + * HiDiffusion: Window Attention + * @description Apply HiDiffusion window attention blocks + * @default true + */ + hidiffusion_window_attn?: boolean; + /** + * HiDiffusion: T1 Ratio + * @description Override HiDiffusion early switch threshold (T1 ratio) + * @default 0.4 + */ + hidiffusion_t1_ratio?: number; + /** + * HiDiffusion: T2 Ratio + * @description Override HiDiffusion late switch threshold (T2 ratio) + * @default 0 + */ + hidiffusion_t2_ratio?: number; /** * @description Latents tensor * @default null diff --git a/pyproject.toml b/pyproject.toml index 5890ad3b37d..20087cc0955 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -239,6 +239,7 @@ version = { attr = "invokeai.version.__version__" } "invokeai.backend.qwen3" = ["tokenizer/*.json", "tokenizer/*.json.gz"] "invokeai.backend.t5" = ["tokenizer/*.json"] "invokeai.app.services.workflow_records.default_workflows" = ["*.json"] +"invokeai.backend.hidiffusion" = ["sd_module_key/*.txt"] "invokeai.app.services.style_preset_records" = ["*.json"] "invokeai.app.services.style_preset_images.default_style_preset_images" = [ "*.png", diff --git a/tests/app/invocations/test_denoise_noise_inputs.py b/tests/app/invocations/test_denoise_noise_inputs.py index 611d28226a6..eb11568f961 100644 --- a/tests/app/invocations/test_denoise_noise_inputs.py +++ b/tests/app/invocations/test_denoise_noise_inputs.py @@ -9,7 +9,11 @@ from invokeai.app.invocations.cogview4_denoise import CogView4DenoiseInvocation from invokeai.app.invocations.flux2_denoise import Flux2DenoiseInvocation from invokeai.app.invocations.flux_denoise import FluxDenoiseInvocation -from invokeai.app.invocations.metadata_linked import FluxDenoiseLatentsMetaInvocation, ZImageDenoiseMetaInvocation +from invokeai.app.invocations.metadata_linked import ( + DenoiseLatentsMetaInvocation, + FluxDenoiseLatentsMetaInvocation, + ZImageDenoiseMetaInvocation, +) from invokeai.app.invocations.primitives import LatentsOutput from invokeai.app.invocations.sd3_denoise import SD3DenoiseInvocation from invokeai.app.invocations.z_image_denoise import ZImageDenoiseInvocation @@ -392,6 +396,36 @@ def test_z_image_metadata_ignores_external_noise_seed_when_noise_not_used(): assert result.metadata.root["seed"] == 999 +def test_denoise_metadata_persists_hidiffusion_fields_when_disabled(): + invocation = DenoiseLatentsMetaInvocation.model_construct( + width=64, + height=64, + steps=10, + cfg_scale=7.5, + cfg_rescale_multiplier=0, + denoising_start=0.0, + denoising_end=1.0, + scheduler="euler", + unet=SimpleNamespace(unet="main-model", loras=[]), + hidiffusion=False, + hidiffusion_raunet=False, + hidiffusion_window_attn=False, + hidiffusion_t1_ratio=0.25, + hidiffusion_t2_ratio=0.1, + ) + mock_context = MagicMock() + output = LatentsOutput.build("latents", torch.zeros(1, 4, 8, 8), seed=None) + + with patch("invokeai.app.invocations.metadata_linked.DenoiseLatentsInvocation.invoke", return_value=output): + result = invocation.invoke(mock_context) + + assert result.metadata.root["hidiffusion"] is False + assert result.metadata.root["hidiffusion_raunet"] is False + assert result.metadata.root["hidiffusion_window_attn"] is False + assert result.metadata.root["hidiffusion_t1_ratio"] == 0.25 + assert result.metadata.root["hidiffusion_t2_ratio"] == 0.1 + + def _get_first_scheduler_sigma( scheduler, *, scheduler_name: str, sigmas: list[float], mu: float | None = None ) -> float: diff --git a/tests/backend/stable_diffusion/test_hidiffusion_utils.py b/tests/backend/stable_diffusion/test_hidiffusion_utils.py new file mode 100644 index 00000000000..74c92ab9604 --- /dev/null +++ b/tests/backend/stable_diffusion/test_hidiffusion_utils.py @@ -0,0 +1,324 @@ +import copy +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from invokeai.backend.hidiffusion.hidiffusion import ( + remove_hidiffusion as real_remove_hidiffusion, +) +from invokeai.backend.hidiffusion.hidiffusion import ( + switching_threshold_ratio_dict, + text_to_img_controlnet_switching_threshold_ratio_dict, +) +from invokeai.backend.stable_diffusion.hidiffusion_utils import hidiffusion_patch + + +class DummySubmodule: + pass + + +class PatchedSubmodule(DummySubmodule): + _parent = DummySubmodule + + +class DummyUNet: + def __init__(self): + self.num_upsamplers = 3 + self.layer = DummySubmodule() + + def named_modules(self): + return [("", self), ("layer", self.layer)] + + +class ModelMixin(torch.nn.Module): + """Minimal diffusers-like UNet accepted by the vendored HiDiffusion type check.""" + + def __init__(self): + super().__init__() + self.num_upsamplers = 3 + + +class WindowMeanAttention(torch.nn.Module): + def forward(self, hidden_states: torch.Tensor, **_kwargs): + return hidden_states.mean(dim=1, keepdim=True).expand_as(hidden_states) + + +class WindowAttentionBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.use_ada_layer_norm = False + self.use_ada_layer_norm_zero = False + self.use_layer_norm = True + self.use_ada_layer_norm_continuous = False + self.use_ada_layer_norm_single = False + self.pos_embed = None + self.norm1 = torch.nn.Identity() + self.attn1 = WindowMeanAttention() + self.only_cross_attention = False + self.attn2 = None + self.norm3 = torch.nn.Identity() + self.ff = torch.nn.Identity() + self._chunk_size = None + + +class WindowAttentionModelMixin(ModelMixin): + def __init__(self): + super().__init__() + self.transformer = WindowAttentionBlock() + + +class CachedHiDiffusionModelMixin(ModelMixin): + def __init__(self): + super().__init__() + self.block = torch.nn.Module() + + +def test_hidiffusion_patch_supports_bare_model_mixin_without_public_name_or_path(): + model = ModelMixin() + + assert not hasattr(model, "name_or_path") + assert not hasattr(model, "_name_or_path") + + with hidiffusion_patch(model, name_or_path="runwayml/stable-diffusion-v1-5"): + assert model.info["pipeline"] is model + assert model.num_upsamplers == 15 + + assert model.num_upsamplers == 3 + assert not hasattr(model, "_name_or_path") + assert not hasattr(model, "info") + + +def test_hidiffusion_window_attention_uses_seeded_generator_instead_of_global_rng(): + module_keys = { + "down_module_key": [], + "down_module_key_extra": [], + "up_module_key": [], + "up_module_key_extra": [], + "windown_attn_module_key": ["transformer"], + } + hidden_states = torch.arange(64, dtype=torch.float32).reshape(1, 64, 1) + + def run_with_global_seed(global_seed: int) -> torch.Tensor: + torch.manual_seed(global_seed) + model = WindowAttentionModelMixin() + generator = torch.Generator(device="cpu").manual_seed(1234) + + with ( + patch("invokeai.backend.hidiffusion.hidiffusion.sd15_hidiffusion_key", return_value=module_keys), + hidiffusion_patch( + model, + name_or_path="runwayml/stable-diffusion-v1-5", + apply_raunet=False, + apply_window_attn=True, + generator=generator, + ), + ): + model.info["size"] = (8, 8) + return model.transformer(hidden_states).clone() + + first = run_with_global_seed(0) + second = run_with_global_seed(1) + + torch.testing.assert_close(first, second) + + +def test_hidiffusion_patch_resets_cached_runtime_state_when_reenabled(): + module_keys = { + "down_module_key": [], + "down_module_key_extra": ["block"], + "up_module_key": [], + "up_module_key_extra": [], + "windown_attn_module_key": [], + } + model = CachedHiDiffusionModelMixin() + + with patch("invokeai.backend.hidiffusion.hidiffusion.sd15_hidiffusion_key", return_value=module_keys): + with hidiffusion_patch(model, name_or_path="runwayml/stable-diffusion-v1-5"): + model.block.timestep = 7 + model.block.aggressive_raunet = True + model.block.T1_ratio = 0.9 + model.block.T1 = 9 + model.block.T1_start = 2 + model.block.T1_end = 8 + model.block.max_timestep = 99 + + assert "timestep" not in model.block.__dict__ + + with hidiffusion_patch(model, name_or_path="runwayml/stable-diffusion-v1-5"): + assert model.block.timestep == 0 + assert model.block.aggressive_raunet is False + assert model.block.T1_ratio == 0 + assert model.block.T1 == 0 + assert model.block.T1_start == 0 + assert model.block.T1_end == 0 + assert model.block.max_timestep == 50 + + +def test_hidiffusion_teardown_restores_downsampler_geometry_after_forward_error(): + module_keys = { + "down_module_key": ["block"], + "down_module_key_extra": [], + "up_module_key": [], + "up_module_key_extra": [], + "windown_attn_module_key": [], + } + model = ModelMixin() + model._num_timesteps = 10 + model.block = torch.nn.Conv2d(1, 1, kernel_size=3, stride=2, padding=1) + original_stride = model.block.stride + original_padding = model.block.padding + original_dilation = model.block.dilation + + with patch("invokeai.backend.hidiffusion.hidiffusion.sd15_hidiffusion_key", return_value=module_keys): + with hidiffusion_patch( + model, + name_or_path="runwayml/stable-diffusion-v1-5", + apply_window_attn=False, + ): + model.info["size"] = (64, 64) + with ( + patch( + "invokeai.backend.hidiffusion.hidiffusion.F.conv2d", + side_effect=RuntimeError("injected convolution failure"), + ), + pytest.raises(RuntimeError, match="injected convolution failure"), + ): + model.block(torch.zeros(1, 1, 16, 16)) + + # Temporary geometry is passed directly to conv2d and never written + # to the cached module, even before teardown runs. + assert model.block.stride == original_stride + assert model.block.padding == original_padding + assert model.block.dilation == original_dilation + + assert model.block.stride == original_stride + assert model.block.padding == original_padding + assert model.block.dilation == original_dilation + + +def test_hidiffusion_patch_restores_state_when_apply_hidiffusion_raises(): + original_switching = copy.deepcopy(switching_threshold_ratio_dict) + original_controlnet = copy.deepcopy(text_to_img_controlnet_switching_threshold_ratio_dict) + + model = SimpleNamespace( + unet=DummyUNet(), + _name_or_path="original-model-name", + config=SimpleNamespace(_name_or_path="original-config-name"), + ) + hook = MagicMock() + + def fake_apply_hidiffusion(patched_model, **_kwargs): + assert patched_model._name_or_path == "patched-model-name" + assert patched_model.config._name_or_path == "patched-model-name" + + first_switching_entry = next(iter(switching_threshold_ratio_dict.values())) + first_controlnet_entry = next(iter(text_to_img_controlnet_switching_threshold_ratio_dict.values())) + assert first_switching_entry["T1_ratio"] == 0.25 + assert first_switching_entry["T2_ratio"] == 0.1 + assert first_controlnet_entry["T1_ratio"] == 0.25 + assert first_controlnet_entry["T2_ratio"] == 0.1 + + patched_model.unet.num_upsamplers = 99 + patched_model.unet.layer.info = {"hooks": [hook]} + patched_model.unet.layer.__class__ = PatchedSubmodule + raise RuntimeError("hidiffusion boom") + + try: + with ( + patch("invokeai.backend.hidiffusion.hidiffusion.apply_hidiffusion", side_effect=fake_apply_hidiffusion), + patch( + "invokeai.backend.hidiffusion.hidiffusion.remove_hidiffusion", + wraps=real_remove_hidiffusion, + ) as mock_remove_hidiffusion, + ): + with pytest.raises(RuntimeError, match="hidiffusion boom"): + with hidiffusion_patch( + model, + name_or_path="patched-model-name", + t1_ratio=0.25, + t2_ratio=0.1, + ): + pass + + assert mock_remove_hidiffusion.call_count == 1 + assert switching_threshold_ratio_dict == original_switching + assert text_to_img_controlnet_switching_threshold_ratio_dict == original_controlnet + assert model.unet.num_upsamplers == 3 + assert model.unet.layer.__class__ is DummySubmodule + assert model.unet.layer.info["hooks"] == [] + hook.remove.assert_called_once() + assert model._name_or_path == "original-model-name" + assert model.config._name_or_path == "original-config-name" + finally: + switching_threshold_ratio_dict.clear() + switching_threshold_ratio_dict.update(original_switching) + text_to_img_controlnet_switching_threshold_ratio_dict.clear() + text_to_img_controlnet_switching_threshold_ratio_dict.update(original_controlnet) + + +def test_hidiffusion_patch_restores_state_before_propagating_remove_error(): + original_switching = copy.deepcopy(switching_threshold_ratio_dict) + original_controlnet = copy.deepcopy(text_to_img_controlnet_switching_threshold_ratio_dict) + + model = SimpleNamespace( + unet=DummyUNet(), + _name_or_path="original-model-name", + config=SimpleNamespace(_name_or_path="original-config-name"), + ) + + def fake_apply_hidiffusion(patched_model, **_kwargs): + patched_model.unet.num_upsamplers = 99 + + try: + with ( + patch("invokeai.backend.hidiffusion.hidiffusion.apply_hidiffusion", side_effect=fake_apply_hidiffusion), + patch( + "invokeai.backend.hidiffusion.hidiffusion.remove_hidiffusion", + side_effect=RuntimeError("remove boom"), + ), + ): + with pytest.raises(RuntimeError, match="remove boom"): + with hidiffusion_patch( + model, + name_or_path="patched-model-name", + t1_ratio=0.25, + t2_ratio=0.1, + ): + pass + + assert switching_threshold_ratio_dict == original_switching + assert text_to_img_controlnet_switching_threshold_ratio_dict == original_controlnet + assert model.unet.num_upsamplers == 3 + assert model._name_or_path == "original-model-name" + assert model.config._name_or_path == "original-config-name" + finally: + switching_threshold_ratio_dict.clear() + switching_threshold_ratio_dict.update(original_switching) + text_to_img_controlnet_switching_threshold_ratio_dict.clear() + text_to_img_controlnet_switching_threshold_ratio_dict.update(original_controlnet) + + +def test_hidiffusion_patch_removes_spoofed_name_from_config_internal_dict(): + class InternalDictConfig: + def __init__(self): + self._internal_dict = {} + + def __getattr__(self, name): + try: + return self._internal_dict[name] + except KeyError as error: + raise AttributeError(name) from error + + config = InternalDictConfig() + model = SimpleNamespace(unet=DummyUNet(), config=config) + + with ( + patch("invokeai.backend.hidiffusion.hidiffusion.apply_hidiffusion"), + patch("invokeai.backend.hidiffusion.hidiffusion.remove_hidiffusion"), + ): + with hidiffusion_patch(model, name_or_path="patched-model-name"): + assert config._internal_dict["_name_or_path"] == "patched-model-name" + + assert "_name_or_path" not in config._internal_dict