From ec39445a270d788406e35e99f0ef80b92ac56c7e Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Sun, 25 Jan 2026 18:40:37 +0200 Subject: [PATCH 01/22] Initial implementation of HiDiffusion pipeline --- docs/features/hidiffusion.md | 32 +++++ invokeai/app/invocations/denoise_latents.py | 97 +++++++++++---- invokeai/app/invocations/fields.py | 3 + invokeai/app/invocations/metadata_linked.py | 4 + .../extensions/hidiffusion.py | 33 ++++++ .../stable_diffusion/hidiffusion_utils.py | 111 ++++++++++++++++++ invokeai/frontend/web/openapi.json | 60 ++++++++++ invokeai/frontend/web/public/locales/en.json | 27 +++++ .../InformationalPopover/constants.ts | 6 + .../controlLayers/store/paramsSlice.ts | 25 ++++ .../src/features/controlLayers/store/types.ts | 6 + .../web/src/features/metadata/parsing.tsx | 60 ++++++++++ .../util/graph/generation/addSDXLRefiner.ts | 6 + .../util/graph/generation/buildSD1Graph.ts | 9 ++ .../util/graph/generation/buildSDXLGraph.ts | 9 ++ .../Advanced/ParamHiDiffusionToggle.tsx | 92 +++++++++++++++ .../AdvancedSettingsAccordion.tsx | 21 +++- .../frontend/web/src/services/api/schema.ts | 36 ++++++ mkdocs.yml | 1 + pyproject.toml | 1 + uv.lock | 15 ++- 21 files changed, 622 insertions(+), 32 deletions(-) create mode 100644 docs/features/hidiffusion.md create mode 100644 invokeai/backend/stable_diffusion/extensions/hidiffusion.py create mode 100644 invokeai/backend/stable_diffusion/hidiffusion_utils.py create mode 100644 invokeai/frontend/web/src/features/parameters/components/Advanced/ParamHiDiffusionToggle.tsx diff --git a/docs/features/hidiffusion.md b/docs/features/hidiffusion.md new file mode 100644 index 00000000000..63141ecb46f --- /dev/null +++ b/docs/features/hidiffusion.md @@ -0,0 +1,32 @@ +--- +title: HiDiffusion +--- + +# HiDiffusion + +HiDiffusion is an optional denoising enhancement that can improve detail and structure at higher resolutions for SD 1.5 and SDXL. It modifies the UNet during denoising and is most noticeable at 1536px and above. + +## 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: + - **HiDiffusion: RAU‑Net** + - **HiDiffusion: Window Attention** + +## 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. + +## 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 bb114263e23..7e591634e23 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 @@ -75,6 +76,7 @@ from invokeai.backend.stable_diffusion.extensions_manager import ExtensionsManager from invokeai.backend.stable_diffusion.schedulers import SCHEDULER_MAP from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES +from invokeai.backend.stable_diffusion.hidiffusion_utils import hidiffusion_patch from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.hotfixes import ControlNetModel from invokeai.backend.util.mask import to_standard_float_mask @@ -191,6 +193,21 @@ 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", + ) latents: Optional[LatentsField] = InputField( default=None, description=FieldDescriptions.latents, @@ -486,6 +503,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 +862,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 +900,14 @@ 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, + ) + ) ### cfg rescale if self.cfg_rescale_multiplier > 0: @@ -940,14 +974,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 +1034,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) @@ -1084,23 +1122,34 @@ def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]: 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, + ) + 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 6137420c873..feb377d77d4 100644 --- a/invokeai/app/invocations/fields.py +++ b/invokeai/app/invocations/fields.py @@ -147,6 +147,9 @@ 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" 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 6a9db3e5890..e6dc51c8670 100644 --- a/invokeai/app/invocations/metadata_linked.py +++ b/invokeai/app/invocations/metadata_linked.py @@ -661,6 +661,10 @@ 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}) + if self.hidiffusion: + md.update({"hidiffusion": self.hidiffusion}) + md.update({"hidiffusion_raunet": self.hidiffusion_raunet}) + md.update({"hidiffusion_window_attn": self.hidiffusion_window_attn}) 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/stable_diffusion/extensions/hidiffusion.py b/invokeai/backend/stable_diffusion/extensions/hidiffusion.py new file mode 100644 index 00000000000..6a743febde3 --- /dev/null +++ b/invokeai/backend/stable_diffusion/extensions/hidiffusion.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from contextlib import contextmanager +from typing import Optional + +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, + ): + super().__init__() + self._name_or_path = name_or_path + self._apply_raunet = apply_raunet + self._apply_window_attn = apply_window_attn + + @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, + ): + 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..295b52cdb15 --- /dev/null +++ b/invokeai/backend/stable_diffusion/hidiffusion_utils.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +from contextlib import contextmanager +from typing import Any, Optional + + +@contextmanager +def hidiffusion_patch( + model: Any, + name_or_path: Optional[str], + apply_raunet: bool = True, + apply_window_attn: bool = True, +): + """Context manager that applies HiDiffusion and restores the model on exit.""" + try: + from hidiffusion import apply_hidiffusion, remove_hidiffusion + except ImportError as exc: + raise ImportError( + "HiDiffusion is not installed. Install it with `pip install hidiffusion` to enable this option." + ) from exc + + 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 + + set_model_name_or_path = False + 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 + + original_num_upsamplers = getattr(target, "num_upsamplers", None) + + apply_hidiffusion(model, apply_raunet=apply_raunet, apply_window_attn=apply_window_attn) + try: + yield + finally: + remove_hidiffusion(model) + 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) + else: + if config is not None: + try: + delattr(config, "_name_or_path") + except AttributeError: + pass diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 1be662f85e8..b17ed3fd9d0 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -15445,6 +15445,36 @@ "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" + }, "latents": { "anyOf": [ { @@ -15803,6 +15833,36 @@ "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" + }, "latents": { "anyOf": [ { diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json index 5bdd7ba08ee..5bcb7550f55 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -838,6 +838,9 @@ "generationMode": "Generation Mode", "guidance": "Guidance", "height": "Height", + "hiDiffusion": "HiDiffusion", + "hiDiffusionRauNet": "HiDiffusion: RAU-Net", + "hiDiffusionWindowAttn": "HiDiffusion: Window Attention", "imageDetails": "Image Details", "imageDimensions": "Image Dimensions", "metadata": "Metadata", @@ -1293,6 +1296,9 @@ "cfgScale": "CFG Scale", "cfgRescaleMultiplier": "CFG Rescale Multiplier", "clipSkip": "CLIP Skip", + "hiDiffusion": "HiDiffusion", + "hiDiffusionRauNet": "HiDiffusion: RAU-Net", + "hiDiffusionWindowAttn": "HiDiffusion: Window Attention", "coherenceMode": "Mode", "coherenceEdgeSize": "Edge Size", "coherenceMinDenoise": "Min Denoise", @@ -1582,6 +1588,27 @@ "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." + ] + }, "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 fc045356e0c..d8eb841c417 100644 --- a/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts +++ b/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts @@ -65,6 +65,9 @@ export type Feature = | 'scaleBeforeProcessing' | 'seamlessTilingXAxis' | 'seamlessTilingYAxis' + | 'hidiffusion' + | 'hidiffusionRauNet' + | 'hidiffusionWindowAttn' | 'colorCompensation' | 'upscaleModel' | 'scale' @@ -220,6 +223,9 @@ export const POPOVER_DATA: { [key in Feature]?: PopoverData } = { seamlessTilingYAxis: { href: 'https://support.invoke.ai/support/solutions/articles/151000178161-advanced-settings', }, + hidiffusion: {}, + hidiffusionRauNet: {}, + hidiffusionWindowAttn: {}, } as const; export const OPEN_DELAY = 1000; // in milliseconds diff --git a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts index 0190aba602b..20c2ee91cc5 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts @@ -98,6 +98,15 @@ 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; + }, setSeamlessXAxis: (state, action: PayloadAction) => { state.seamlessXAxis = action.payload; }, @@ -474,6 +483,9 @@ export const { setSeed, setImg2imgStrength, setOptimizedDenoisingEnabled, + setHiDiffusionEnabled, + setHiDiffusionRauNetEnabled, + setHiDiffusionWindowAttnEnabled, setSeamlessXAxis, setSeamlessYAxis, setShouldRandomizeSeed, @@ -538,6 +550,16 @@ export const paramsSliceConfig: SliceConfig = { state.positivePromptHistory = []; } + if (!('hiDiffusionEnabled' in state)) { + state.hiDiffusionEnabled = false; + } + if (!('hiDiffusionRauNetEnabled' in state)) { + state.hiDiffusionRauNetEnabled = true; + } + if (!('hiDiffusionWindowAttnEnabled' in state)) { + state.hiDiffusionWindowAttnEnabled = true; + } + return zParamsState.parse(state); }, }, @@ -591,6 +613,9 @@ 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 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 fbd3d415b79..0d90a5e705b 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/types.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/types.ts @@ -596,6 +596,9 @@ export const zParamsState = z.object({ guidance: zParameterGuidance, img2imgStrength: zParameterStrength, optimizedDenoisingEnabled: z.boolean(), + hiDiffusionEnabled: z.boolean(), + hiDiffusionRauNetEnabled: z.boolean(), + hiDiffusionWindowAttnEnabled: z.boolean(), iterations: z.number(), scheduler: zParameterScheduler, fluxScheduler: zParameterFluxScheduler, @@ -656,6 +659,9 @@ export const getInitialParamsState = (): ParamsState => ({ guidance: 4, img2imgStrength: 0.75, optimizedDenoisingEnabled: true, + hiDiffusionEnabled: false, + hiDiffusionRauNetEnabled: true, + hiDiffusionWindowAttnEnabled: true, iterations: 1, scheduler: 'dpmpp_3m_k', fluxScheduler: 'euler', diff --git a/invokeai/frontend/web/src/features/metadata/parsing.tsx b/invokeai/frontend/web/src/features/metadata/parsing.tsx index d9201f15ffa..1ed7d4505c1 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.tsx +++ b/invokeai/frontend/web/src/features/metadata/parsing.tsx @@ -18,6 +18,9 @@ import { setClipSkip, setFluxScheduler, setGuidance, + setHiDiffusionEnabled, + setHiDiffusionRauNetEnabled, + setHiDiffusionWindowAttnEnabled, setImg2imgStrength, setRefinerCFGScale, setRefinerNegativeAestheticScore, @@ -537,6 +540,60 @@ const SeamlessY: SingleMetadataHandler = { }; //#endregion SeamlessY +//#region HiDiffusion +const HiDiffusion: SingleMetadataHandler = { + [SingleMetadataKey]: true, + type: 'HiDiffusion', + parse: (metadata, _store) => { + const raw = getProperty(metadata, 'hidiffusion'); + const parsed = 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 ZImageSeedVarianceEnabled const ZImageSeedVarianceEnabled: SingleMetadataHandler = { [SingleMetadataKey]: true, @@ -1071,6 +1128,9 @@ export const ImageMetadataHandlers = { DenoisingStrength, SeamlessX, SeamlessY, + HiDiffusion, + HiDiffusionRauNet, + HiDiffusionWindowAttn, 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..757dc638c0d 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,9 @@ export const addSDXLRefiner = async ( refinerScheduler, refinerCFGScale, refinerStart, + hiDiffusionEnabled, + hiDiffusionRauNetEnabled, + hiDiffusionWindowAttnEnabled, } = state.params; assert(refinerModel, 'No refiner model found in state'); @@ -57,6 +60,9 @@ export const addSDXLRefiner = async ( cfg_scale: refinerCFGScale, steps: refinerSteps, scheduler: refinerScheduler, + hidiffusion: hiDiffusionEnabled, + hidiffusion_raunet: hiDiffusionRauNetEnabled, + hidiffusion_window_attn: hiDiffusionWindowAttnEnabled, denoising_start: refinerStart, denoising_end: 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 eae07532011..95ab2973ba9 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,9 @@ 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(); + + return ( + + + + {t('parameters.hiDiffusionRauNet')} + + + ) => + dispatch(setHiDiffusionRauNetEnabled(event.target.checked)) + } + /> + + ); +}); + +ParamHiDiffusionRauNetToggle.displayName = 'ParamHiDiffusionRauNetToggle'; + +export const ParamHiDiffusionWindowAttnToggle = memo(() => { + const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled); + const hiDiffusionWindowAttnEnabled = useAppSelector(selectHiDiffusionWindowAttnEnabled); + const dispatch = useAppDispatch(); + const { t } = useTranslation(); + + return ( + + + + {t('parameters.hiDiffusionWindowAttn')} + + + ) => + dispatch(setHiDiffusionWindowAttnEnabled(event.target.checked)) + } + /> + + ); +}); + +ParamHiDiffusionWindowAttnToggle.displayName = 'ParamHiDiffusionWindowAttnToggle'; 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 e52d3ad40d0..ee47fb4797e 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'; @@ -15,6 +15,11 @@ import ParamCLIPEmbedModelSelect from 'features/parameters/components/Advanced/P import ParamCLIPGEmbedModelSelect from 'features/parameters/components/Advanced/ParamCLIPGEmbedModelSelect'; import ParamCLIPLEmbedModelSelect from 'features/parameters/components/Advanced/ParamCLIPLEmbedModelSelect'; import ParamClipSkip from 'features/parameters/components/Advanced/ParamClipSkip'; +import { + ParamHiDiffusionRauNetToggle, + ParamHiDiffusionToggle, + ParamHiDiffusionWindowAttnToggle, +} from 'features/parameters/components/Advanced/ParamHiDiffusionToggle'; import ParamT5EncoderModelSelect from 'features/parameters/components/Advanced/ParamT5EncoderModelSelect'; import ParamZImageQwen3VaeModelSelect from 'features/parameters/components/Advanced/ParamZImageQwen3VaeModelSelect'; import ParamSeamlessXAxis from 'features/parameters/components/Seamless/ParamSeamlessXAxis'; @@ -71,9 +76,12 @@ export const AdvancedSettingsAccordion = memo(() => { if (params.cfgRescaleMultiplier) { badges.push(`Rescale ${params.cfgRescaleMultiplier}`); } - if (params.seamlessXAxis || params.seamlessYAxis) { - badges.push('seamless'); - } + if (params.seamlessXAxis || params.seamlessYAxis) { + badges.push('seamless'); + } + if (params.hiDiffusionEnabled) { + badges.push('HiDiffusion'); + } } return badges; @@ -107,9 +115,10 @@ 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 829c5d435cd..22d9fcb4652 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -6518,6 +6518,24 @@ 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; /** * @description Latents tensor * @default null @@ -6636,6 +6654,24 @@ 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; /** * @description Latents tensor * @default null diff --git a/mkdocs.yml b/mkdocs.yml index 656baec9c3d..6291411cf2b 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -139,6 +139,7 @@ nav: - Features: - Low VRAM mode: 'features/low-vram.md' - Database: 'features/database.md' + - HiDiffusion: 'features/hidiffusion.md' - New to InvokeAI?: 'help/gettingStartedWithAI.md' - Contributing: - Overview: 'contributing/index.md' diff --git a/pyproject.toml b/pyproject.toml index adfe5982baf..5039f0445ad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ dependencies = [ "compel==2.1.1", "diffusers[torch]==0.36.0", "gguf", + "hidiffusion", "mediapipe==0.10.14", # needed for "mediapipeface" controlnet model "numpy<2.0.0", "onnx==1.16.1", diff --git a/uv.lock b/uv.lock index f6841cb6e71..60d18aeec0c 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = ">=3.11, <3.13" resolution-markers = [ "python_full_version >= '3.12' and platform_machine == 'aarch64' and sys_platform == 'linux' and extra != 'extra-8-invokeai-cpu' and extra != 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-rocm'", @@ -815,6 +815,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f0/55/ef77a85ee443ae05a9e9cba1c9f0dd9241eb42da2aeba1dc50f51154c81a/hf_xet-1.1.5-cp37-abi3-win_amd64.whl", hash = "sha256:73e167d9807d166596b4b2f0b585c6d5bd84a26dea32843665a8b58f6edba245", size = 2738931, upload-time = "2025-06-20T21:48:39.482Z" }, ] +[[package]] +name = "hidiffusion" +version = "0.1.10" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/2b/422acb8106f1295e94a8f98363e2b9b0d6be6f90f672d353f8b2c38654fe/hidiffusion-0.1.10.tar.gz", hash = "sha256:e8a7cfab49a3a29d83b33f3c6c4adcaf32377cb535d3a279bbf53125e5a57bcd", size = 41303, upload-time = "2024-06-16T17:23:58.739Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/64/fb/ebee9c90768db7c8368922d1387c64311051c4423de88aecfe1dcb8387ad/hidiffusion-0.1.10-py3-none-any.whl", hash = "sha256:9498c8fcc07f60645e533feeaaf66aeb6c74b289788be333f59507980556e9e4", size = 37727, upload-time = "2024-06-16T17:23:56.478Z" }, +] + [[package]] name = "httpcore" version = "1.0.9" @@ -972,6 +981,7 @@ dependencies = [ { name = "fastapi" }, { name = "fastapi-events" }, { name = "gguf" }, + { name = "hidiffusion" }, { name = "huggingface-hub" }, { name = "mediapipe" }, { name = "numpy" }, @@ -1079,6 +1089,7 @@ requires-dist = [ { name = "fastapi-events" }, { name = "gguf" }, { name = "gprof2dot", marker = "extra == 'dev'" }, + { name = "hidiffusion" }, { name = "httpx", marker = "extra == 'test'" }, { name = "huggingface-hub" }, { name = "humanize", marker = "extra == 'test'", specifier = "==4.12.1" }, @@ -3650,7 +3661,7 @@ name = "triton" version = "3.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "setuptools", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux' and extra != 'extra-8-invokeai-cpu' and extra != 'extra-8-invokeai-rocm') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-8-invokeai-cuda') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-8-invokeai-cpu' and extra != 'extra-8-invokeai-rocm') or (sys_platform == 'linux' and extra == 'extra-8-invokeai-cuda') or (sys_platform == 'darwin' and extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-rocm') or (extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-cuda') or (extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-rocm')" }, + { name = "setuptools", marker = "sys_platform != 'darwin'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/21/2f/3e56ea7b58f80ff68899b1dbe810ff257c9d177d288c6b0f55bf2fe4eb50/triton-3.3.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b31e3aa26f8cb3cc5bf4e187bf737cbacf17311e1112b781d4a059353dfd731b", size = 155689937, upload-time = "2025-05-29T23:39:44.182Z" }, From 4d52b493a6481ac2d31721b15c69d3145edc7048 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Sun, 25 Jan 2026 18:42:01 +0200 Subject: [PATCH 02/22] Added URLs to original HiDiffusion repository in documentation and frontend popover. --- docs/features/hidiffusion.md | 2 ++ .../src/common/components/InformationalPopover/constants.ts | 4 +++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/docs/features/hidiffusion.md b/docs/features/hidiffusion.md index 63141ecb46f..ff9fae6add1 100644 --- a/docs/features/hidiffusion.md +++ b/docs/features/hidiffusion.md @@ -6,6 +6,8 @@ title: HiDiffusion HiDiffusion is an optional denoising enhancement that can improve detail and structure at higher resolutions for SD 1.5 and SDXL. It modifies the UNet during denoising and is most noticeable at 1536px and above. +Learn more: https://github.com/megvii-research/HiDiffusion + ## Where to find the switches 1. Open the **Canvas** tab. diff --git a/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts b/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts index d8eb841c417..bdc130b63f3 100644 --- a/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts +++ b/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts @@ -223,7 +223,9 @@ export const POPOVER_DATA: { [key in Feature]?: PopoverData } = { seamlessTilingYAxis: { href: 'https://support.invoke.ai/support/solutions/articles/151000178161-advanced-settings', }, - hidiffusion: {}, + hidiffusion: { + href: 'https://github.com/megvii-research/HiDiffusion', + }, hidiffusionRauNet: {}, hidiffusionWindowAttn: {}, } as const; From 8c1b8e2ef1c886e49db951bf047a7ae60b32edd4 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Sun, 25 Jan 2026 19:06:32 +0200 Subject: [PATCH 03/22] Added comment next to HiDiffusion in pyproject.toml to clarify its purpose. --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 5039f0445ad..17e46baa27f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,7 +38,7 @@ dependencies = [ "compel==2.1.1", "diffusers[torch]==0.36.0", "gguf", - "hidiffusion", + "hidiffusion", # optional high-res denoise patch (HiDiffusion) "mediapipe==0.10.14", # needed for "mediapipeface" controlnet model "numpy<2.0.0", "onnx==1.16.1", From 514fe8a5c43fcec97d8ff1784f5e0b7b03802c17 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Sun, 25 Jan 2026 19:17:14 +0200 Subject: [PATCH 04/22] Formatting --- .../Advanced/ParamHiDiffusionToggle.tsx | 30 ++++++++++--------- .../AdvancedSettingsAccordion.tsx | 12 ++++---- 2 files changed, 22 insertions(+), 20 deletions(-) diff --git a/invokeai/frontend/web/src/features/parameters/components/Advanced/ParamHiDiffusionToggle.tsx b/invokeai/frontend/web/src/features/parameters/components/Advanced/ParamHiDiffusionToggle.tsx index cf396b946cb..de6f7e811c9 100644 --- a/invokeai/frontend/web/src/features/parameters/components/Advanced/ParamHiDiffusionToggle.tsx +++ b/invokeai/frontend/web/src/features/parameters/components/Advanced/ParamHiDiffusionToggle.tsx @@ -45,6 +45,13 @@ export const ParamHiDiffusionRauNetToggle = memo(() => { const dispatch = useAppDispatch(); const { t } = useTranslation(); + const onChange = useCallback( + (event: ChangeEvent) => { + dispatch(setHiDiffusionRauNetEnabled(event.target.checked)); + }, + [dispatch] + ); + return ( @@ -52,13 +59,7 @@ export const ParamHiDiffusionRauNetToggle = memo(() => { {t('parameters.hiDiffusionRauNet')} - ) => - dispatch(setHiDiffusionRauNetEnabled(event.target.checked)) - } - /> + ); }); @@ -71,6 +72,13 @@ export const ParamHiDiffusionWindowAttnToggle = memo(() => { const dispatch = useAppDispatch(); const { t } = useTranslation(); + const onChange = useCallback( + (event: ChangeEvent) => { + dispatch(setHiDiffusionWindowAttnEnabled(event.target.checked)); + }, + [dispatch] + ); + return ( @@ -78,13 +86,7 @@ export const ParamHiDiffusionWindowAttnToggle = memo(() => { {t('parameters.hiDiffusionWindowAttn')} - ) => - dispatch(setHiDiffusionWindowAttnEnabled(event.target.checked)) - } - /> + ); }); 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 ee47fb4797e..878200fdf19 100644 --- a/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx +++ b/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx @@ -76,12 +76,12 @@ export const AdvancedSettingsAccordion = memo(() => { if (params.cfgRescaleMultiplier) { badges.push(`Rescale ${params.cfgRescaleMultiplier}`); } - if (params.seamlessXAxis || params.seamlessYAxis) { - badges.push('seamless'); - } - if (params.hiDiffusionEnabled) { - badges.push('HiDiffusion'); - } + if (params.seamlessXAxis || params.seamlessYAxis) { + badges.push('seamless'); + } + if (params.hiDiffusionEnabled) { + badges.push('HiDiffusion'); + } } return badges; From 979678c24dbb5c6fb256e1b2e5c6bdfe29ddf2b1 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Sun, 25 Jan 2026 19:20:38 +0200 Subject: [PATCH 05/22] Import order fix --- invokeai/app/invocations/denoise_latents.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index 7e591634e23..2ddc8f5ffe0 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -74,9 +74,9 @@ 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.stable_diffusion.hidiffusion_utils import hidiffusion_patch from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.hotfixes import ControlNetModel from invokeai.backend.util.mask import to_standard_float_mask From 9c32ada18e2c49cd5e6e3566cde405c45f454c59 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Mon, 26 Jan 2026 00:45:02 +0200 Subject: [PATCH 06/22] Add HiDiffusion T1/T2 ratio controls and docs updates --- docs/features/hidiffusion.md | 14 ++- invokeai/app/invocations/denoise_latents.py | 18 +++ invokeai/app/invocations/fields.py | 2 + invokeai/app/invocations/metadata_linked.py | 2 + .../extensions/hidiffusion.py | 6 + .../stable_diffusion/hidiffusion_utils.py | 39 +++++++ invokeai/frontend/web/openapi.json | 48 ++++++++ invokeai/frontend/web/public/locales/en.json | 18 +++ .../InformationalPopover/constants.ts | 4 + .../controlLayers/store/paramsSlice.ts | 16 +++ .../src/features/controlLayers/store/types.ts | 4 + .../web/src/features/metadata/parsing.tsx | 40 +++++++ .../util/graph/generation/addSDXLRefiner.ts | 4 + .../util/graph/generation/buildSD1Graph.ts | 6 + .../util/graph/generation/buildSDXLGraph.ts | 6 + .../Advanced/ParamHiDiffusionToggle.tsx | 107 +++++++++++++++++- .../AdvancedSettingsAccordion.tsx | 4 + .../frontend/web/src/services/api/schema.ts | 24 ++++ 18 files changed, 360 insertions(+), 2 deletions(-) diff --git a/docs/features/hidiffusion.md b/docs/features/hidiffusion.md index ff9fae6add1..4c3f3346ec8 100644 --- a/docs/features/hidiffusion.md +++ b/docs/features/hidiffusion.md @@ -12,9 +12,11 @@ Learn more: https://github.com/megvii-research/HiDiffusion 1. Open the **Canvas** tab. 2. Expand **Advanced Settings**. -3. In the **Advanced** grid, enable **HiDiffusion** and optionally adjust the two sub‑toggles: +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 @@ -27,8 +29,18 @@ Learn more: https://github.com/megvii-research/HiDiffusion - **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. + +--- + +TODO: Decide whether HiDiffusion toggles and ratios should always be emitted in metadata for recall. diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index 2ddc8f5ffe0..676513ea9a8 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -208,6 +208,20 @@ class DenoiseLatentsInvocation(BaseInvocation): 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, @@ -906,6 +920,8 @@ def step_callback(state: PipelineIntermediateState) -> None: 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, ) ) @@ -1130,6 +1146,8 @@ def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]: 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, ) if self.hidiffusion else nullcontext() diff --git a/invokeai/app/invocations/fields.py b/invokeai/app/invocations/fields.py index feb377d77d4..b92d4991f52 100644 --- a/invokeai/app/invocations/fields.py +++ b/invokeai/app/invocations/fields.py @@ -150,6 +150,8 @@ class FieldDescriptions: 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 e6dc51c8670..c02e1c4c075 100644 --- a/invokeai/app/invocations/metadata_linked.py +++ b/invokeai/app/invocations/metadata_linked.py @@ -665,6 +665,8 @@ def _loras_to_json(obj: Union[Any, list[Any]]): 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/stable_diffusion/extensions/hidiffusion.py b/invokeai/backend/stable_diffusion/extensions/hidiffusion.py index 6a743febde3..0dc063b10d0 100644 --- a/invokeai/backend/stable_diffusion/extensions/hidiffusion.py +++ b/invokeai/backend/stable_diffusion/extensions/hidiffusion.py @@ -16,11 +16,15 @@ def __init__( name_or_path: Optional[str], apply_raunet: bool = True, apply_window_attn: bool = True, + t1_ratio: Optional[float] = None, + t2_ratio: Optional[float] = 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 @contextmanager def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeightsStorage): @@ -29,5 +33,7 @@ def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeigh 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, ): yield None diff --git a/invokeai/backend/stable_diffusion/hidiffusion_utils.py b/invokeai/backend/stable_diffusion/hidiffusion_utils.py index 295b52cdb15..ffcb0291c36 100644 --- a/invokeai/backend/stable_diffusion/hidiffusion_utils.py +++ b/invokeai/backend/stable_diffusion/hidiffusion_utils.py @@ -3,6 +3,8 @@ from contextlib import contextmanager from typing import Any, Optional +import copy + @contextmanager def hidiffusion_patch( @@ -10,9 +12,12 @@ def hidiffusion_patch( name_or_path: Optional[str], apply_raunet: bool = True, apply_window_attn: bool = True, + t1_ratio: Optional[float] = None, + t2_ratio: Optional[float] = None, ): """Context manager that applies HiDiffusion and restores the model on exit.""" try: + import hidiffusion as hd from hidiffusion import apply_hidiffusion, remove_hidiffusion except ImportError as exc: raise ImportError( @@ -82,11 +87,45 @@ def _set_name_or_path_on_config(cfg, value: str) -> bool: 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: + try: + from hidiffusion.hidiffusion import ( + switching_threshold_ratio_dict as _switching_threshold_ratio_dict, + text_to_img_controlnet_switching_threshold_ratio_dict as _text_to_img_controlnet_switching_threshold_ratio_dict, + ) + + ratio_dicts = ( + _switching_threshold_ratio_dict, + _text_to_img_controlnet_switching_threshold_ratio_dict, + ) + except Exception: + ratio_dicts = None + + 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) try: yield finally: remove_hidiffusion(model) + 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: diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index b17ed3fd9d0..8f0405dc7de 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -15475,6 +15475,54 @@ "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, + "description": "Override HiDiffusion late switch threshold (T2 ratio)", + "field_kind": "input", + "input": "any", + "maximum": 1, + "minimum": 0, + "orig_default": 0, + "orig_required": false, + "title": "HiDiffusion: T2 Ratio", + "type": "number" + }, + "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, + "description": "Override HiDiffusion late switch threshold (T2 ratio)", + "field_kind": "input", + "input": "any", + "maximum": 1, + "minimum": 0, + "orig_default": 0, + "orig_required": false, + "title": "HiDiffusion: T2 Ratio", + "type": "number" + }, "latents": { "anyOf": [ { diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json index 5bcb7550f55..915ce4f28ad 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -841,6 +841,8 @@ "hiDiffusion": "HiDiffusion", "hiDiffusionRauNet": "HiDiffusion: RAU-Net", "hiDiffusionWindowAttn": "HiDiffusion: Window Attention", + "hiDiffusionT1Ratio": "HiDiffusion: T1 Ratio", + "hiDiffusionT2Ratio": "HiDiffusion: T2 Ratio", "imageDetails": "Image Details", "imageDimensions": "Image Dimensions", "metadata": "Metadata", @@ -1299,6 +1301,8 @@ "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", @@ -1609,6 +1613,20 @@ "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 bdc130b63f3..b898e550f7c 100644 --- a/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts +++ b/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts @@ -68,6 +68,8 @@ export type Feature = | 'hidiffusion' | 'hidiffusionRauNet' | 'hidiffusionWindowAttn' + | 'hidiffusionT1Ratio' + | 'hidiffusionT2Ratio' | 'colorCompensation' | 'upscaleModel' | 'scale' @@ -228,6 +230,8 @@ export const POPOVER_DATA: { [key in Feature]?: PopoverData } = { }, hidiffusionRauNet: {}, hidiffusionWindowAttn: {}, + hidiffusionT1Ratio: {}, + hidiffusionT2Ratio: {}, } as const; export const OPEN_DELAY = 1000; // in milliseconds diff --git a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts index 20c2ee91cc5..30037dd0ea5 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts @@ -107,6 +107,12 @@ const slice = createSlice({ 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; }, @@ -486,6 +492,8 @@ export const { setHiDiffusionEnabled, setHiDiffusionRauNetEnabled, setHiDiffusionWindowAttnEnabled, + setHiDiffusionT1Ratio, + setHiDiffusionT2Ratio, setSeamlessXAxis, setSeamlessYAxis, setShouldRandomizeSeed, @@ -559,6 +567,12 @@ export const paramsSliceConfig: SliceConfig = { 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); }, @@ -616,6 +630,8 @@ export const selectOptimizedDenoisingEnabled = createParamsSelector((params) => 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 0d90a5e705b..8e6a99373bd 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/types.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/types.ts @@ -599,6 +599,8 @@ export const zParamsState = z.object({ hiDiffusionEnabled: z.boolean(), hiDiffusionRauNetEnabled: z.boolean(), hiDiffusionWindowAttnEnabled: z.boolean(), + hiDiffusionT1Ratio: z.number(), + hiDiffusionT2Ratio: z.number(), iterations: z.number(), scheduler: zParameterScheduler, fluxScheduler: zParameterFluxScheduler, @@ -662,6 +664,8 @@ export const getInitialParamsState = (): ParamsState => ({ 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.tsx b/invokeai/frontend/web/src/features/metadata/parsing.tsx index 1ed7d4505c1..23d0fa4c067 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.tsx +++ b/invokeai/frontend/web/src/features/metadata/parsing.tsx @@ -20,6 +20,8 @@ import { setGuidance, setHiDiffusionEnabled, setHiDiffusionRauNetEnabled, + setHiDiffusionT1Ratio, + setHiDiffusionT2Ratio, setHiDiffusionWindowAttnEnabled, setImg2imgStrength, setRefinerCFGScale, @@ -594,6 +596,42 @@ const HiDiffusionWindowAttn: SingleMetadataHandler = { }; //#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, @@ -1131,6 +1169,8 @@ export const ImageMetadataHandlers = { 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 757dc638c0d..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 @@ -25,6 +25,8 @@ export const addSDXLRefiner = async ( refinerStart, hiDiffusionEnabled, hiDiffusionRauNetEnabled, + hiDiffusionT1Ratio, + hiDiffusionT2Ratio, hiDiffusionWindowAttnEnabled, } = state.params; @@ -63,6 +65,8 @@ export const addSDXLRefiner = async ( 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/buildSD1Graph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts index 95ab2973ba9..a5cebc8c04e 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 @@ -44,6 +44,8 @@ export const buildSD1Graph = async (arg: GraphBuilderArg): Promise { }); 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 878200fdf19..00f55bfda45 100644 --- a/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx +++ b/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx @@ -17,6 +17,8 @@ import ParamCLIPLEmbedModelSelect from 'features/parameters/components/Advanced/ import ParamClipSkip from 'features/parameters/components/Advanced/ParamClipSkip'; import { ParamHiDiffusionRauNetToggle, + ParamHiDiffusionT1Ratio, + ParamHiDiffusionT2Ratio, ParamHiDiffusionToggle, ParamHiDiffusionWindowAttnToggle, } from 'features/parameters/components/Advanced/ParamHiDiffusionToggle'; @@ -119,6 +121,8 @@ export const AdvancedSettingsAccordion = memo(() => { + + diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 22d9fcb4652..3e615143636 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -6536,6 +6536,30 @@ export type components = { * @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.0 + */ + hidiffusion_t2_ratio?: number; + /** + * 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.0 + */ + hidiffusion_t2_ratio?: number; /** * @description Latents tensor * @default null From f27143bded219ca06c2b981447d0771eaa2dff07 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Mon, 26 Jan 2026 04:21:00 +0200 Subject: [PATCH 07/22] Checks appeasing --- .../stable_diffusion/hidiffusion_utils.py | 13 +++++----- .../frontend/web/src/services/api/schema.ts | 24 +++++++++---------- 2 files changed, 18 insertions(+), 19 deletions(-) diff --git a/invokeai/backend/stable_diffusion/hidiffusion_utils.py b/invokeai/backend/stable_diffusion/hidiffusion_utils.py index ffcb0291c36..436a5c4e566 100644 --- a/invokeai/backend/stable_diffusion/hidiffusion_utils.py +++ b/invokeai/backend/stable_diffusion/hidiffusion_utils.py @@ -1,10 +1,10 @@ from __future__ import annotations +import copy +import importlib.util from contextlib import contextmanager from typing import Any, Optional -import copy - @contextmanager def hidiffusion_patch( @@ -16,13 +16,12 @@ def hidiffusion_patch( t2_ratio: Optional[float] = None, ): """Context manager that applies HiDiffusion and restores the model on exit.""" - try: - import hidiffusion as hd - from hidiffusion import apply_hidiffusion, remove_hidiffusion - except ImportError as exc: + if importlib.util.find_spec("hidiffusion") is None: raise ImportError( "HiDiffusion is not installed. Install it with `pip install hidiffusion` to enable this option." - ) from exc + ) + + from hidiffusion import apply_hidiffusion, remove_hidiffusion target = model.unet if hasattr(model, "unet") else model diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 3e615143636..45e2d4540fb 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -6548,18 +6548,6 @@ export type components = { * @default 0.0 */ hidiffusion_t2_ratio?: number; - /** - * 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.0 - */ - hidiffusion_t2_ratio?: number; /** * @description Latents tensor * @default null @@ -6696,6 +6684,18 @@ export type components = { * @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.0 + */ + hidiffusion_t2_ratio?: number; /** * @description Latents tensor * @default null From 6f3aa5a11325b83c4b719ab330c7d98345e37455 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Mon, 26 Jan 2026 04:27:44 +0200 Subject: [PATCH 08/22] Refactor HiDiffusion import and update T2 ratio default value to 0 --- invokeai/backend/stable_diffusion/hidiffusion_utils.py | 2 +- invokeai/frontend/web/src/services/api/schema.ts | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/invokeai/backend/stable_diffusion/hidiffusion_utils.py b/invokeai/backend/stable_diffusion/hidiffusion_utils.py index 436a5c4e566..1672099cf08 100644 --- a/invokeai/backend/stable_diffusion/hidiffusion_utils.py +++ b/invokeai/backend/stable_diffusion/hidiffusion_utils.py @@ -90,8 +90,8 @@ def _set_name_or_path_on_config(cfg, value: str) -> bool: ratio_dicts = None if t1_ratio is not None or t2_ratio is not None: try: + from hidiffusion.hidiffusion import switching_threshold_ratio_dict as _switching_threshold_ratio_dict from hidiffusion.hidiffusion import ( - switching_threshold_ratio_dict as _switching_threshold_ratio_dict, text_to_img_controlnet_switching_threshold_ratio_dict as _text_to_img_controlnet_switching_threshold_ratio_dict, ) diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 45e2d4540fb..3ce93af6896 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -6545,7 +6545,7 @@ export type components = { /** * HiDiffusion: T2 Ratio * @description Override HiDiffusion late switch threshold (T2 ratio) - * @default 0.0 + * @default 0 */ hidiffusion_t2_ratio?: number; /** @@ -6693,7 +6693,7 @@ export type components = { /** * HiDiffusion: T2 Ratio * @description Override HiDiffusion late switch threshold (T2 ratio) - * @default 0.0 + * @default 0 */ hidiffusion_t2_ratio?: number; /** From 0fdc6d9497dd90554169344d04ed4698d63fdd1d Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Tue, 27 Jan 2026 20:11:09 +0200 Subject: [PATCH 09/22] Changed to vendoring an updated version of HiDiffusion and removed deps for external dependency --- invokeai/backend/hidiffusion/LICENSE | 201 ++ invokeai/backend/hidiffusion/NOTICE | 9 + invokeai/backend/hidiffusion/__init__.py | 3 + invokeai/backend/hidiffusion/hidiffusion.py | 2112 +++++++++++++++++ .../sd_module_key/sd15_module_key.txt | 685 ++++++ .../sd_module_key/sdxl_module_key.txt | 1910 +++++++++++++++ invokeai/backend/hidiffusion/utils.py | 36 + .../stable_diffusion/hidiffusion_utils.py | 12 +- pyproject.toml | 2 +- uv.lock | 11 - 10 files changed, 4960 insertions(+), 21 deletions(-) create mode 100644 invokeai/backend/hidiffusion/LICENSE create mode 100644 invokeai/backend/hidiffusion/NOTICE create mode 100644 invokeai/backend/hidiffusion/__init__.py create mode 100644 invokeai/backend/hidiffusion/hidiffusion.py create mode 100644 invokeai/backend/hidiffusion/sd_module_key/sd15_module_key.txt create mode 100644 invokeai/backend/hidiffusion/sd_module_key/sdxl_module_key.txt create mode 100644 invokeai/backend/hidiffusion/utils.py diff --git a/invokeai/backend/hidiffusion/LICENSE b/invokeai/backend/hidiffusion/LICENSE new file mode 100644 index 00000000000..261eeb9e9f8 --- /dev/null +++ b/invokeai/backend/hidiffusion/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/invokeai/backend/hidiffusion/NOTICE b/invokeai/backend/hidiffusion/NOTICE new file mode 100644 index 00000000000..9d6d1178ca1 --- /dev/null +++ b/invokeai/backend/hidiffusion/NOTICE @@ -0,0 +1,9 @@ +This code is vendored from: https://github.com/monofy-org/HiDiffusion + +This repository is a fork which implements some fixes to HiDiffusion + +I have further modified HiDiffusion so that apply_hidiffusion accepts +a torch.Generator object, so that deterministic images based on a seed +can be achieved. + +HiDiffusion is here: https://github.com/megvii-research/HiDiffusion \ No newline at end of file diff --git a/invokeai/backend/hidiffusion/__init__.py b/invokeai/backend/hidiffusion/__init__.py new file mode 100644 index 00000000000..9c9e6a3ed58 --- /dev/null +++ b/invokeai/backend/hidiffusion/__init__.py @@ -0,0 +1,3 @@ +from .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..e44a6fa1f13 --- /dev/null +++ b/invokeai/backend/hidiffusion/hidiffusion.py @@ -0,0 +1,2112 @@ +import importlib.resources + +import torch +import math +from typing import Type, Dict, Any, Tuple, Callable, Optional, Union, List +import torch.nn.functional as F +from .utils import isinstance_str +import diffusers +from diffusers.utils import USE_PEFT_BACKEND, scale_lora_layers, unscale_lora_layers, deprecate +from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput +from diffusers.image_processor import PipelineImageInput +from diffusers.utils.torch_utils import is_compiled_module, is_torch_version, apply_freeu +from diffusers.pipelines import auto_pipeline +from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel +from diffusers.models import ControlNetModel +from diffusers.models.attention import _chunked_feed_forward +import warnings + +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 = dict() + 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 = dict() + 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 = dict() + 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: List[str] = ["latents"], + **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. + """ + + # 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: + assert False + 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: + assert False + # 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) + ] + 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])] + 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 + ): + _, _, 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." + ) + 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) + # 1. Retrieve lora scale. + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + # 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 = () + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + blocks = list(zip(self.resnets, self.attentions)) + + 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) + + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + 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)): + # 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) + if self.timestep < self.T1: + self.ori_stride = self.stride + self.ori_padding = self.padding + self.ori_dilation = self.dilation + + self.stride = (4,4) + self.padding = (2,2) + self.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, self.stride, self.padding, self.dilation, self.groups + ) + if self.timestep < self.T1: + self.stride = self.ori_stride + self.padding = self.ori_padding + self.dilation = self.ori_dilation + 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, self.stride, self.padding, self.dilation, self.groups + ) + return original_outputs + (scale * self.lora_layer(hidden_states)) + else: + hidden_states = F.conv2d( + hidden_states, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups + ) + if self.timestep < self.T1: + self.stride = self.ori_stride + self.padding = self.ori_padding + self.dilation = self.ori_dilation + 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)) + + + +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 + + # Hack, avoid non-square problem. See unet_2d_condition.py in diffusers + diffusion_model.num_upsamplers += 12 + + name_or_path = 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']: + 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']: + 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']: + 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']: + 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']: + 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']: + 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']: + 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']: + 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']: + 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']: + 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'{model.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() + + if hasattr(module, "_parent"): + 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..9ccda215833 --- /dev/null +++ b/invokeai/backend/hidiffusion/utils.py @@ -0,0 +1,36 @@ +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): + """ + 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 + \ No newline at end of file diff --git a/invokeai/backend/stable_diffusion/hidiffusion_utils.py b/invokeai/backend/stable_diffusion/hidiffusion_utils.py index 1672099cf08..b3772598c8d 100644 --- a/invokeai/backend/stable_diffusion/hidiffusion_utils.py +++ b/invokeai/backend/stable_diffusion/hidiffusion_utils.py @@ -1,7 +1,6 @@ from __future__ import annotations import copy -import importlib.util from contextlib import contextmanager from typing import Any, Optional @@ -16,12 +15,7 @@ def hidiffusion_patch( t2_ratio: Optional[float] = None, ): """Context manager that applies HiDiffusion and restores the model on exit.""" - if importlib.util.find_spec("hidiffusion") is None: - raise ImportError( - "HiDiffusion is not installed. Install it with `pip install hidiffusion` to enable this option." - ) - - from hidiffusion import apply_hidiffusion, remove_hidiffusion + from invokeai.backend.hidiffusion.hidiffusion import apply_hidiffusion, remove_hidiffusion target = model.unet if hasattr(model, "unet") else model @@ -90,8 +84,8 @@ def _set_name_or_path_on_config(cfg, value: str) -> bool: ratio_dicts = None if t1_ratio is not None or t2_ratio is not None: try: - from hidiffusion.hidiffusion import switching_threshold_ratio_dict as _switching_threshold_ratio_dict - from hidiffusion.hidiffusion import ( + from invokeai.backend.hidiffusion.hidiffusion import ( + switching_threshold_ratio_dict as _switching_threshold_ratio_dict, text_to_img_controlnet_switching_threshold_ratio_dict as _text_to_img_controlnet_switching_threshold_ratio_dict, ) diff --git a/pyproject.toml b/pyproject.toml index 17e46baa27f..6310835de42 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,7 +38,6 @@ dependencies = [ "compel==2.1.1", "diffusers[torch]==0.36.0", "gguf", - "hidiffusion", # optional high-res denoise patch (HiDiffusion) "mediapipe==0.10.14", # needed for "mediapipeface" controlnet model "numpy<2.0.0", "onnx==1.16.1", @@ -189,6 +188,7 @@ version = { attr = "invokeai.version.__version__" } [tool.setuptools.package-data] "invokeai.app.assets" = ["**/*.png"] "invokeai.app.services.workflow_records.default_workflows" = ["*.json"] +"invokeai.backend.hidiffusion" = ["LICENSE", "NOTICE", "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/uv.lock b/uv.lock index 60d18aeec0c..3aa32aa58ad 100644 --- a/uv.lock +++ b/uv.lock @@ -815,15 +815,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f0/55/ef77a85ee443ae05a9e9cba1c9f0dd9241eb42da2aeba1dc50f51154c81a/hf_xet-1.1.5-cp37-abi3-win_amd64.whl", hash = "sha256:73e167d9807d166596b4b2f0b585c6d5bd84a26dea32843665a8b58f6edba245", size = 2738931, upload-time = "2025-06-20T21:48:39.482Z" }, ] -[[package]] -name = "hidiffusion" -version = "0.1.10" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/fd/2b/422acb8106f1295e94a8f98363e2b9b0d6be6f90f672d353f8b2c38654fe/hidiffusion-0.1.10.tar.gz", hash = "sha256:e8a7cfab49a3a29d83b33f3c6c4adcaf32377cb535d3a279bbf53125e5a57bcd", size = 41303, upload-time = "2024-06-16T17:23:58.739Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/64/fb/ebee9c90768db7c8368922d1387c64311051c4423de88aecfe1dcb8387ad/hidiffusion-0.1.10-py3-none-any.whl", hash = "sha256:9498c8fcc07f60645e533feeaaf66aeb6c74b289788be333f59507980556e9e4", size = 37727, upload-time = "2024-06-16T17:23:56.478Z" }, -] - [[package]] name = "httpcore" version = "1.0.9" @@ -981,7 +972,6 @@ dependencies = [ { name = "fastapi" }, { name = "fastapi-events" }, { name = "gguf" }, - { name = "hidiffusion" }, { name = "huggingface-hub" }, { name = "mediapipe" }, { name = "numpy" }, @@ -1089,7 +1079,6 @@ requires-dist = [ { name = "fastapi-events" }, { name = "gguf" }, { name = "gprof2dot", marker = "extra == 'dev'" }, - { name = "hidiffusion" }, { name = "httpx", marker = "extra == 'test'" }, { name = "huggingface-hub" }, { name = "humanize", marker = "extra == 'test'", specifier = "==4.12.1" }, From 4d8cc51149b6561104ec79b2ea48635339ac2cd0 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Wed, 4 Feb 2026 18:44:37 +0200 Subject: [PATCH 10/22] Ruff --- invokeai/backend/hidiffusion/__init__.py | 2 +- invokeai/backend/hidiffusion/hidiffusion.py | 58 ++++++++++--------- invokeai/backend/hidiffusion/utils.py | 7 +-- .../stable_diffusion/hidiffusion_utils.py | 23 ++++---- .../AdvancedSettingsAccordion.tsx | 2 +- 5 files changed, 46 insertions(+), 46 deletions(-) diff --git a/invokeai/backend/hidiffusion/__init__.py b/invokeai/backend/hidiffusion/__init__.py index 9c9e6a3ed58..c4b4d23d94c 100644 --- a/invokeai/backend/hidiffusion/__init__.py +++ b/invokeai/backend/hidiffusion/__init__.py @@ -1,3 +1,3 @@ -from .hidiffusion import apply_hidiffusion, remove_hidiffusion +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 index e44a6fa1f13..983c5eef2b9 100644 --- a/invokeai/backend/hidiffusion/hidiffusion.py +++ b/invokeai/backend/hidiffusion/hidiffusion.py @@ -1,20 +1,21 @@ import importlib.resources +import math +import warnings +from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union +import diffusers import torch -import math -from typing import Type, Dict, Any, Tuple, Callable, Optional, Union, List import torch.nn.functional as F -from .utils import isinstance_str -import diffusers -from diffusers.utils import USE_PEFT_BACKEND, scale_lora_layers, unscale_lora_layers, deprecate -from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput from diffusers.image_processor import PipelineImageInput -from diffusers.utils.torch_utils import is_compiled_module, is_torch_version, apply_freeu -from diffusers.pipelines import auto_pipeline -from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel from diffusers.models import ControlNetModel from diffusers.models.attention import _chunked_feed_forward -import warnings +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": @@ -25,7 +26,7 @@ old_diffusers = False def sd15_hidiffusion_key(): - modified_key = dict() + 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'] @@ -38,7 +39,7 @@ def sd15_hidiffusion_key(): return modified_key def sdxl_hidiffusion_key(): - modified_key = dict() + 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'] @@ -58,7 +59,7 @@ def sdxl_hidiffusion_key(): def sdxl_turbo_hidiffusion_key(): - modified_key = dict() + 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', @@ -169,7 +170,7 @@ def __call__( 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: List[str] = ["latents"], + callback_on_step_end_tensor_inputs: Optional[List[str]] = None, **kwargs, ): r""" @@ -328,6 +329,8 @@ def __call__( [`~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: @@ -535,7 +538,7 @@ def __call__( control_image = control_images height, width = control_image[0].shape[-2:] else: - assert False + raise AssertionError("Unsupported controlnet type for control image preprocessing.") else: if isinstance(controlnet, ControlNetModel): control_image = self.prepare_image( @@ -571,7 +574,7 @@ def __call__( control_image = images height, width = image[0].shape[-2:] else: - assert False + 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: @@ -631,7 +634,7 @@ def __call__( 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) + for s, e in zip(control_guidance_start, control_guidance_end, strict=False) ] controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) @@ -744,7 +747,10 @@ def __call__( 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])] + 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): @@ -1185,7 +1191,7 @@ def forward( new_down_block_res_samples = () for down_block_res_sample, down_block_additional_residual in zip( - down_block_res_samples, down_block_additional_residuals + 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') @@ -1311,7 +1317,8 @@ def window_partition(x, window_size, shift_size, H, W): 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." + 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: @@ -1399,9 +1406,6 @@ def window_reverse(windows, window_size, H, W, shift_size): 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) - # 1. Retrieve lora scale. - lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 - # 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) @@ -1555,9 +1559,8 @@ def forward( self.T1 = int(self.max_timestep * self.T1_ratio) output_states = () - lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 - blocks = list(zip(self.resnets, self.attentions)) + blocks = list(zip(self.resnets, self.attentions, strict=False)) for i, (resnet, attn) in enumerate(blocks): if self.training and self.gradient_checkpointing: @@ -1687,7 +1690,6 @@ def forward( else: self.T1 = int(self.max_timestep * self.T1_ratio) - lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 is_freeu_enabled = ( getattr(self, "s1", None) and getattr(self, "s2", None) @@ -1695,7 +1697,7 @@ def forward( and getattr(self, "b2", None) ) - for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)): + 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] @@ -1995,7 +1997,7 @@ def apply_hidiffusion( name_or_path = 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(): + 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' diff --git a/invokeai/backend/hidiffusion/utils.py b/invokeai/backend/hidiffusion/utils.py index 9ccda215833..b8db24238b6 100644 --- a/invokeai/backend/hidiffusion/utils.py +++ b/invokeai/backend/hidiffusion/utils.py @@ -5,7 +5,7 @@ def isinstance_str(x: object, cls_name: str, prefix: bool = False, contains: boo """ 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! """ @@ -16,11 +16,11 @@ def isinstance_str(x: object, cls_name: str, prefix: bool = False, contains: boo return True if _cls.__name__ == cls_name: return True - + return False -def init_generator(device: torch.device, fallback: torch.Generator=None): +def init_generator(device: torch.device, fallback: torch.Generator | None = None): """ Forks the current default random generator given device. """ @@ -33,4 +33,3 @@ def init_generator(device: torch.device, fallback: torch.Generator=None): return init_generator(torch.device("cpu")) else: return fallback - \ No newline at end of file diff --git a/invokeai/backend/stable_diffusion/hidiffusion_utils.py b/invokeai/backend/stable_diffusion/hidiffusion_utils.py index b3772598c8d..50c257a3398 100644 --- a/invokeai/backend/stable_diffusion/hidiffusion_utils.py +++ b/invokeai/backend/stable_diffusion/hidiffusion_utils.py @@ -4,6 +4,13 @@ from contextlib import contextmanager from typing import Any, Optional +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( @@ -83,18 +90,10 @@ def _set_name_or_path_on_config(cfg, value: str) -> bool: ratio_overrides = None ratio_dicts = None if t1_ratio is not None or t2_ratio is not None: - try: - from invokeai.backend.hidiffusion.hidiffusion import ( - switching_threshold_ratio_dict as _switching_threshold_ratio_dict, - text_to_img_controlnet_switching_threshold_ratio_dict as _text_to_img_controlnet_switching_threshold_ratio_dict, - ) - - ratio_dicts = ( - _switching_threshold_ratio_dict, - _text_to_img_controlnet_switching_threshold_ratio_dict, - ) - except Exception: - ratio_dicts = None + ratio_dicts = ( + _switching_threshold_ratio_dict, + _text_to_img_controlnet_switching_threshold_ratio_dict, + ) if ratio_dicts is not None: ratio_overrides = (copy.deepcopy(ratio_dicts[0]), copy.deepcopy(ratio_dicts[1])) 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 9980cf34868..515c22b0125 100644 --- a/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx +++ b/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx @@ -16,6 +16,7 @@ import ParamCLIPEmbedModelSelect from 'features/parameters/components/Advanced/P import ParamCLIPGEmbedModelSelect from 'features/parameters/components/Advanced/ParamCLIPGEmbedModelSelect'; 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, @@ -23,7 +24,6 @@ import { ParamHiDiffusionToggle, ParamHiDiffusionWindowAttnToggle, } from 'features/parameters/components/Advanced/ParamHiDiffusionToggle'; -import ParamFlux2KleinModelSelect from 'features/parameters/components/Advanced/ParamFlux2KleinModelSelect'; import ParamT5EncoderModelSelect from 'features/parameters/components/Advanced/ParamT5EncoderModelSelect'; import ParamZImageQwen3VaeModelSelect from 'features/parameters/components/Advanced/ParamZImageQwen3VaeModelSelect'; import ParamSeamlessXAxis from 'features/parameters/components/Seamless/ParamSeamlessXAxis'; From 0ba1b665ce52d1e3b1f7374f52bbb3a209af5dd5 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Wed, 4 Feb 2026 18:50:26 +0200 Subject: [PATCH 11/22] Ruff again --- invokeai/backend/hidiffusion/hidiffusion.py | 470 +++++++++++--------- 1 file changed, 257 insertions(+), 213 deletions(-) diff --git a/invokeai/backend/hidiffusion/hidiffusion.py b/invokeai/backend/hidiffusion/hidiffusion.py index 983c5eef2b9..9680459e31c 100644 --- a/invokeai/backend/hidiffusion/hidiffusion.py +++ b/invokeai/backend/hidiffusion/hidiffusion.py @@ -20,81 +20,94 @@ 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'] + 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'] + 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'] + 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' + "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}, + "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}, + "sdxl_2048": {"T1_ratio": 0.5, "T2_ratio": 0.0}, } controlnet_apply_steps_rate = 0.6 @@ -105,12 +118,10 @@ def sdxl_turbo_hidiffusion_key(): playground_is_aggressive_raunet = False -with importlib.resources.open_text( - f"{__package__}.sd_module_key", "sd15_module_key.txt", encoding="utf-8") as f: +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: +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() @@ -118,15 +129,14 @@ 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'): + 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 @@ -336,9 +346,9 @@ def __call__( if control_image is None: control_image = image image = None - self.info['text_to_img_controlnet'] = True + self.info["text_to_img_controlnet"] = True else: - self.info['text_to_img_controlnet'] = False + self.info["text_to_img_controlnet"] = False callback = kwargs.pop("callback", None) callback_steps = kwargs.pop("callback_steps", None) @@ -748,8 +758,7 @@ def __call__( if isinstance(controlnet_keep[i], list): cond_scale = [ - c * s - for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i], strict=False) + c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i], strict=False) ] else: controlnet_cond_scale = controlnet_conditioning_scale @@ -757,19 +766,20 @@ def __call__( 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) + 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) + 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] + 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), @@ -888,10 +898,10 @@ def __call__( 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, @@ -1194,7 +1204,9 @@ def forward( 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_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,) @@ -1224,7 +1236,9 @@ def forward( 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') + mid_block_additional_residual = F.interpolate( + mid_block_additional_residual, (ori_H, ori_W), mode="bicubic" + ) sample = sample + mid_block_additional_residual # 5. up @@ -1280,10 +1294,13 @@ def forward( 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]: +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 @@ -1300,7 +1317,6 @@ def forward( 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): """ @@ -1312,15 +1328,23 @@ def window_partition(x, window_size, shift_size, H, W): windows: (num_windows*B, window_size, window_size, C) """ B, N, C = x.shape - x = x.view(B,H,W,C) + 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"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() + 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)) @@ -1332,7 +1356,6 @@ def window_partition(x, window_size, shift_size, H, W): windows = windows.view(-1, window_size[0] * window_size[1], C) return windows - def window_reverse(windows, window_size, H, W, shift_size): """ Args: @@ -1346,9 +1369,9 @@ def window_reverse(windows, window_size, H, W, shift_size): """ B, N, C = windows.shape windows = windows.view(-1, window_size[0], window_size[1], C) - B = int(windows.shape[0] / 4) # 2x2 + 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) + 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)) @@ -1356,8 +1379,12 @@ def window_reverse(windows, window_size, H, W, shift_size): 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) + 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. @@ -1393,18 +1420,18 @@ def window_reverse(windows, window_size, H, W, shift_size): 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)) + 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) + shift_size = (0, 0) if rand_num > 0.25 and rand_num <= 0.5: - shift_size = (widow_size[0]//4, widow_size[1]//4) + 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) + 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) + 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 {} @@ -1472,9 +1499,7 @@ def window_reverse(windows, window_size, H, W, shift_size): 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) # ff_output = _chunked_feed_forward( # self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size, lora_scale=lora_scale # ) @@ -1507,7 +1532,7 @@ class cross_attn_down_block(block_class): T1_start = 0 T1_end = 0 aggressive_raunet = False - T1 = 0 # to avoid confict with sdxl-turbo + T1 = 0 # to avoid confict with sdxl-turbo max_timestep = 50 info: dict = None model: str = None @@ -1522,39 +1547,40 @@ def forward( 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': + 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] + 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': + 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] + 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] + self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] - if self.info['is_inpainting_task']: + if self.info["is_inpainting_task"]: self.aggressive_raunet = inpainting_is_aggressive_raunet - elif self.info['is_playground']: + 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] + 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_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 + self.T1 = 0 # to avoid confict with sdxl-turbo else: self.T1 = int(self.max_timestep * self.T1_ratio) @@ -1608,10 +1634,10 @@ def custom_forward(*inputs): 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) + 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) + 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: @@ -1626,8 +1652,10 @@ def custom_forward(*inputs): 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): @@ -1639,7 +1667,7 @@ class cross_attn_up_block(block_class): T1_start = 0 T1_end = 0 aggressive_raunet = False - T1 = 0 # to avoid confict with sdxl-turbo + T1 = 0 # to avoid confict with sdxl-turbo max_timestep = 50 def forward( @@ -1653,40 +1681,41 @@ def forward( 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': + 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] + 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': + 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] + 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] + self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] - if self.info['is_inpainting_task']: + if self.info["is_inpainting_task"]: self.aggressive_raunet = inpainting_is_aggressive_raunet - elif self.info['is_playground']: + 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] + 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_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 + self.T1 = 0 # to avoid confict with sdxl-turbo else: self.T1 = int(self.max_timestep * self.T1_ratio) @@ -1756,9 +1785,13 @@ def custom_forward(*inputs): 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') + 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') + 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) @@ -1769,8 +1802,8 @@ def custom_forward(*inputs): self.timestep = 0 return hidden_states - return cross_attn_up_block + return cross_attn_up_block def make_diffusers_downsampler_block(block_class: Type[torch.nn.Module]) -> Type[torch.nn.Module]: @@ -1784,37 +1817,39 @@ class downsampler_block(block_class): aggressive_raunet = False max_timestep = 50 - def forward(self, hidden_states: torch.Tensor, scale = 1.0) -> torch.Tensor: + 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': + 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] + 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': + 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] + 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] + self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] - if self.info['is_inpainting_task']: + if self.info["is_inpainting_task"]: self.aggressive_raunet = inpainting_is_aggressive_raunet - elif self.info['is_playground']: + 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] + 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) + self.T1 = int(aggressive_step / 50 * self.max_timestep) else: self.T1 = int(self.max_timestep * self.T1_ratio) if self.timestep < self.T1: @@ -1822,9 +1857,9 @@ def forward(self, hidden_states: torch.Tensor, scale = 1.0) -> torch.Tensor: self.ori_padding = self.padding self.ori_dilation = self.dilation - self.stride = (4,4) - self.padding = (2,2) - self.dilation = (2,2) + self.stride = (4, 4) + self.padding = (2, 2) + self.dilation = (2, 2) if old_diffusers: if self.lora_layer is None: @@ -1858,6 +1893,7 @@ def forward(self, hidden_states: torch.Tensor, scale = 1.0) -> torch.Tensor: if self.timestep == self.max_timestep: self.timestep = 0 return hidden_states + return downsampler_block @@ -1873,38 +1909,39 @@ class upsampler_block(block_class): max_timestep = 50 info: dict = None - def forward(self, hidden_states: torch.Tensor, scale = 1.0) -> torch.Tensor: + 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': + 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] + 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': + 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] + 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] + self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] - if self.info['is_inpainting_task']: + if self.info["is_inpainting_task"]: self.aggressive_raunet = inpainting_is_aggressive_raunet - elif self.info['is_playground']: + 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] + 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) + self.T1 = int(aggressive_step / 50 * self.max_timestep) else: self.T1 = int(self.max_timestep * self.T1_ratio) self.timestep += 1 @@ -1927,12 +1964,13 @@ def forward(self, hidden_states: torch.Tensor, scale = 1.0) -> torch.Tensor: return F.conv2d( hidden_states, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups ) - return upsampler_block + 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. """ + """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 @@ -1940,13 +1978,13 @@ def hook(module, args): model.info["hooks"].append(model.register_forward_pre_hook(hook)) - 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: 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. @@ -1964,12 +2002,16 @@ def apply_hidiffusion( 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) + 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) + 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 @@ -2000,105 +2042,107 @@ def apply_hidiffusion( 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' + 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' + 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 + "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']: + 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']: + if apply_raunet and key in modified_key["down_module_key"]: 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']: + module.switching_threshold_ratio = "T1_ratio" + if apply_raunet and key in modified_key["down_module_key_extra"]: 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']: + module.switching_threshold_ratio = "T2_ratio" + if apply_raunet and key in modified_key["up_module_key"]: 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']: + module.switching_threshold_ratio = "T1_ratio" + if apply_raunet and key in modified_key["up_module_key_extra"]: 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']: + 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.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']: + 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']: + if apply_raunet and key in modified_key["down_module_key"]: make_block_fn = make_diffusers_cross_attn_down_block module.__class__ = make_block_fn(module.__class__) - module.switching_threshold_ratio = 'T1_ratio' + module.switching_threshold_ratio = "T1_ratio" - if apply_raunet and key in modified_key['down_module_key_extra']: + if apply_raunet and key in modified_key["down_module_key_extra"]: make_block_fn = make_diffusers_downsampler_block module.__class__ = make_block_fn(module.__class__) - module.switching_threshold_ratio = 'T2_ratio' + module.switching_threshold_ratio = "T2_ratio" - if apply_raunet and key in modified_key['up_module_key']: + if apply_raunet and key in modified_key["up_module_key"]: make_block_fn = make_diffusers_cross_attn_up_block module.__class__ = make_block_fn(module.__class__) - module.switching_threshold_ratio = 'T1_ratio' + module.switching_threshold_ratio = "T1_ratio" - if apply_raunet and key in modified_key['up_module_key_extra']: + if apply_raunet and key in modified_key["up_module_key_extra"]: make_block_fn = make_diffusers_upsampler_block module.__class__ = make_block_fn(module.__class__) - module.switching_threshold_ratio = 'T2_ratio' + module.switching_threshold_ratio = "T2_ratio" - if apply_window_attn and key in modified_key['windown_attn_module_key']: + 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.model = "sdxl" module.info = diffusion_model.info - elif name_or_path == 'stabilityai/sdxl-turbo': + 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']: + if apply_raunet and key in modified_key["down_module_key"]: make_block_fn = make_diffusers_cross_attn_down_block module.__class__ = make_block_fn(module.__class__) - module.switching_threshold_ratio = 'T1_ratio' + module.switching_threshold_ratio = "T1_ratio" - if apply_raunet and key in modified_key['up_module_key']: + if apply_raunet and key in modified_key["up_module_key"]: make_block_fn = make_diffusers_cross_attn_up_block module.__class__ = make_block_fn(module.__class__) - module.switching_threshold_ratio = 'T1_ratio' + module.switching_threshold_ratio = "T1_ratio" - if apply_window_attn and key in modified_key['windown_attn_module_key']: + 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.model = "sdxl_turbo" module.info = diffusion_model.info else: - raise Exception(f'{model.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.') + raise Exception( + f"{model.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. """ + """Removes hidiffusion from a Diffusion module if it was already patched.""" # For diffusers model = model.unet if hasattr(model, "unet") else model From 05f7f1ab991db7e6c76010caa33e5e5ba7d8ff4c Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Tue, 19 May 2026 22:28:29 +0300 Subject: [PATCH 12/22] chore(api): refresh hidiffusion openapi schema --- invokeai/frontend/web/openapi.json | 52 +++++++++++++++--------------- 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 900a9c6b15a..2e041e0c82d 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -20657,37 +20657,13 @@ "type": "number" }, "hidiffusion_t2_ratio": { - "default": 0, - "description": "Override HiDiffusion late switch threshold (T2 ratio)", - "field_kind": "input", - "input": "any", - "maximum": 1, - "minimum": 0, - "orig_default": 0, - "orig_required": false, - "title": "HiDiffusion: T2 Ratio", - "type": "number" - }, - "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, + "default": 0.0, "description": "Override HiDiffusion late switch threshold (T2 ratio)", "field_kind": "input", "input": "any", "maximum": 1, "minimum": 0, - "orig_default": 0, + "orig_default": 0.0, "orig_required": false, "title": "HiDiffusion: T2 Ratio", "type": "number" @@ -21081,6 +21057,30 @@ "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": [ { From 72b369583c33e095cd59c7bd0769ef0b7e961695 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Tue, 19 May 2026 22:40:13 +0300 Subject: [PATCH 13/22] Fix HiDiffusion cleanup and metadata recall --- invokeai/app/invocations/metadata_linked.py | 11 +- .../stable_diffusion/hidiffusion_utils.py | 79 ++++---- .../generation/buildHiDiffusionGraph.test.ts | 190 ++++++++++++++++++ .../util/graph/generation/buildSD1Graph.ts | 10 +- .../util/graph/generation/buildSDXLGraph.ts | 10 +- .../invocations/test_denoise_noise_inputs.py | 36 +++- .../test_hidiffusion_utils.py | 93 +++++++++ 7 files changed, 376 insertions(+), 53 deletions(-) create mode 100644 invokeai/frontend/web/src/features/nodes/util/graph/generation/buildHiDiffusionGraph.test.ts create mode 100644 tests/backend/stable_diffusion/test_hidiffusion_utils.py diff --git a/invokeai/app/invocations/metadata_linked.py b/invokeai/app/invocations/metadata_linked.py index 1aa154c46ac..eefba505601 100644 --- a/invokeai/app/invocations/metadata_linked.py +++ b/invokeai/app/invocations/metadata_linked.py @@ -661,12 +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}) - if self.hidiffusion: - 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}) + 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/stable_diffusion/hidiffusion_utils.py b/invokeai/backend/stable_diffusion/hidiffusion_utils.py index 50c257a3398..580a2672b77 100644 --- a/invokeai/backend/stable_diffusion/hidiffusion_utils.py +++ b/invokeai/backend/stable_diffusion/hidiffusion_utils.py @@ -1,6 +1,7 @@ from __future__ import annotations import copy +import sys from contextlib import contextmanager from typing import Any, Optional @@ -62,29 +63,6 @@ def _set_name_or_path_on_config(cfg, value: str) -> bool: except Exception: return False - set_model_name_or_path = False - 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 - original_num_upsamplers = getattr(target, "num_upsamplers", None) ratio_overrides = None @@ -95,24 +73,53 @@ def _set_name_or_path_on_config(cfg, value: str) -> bool: _text_to_img_controlnet_switching_threshold_ratio_dict, ) - if ratio_dicts is not None: - ratio_overrides = (copy.deepcopy(ratio_dicts[0]), copy.deepcopy(ratio_dicts[1])) + 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 + 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_ratio_overrides(ratio_dicts[0]) + _apply_ratio_overrides(ratio_dicts[1]) - apply_hidiffusion(model, apply_raunet=apply_raunet, apply_window_attn=apply_window_attn) - try: + apply_hidiffusion(model, apply_raunet=apply_raunet, apply_window_attn=apply_window_attn) yield finally: - remove_hidiffusion(model) + had_active_exception = sys.exc_info()[0] is not None + try: + remove_hidiffusion(model) + except Exception: + if not had_active_exception: + raise if ratio_overrides is not None and ratio_dicts is not None: ratio_dicts[0].clear() ratio_dicts[0].update(ratio_overrides[0]) 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..03956bd0ef8 --- /dev/null +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildHiDiffusionGraph.test.ts @@ -0,0 +1,190 @@ +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, +}; + +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', () => ({ + 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 type { Graph } from './Graph'; +import { buildSDXLGraph } from './buildSDXLGraph'; + +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 a5cebc8c04e..b10a92decb3 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 @@ -147,11 +147,11 @@ export const buildSD1Graph = async (arg: GraphBuilderArg): Promise 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..35f2f7679a0 --- /dev/null +++ b/tests/backend/stable_diffusion/test_hidiffusion_utils.py @@ -0,0 +1,93 @@ +import copy +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from invokeai.backend.hidiffusion.hidiffusion import ( + remove_hidiffusion as real_remove_hidiffusion, +) +from invokeai.backend.hidiffusion.hidiffusion import ( + switching_threshold_ratio_dict, +) +from invokeai.backend.hidiffusion.hidiffusion import ( + 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)] + + +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) From 3cf086339e41bacebc13b7fe1f4516f6743fa236 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Tue, 19 May 2026 22:51:33 +0300 Subject: [PATCH 14/22] chore: fix lint import ordering --- .../nodes/util/graph/generation/buildHiDiffusionGraph.test.ts | 2 +- tests/backend/stable_diffusion/test_hidiffusion_utils.py | 2 -- 2 files changed, 1 insertion(+), 3 deletions(-) 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 index 03956bd0ef8..870f05978f8 100644 --- 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 @@ -134,8 +134,8 @@ vi.mock('./addRegions', () => ({ import type { GraphBuilderArg } from 'features/nodes/util/graph/types'; import { buildSD1Graph } from './buildSD1Graph'; -import type { Graph } from './Graph'; import { buildSDXLGraph } from './buildSDXLGraph'; +import type { Graph } from './Graph'; const buildGraphArg = (): GraphBuilderArg => ({ diff --git a/tests/backend/stable_diffusion/test_hidiffusion_utils.py b/tests/backend/stable_diffusion/test_hidiffusion_utils.py index 35f2f7679a0..2ac967987fe 100644 --- a/tests/backend/stable_diffusion/test_hidiffusion_utils.py +++ b/tests/backend/stable_diffusion/test_hidiffusion_utils.py @@ -9,8 +9,6 @@ ) from invokeai.backend.hidiffusion.hidiffusion import ( switching_threshold_ratio_dict, -) -from invokeai.backend.hidiffusion.hidiffusion import ( text_to_img_controlnet_switching_threshold_ratio_dict, ) from invokeai.backend.stable_diffusion.hidiffusion_utils import hidiffusion_patch From 3fe3aa2789791e83292418b73228e208c434ecd7 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Tue, 28 Jul 2026 22:47:25 +0300 Subject: [PATCH 15/22] Fix HiDiffusion cleanup and metadata recall --- .../src/content/docs/features/hidiffusion.mdx | 8 +-- .../stable_diffusion/hidiffusion_utils.py | 20 ++++-- .../web/src/features/metadata/parsing.test.ts | 35 ++++++++++ .../web/src/features/metadata/parsing.tsx | 2 +- .../test_hidiffusion_utils.py | 66 +++++++++++++++++++ 5 files changed, 119 insertions(+), 12 deletions(-) diff --git a/docs/src/content/docs/features/hidiffusion.mdx b/docs/src/content/docs/features/hidiffusion.mdx index 4c3f3346ec8..5aee05985c7 100644 --- a/docs/src/content/docs/features/hidiffusion.mdx +++ b/docs/src/content/docs/features/hidiffusion.mdx @@ -4,7 +4,9 @@ title: HiDiffusion # HiDiffusion -HiDiffusion is an optional denoising enhancement that can improve detail and structure at higher resolutions for SD 1.5 and SDXL. It modifies the UNet during denoising and is most noticeable at 1536px and above. +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 @@ -40,7 +42,3 @@ Learn more: https://github.com/megvii-research/HiDiffusion - 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. - ---- - -TODO: Decide whether HiDiffusion toggles and ratios should always be emitted in metadata for recall. diff --git a/invokeai/backend/stable_diffusion/hidiffusion_utils.py b/invokeai/backend/stable_diffusion/hidiffusion_utils.py index 580a2672b77..baf0a779125 100644 --- a/invokeai/backend/stable_diffusion/hidiffusion_utils.py +++ b/invokeai/backend/stable_diffusion/hidiffusion_utils.py @@ -115,11 +115,12 @@ def _apply_ratio_overrides(ratio_dict: dict) -> None: yield finally: had_active_exception = sys.exc_info()[0] is not None + teardown_error: Exception | None = None try: remove_hidiffusion(model) - except Exception: + except Exception as error: if not had_active_exception: - raise + 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]) @@ -141,9 +142,16 @@ def _apply_ratio_overrides(ratio_dict: dict) -> None: 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) - else: - if config is not None: + elif config is not None: + internal_dict = getattr(config, "_internal_dict", None) + if internal_dict is not None: try: - delattr(config, "_name_or_path") - except AttributeError: + 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/src/features/metadata/parsing.test.ts b/invokeai/frontend/web/src/features/metadata/parsing.test.ts index f4a542f658c..cc54ab304c0 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,36 @@ 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(hiDiffusionEnabled).toBe(false); + }); +}); diff --git a/invokeai/frontend/web/src/features/metadata/parsing.tsx b/invokeai/frontend/web/src/features/metadata/parsing.tsx index 6ab04f14628..e2064c33535 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.tsx +++ b/invokeai/frontend/web/src/features/metadata/parsing.tsx @@ -694,7 +694,7 @@ const HiDiffusion: SingleMetadataHandler = { type: 'HiDiffusion', parse: (metadata, _store) => { const raw = getProperty(metadata, 'hidiffusion'); - const parsed = z.boolean().parse(raw); + const parsed = raw === undefined ? false : z.boolean().parse(raw); return Promise.resolve(parsed); }, recall: (value, store) => { diff --git a/tests/backend/stable_diffusion/test_hidiffusion_utils.py b/tests/backend/stable_diffusion/test_hidiffusion_utils.py index 2ac967987fe..3f8f0494dc1 100644 --- a/tests/backend/stable_diffusion/test_hidiffusion_utils.py +++ b/tests/backend/stable_diffusion/test_hidiffusion_utils.py @@ -89,3 +89,69 @@ def fake_apply_hidiffusion(patched_model, **_kwargs): 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 From 80c457d4bd1d42ee85624f95659e46d49c98c2f4 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Thu, 30 Jul 2026 08:27:39 +0300 Subject: [PATCH 16/22] Move HiDiffusion license notice to repository root --- LICENSE-HiDiffusion.txt | 65 +++++++++ invokeai/backend/hidiffusion/LICENSE | 201 --------------------------- invokeai/backend/hidiffusion/NOTICE | 9 -- pyproject.toml | 2 +- 4 files changed, 66 insertions(+), 211 deletions(-) create mode 100644 LICENSE-HiDiffusion.txt delete mode 100644 invokeai/backend/hidiffusion/LICENSE delete mode 100644 invokeai/backend/hidiffusion/NOTICE diff --git a/LICENSE-HiDiffusion.txt b/LICENSE-HiDiffusion.txt new file mode 100644 index 00000000000..f8c5a3567d8 --- /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. \ No newline at end of file diff --git a/invokeai/backend/hidiffusion/LICENSE b/invokeai/backend/hidiffusion/LICENSE deleted file mode 100644 index 261eeb9e9f8..00000000000 --- a/invokeai/backend/hidiffusion/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/invokeai/backend/hidiffusion/NOTICE b/invokeai/backend/hidiffusion/NOTICE deleted file mode 100644 index 9d6d1178ca1..00000000000 --- a/invokeai/backend/hidiffusion/NOTICE +++ /dev/null @@ -1,9 +0,0 @@ -This code is vendored from: https://github.com/monofy-org/HiDiffusion - -This repository is a fork which implements some fixes to HiDiffusion - -I have further modified HiDiffusion so that apply_hidiffusion accepts -a torch.Generator object, so that deterministic images based on a seed -can be achieved. - -HiDiffusion is here: https://github.com/megvii-research/HiDiffusion \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 40666c67395..20087cc0955 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -239,7 +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" = ["LICENSE", "NOTICE", "sd_module_key/*.txt"] +"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", From b49af9316c9dd1c4ae3d6a97b028c84145fafa9e Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Fri, 31 Jul 2026 04:23:44 +0300 Subject: [PATCH 17/22] Fix HiDiffusion modular denoise and docs formatting --- .../src/content/docs/features/hidiffusion.mdx | 15 ++++-------- invokeai/backend/hidiffusion/hidiffusion.py | 4 ++-- .../test_hidiffusion_utils.py | 24 +++++++++++++++++++ 3 files changed, 31 insertions(+), 12 deletions(-) diff --git a/docs/src/content/docs/features/hidiffusion.mdx b/docs/src/content/docs/features/hidiffusion.mdx index 5aee05985c7..8158f4980a1 100644 --- a/docs/src/content/docs/features/hidiffusion.mdx +++ b/docs/src/content/docs/features/hidiffusion.mdx @@ -22,20 +22,15 @@ Learn more: https://github.com/megvii-research/HiDiffusion ## 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**: 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: 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: 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: 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. +- **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 diff --git a/invokeai/backend/hidiffusion/hidiffusion.py b/invokeai/backend/hidiffusion/hidiffusion.py index 9680459e31c..b06ae5345e0 100644 --- a/invokeai/backend/hidiffusion/hidiffusion.py +++ b/invokeai/backend/hidiffusion/hidiffusion.py @@ -2036,7 +2036,7 @@ def apply_hidiffusion( # Hack, avoid non-square problem. See unet_2d_condition.py in diffusers diffusion_model.num_upsamplers += 12 - name_or_path = model.name_or_path + 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(): @@ -2136,7 +2136,7 @@ def apply_hidiffusion( module.info = diffusion_model.info else: raise Exception( - f"{model.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." + 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 diff --git a/tests/backend/stable_diffusion/test_hidiffusion_utils.py b/tests/backend/stable_diffusion/test_hidiffusion_utils.py index 3f8f0494dc1..1859a7c4093 100644 --- a/tests/backend/stable_diffusion/test_hidiffusion_utils.py +++ b/tests/backend/stable_diffusion/test_hidiffusion_utils.py @@ -3,6 +3,7 @@ from unittest.mock import MagicMock, patch import pytest +import torch from invokeai.backend.hidiffusion.hidiffusion import ( remove_hidiffusion as real_remove_hidiffusion, @@ -31,6 +32,29 @@ 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 + + +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 model.info["hooks"] == [] + + 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) From d7a7449b60e60db2be16f2c21d1aa46546bce47e Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Sat, 1 Aug 2026 08:50:41 +0300 Subject: [PATCH 18/22] Bump HiDiffusion node versions and refresh schema --- LICENSE-HiDiffusion.txt | 2 +- invokeai/app/invocations/denoise_latents.py | 2 +- invokeai/app/invocations/metadata_linked.py | 2 +- invokeai/frontend/web/openapi.json | 4 ++-- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/LICENSE-HiDiffusion.txt b/LICENSE-HiDiffusion.txt index f8c5a3567d8..73095184ff3 100644 --- a/LICENSE-HiDiffusion.txt +++ b/LICENSE-HiDiffusion.txt @@ -62,4 +62,4 @@ The following changes were applied when integrating HiDiffusion into InvokeAI: * 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. \ No newline at end of file + generation lifecycle. diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index f50cf6ca94c..93ef5b53f9a 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -132,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""" diff --git a/invokeai/app/invocations/metadata_linked.py b/invokeai/app/invocations/metadata_linked.py index eefba505601..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: diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index ec33b529884..4090bd23846 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -24136,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" } @@ -24549,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" } From f94e1f49075dd52725c00a406e19c7b8cff7948c Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Sun, 2 Aug 2026 02:37:28 +0300 Subject: [PATCH 19/22] Seed HiDiffusion window attention deterministically --- invokeai/app/invocations/denoise_latents.py | 2 + .../extensions/hidiffusion.py | 4 ++ .../stable_diffusion/hidiffusion_utils.py | 10 ++- .../test_hidiffusion_utils.py | 63 +++++++++++++++++++ 4 files changed, 78 insertions(+), 1 deletion(-) diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index 93ef5b53f9a..68e8dfc0dcd 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -922,6 +922,7 @@ def step_callback(state: PipelineIntermediateState) -> None: 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), ) ) @@ -1147,6 +1148,7 @@ def _lora_loader() -> Iterator[PatchSpec]: 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() diff --git a/invokeai/backend/stable_diffusion/extensions/hidiffusion.py b/invokeai/backend/stable_diffusion/extensions/hidiffusion.py index 0dc063b10d0..444c90f5480 100644 --- a/invokeai/backend/stable_diffusion/extensions/hidiffusion.py +++ b/invokeai/backend/stable_diffusion/extensions/hidiffusion.py @@ -3,6 +3,7 @@ from contextlib import contextmanager from typing import Optional +import torch from diffusers import UNet2DConditionModel from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase @@ -18,6 +19,7 @@ def __init__( 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 @@ -25,6 +27,7 @@ def __init__( 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): @@ -35,5 +38,6 @@ def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeigh 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 index baf0a779125..327d7d083b1 100644 --- a/invokeai/backend/stable_diffusion/hidiffusion_utils.py +++ b/invokeai/backend/stable_diffusion/hidiffusion_utils.py @@ -5,6 +5,8 @@ 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, ) @@ -21,6 +23,7 @@ def hidiffusion_patch( 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 @@ -111,7 +114,12 @@ def _apply_ratio_overrides(ratio_dict: dict) -> None: _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) + 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 diff --git a/tests/backend/stable_diffusion/test_hidiffusion_utils.py b/tests/backend/stable_diffusion/test_hidiffusion_utils.py index 1859a7c4093..98d207b5e66 100644 --- a/tests/backend/stable_diffusion/test_hidiffusion_utils.py +++ b/tests/backend/stable_diffusion/test_hidiffusion_utils.py @@ -40,6 +40,35 @@ def __init__(self): 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() + + def test_hidiffusion_patch_supports_bare_model_mixin_without_public_name_or_path(): model = ModelMixin() @@ -55,6 +84,40 @@ def test_hidiffusion_patch_supports_bare_model_mixin_without_public_name_or_path assert model.info["hooks"] == [] +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_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) From 3e04bed199bf389044f30e7aeed731cfe4d23a63 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Sun, 2 Aug 2026 14:53:11 +0300 Subject: [PATCH 20/22] fix: stale HiDiffusion state in cached UNet --- invokeai/backend/hidiffusion/hidiffusion.py | 38 ++++++++++++++++++ .../test_hidiffusion_utils.py | 40 +++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/invokeai/backend/hidiffusion/hidiffusion.py b/invokeai/backend/hidiffusion/hidiffusion.py index b06ae5345e0..a9cb7f774ae 100644 --- a/invokeai/backend/hidiffusion/hidiffusion.py +++ b/invokeai/backend/hidiffusion/hidiffusion.py @@ -1978,6 +1978,34 @@ def hook(module, args): model.info["hooks"].append(model.register_forward_pre_hook(hook)) +def _reset_hidiffusion_runtime_state(module: torch.nn.Module) -> None: + """Clear runtime state left on a cached module by a previous HiDiffusion patch.""" + # The downsampler temporarily overrides these attributes while resizing. If its + # forward pass was interrupted, restore the originals before dropping the + # bookkeeping attributes. + 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 ( + "timestep", + "aggressive_raunet", + "T1_ratio", + "T1", + "T1_start", + "T1_end", + "max_timestep", + "ori_stride", + "ori_padding", + "ori_dilation", + ): + module.__dict__.pop(attribute, None) + + def apply_hidiffusion( model: torch.nn.Module, apply_raunet: bool = True, @@ -2062,18 +2090,22 @@ def apply_hidiffusion( 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" @@ -2090,21 +2122,25 @@ def apply_hidiffusion( 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" @@ -2119,11 +2155,13 @@ def apply_hidiffusion( 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" diff --git a/tests/backend/stable_diffusion/test_hidiffusion_utils.py b/tests/backend/stable_diffusion/test_hidiffusion_utils.py index 98d207b5e66..1c7efcb5069 100644 --- a/tests/backend/stable_diffusion/test_hidiffusion_utils.py +++ b/tests/backend/stable_diffusion/test_hidiffusion_utils.py @@ -69,6 +69,12 @@ def __init__(self): 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() @@ -118,6 +124,40 @@ def run_with_global_seed(global_seed: int) -> torch.Tensor: 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 + + # A normal generation while HiDiffusion is disabled reuses the cached + # module without exercising its dormant HiDiffusion attributes. + assert model.block.timestep == 7 + + 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_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) From 2cf2e0b48e21a1358271ca027619a34b206a724c Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Sun, 2 Aug 2026 15:55:24 +0300 Subject: [PATCH 21/22] Make HiDiffusion teardown transactional --- invokeai/backend/hidiffusion/hidiffusion.py | 102 ++++++++++++------ .../test_hidiffusion_utils.py | 48 ++++++++- 2 files changed, 112 insertions(+), 38 deletions(-) diff --git a/invokeai/backend/hidiffusion/hidiffusion.py b/invokeai/backend/hidiffusion/hidiffusion.py index a9cb7f774ae..cbf156f5ade 100644 --- a/invokeai/backend/hidiffusion/hidiffusion.py +++ b/invokeai/backend/hidiffusion/hidiffusion.py @@ -1852,43 +1852,34 @@ def forward(self, hidden_states: torch.Tensor, scale=1.0) -> torch.Tensor: 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: - self.ori_stride = self.stride - self.ori_padding = self.padding - self.ori_dilation = self.dilation - - self.stride = (4, 4) - self.padding = (2, 2) - self.dilation = (2, 2) + 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, self.stride, self.padding, self.dilation, self.groups + hidden_states, self.weight, self.bias, stride, padding, dilation, self.groups ) - if self.timestep < self.T1: - self.stride = self.ori_stride - self.padding = self.ori_padding - self.dilation = self.ori_dilation 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, self.stride, self.padding, self.dilation, self.groups + 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, self.stride, self.padding, self.dilation, self.groups + hidden_states, self.weight, self.bias, stride, padding, dilation, self.groups ) - if self.timestep < self.T1: - self.stride = self.ori_stride - self.padding = self.ori_padding - self.dilation = self.ori_dilation self.timestep += 1 if self.timestep == self.max_timestep: self.timestep = 0 @@ -1978,11 +1969,42 @@ def hook(module, args): 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.""" - # The downsampler temporarily overrides these attributes while resizing. If its - # forward pass was interrupted, restore the originals before dropping the - # bookkeeping attributes. for attribute, original_attribute in ( ("stride", "ori_stride"), ("padding", "ori_padding"), @@ -1991,21 +2013,24 @@ def _reset_hidiffusion_runtime_state(module: torch.nn.Module) -> None: if original_attribute in module.__dict__: setattr(module, attribute, module.__dict__[original_attribute]) - for attribute in ( - "timestep", - "aggressive_raunet", - "T1_ratio", - "T1", - "T1_start", - "T1_end", - "max_timestep", - "ori_stride", - "ori_padding", - "ori_dilation", - ): + 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, @@ -2061,6 +2086,9 @@ def apply_hidiffusion( 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 @@ -2190,7 +2218,13 @@ def remove_hidiffusion(model: torch.nn.Module): hook.remove() module.info["hooks"].clear() - if hasattr(module, "_parent"): + 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/tests/backend/stable_diffusion/test_hidiffusion_utils.py b/tests/backend/stable_diffusion/test_hidiffusion_utils.py index 1c7efcb5069..74c92ab9604 100644 --- a/tests/backend/stable_diffusion/test_hidiffusion_utils.py +++ b/tests/backend/stable_diffusion/test_hidiffusion_utils.py @@ -87,7 +87,7 @@ def test_hidiffusion_patch_supports_bare_model_mixin_without_public_name_or_path assert model.num_upsamplers == 3 assert not hasattr(model, "_name_or_path") - assert model.info["hooks"] == [] + assert not hasattr(model, "info") def test_hidiffusion_window_attention_uses_seeded_generator_instead_of_global_rng(): @@ -144,9 +144,7 @@ def test_hidiffusion_patch_resets_cached_runtime_state_when_reenabled(): model.block.T1_end = 8 model.block.max_timestep = 99 - # A normal generation while HiDiffusion is disabled reuses the cached - # module without exercising its dormant HiDiffusion attributes. - assert model.block.timestep == 7 + assert "timestep" not in model.block.__dict__ with hidiffusion_patch(model, name_or_path="runwayml/stable-diffusion-v1-5"): assert model.block.timestep == 0 @@ -158,6 +156,48 @@ def test_hidiffusion_patch_resets_cached_runtime_state_when_reenabled(): 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) From 013a17b50af91a0249e3b705a67411da97ff91d9 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Sun, 2 Aug 2026 15:58:24 +0300 Subject: [PATCH 22/22] Format HiDiffusion teardown with Ruff --- invokeai/backend/hidiffusion/hidiffusion.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/invokeai/backend/hidiffusion/hidiffusion.py b/invokeai/backend/hidiffusion/hidiffusion.py index cbf156f5ade..5d67f2554ec 100644 --- a/invokeai/backend/hidiffusion/hidiffusion.py +++ b/invokeai/backend/hidiffusion/hidiffusion.py @@ -1877,9 +1877,7 @@ def forward(self, hidden_states: torch.Tensor, scale=1.0) -> torch.Tensor: ) 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 - ) + 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