invoke-ai/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI.git synced 2026-10-12 02:16:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Feature: HiDiffusion integration (#8787)

Browse Source 
* Initial implementation of HiDiffusion pipeline

* Added URLs to original HiDiffusion repository in documentation and frontend popover.

* Added comment next to HiDiffusion in pyproject.toml to clarify its purpose.

* Formatting

* Import order fix

* Add HiDiffusion T1/T2 ratio controls and docs updates

* Checks appeasing

* Refactor HiDiffusion import and update T2 ratio default value to 0

* Changed to vendoring an updated version of HiDiffusion and removed deps for external dependency

* Ruff

* Ruff again

* chore(api): refresh hidiffusion openapi schema

* Fix HiDiffusion cleanup and metadata recall

* chore: fix lint import ordering

* Fix HiDiffusion cleanup and metadata recall

* Move HiDiffusion license notice to repository root

* Fix HiDiffusion modular denoise and docs formatting

* Bump HiDiffusion node versions and refresh schema

* Seed HiDiffusion window attention deterministically

* fix: stale HiDiffusion state in cached UNet

* Make HiDiffusion teardown transactional

* Format HiDiffusion teardown with Ruff

---------

Co-authored-by: Jonathan <[email protected]>
This commit is contained in:
Valeri Che
2026-08-03 10:25:51 -04:00
committed by GitHub
co-authored by Jonathan
parent 9f2b2be350
commit 2f470d36f6
30 changed files with 6515 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
LICENSE-HiDiffusion.txt
+65
View File
@@ -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.
docs/src/content/docs/features/hidiffusion.mdx
+39
View File
@@ -0,0 +1,39 @@
---
title: HiDiffusion
---
# HiDiffusion
HiDiffusion is an optional denoising enhancement that can improve detail and structure at higher resolutions for SD 1.5, SD 2.1, and SDXL. It modifies the UNet during denoising and is most noticeable at 1536px and above.
The vendored implementation includes module-key definitions for SD 1.5 and SDXL. SD 2.1 intentionally reuses the SD 1.5 module-key mapping.
Learn more: https://github.com/megvii-research/HiDiffusion
## Where to find the switches
1. Open the **Canvas** tab.
2. Expand **Advanced Settings**.
3. In the **Advanced** grid, enable **HiDiffusion** and optionally adjust the two sub‑toggles and ratios:
- **HiDiffusion: RAU‑Net**
- **HiDiffusion: Window Attention**
- **HiDiffusion: T1 Ratio**
- **HiDiffusion: T2 Ratio**
## What the switches do
- **HiDiffusion**: Enables the HiDiffusion patch for denoising. Use this for high‑resolution generations; the effect is subtle at lower sizes.
- **HiDiffusion: RAU‑Net**: Enables RAU‑Net blocks. This typically improves structure and mid‑frequency detail, especially at larger resolutions.
- **HiDiffusion: Window Attention**: Enables windowed attention blocks. This can boost local texture/detail, but may slightly affect global coherence in some prompts.
- **HiDiffusion: T1 Ratio**: Controls when HiDiffusion switches into its mid‑stage behavior. Lower values switch earlier; higher values preserve global structure longer.
- **HiDiffusion: T2 Ratio**: Controls when HiDiffusion switches into its late‑stage behavior. Higher values keep window attention active longer and can sharpen local detail.
## Tips
- Try **1536–2048 px** for the clearest benefits (SDXL).
- If results look worse, disable **Window Attention** first, then RAU‑Net.
- Effects vary by scheduler and model; compare with the same seed for a fair test.
invokeai/app/invocations/denoise_latents.py
+94 -25
View File
@@ -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.diffusion_backend import StableDiffusionB
from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType
from invokeai.backend.stable_diffusion.extensions.controlnet import ControlNetExt
from invokeai.backend.stable_diffusion.extensions.freeu import FreeUExt
from invokeai.backend.stable_diffusion.extensions.hidiffusion import HiDiffusionExt
from invokeai.backend.stable_diffusion.extensions.inpaint import InpaintExt
from invokeai.backend.stable_diffusion.extensions.inpaint_model import InpaintModelExt
from invokeai.backend.stable_diffusion.extensions.lora import LoRAExt
@@ -73,6 +74,7 @@ from invokeai.backend.stable_diffusion.extensions.rescale_cfg import RescaleCFGE
from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt
from invokeai.backend.stable_diffusion.extensions.t2i_adapter import T2IAdapterExt
from invokeai.backend.stable_diffusion.extensions_manager import ExtensionsManager
from invokeai.backend.stable_diffusion.hidiffusion_utils import hidiffusion_patch
from invokeai.backend.stable_diffusion.schedulers import SCHEDULER_MAP
from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES
from invokeai.backend.util.devices import TorchDevice
@@ -130,7 +132,7 @@ def get_scheduler(
title="Denoise - SD1.5, SDXL",
tags=["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"],
category="latents",
version="1.5.4",
version="1.6.0",
)
class DenoiseLatentsInvocation(BaseInvocation):
"""Denoises noisy latents to decodable images"""
@@ -191,6 +193,35 @@ class DenoiseLatentsInvocation(BaseInvocation):
cfg_rescale_multiplier: float = InputField(
title="CFG Rescale Multiplier", default=0, ge=0, lt=1, description=FieldDescriptions.cfg_rescale_multiplier
)
hidiffusion: bool = InputField(
default=False,
description=FieldDescriptions.hidiffusion,
title="HiDiffusion",
)
hidiffusion_raunet: bool = InputField(
default=True,
description=FieldDescriptions.hidiffusion_raunet,
title="HiDiffusion: RAU-Net",
)
hidiffusion_window_attn: bool = InputField(
default=True,
description=FieldDescriptions.hidiffusion_window_attn,
title="HiDiffusion: Window Attention",
)
hidiffusion_t1_ratio: float = InputField(
default=0.4,
ge=0,
le=1,
description=FieldDescriptions.hidiffusion_t1_ratio,
title="HiDiffusion: T1 Ratio",
)
hidiffusion_t2_ratio: float = InputField(
default=0.0,
ge=0,
le=1,
description=FieldDescriptions.hidiffusion_t2_ratio,
title="HiDiffusion: T2 Ratio",
)
latents: Optional[LatentsField] = InputField(
default=None,
description=FieldDescriptions.latents,
@@ -486,6 +517,14 @@ class DenoiseLatentsInvocation(BaseInvocation):
return controlnet_data
@staticmethod
def _get_hidiffusion_name_or_path(unet_config: AnyModelConfig) -> Optional[str]:
return (
getattr(unet_config, "source", None)
or getattr(unet_config, "path", None)
or getattr(unet_config, "name", None)
)
@staticmethod
def parse_controlnet_field(
exit_stack: ExitStack,
@@ -837,6 +876,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
# get the unet's config so that we can pass the base to sd_step_callback()
unet_config = context.models.get_config(self.unet.unet.key)
hidiffusion_name_or_path = self._get_hidiffusion_name_or_path(unet_config)
conditioning_data = self.get_conditioning_data(
context=context,
@@ -874,6 +914,17 @@ class DenoiseLatentsInvocation(BaseInvocation):
context.util.sd_step_callback(state, unet_config.base)
ext_manager.add_extension(PreviewExt(step_callback))
if self.hidiffusion:
ext_manager.add_extension(
HiDiffusionExt(
name_or_path=hidiffusion_name_or_path,
apply_raunet=self.hidiffusion_raunet,
apply_window_attn=self.hidiffusion_window_attn,
t1_ratio=self.hidiffusion_t1_ratio,
t2_ratio=self.hidiffusion_t2_ratio,
generator=torch.Generator(device="cpu").manual_seed(seed),
)
)
### cfg rescale
if self.cfg_rescale_multiplier > 0:
@@ -940,14 +991,17 @@ class DenoiseLatentsInvocation(BaseInvocation):
# ext: t2i/ip adapter
ext_manager.run_callback(ExtensionCallbackType.SETUP, denoise_ctx)
with (
context.models.load(self.unet.unet).model_on_device() as (cached_weights, unet),
ModelPatcher.patch_unet_attention_processor(unet, denoise_ctx.inputs.attention_processor_cls),
with ExitStack() as unet_stack:
cached_weights, unet = unet_stack.enter_context(context.models.load(self.unet.unet).model_on_device())
unet._num_timesteps = timesteps.shape[0]
unet_stack.enter_context(
ModelPatcher.patch_unet_attention_processor(unet, denoise_ctx.inputs.attention_processor_cls)
)
# ext: controlnet
ext_manager.patch_extensions(denoise_ctx),
# ext: freeu, seamless, ip adapter, lora
ext_manager.patch_unet(unet, cached_weights),
):
unet_stack.enter_context(ext_manager.patch_extensions(denoise_ctx))
# ext: freeu, seamless, ip adapter, lora, hidiffusion
unet_stack.enter_context(ext_manager.patch_unet(unet, cached_weights))
sd_backend = StableDiffusionBackend(unet, scheduler)
denoise_ctx.unet = unet
result_latents = sd_backend.latents_from_embeddings(denoise_ctx, ext_manager)
@@ -997,6 +1051,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
# get the unet's config so that we can pass the base to sd_step_callback()
unet_config = context.models.get_config(self.unet.unet.key)
hidiffusion_name_or_path = self._get_hidiffusion_name_or_path(unet_config)
def step_callback(state: PipelineIntermediateState) -> None:
context.util.sd_step_callback(state, unet_config.base)
@@ -1083,23 +1138,37 @@ class DenoiseLatentsInvocation(BaseInvocation):
denoising_end=self.denoising_end,
seed=seed,
)
pipeline._num_timesteps = timesteps.shape[0]
result_latents = pipeline.latents_from_embeddings(
latents=latents,
timesteps=timesteps,
init_timestep=init_timestep,
noise=noise,
seed=seed,
mask=mask,
masked_latents=masked_latents,
is_gradient_mask=gradient_mask,
scheduler_step_kwargs=scheduler_step_kwargs,
conditioning_data=conditioning_data,
control_data=controlnet_data,
ip_adapter_data=ip_adapter_data,
t2i_adapter_data=t2i_adapter_data,
callback=step_callback,
)
with (
hidiffusion_patch(
pipeline,
name_or_path=hidiffusion_name_or_path,
apply_raunet=self.hidiffusion_raunet,
apply_window_attn=self.hidiffusion_window_attn,
t1_ratio=self.hidiffusion_t1_ratio,
t2_ratio=self.hidiffusion_t2_ratio,
generator=torch.Generator(device="cpu").manual_seed(seed),
)
if self.hidiffusion
else nullcontext()
):
result_latents = pipeline.latents_from_embeddings(
latents=latents,
timesteps=timesteps,
init_timestep=init_timestep,
noise=noise,
seed=seed,
mask=mask,
masked_latents=masked_latents,
is_gradient_mask=gradient_mask,
scheduler_step_kwargs=scheduler_step_kwargs,
conditioning_data=conditioning_data,
control_data=controlnet_data,
ip_adapter_data=ip_adapter_data,
t2i_adapter_data=t2i_adapter_data,
callback=step_callback,
)
# https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
result_latents = result_latents.to("cpu")
invokeai/app/invocations/fields.py
+5
View File
@@ -149,6 +149,11 @@ class FieldDescriptions:
denoising_end = "When to stop denoising, expressed a percentage of total steps"
cfg_scale = "Classifier-Free Guidance scale"
cfg_rescale_multiplier = "Rescale multiplier for CFG guidance, used for models trained with zero-terminal SNR"
hidiffusion = "Apply HiDiffusion (RAU-Net + MSW-MSA) for higher-resolution denoising"
hidiffusion_raunet = "Apply HiDiffusion RAU-Net blocks"
hidiffusion_window_attn = "Apply HiDiffusion window attention blocks"
hidiffusion_t1_ratio = "Override HiDiffusion early switch threshold (T1 ratio)"
hidiffusion_t2_ratio = "Override HiDiffusion late switch threshold (T2 ratio)"
scheduler = "Scheduler to use during inference"
positive_cond = "Positive conditioning tensor"
negative_cond = "Negative conditioning tensor"
invokeai/app/invocations/metadata_linked.py
+6 -1
View File
@@ -622,7 +622,7 @@ class LatentsMetaOutput(LatentsOutput, MetadataOutput):
title=f"{DenoiseLatentsInvocation.UIConfig.title} + Metadata",
tags=["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"],
category="metadata",
version="1.1.1",
version="1.2.0",
)
class DenoiseLatentsMetaInvocation(DenoiseLatentsInvocation, WithMetadata):
def invoke(self, context: InvocationContext) -> LatentsMetaOutput:
@@ -661,6 +661,11 @@ class DenoiseLatentsMetaInvocation(DenoiseLatentsInvocation, WithMetadata):
md.update({"denoising_end": self.denoising_end})
md.update({"scheduler": self.scheduler})
md.update({"model": self.unet.unet})
md.update({"hidiffusion": self.hidiffusion})
md.update({"hidiffusion_raunet": self.hidiffusion_raunet})
md.update({"hidiffusion_window_attn": self.hidiffusion_window_attn})
md.update({"hidiffusion_t1_ratio": self.hidiffusion_t1_ratio})
md.update({"hidiffusion_t2_ratio": self.hidiffusion_t2_ratio})
if isinstance(self.control, ControlField) or (isinstance(self.control, list) and len(self.control) > 0):
md.update({"controlnets": _to_json(self.control)})
if isinstance(self.ip_adapter, IPAdapterField) or (
invokeai/backend/hidiffusion/__init__.py
+3
View File
@@ -0,0 +1,3 @@
from invokeai.backend.hidiffusion.hidiffusion import apply_hidiffusion, remove_hidiffusion
__all__ = ["apply_hidiffusion", "remove_hidiffusion"]
invokeai/backend/hidiffusion/hidiffusion.py
+2228
View File
File diff suppressed because it is too large Load Diff
invokeai/backend/hidiffusion/sd_module_key/sd15_module_key.txt
+685
View File
@@ -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
invokeai/backend/hidiffusion/sd_module_key/sdxl_module_key.txt
+1910
View File
File diff suppressed because it is too large Load Diff
invokeai/backend/hidiffusion/utils.py
+35
View File
@@ -0,0 +1,35 @@
import torch
def isinstance_str(x: object, cls_name: str, prefix: bool = False, contains: bool = False):
"""
Checks whether x has any class equal to, prefixed with, or contains (cls_name) in its ancestry.
Doesn't require access to the class's implementation.
Useful for patching!
"""
for _cls in x.__class__.__mro__:
if _cls.__name__.startswith(cls_name) and prefix:
return True
if contains and cls_name in _cls.__name__:
return True
if _cls.__name__ == cls_name:
return True
return False
def init_generator(device: torch.device, fallback: torch.Generator | None = None):
"""
Forks the current default random generator given device.
"""
if device.type == "cpu":
return torch.Generator(device="cpu").set_state(torch.get_rng_state())
elif device.type == "cuda":
return torch.Generator(device=device).set_state(torch.cuda.get_rng_state())
else:
if fallback is None:
return init_generator(torch.device("cpu"))
else:
return fallback
invokeai/backend/stable_diffusion/extensions/hidiffusion.py
+43
View File
@@ -0,0 +1,43 @@
from __future__ import annotations
from contextlib import contextmanager
from typing import Optional
import torch
from diffusers import UNet2DConditionModel
from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase
from invokeai.backend.stable_diffusion.hidiffusion_utils import hidiffusion_patch
from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage
class HiDiffusionExt(ExtensionBase):
def __init__(
self,
name_or_path: Optional[str],
apply_raunet: bool = True,
apply_window_attn: bool = True,
t1_ratio: Optional[float] = None,
t2_ratio: Optional[float] = None,
generator: torch.Generator | None = None,
):
super().__init__()
self._name_or_path = name_or_path
self._apply_raunet = apply_raunet
self._apply_window_attn = apply_window_attn
self._t1_ratio = t1_ratio
self._t2_ratio = t2_ratio
self._generator = generator
@contextmanager
def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeightsStorage):
with hidiffusion_patch(
unet,
name_or_path=self._name_or_path,
apply_raunet=self._apply_raunet,
apply_window_attn=self._apply_window_attn,
t1_ratio=self._t1_ratio,
t2_ratio=self._t2_ratio,
generator=self._generator,
):
yield None
invokeai/backend/stable_diffusion/hidiffusion_utils.py
+165
View File
@@ -0,0 +1,165 @@
from __future__ import annotations
import copy
import sys
from contextlib import contextmanager
from typing import Any, Optional
import torch
from invokeai.backend.hidiffusion.hidiffusion import (
switching_threshold_ratio_dict as _switching_threshold_ratio_dict,
)
from invokeai.backend.hidiffusion.hidiffusion import (
text_to_img_controlnet_switching_threshold_ratio_dict as _text_to_img_controlnet_switching_threshold_ratio_dict,
)
@contextmanager
def hidiffusion_patch(
model: Any,
name_or_path: Optional[str],
apply_raunet: bool = True,
apply_window_attn: bool = True,
t1_ratio: Optional[float] = None,
t2_ratio: Optional[float] = None,
generator: torch.Generator | None = None,
):
"""Context manager that applies HiDiffusion and restores the model on exit."""
from invokeai.backend.hidiffusion.hidiffusion import apply_hidiffusion, remove_hidiffusion
target = model.unet if hasattr(model, "unet") else model
had_model_name_or_path = hasattr(model, "_name_or_path")
had_config = hasattr(model, "config")
config = model.config if had_config else None
had_config_name_or_path = bool(config) and hasattr(config, "_name_or_path")
original_model_name_or_path = model._name_or_path if had_model_name_or_path else None
original_config_name_or_path = config._name_or_path if had_config_name_or_path else None
effective_name_or_path = (
name_or_path
or getattr(model, "name_or_path", None)
or original_model_name_or_path
or original_config_name_or_path
or ""
)
def _set_name_or_path_on_config(cfg, value: str) -> bool:
if cfg is None:
return False
if hasattr(cfg, "_internal_dict"):
try:
cfg._internal_dict["_name_or_path"] = value
return True
except Exception:
pass
try:
object.__setattr__(cfg, "_name_or_path", value)
return True
except Exception:
pass
try:
cfg.__dict__["_name_or_path"] = value
return True
except Exception:
return False
original_num_upsamplers = getattr(target, "num_upsamplers", None)
ratio_overrides = None
ratio_dicts = None
if t1_ratio is not None or t2_ratio is not None:
ratio_dicts = (
_switching_threshold_ratio_dict,
_text_to_img_controlnet_switching_threshold_ratio_dict,
)
set_model_name_or_path = False
set_config_name_or_path = False
try:
try:
object.__setattr__(model, "_name_or_path", effective_name_or_path)
set_model_name_or_path = True
except Exception:
set_model_name_or_path = False
set_config_name_or_path = _set_name_or_path_on_config(config, effective_name_or_path)
# Ensure the property resolves to a non-None value before calling HiDiffusion.
try:
if getattr(model, "name_or_path", None) is None:
if not set_model_name_or_path:
try:
object.__setattr__(model, "_name_or_path", effective_name_or_path)
set_model_name_or_path = True
except Exception:
pass
if not set_config_name_or_path:
set_config_name_or_path = _set_name_or_path_on_config(config, effective_name_or_path)
except Exception:
pass
if ratio_dicts is not None:
ratio_overrides = (copy.deepcopy(ratio_dicts[0]), copy.deepcopy(ratio_dicts[1]))
def _apply_ratio_overrides(ratio_dict: dict) -> None:
for _, entry in ratio_dict.items():
if t1_ratio is not None:
entry["T1_ratio"] = t1_ratio
if t2_ratio is not None and "T2_ratio" in entry:
entry["T2_ratio"] = t2_ratio
_apply_ratio_overrides(ratio_dicts[0])
_apply_ratio_overrides(ratio_dicts[1])
apply_hidiffusion(
model,
apply_raunet=apply_raunet,
apply_window_attn=apply_window_attn,
generator=generator,
)
yield
finally:
had_active_exception = sys.exc_info()[0] is not None
teardown_error: Exception | None = None
try:
remove_hidiffusion(model)
except Exception as error:
if not had_active_exception:
teardown_error = error
if ratio_overrides is not None and ratio_dicts is not None:
ratio_dicts[0].clear()
ratio_dicts[0].update(ratio_overrides[0])
ratio_dicts[1].clear()
ratio_dicts[1].update(ratio_overrides[1])
if original_num_upsamplers is not None:
target.num_upsamplers = original_num_upsamplers
if set_model_name_or_path:
if had_model_name_or_path:
try:
object.__setattr__(model, "_name_or_path", original_model_name_or_path)
except Exception:
pass
else:
try:
delattr(model, "_name_or_path")
except AttributeError:
pass
if set_config_name_or_path and had_config:
if had_config_name_or_path:
_set_name_or_path_on_config(config, original_config_name_or_path)
elif config is not None:
internal_dict = getattr(config, "_internal_dict", None)
if internal_dict is not None:
try:
internal_dict.pop("_name_or_path", None)
except Exception:
pass
try:
delattr(config, "_name_or_path")
except AttributeError:
pass
if teardown_error is not None:
raise teardown_error
invokeai/frontend/web/openapi.json
+110 -2
View File
@@ -24036,6 +24036,60 @@
"title": "CFG Rescale Multiplier",
"type": "number"
},
"hidiffusion": {
"default": false,
"description": "Apply HiDiffusion (RAU-Net + MSW-MSA) for higher-resolution denoising",
"field_kind": "input",
"input": "any",
"orig_default": false,
"orig_required": false,
"title": "HiDiffusion",
"type": "boolean"
},
"hidiffusion_raunet": {
"default": true,
"description": "Apply HiDiffusion RAU-Net blocks",
"field_kind": "input",
"input": "any",
"orig_default": true,
"orig_required": false,
"title": "HiDiffusion: RAU-Net",
"type": "boolean"
},
"hidiffusion_window_attn": {
"default": true,
"description": "Apply HiDiffusion window attention blocks",
"field_kind": "input",
"input": "any",
"orig_default": true,
"orig_required": false,
"title": "HiDiffusion: Window Attention",
"type": "boolean"
},
"hidiffusion_t1_ratio": {
"default": 0.4,
"description": "Override HiDiffusion early switch threshold (T1 ratio)",
"field_kind": "input",
"input": "any",
"maximum": 1,
"minimum": 0,
"orig_default": 0.4,
"orig_required": false,
"title": "HiDiffusion: T1 Ratio",
"type": "number"
},
"hidiffusion_t2_ratio": {
"default": 0.0,
"description": "Override HiDiffusion late switch threshold (T2 ratio)",
"field_kind": "input",
"input": "any",
"maximum": 1,
"minimum": 0,
"orig_default": 0.0,
"orig_required": false,
"title": "HiDiffusion: T2 Ratio",
"type": "number"
},
"latents": {
"anyOf": [
{
@@ -24082,7 +24136,7 @@
"tags": ["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"],
"title": "Denoise - SD1.5, SDXL",
"type": "object",
"version": "1.5.4",
"version": "1.6.0",
"output": {
"$ref": "#/components/schemas/LatentsOutput"
}
@@ -24395,6 +24449,60 @@
"title": "CFG Rescale Multiplier",
"type": "number"
},
"hidiffusion": {
"default": false,
"description": "Apply HiDiffusion (RAU-Net + MSW-MSA) for higher-resolution denoising",
"field_kind": "input",
"input": "any",
"orig_default": false,
"orig_required": false,
"title": "HiDiffusion",
"type": "boolean"
},
"hidiffusion_raunet": {
"default": true,
"description": "Apply HiDiffusion RAU-Net blocks",
"field_kind": "input",
"input": "any",
"orig_default": true,
"orig_required": false,
"title": "HiDiffusion: RAU-Net",
"type": "boolean"
},
"hidiffusion_window_attn": {
"default": true,
"description": "Apply HiDiffusion window attention blocks",
"field_kind": "input",
"input": "any",
"orig_default": true,
"orig_required": false,
"title": "HiDiffusion: Window Attention",
"type": "boolean"
},
"hidiffusion_t1_ratio": {
"default": 0.4,
"description": "Override HiDiffusion early switch threshold (T1 ratio)",
"field_kind": "input",
"input": "any",
"maximum": 1,
"minimum": 0,
"orig_default": 0.4,
"orig_required": false,
"title": "HiDiffusion: T1 Ratio",
"type": "number"
},
"hidiffusion_t2_ratio": {
"default": 0.0,
"description": "Override HiDiffusion late switch threshold (T2 ratio)",
"field_kind": "input",
"input": "any",
"maximum": 1,
"minimum": 0,
"orig_default": 0.0,
"orig_required": false,
"title": "HiDiffusion: T2 Ratio",
"type": "number"
},
"latents": {
"anyOf": [
{
@@ -24441,7 +24549,7 @@
"tags": ["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"],
"title": "Denoise - SD1.5, SDXL + Metadata",
"type": "object",
"version": "1.1.1",
"version": "1.2.0",
"output": {
"$ref": "#/components/schemas/LatentsMetaOutput"
}
invokeai/frontend/web/public/locales/en.json
+45
View File
@@ -1068,6 +1068,11 @@
"seedreamOptimizePrompt": "Seedream Optimize Prompt",
"guidance": "Guidance",
"height": "Height",
"hiDiffusion": "HiDiffusion",
"hiDiffusionRauNet": "HiDiffusion: RAU-Net",
"hiDiffusionWindowAttn": "HiDiffusion: Window Attention",
"hiDiffusionT1Ratio": "HiDiffusion: T1 Ratio",
"hiDiffusionT2Ratio": "HiDiffusion: T2 Ratio",
"imageDetails": "Image Details",
"imageDimensions": "Image Dimensions",
"imageSize": "Image Size",
@@ -1717,6 +1722,11 @@
"cfgScale": "CFG Scale",
"cfgRescaleMultiplier": "CFG Rescale Multiplier",
"clipSkip": "CLIP Skip",
"hiDiffusion": "HiDiffusion",
"hiDiffusionRauNet": "HiDiffusion: RAU-Net",
"hiDiffusionWindowAttn": "HiDiffusion: Window Attention",
"hiDiffusionT1Ratio": "HiDiffusion: T1 Ratio",
"hiDiffusionT2Ratio": "HiDiffusion: T2 Ratio",
"coherenceMode": "Mode",
"coherenceEdgeSize": "Edge Size",
"coherenceMinDenoise": "Min Denoise",
@@ -2131,6 +2141,41 @@
"invalidBboxDesc": "The bounding box has no valid dimensions"
},
"popovers": {
"hidiffusion": {
"heading": "HiDiffusion",
"paragraphs": [
"Enables HiDiffusion denoising for higher-resolution generation.",
"Try at 1536px or higher for the most visible effect."
]
},
"hidiffusionRauNet": {
"heading": "HiDiffusion: RAU-Net",
"paragraphs": [
"Enables RAU-Net blocks used by HiDiffusion.",
"Can improve structure and detail at higher resolutions."
]
},
"hidiffusionWindowAttn": {
"heading": "HiDiffusion: Window Attention",
"paragraphs": [
"Enables windowed attention used by HiDiffusion.",
"Can boost local detail, but may affect global coherence."
]
},
"hidiffusionT1Ratio": {
"heading": "HiDiffusion: T1 Ratio",
"paragraphs": [
"Controls the early switch point for HiDiffusion (T1).",
"Lower values switch earlier; higher values preserve global structure longer."
]
},
"hidiffusionT2Ratio": {
"heading": "HiDiffusion: T2 Ratio",
"paragraphs": [
"Controls the late switch point for HiDiffusion (T2).",
"Higher values keep window attention active longer."
]
},
"clipSkip": {
"heading": "CLIP Skip",
"paragraphs": [
invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts
+12
View File
@@ -75,6 +75,11 @@ export type Feature =
| 'scaleBeforeProcessing'
| 'seamlessTilingXAxis'
| 'seamlessTilingYAxis'
| 'hidiffusion'
| 'hidiffusionRauNet'
| 'hidiffusionWindowAttn'
| 'hidiffusionT1Ratio'
| 'hidiffusionT2Ratio'
| 'colorCompensation'
| 'upscaleModel'
| 'scale'
@@ -246,6 +251,13 @@ export const POPOVER_DATA: { [key in Feature]?: PopoverData } = {
seamlessTilingYAxis: {
href: 'https://support.invoke.ai/support/solutions/articles/151000178161-advanced-settings',
},
hidiffusion: {
href: 'https://github.com/megvii-research/HiDiffusion',
},
hidiffusionRauNet: {},
hidiffusionWindowAttn: {},
hidiffusionT1Ratio: {},
hidiffusionT2Ratio: {},
} as const;
export const OPEN_DELAY = 1000; // in milliseconds
invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.test.ts
+12 -2
View File
@@ -144,10 +144,10 @@ describe('paramsSlice selectors for external models', () => {
describe('paramsSliceConfig persisted state migration', () => {
const migrate = paramsSliceConfig.persistConfig?.migrate;
it('backfills new Qwen Image fields when migrating from v2 and preserves existing params', () => {
it('backfills Qwen Image and HiDiffusion fields when migrating from v2 and preserves existing params', () => {
expect(migrate).toBeDefined();
// Build a valid pre-PR v2 persisted state by removing the fields that were added in v3
// Build a valid pre-PR v2 persisted state by removing the fields that were added later.
const initial = getInitialParamsState();
const v2State: Record<string, unknown> = {
...initial,
@@ -159,6 +159,11 @@ describe('paramsSliceConfig persisted state migration', () => {
};
delete v2State.qwenImageVaeModel;
delete v2State.qwenImageQwenVLEncoderModel;
delete v2State.hiDiffusionEnabled;
delete v2State.hiDiffusionRauNetEnabled;
delete v2State.hiDiffusionWindowAttnEnabled;
delete v2State.hiDiffusionT1Ratio;
delete v2State.hiDiffusionT2Ratio;
const result = migrate?.(v2State) as ReturnType<typeof getInitialParamsState>;
@@ -167,6 +172,11 @@ describe('paramsSliceConfig persisted state migration', () => {
expect(result._version).toBe(4);
expect(result.qwenImageVaeModel).toBeNull();
expect(result.qwenImageQwenVLEncoderModel).toBeNull();
expect(result.hiDiffusionEnabled).toBe(false);
expect(result.hiDiffusionRauNetEnabled).toBe(true);
expect(result.hiDiffusionWindowAttnEnabled).toBe(true);
expect(result.hiDiffusionT1Ratio).toBe(0.4);
expect(result.hiDiffusionT2Ratio).toBe(0.0);
// Existing params should be preserved
expect(result.positivePrompt).toBe('a fluffy cat');
expect(result.seed).toBe(42);
invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts
+41
View File
@@ -170,6 +170,21 @@ const slice = createSlice({
setOptimizedDenoisingEnabled: (state, action: PayloadAction<boolean>) => {
state.optimizedDenoisingEnabled = action.payload;
},
setHiDiffusionEnabled: (state, action: PayloadAction<boolean>) => {
state.hiDiffusionEnabled = action.payload;
},
setHiDiffusionRauNetEnabled: (state, action: PayloadAction<boolean>) => {
state.hiDiffusionRauNetEnabled = action.payload;
},
setHiDiffusionWindowAttnEnabled: (state, action: PayloadAction<boolean>) => {
state.hiDiffusionWindowAttnEnabled = action.payload;
},
setHiDiffusionT1Ratio: (state, action: PayloadAction<number>) => {
state.hiDiffusionT1Ratio = action.payload;
},
setHiDiffusionT2Ratio: (state, action: PayloadAction<number>) => {
state.hiDiffusionT2Ratio = action.payload;
},
setSeamlessXAxis: (state, action: PayloadAction<boolean>) => {
state.seamlessXAxis = action.payload;
},
@@ -863,6 +878,11 @@ export const {
setSeed,
setImg2imgStrength,
setOptimizedDenoisingEnabled,
setHiDiffusionEnabled,
setHiDiffusionRauNetEnabled,
setHiDiffusionWindowAttnEnabled,
setHiDiffusionT1Ratio,
setHiDiffusionT2Ratio,
setSeamlessXAxis,
setSeamlessYAxis,
setShouldRandomizeSeed,
@@ -985,6 +1005,22 @@ export const paramsSliceConfig: SliceConfig<typeof slice> = {
state.pidSteps = 4;
}
if (!('hiDiffusionEnabled' in state)) {
state.hiDiffusionEnabled = false;
}
if (!('hiDiffusionRauNetEnabled' in state)) {
state.hiDiffusionRauNetEnabled = true;
}
if (!('hiDiffusionWindowAttnEnabled' in state)) {
state.hiDiffusionWindowAttnEnabled = true;
}
if (!('hiDiffusionT1Ratio' in state)) {
state.hiDiffusionT1Ratio = 0.4;
}
if (!('hiDiffusionT2Ratio' in state)) {
state.hiDiffusionT2Ratio = 0.0;
}
return zParamsState.parse(state);
},
},
@@ -1069,6 +1105,11 @@ export const selectInfillPatchmatchDownscaleSize = createParamsSelector(
export const selectInfillColorValue = createParamsSelector((params) => params.infillColorValue);
export const selectImg2imgStrength = createParamsSelector((params) => params.img2imgStrength);
export const selectOptimizedDenoisingEnabled = createParamsSelector((params) => params.optimizedDenoisingEnabled);
export const selectHiDiffusionEnabled = createParamsSelector((params) => params.hiDiffusionEnabled);
export const selectHiDiffusionRauNetEnabled = createParamsSelector((params) => params.hiDiffusionRauNetEnabled);
export const selectHiDiffusionWindowAttnEnabled = createParamsSelector((params) => params.hiDiffusionWindowAttnEnabled);
export const selectHiDiffusionT1Ratio = createParamsSelector((params) => params.hiDiffusionT1Ratio);
export const selectHiDiffusionT2Ratio = createParamsSelector((params) => params.hiDiffusionT2Ratio);
export const selectPositivePrompt = createParamsSelector((params) => params.positivePrompt);
export const selectNegativePrompt = createParamsSelector((params) => params.negativePrompt);
export const selectNegativePromptWithFallback = createParamsSelector((params) => params.negativePrompt ?? '');
invokeai/frontend/web/src/features/controlLayers/store/types.ts
+10
View File
@@ -832,6 +832,11 @@ export const zParamsState = z.object({
guidance: zParameterGuidance,
img2imgStrength: zParameterStrength,
optimizedDenoisingEnabled: z.boolean(),
hiDiffusionEnabled: z.boolean(),
hiDiffusionRauNetEnabled: z.boolean(),
hiDiffusionWindowAttnEnabled: z.boolean(),
hiDiffusionT1Ratio: z.number(),
hiDiffusionT2Ratio: z.number(),
iterations: z.number(),
scheduler: zParameterScheduler,
fluxScheduler: zParameterFluxScheduler,
@@ -961,6 +966,11 @@ export const getInitialParamsState = (): ParamsState => ({
guidance: 4,
img2imgStrength: 0.75,
optimizedDenoisingEnabled: true,
hiDiffusionEnabled: false,
hiDiffusionRauNetEnabled: true,
hiDiffusionWindowAttnEnabled: true,
hiDiffusionT1Ratio: 0.4,
hiDiffusionT2Ratio: 0.0,
iterations: 1,
scheduler: 'dpmpp_3m_k',
fluxScheduler: 'euler',
invokeai/frontend/web/src/features/metadata/parsing.test.ts
+36
View File
@@ -1,3 +1,5 @@
import type { AppStore } from 'app/store/store';
import { setHiDiffusionEnabled } from 'features/controlLayers/store/paramsSlice';
import { describe, expect, it, vi } from 'vitest';
import { ImageMetadataHandlers, MetadataUtils, parseMetadataHandler } from './parsing';
@@ -159,3 +161,37 @@ describe('Qwen metadata parsing', () => {
expect(mockStore.dispatch).toHaveBeenCalledTimes(1);
});
});
describe('HiDiffusion metadata parsing', () => {
it('disables HiDiffusion when recalling all metadata from an older image', async () => {
let hiDiffusionEnabled = true;
const store = {
dispatch: vi.fn((action) => {
if (action.type === setHiDiffusionEnabled.type) {
hiDiffusionEnabled = action.payload;
}
return action;
}),
getState: vi.fn(() => ({
params: { model: null },
})),
} as unknown as AppStore;
await MetadataUtils.recallAllImageMetadata(
{
generation_mode: 'txt2img',
width: 512,
height: 512,
steps: 20,
cfg_scale: 7.5,
scheduler: 'euler',
positive_prompt: 'an older image',
negative_prompt: '',
},
store
);
expect(store.dispatch).toHaveBeenCalledWith(setHiDiffusionEnabled(false));
expect(hiDiffusionEnabled).toBe(false);
});
});
invokeai/frontend/web/src/features/metadata/parsing.tsx
+100
View File
@@ -42,6 +42,11 @@ import {
setFluxDypeScale,
setFluxScheduler,
setGuidance,
setHiDiffusionEnabled,
setHiDiffusionRauNetEnabled,
setHiDiffusionT1Ratio,
setHiDiffusionT2Ratio,
setHiDiffusionWindowAttnEnabled,
setIdeogram4ColorPalette,
setIdeogram4GuidanceScale,
setIdeogram4Mu,
@@ -698,6 +703,96 @@ const SeamlessY: SingleMetadataHandler<ParameterSeamlessY> = {
};
//#endregion SeamlessY
//#region HiDiffusion
const HiDiffusion: SingleMetadataHandler<boolean> = {
[SingleMetadataKey]: true,
type: 'HiDiffusion',
parse: (metadata, _store) => {
const raw = getProperty(metadata, 'hidiffusion');
const parsed = raw === undefined ? false : z.boolean().parse(raw);
return Promise.resolve(parsed);
},
recall: (value, store) => {
store.dispatch(setHiDiffusionEnabled(value));
},
i18nKey: 'metadata.hiDiffusion',
LabelComponent: MetadataLabel,
ValueComponent: ({ value }: SingleMetadataValueProps<boolean>) => <MetadataPrimitiveValue value={value} />,
};
//#endregion HiDiffusion
//#region HiDiffusionRAUNet
const HiDiffusionRauNet: SingleMetadataHandler<boolean> = {
[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<boolean>) => <MetadataPrimitiveValue value={value} />,
};
//#endregion HiDiffusionRAUNet
//#region HiDiffusionWindowAttn
const HiDiffusionWindowAttn: SingleMetadataHandler<boolean> = {
[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<boolean>) => <MetadataPrimitiveValue value={value} />,
};
//#endregion HiDiffusionWindowAttn
//#region HiDiffusionT1Ratio
const HiDiffusionT1Ratio: SingleMetadataHandler<number> = {
[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<number>) => <MetadataPrimitiveValue value={value} />,
};
//#endregion HiDiffusionT1Ratio
//#region HiDiffusionT2Ratio
const HiDiffusionT2Ratio: SingleMetadataHandler<number> = {
[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<number>) => <MetadataPrimitiveValue value={value} />,
};
//#endregion HiDiffusionT2Ratio
//#region ZImageSeedVarianceEnabled
const ZImageSeedVarianceEnabled: SingleMetadataHandler<boolean> = {
[SingleMetadataKey]: true,
@@ -2153,6 +2248,11 @@ export const ImageMetadataHandlers = {
DenoisingStrength,
SeamlessX,
SeamlessY,
HiDiffusion,
HiDiffusionRauNet,
HiDiffusionWindowAttn,
HiDiffusionT1Ratio,
HiDiffusionT2Ratio,
RefinerModel,
RefinerSteps,
RefinerCFGScale,
invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.ts
+10
View File
@@ -23,6 +23,11 @@ export const addSDXLRefiner = async (
refinerScheduler,
refinerCFGScale,
refinerStart,
hiDiffusionEnabled,
hiDiffusionRauNetEnabled,
hiDiffusionT1Ratio,
hiDiffusionT2Ratio,
hiDiffusionWindowAttnEnabled,
} = state.params;
assert(refinerModel, 'No refiner model found in state');
@@ -57,6 +62,11 @@ export const addSDXLRefiner = async (
cfg_scale: refinerCFGScale,
steps: refinerSteps,
scheduler: refinerScheduler,
hidiffusion: hiDiffusionEnabled,
hidiffusion_raunet: hiDiffusionRauNetEnabled,
hidiffusion_window_attn: hiDiffusionWindowAttnEnabled,
hidiffusion_t1_ratio: hiDiffusionEnabled ? hiDiffusionT1Ratio : undefined,
hidiffusion_t2_ratio: hiDiffusionEnabled ? hiDiffusionT2Ratio : undefined,
denoising_start: refinerStart,
denoising_end: 1,
});
invokeai/frontend/web/src/features/nodes/util/graph/generation/buildHiDiffusionGraph.test.ts
+196
View File
@@ -0,0 +1,196 @@
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
vi.mock('app/logging/logger', () => ({
logger: () => ({
debug: vi.fn(),
}),
}));
let nextId = 0;
vi.mock('features/controlLayers/konva/util', () => ({
getPrefixedId: (prefix: string) => `${prefix}:${nextId++}`,
}));
const sd1Model = {
key: 'sd1-model',
hash: 'sd1-hash',
name: 'SD 1.5',
base: 'sd-1',
type: 'main',
};
const sdxlModel = {
key: 'sdxl-model',
hash: 'sdxl-hash',
name: 'SDXL',
base: 'sdxl',
type: 'main',
};
const defaultParams = {
cfgScale: 7.5,
cfgRescaleMultiplier: 0,
hiDiffusionEnabled: false,
hiDiffusionRauNetEnabled: false,
hiDiffusionT1Ratio: 0.25,
hiDiffusionT2Ratio: 0.1,
hiDiffusionWindowAttnEnabled: false,
scheduler: 'euler',
steps: 20,
clipSkip: 0,
shouldUseCpuNoise: false,
vaePrecision: 'fp16',
vae: null,
colorCompensation: false,
refinerModel: null,
pidMode: 'off',
};
let currentModel: typeof sd1Model | typeof sdxlModel = sd1Model;
let params = { ...defaultParams };
vi.mock('features/controlLayers/store/paramsSlice', () => ({
selectMainModelConfig: vi.fn(() => currentModel),
selectParamsSlice: vi.fn(() => params),
}));
vi.mock('features/controlLayers/store/refImagesSlice', () => ({
selectRefImagesSlice: vi.fn(() => ({ entities: [] })),
}));
vi.mock('features/controlLayers/store/selectors', () => ({
selectCanvasSlice: vi.fn(() => ({
bbox: { rect: { x: 0, y: 0, width: 1024, height: 1024 } },
controlLayers: { entities: [] },
regionalGuidance: { entities: [] },
})),
selectCanvasMetadata: vi.fn(() => ({})),
}));
vi.mock('features/nodes/util/graph/generation/addControlAdapters', () => ({
addControlNets: vi.fn(() => Promise.resolve({ addedControlNets: 0 })),
addT2IAdapters: vi.fn(() => Promise.resolve({ addedT2IAdapters: 0 })),
}));
vi.mock('features/nodes/util/graph/generation/addImageToImage', () => ({
addImageToImage: vi.fn(),
}));
vi.mock('features/nodes/util/graph/generation/addInpaint', () => ({
addInpaint: vi.fn(),
}));
vi.mock('features/nodes/util/graph/generation/addIPAdapters', () => ({
addIPAdapters: vi.fn(() => ({ addedIPAdapters: 0 })),
}));
vi.mock('features/nodes/util/graph/generation/addLoRAs', () => ({
addLoRAs: vi.fn(),
}));
vi.mock('features/nodes/util/graph/generation/addNSFWChecker', () => ({
addNSFWChecker: vi.fn((_g, node) => node),
}));
vi.mock('features/nodes/util/graph/generation/addOutpaint', () => ({
addOutpaint: vi.fn(),
}));
vi.mock('features/nodes/util/graph/generation/addSDXLLoRAs', () => ({
addSDXLLoRAs: vi.fn(),
}));
vi.mock('features/nodes/util/graph/generation/addSDXLRefiner', () => ({
addSDXLRefiner: vi.fn(),
}));
vi.mock('features/nodes/util/graph/generation/addSeamless', () => ({
addSeamless: vi.fn(() => null),
}));
vi.mock('features/nodes/util/graph/generation/addTextToImage', () => ({
addTextToImage: vi.fn(({ l2i }) => l2i),
}));
vi.mock('features/nodes/util/graph/generation/addWatermarker', () => ({
addWatermarker: vi.fn((_g, node) => node),
}));
vi.mock('features/nodes/util/graph/graphBuilderUtils', () => ({
getOriginalAndScaledSizesForTextToImage: vi.fn(() => ({
originalSize: { width: 1024, height: 1024 },
scaledSize: { width: 1024, height: 1024 },
aspectRatio: { id: 'Free', value: 1 },
})),
selectCanvasOutputFields: vi.fn(() => ({})),
selectPresetModifiedPrompts: vi.fn(() => ({
positive: 'a prompt',
negative: 'a negative prompt',
})),
}));
vi.mock('features/ui/store/uiSelectors', () => ({
selectActiveTab: vi.fn(() => 'generation'),
}));
vi.mock('./addRegions', () => ({
addRegions: vi.fn(() => Promise.resolve({ addedRegions: 0 })),
}));
import type { GraphBuilderArg } from 'features/nodes/util/graph/types';
import { buildSD1Graph } from './buildSD1Graph';
import { buildSDXLGraph } from './buildSDXLGraph';
import type { Graph } from './Graph';
const buildGraphArg = (): GraphBuilderArg =>
({
generationMode: 'txt2img',
manager: null,
state: {
system: {
shouldUseNSFWChecker: false,
shouldUseWatermarker: false,
},
},
}) as unknown as GraphBuilderArg;
const getMetadata = (g: Graph): Record<string, unknown> =>
(g as unknown as { getMetadataNode: () => Record<string, unknown> }).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);
});
});
invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts
+15
View File
@@ -42,6 +42,11 @@ export const buildSD1Graph = async (arg: GraphBuilderArg): Promise<GraphBuilderR
const {
cfgScale: cfg_scale,
cfgRescaleMultiplier: cfg_rescale_multiplier,
hiDiffusionEnabled,
hiDiffusionRauNetEnabled,
hiDiffusionT1Ratio,
hiDiffusionT2Ratio,
hiDiffusionWindowAttnEnabled,
scheduler,
steps,
clipSkip: skipped_layers,
@@ -100,6 +105,11 @@ export const buildSD1Graph = async (arg: GraphBuilderArg): Promise<GraphBuilderR
id: getPrefixedId('denoise_latents'),
cfg_scale,
cfg_rescale_multiplier,
hidiffusion: hiDiffusionEnabled,
hidiffusion_raunet: hiDiffusionRauNetEnabled,
hidiffusion_window_attn: hiDiffusionWindowAttnEnabled,
hidiffusion_t1_ratio: hiDiffusionEnabled ? hiDiffusionT1Ratio : undefined,
hidiffusion_t2_ratio: hiDiffusionEnabled ? hiDiffusionT2Ratio : undefined,
scheduler,
steps,
denoising_start: 0,
@@ -139,6 +149,11 @@ export const buildSD1Graph = async (arg: GraphBuilderArg): Promise<GraphBuilderR
g.upsertMetadata({
cfg_scale,
cfg_rescale_multiplier,
hidiffusion: hiDiffusionEnabled,
hidiffusion_raunet: hiDiffusionRauNetEnabled,
hidiffusion_window_attn: hiDiffusionWindowAttnEnabled,
hidiffusion_t1_ratio: hiDiffusionT1Ratio,
hidiffusion_t2_ratio: hiDiffusionT2Ratio,
model: Graph.getModelMetadataField(model),
steps,
rand_device: shouldUseCpuNoise ? 'cpu' : 'cuda',
invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSDXLGraph.ts
+15
View File
@@ -53,6 +53,11 @@ export const buildSDXLGraph = async (arg: GraphBuilderArg): Promise<GraphBuilder
const {
cfgScale: cfg_scale,
cfgRescaleMultiplier: cfg_rescale_multiplier,
hiDiffusionEnabled,
hiDiffusionRauNetEnabled,
hiDiffusionT1Ratio,
hiDiffusionT2Ratio,
hiDiffusionWindowAttnEnabled,
scheduler,
steps,
shouldUseCpuNoise,
@@ -113,6 +118,11 @@ export const buildSDXLGraph = async (arg: GraphBuilderArg): Promise<GraphBuilder
id: getPrefixedId('denoise_latents'),
cfg_scale,
cfg_rescale_multiplier,
hidiffusion: hiDiffusionEnabled,
hidiffusion_raunet: hiDiffusionRauNetEnabled,
hidiffusion_window_attn: hiDiffusionWindowAttnEnabled,
hidiffusion_t1_ratio: hiDiffusionEnabled ? hiDiffusionT1Ratio : undefined,
hidiffusion_t2_ratio: hiDiffusionEnabled ? hiDiffusionT2Ratio : undefined,
scheduler,
steps,
});
@@ -155,6 +165,11 @@ export const buildSDXLGraph = async (arg: GraphBuilderArg): Promise<GraphBuilder
cfg_scale,
cfg_rescale_multiplier,
model: Graph.getModelMetadataField(model),
hidiffusion: hiDiffusionEnabled,
hidiffusion_raunet: hiDiffusionRauNetEnabled,
hidiffusion_window_attn: hiDiffusionWindowAttnEnabled,
hidiffusion_t1_ratio: hiDiffusionT1Ratio,
hidiffusion_t2_ratio: hiDiffusionT2Ratio,
steps,
rand_device: shouldUseCpuNoise ? 'cpu' : 'cuda',
scheduler,
invokeai/frontend/web/src/features/parameters/components/Advanced/ParamHiDiffusionToggle.tsx
+199
View File
@@ -0,0 +1,199 @@
import { CompositeNumberInput, CompositeSlider, FormControl, FormLabel, Switch } from '@invoke-ai/ui-library';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import { InformationalPopover } from 'common/components/InformationalPopover/InformationalPopover';
import {
selectHiDiffusionEnabled,
selectHiDiffusionRauNetEnabled,
selectHiDiffusionT1Ratio,
selectHiDiffusionT2Ratio,
selectHiDiffusionWindowAttnEnabled,
setHiDiffusionEnabled,
setHiDiffusionRauNetEnabled,
setHiDiffusionT1Ratio,
setHiDiffusionT2Ratio,
setHiDiffusionWindowAttnEnabled,
} from 'features/controlLayers/store/paramsSlice';
import type { ChangeEvent } from 'react';
import { memo, useCallback } from 'react';
import { useTranslation } from 'react-i18next';
export const ParamHiDiffusionToggle = memo(() => {
const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled);
const dispatch = useAppDispatch();
const { t } = useTranslation();
const onChange = useCallback(
(event: ChangeEvent<HTMLInputElement>) => {
dispatch(setHiDiffusionEnabled(event.target.checked));
},
[dispatch]
);
return (
<FormControl>
<InformationalPopover feature="hidiffusion">
<FormLabel maxW="100%" whiteSpace="nowrap" overflow="hidden" textOverflow="ellipsis">
{t('parameters.hiDiffusion')}
</FormLabel>
</InformationalPopover>
<Switch isChecked={hiDiffusionEnabled} onChange={onChange} />
</FormControl>
);
});
ParamHiDiffusionToggle.displayName = 'ParamHiDiffusionToggle';
export const ParamHiDiffusionRauNetToggle = memo(() => {
const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled);
const hiDiffusionRauNetEnabled = useAppSelector(selectHiDiffusionRauNetEnabled);
const dispatch = useAppDispatch();
const { t } = useTranslation();
const onChange = useCallback(
(event: ChangeEvent<HTMLInputElement>) => {
dispatch(setHiDiffusionRauNetEnabled(event.target.checked));
},
[dispatch]
);
return (
<FormControl>
<InformationalPopover feature="hidiffusionRauNet">
<FormLabel maxW="100%" whiteSpace="nowrap" overflow="hidden" textOverflow="ellipsis">
{t('parameters.hiDiffusionRauNet')}
</FormLabel>
</InformationalPopover>
<Switch isChecked={hiDiffusionRauNetEnabled} isDisabled={!hiDiffusionEnabled} onChange={onChange} />
</FormControl>
);
});
ParamHiDiffusionRauNetToggle.displayName = 'ParamHiDiffusionRauNetToggle';
export const ParamHiDiffusionWindowAttnToggle = memo(() => {
const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled);
const hiDiffusionWindowAttnEnabled = useAppSelector(selectHiDiffusionWindowAttnEnabled);
const dispatch = useAppDispatch();
const { t } = useTranslation();
const onChange = useCallback(
(event: ChangeEvent<HTMLInputElement>) => {
dispatch(setHiDiffusionWindowAttnEnabled(event.target.checked));
},
[dispatch]
);
return (
<FormControl>
<InformationalPopover feature="hidiffusionWindowAttn">
<FormLabel maxW="100%" whiteSpace="nowrap" overflow="hidden" textOverflow="ellipsis">
{t('parameters.hiDiffusionWindowAttn')}
</FormLabel>
</InformationalPopover>
<Switch isChecked={hiDiffusionWindowAttnEnabled} isDisabled={!hiDiffusionEnabled} onChange={onChange} />
</FormControl>
);
});
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 (
<FormControl isDisabled={!hiDiffusionEnabled} gridColumn="1 / -1">
<InformationalPopover feature="hidiffusionT1Ratio">
<FormLabel minW="9rem" maxW="100%" whiteSpace="nowrap" overflow="hidden" textOverflow="ellipsis">
{t('parameters.hiDiffusionT1Ratio')}
</FormLabel>
</InformationalPopover>
<CompositeSlider
value={hiDiffusionT1Ratio}
defaultValue={RATIO_CONSTRAINTS.t1.initial}
min={RATIO_CONSTRAINTS.t1.sliderMin}
max={RATIO_CONSTRAINTS.t1.sliderMax}
step={RATIO_CONSTRAINTS.t1.coarseStep}
fineStep={RATIO_CONSTRAINTS.t1.fineStep}
onChange={onChange}
marks
/>
<CompositeNumberInput
value={hiDiffusionT1Ratio}
defaultValue={RATIO_CONSTRAINTS.t1.initial}
min={RATIO_CONSTRAINTS.t1.numberInputMin}
max={RATIO_CONSTRAINTS.t1.numberInputMax}
step={RATIO_CONSTRAINTS.t1.coarseStep}
fineStep={RATIO_CONSTRAINTS.t1.fineStep}
onChange={onChange}
/>
</FormControl>
);
});
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 (
<FormControl isDisabled={!hiDiffusionEnabled} gridColumn="1 / -1">
<InformationalPopover feature="hidiffusionT2Ratio">
<FormLabel minW="9rem" maxW="100%" whiteSpace="nowrap" overflow="hidden" textOverflow="ellipsis">
{t('parameters.hiDiffusionT2Ratio')}
</FormLabel>
</InformationalPopover>
<CompositeSlider
value={hiDiffusionT2Ratio}
defaultValue={RATIO_CONSTRAINTS.t2.initial}
min={RATIO_CONSTRAINTS.t2.sliderMin}
max={RATIO_CONSTRAINTS.t2.sliderMax}
step={RATIO_CONSTRAINTS.t2.coarseStep}
fineStep={RATIO_CONSTRAINTS.t2.fineStep}
onChange={onChange}
marks
/>
<CompositeNumberInput
value={hiDiffusionT2Ratio}
defaultValue={RATIO_CONSTRAINTS.t2.initial}
min={RATIO_CONSTRAINTS.t2.numberInputMin}
max={RATIO_CONSTRAINTS.t2.numberInputMax}
step={RATIO_CONSTRAINTS.t2.coarseStep}
fineStep={RATIO_CONSTRAINTS.t2.fineStep}
onChange={onChange}
/>
</FormControl>
);
});
ParamHiDiffusionT2Ratio.displayName = 'ParamHiDiffusionT2Ratio';
invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx
+16 -3
View File
@@ -1,5 +1,5 @@
import type { FormLabelProps } from '@invoke-ai/ui-library';
import { Box, Flex, FormControlGroup, SimpleGrid, StandaloneAccordion } from '@invoke-ai/ui-library';
import { Flex, FormControlGroup, SimpleGrid, StandaloneAccordion } from '@invoke-ai/ui-library';
import { skipToken } from '@reduxjs/toolkit/query';
import { createMemoizedSelector } from 'app/store/createMemoizedSelector';
import { useAppSelector } from 'app/store/storeHooks';
@@ -25,6 +25,13 @@ import ParamCLIPGEmbedModelSelect from 'features/parameters/components/Advanced/
import ParamCLIPLEmbedModelSelect from 'features/parameters/components/Advanced/ParamCLIPLEmbedModelSelect';
import ParamClipSkip from 'features/parameters/components/Advanced/ParamClipSkip';
import ParamFlux2KleinModelSelect from 'features/parameters/components/Advanced/ParamFlux2KleinModelSelect';
import {
ParamHiDiffusionRauNetToggle,
ParamHiDiffusionT1Ratio,
ParamHiDiffusionT2Ratio,
ParamHiDiffusionToggle,
ParamHiDiffusionWindowAttnToggle,
} from 'features/parameters/components/Advanced/ParamHiDiffusionToggle';
import ParamKrea2ModelSelects from 'features/parameters/components/Advanced/ParamKrea2ModelSelects';
import ParamQwenImageComponentSourceSelect from 'features/parameters/components/Advanced/ParamQwenImageComponentSourceSelect';
import ParamQwenImageQuantization from 'features/parameters/components/Advanced/ParamQwenImageQuantization';
@@ -106,6 +113,9 @@ export const AdvancedSettingsAccordion = memo(() => {
if (params.seamlessXAxis || params.seamlessYAxis) {
badges.push('seamless');
}
if (params.hiDiffusionEnabled) {
badges.push('HiDiffusion');
}
}
return badges;
@@ -153,9 +163,12 @@ export const AdvancedSettingsAccordion = memo(() => {
<SimpleGrid columns={2} spacing={4} w="full">
<ParamSeamlessXAxis />
<ParamSeamlessYAxis />
<ParamHiDiffusionToggle />
<ParamColorCompensation />
{/* Empty box for visual alignment. Replace with new option when needed. */}
<Box />
<ParamHiDiffusionRauNetToggle />
<ParamHiDiffusionWindowAttnToggle />
<ParamHiDiffusionT1Ratio />
<ParamHiDiffusionT2Ratio />
</SimpleGrid>
</FormControlGroup>
</Flex>
invokeai/frontend/web/src/services/api/schema.ts
+60
View File
@@ -9204,6 +9204,36 @@ export type components = {
* @default 0
*/
cfg_rescale_multiplier?: number;
/**
* HiDiffusion
* @description Apply HiDiffusion (RAU-Net + MSW-MSA) for higher-resolution denoising
* @default false
*/
hidiffusion?: boolean;
/**
* HiDiffusion: RAU-Net
* @description Apply HiDiffusion RAU-Net blocks
* @default true
*/
hidiffusion_raunet?: boolean;
/**
* HiDiffusion: Window Attention
* @description Apply HiDiffusion window attention blocks
* @default true
*/
hidiffusion_window_attn?: boolean;
/**
* HiDiffusion: T1 Ratio
* @description Override HiDiffusion early switch threshold (T1 ratio)
* @default 0.4
*/
hidiffusion_t1_ratio?: number;
/**
* HiDiffusion: T2 Ratio
* @description Override HiDiffusion late switch threshold (T2 ratio)
* @default 0
*/
hidiffusion_t2_ratio?: number;
/**
* @description Latents tensor
* @default null
@@ -9322,6 +9352,36 @@ export type components = {
* @default 0
*/
cfg_rescale_multiplier?: number;
/**
* HiDiffusion
* @description Apply HiDiffusion (RAU-Net + MSW-MSA) for higher-resolution denoising
* @default false
*/
hidiffusion?: boolean;
/**
* HiDiffusion: RAU-Net
* @description Apply HiDiffusion RAU-Net blocks
* @default true
*/
hidiffusion_raunet?: boolean;
/**
* HiDiffusion: Window Attention
* @description Apply HiDiffusion window attention blocks
* @default true
*/
hidiffusion_window_attn?: boolean;
/**
* HiDiffusion: T1 Ratio
* @description Override HiDiffusion early switch threshold (T1 ratio)
* @default 0.4
*/
hidiffusion_t1_ratio?: number;
/**
* HiDiffusion: T2 Ratio
* @description Override HiDiffusion late switch threshold (T2 ratio)
* @default 0
*/
hidiffusion_t2_ratio?: number;
/**
* @description Latents tensor
* @default null
pyproject.toml
+1
View File
@@ -239,6 +239,7 @@ version = { attr = "invokeai.version.__version__" }
"invokeai.backend.qwen3" = ["tokenizer/*.json", "tokenizer/*.json.gz"]
"invokeai.backend.t5" = ["tokenizer/*.json"]
"invokeai.app.services.workflow_records.default_workflows" = ["*.json"]
"invokeai.backend.hidiffusion" = ["sd_module_key/*.txt"]
"invokeai.app.services.style_preset_records" = ["*.json"]
"invokeai.app.services.style_preset_images.default_style_preset_images" = [
"*.png",
tests/app/invocations/test_denoise_noise_inputs.py
+35 -1
View File
@@ -9,7 +9,11 @@ from invokeai.app.invocations.anima_denoise import AnimaDenoiseInvocation
from invokeai.app.invocations.cogview4_denoise import CogView4DenoiseInvocation
from invokeai.app.invocations.flux2_denoise import Flux2DenoiseInvocation
from invokeai.app.invocations.flux_denoise import FluxDenoiseInvocation
from invokeai.app.invocations.metadata_linked import FluxDenoiseLatentsMetaInvocation, ZImageDenoiseMetaInvocation
from invokeai.app.invocations.metadata_linked import (
DenoiseLatentsMetaInvocation,
FluxDenoiseLatentsMetaInvocation,
ZImageDenoiseMetaInvocation,
)
from invokeai.app.invocations.primitives import LatentsOutput
from invokeai.app.invocations.sd3_denoise import SD3DenoiseInvocation
from invokeai.app.invocations.z_image_denoise import ZImageDenoiseInvocation
@@ -392,6 +396,36 @@ def test_z_image_metadata_ignores_external_noise_seed_when_noise_not_used():
assert result.metadata.root["seed"] == 999
def test_denoise_metadata_persists_hidiffusion_fields_when_disabled():
invocation = DenoiseLatentsMetaInvocation.model_construct(
width=64,
height=64,
steps=10,
cfg_scale=7.5,
cfg_rescale_multiplier=0,
denoising_start=0.0,
denoising_end=1.0,
scheduler="euler",
unet=SimpleNamespace(unet="main-model", loras=[]),
hidiffusion=False,
hidiffusion_raunet=False,
hidiffusion_window_attn=False,
hidiffusion_t1_ratio=0.25,
hidiffusion_t2_ratio=0.1,
)
mock_context = MagicMock()
output = LatentsOutput.build("latents", torch.zeros(1, 4, 8, 8), seed=None)
with patch("invokeai.app.invocations.metadata_linked.DenoiseLatentsInvocation.invoke", return_value=output):
result = invocation.invoke(mock_context)
assert result.metadata.root["hidiffusion"] is False
assert result.metadata.root["hidiffusion_raunet"] is False
assert result.metadata.root["hidiffusion_window_attn"] is False
assert result.metadata.root["hidiffusion_t1_ratio"] == 0.25
assert result.metadata.root["hidiffusion_t2_ratio"] == 0.1
def _get_first_scheduler_sigma(
scheduler, *, scheduler_name: str, sigmas: list[float], mu: float | None = None
) -> float:
tests/backend/stable_diffusion/test_hidiffusion_utils.py
+324
View File
@@ -0,0 +1,324 @@
import copy
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
import torch
from invokeai.backend.hidiffusion.hidiffusion import (
remove_hidiffusion as real_remove_hidiffusion,
)
from invokeai.backend.hidiffusion.hidiffusion import (
switching_threshold_ratio_dict,
text_to_img_controlnet_switching_threshold_ratio_dict,
)
from invokeai.backend.stable_diffusion.hidiffusion_utils import hidiffusion_patch
class DummySubmodule:
pass
class PatchedSubmodule(DummySubmodule):
_parent = DummySubmodule
class DummyUNet:
def __init__(self):
self.num_upsamplers = 3
self.layer = DummySubmodule()
def named_modules(self):
return [("", self), ("layer", self.layer)]
class ModelMixin(torch.nn.Module):
"""Minimal diffusers-like UNet accepted by the vendored HiDiffusion type check."""
def __init__(self):
super().__init__()
self.num_upsamplers = 3
class WindowMeanAttention(torch.nn.Module):
def forward(self, hidden_states: torch.Tensor, **_kwargs):
return hidden_states.mean(dim=1, keepdim=True).expand_as(hidden_states)
class WindowAttentionBlock(torch.nn.Module):
def __init__(self):
super().__init__()
self.use_ada_layer_norm = False
self.use_ada_layer_norm_zero = False
self.use_layer_norm = True
self.use_ada_layer_norm_continuous = False
self.use_ada_layer_norm_single = False
self.pos_embed = None
self.norm1 = torch.nn.Identity()
self.attn1 = WindowMeanAttention()
self.only_cross_attention = False
self.attn2 = None
self.norm3 = torch.nn.Identity()
self.ff = torch.nn.Identity()
self._chunk_size = None
class WindowAttentionModelMixin(ModelMixin):
def __init__(self):
super().__init__()
self.transformer = WindowAttentionBlock()
class CachedHiDiffusionModelMixin(ModelMixin):
def __init__(self):
super().__init__()
self.block = torch.nn.Module()
def test_hidiffusion_patch_supports_bare_model_mixin_without_public_name_or_path():
model = ModelMixin()
assert not hasattr(model, "name_or_path")
assert not hasattr(model, "_name_or_path")
with hidiffusion_patch(model, name_or_path="runwayml/stable-diffusion-v1-5"):
assert model.info["pipeline"] is model
assert model.num_upsamplers == 15
assert model.num_upsamplers == 3
assert not hasattr(model, "_name_or_path")
assert not hasattr(model, "info")
def test_hidiffusion_window_attention_uses_seeded_generator_instead_of_global_rng():
module_keys = {
"down_module_key": [],
"down_module_key_extra": [],
"up_module_key": [],
"up_module_key_extra": [],
"windown_attn_module_key": ["transformer"],
}
hidden_states = torch.arange(64, dtype=torch.float32).reshape(1, 64, 1)
def run_with_global_seed(global_seed: int) -> torch.Tensor:
torch.manual_seed(global_seed)
model = WindowAttentionModelMixin()
generator = torch.Generator(device="cpu").manual_seed(1234)
with (
patch("invokeai.backend.hidiffusion.hidiffusion.sd15_hidiffusion_key", return_value=module_keys),
hidiffusion_patch(
model,
name_or_path="runwayml/stable-diffusion-v1-5",
apply_raunet=False,
apply_window_attn=True,
generator=generator,
),
):
model.info["size"] = (8, 8)
return model.transformer(hidden_states).clone()
first = run_with_global_seed(0)
second = run_with_global_seed(1)
torch.testing.assert_close(first, second)
def test_hidiffusion_patch_resets_cached_runtime_state_when_reenabled():
module_keys = {
"down_module_key": [],
"down_module_key_extra": ["block"],
"up_module_key": [],
"up_module_key_extra": [],
"windown_attn_module_key": [],
}
model = CachedHiDiffusionModelMixin()
with patch("invokeai.backend.hidiffusion.hidiffusion.sd15_hidiffusion_key", return_value=module_keys):
with hidiffusion_patch(model, name_or_path="runwayml/stable-diffusion-v1-5"):
model.block.timestep = 7
model.block.aggressive_raunet = True
model.block.T1_ratio = 0.9
model.block.T1 = 9
model.block.T1_start = 2
model.block.T1_end = 8
model.block.max_timestep = 99
assert "timestep" not in model.block.__dict__
with hidiffusion_patch(model, name_or_path="runwayml/stable-diffusion-v1-5"):
assert model.block.timestep == 0
assert model.block.aggressive_raunet is False
assert model.block.T1_ratio == 0
assert model.block.T1 == 0
assert model.block.T1_start == 0
assert model.block.T1_end == 0
assert model.block.max_timestep == 50
def test_hidiffusion_teardown_restores_downsampler_geometry_after_forward_error():
module_keys = {
"down_module_key": ["block"],
"down_module_key_extra": [],
"up_module_key": [],
"up_module_key_extra": [],
"windown_attn_module_key": [],
}
model = ModelMixin()
model._num_timesteps = 10
model.block = torch.nn.Conv2d(1, 1, kernel_size=3, stride=2, padding=1)
original_stride = model.block.stride
original_padding = model.block.padding
original_dilation = model.block.dilation
with patch("invokeai.backend.hidiffusion.hidiffusion.sd15_hidiffusion_key", return_value=module_keys):
with hidiffusion_patch(
model,
name_or_path="runwayml/stable-diffusion-v1-5",
apply_window_attn=False,
):
model.info["size"] = (64, 64)
with (
patch(
"invokeai.backend.hidiffusion.hidiffusion.F.conv2d",
side_effect=RuntimeError("injected convolution failure"),
),
pytest.raises(RuntimeError, match="injected convolution failure"),
):
model.block(torch.zeros(1, 1, 16, 16))
# Temporary geometry is passed directly to conv2d and never written
# to the cached module, even before teardown runs.
assert model.block.stride == original_stride
assert model.block.padding == original_padding
assert model.block.dilation == original_dilation
assert model.block.stride == original_stride
assert model.block.padding == original_padding
assert model.block.dilation == original_dilation
def test_hidiffusion_patch_restores_state_when_apply_hidiffusion_raises():
original_switching = copy.deepcopy(switching_threshold_ratio_dict)
original_controlnet = copy.deepcopy(text_to_img_controlnet_switching_threshold_ratio_dict)
model = SimpleNamespace(
unet=DummyUNet(),
_name_or_path="original-model-name",
config=SimpleNamespace(_name_or_path="original-config-name"),
)
hook = MagicMock()
def fake_apply_hidiffusion(patched_model, **_kwargs):
assert patched_model._name_or_path == "patched-model-name"
assert patched_model.config._name_or_path == "patched-model-name"
first_switching_entry = next(iter(switching_threshold_ratio_dict.values()))
first_controlnet_entry = next(iter(text_to_img_controlnet_switching_threshold_ratio_dict.values()))
assert first_switching_entry["T1_ratio"] == 0.25
assert first_switching_entry["T2_ratio"] == 0.1
assert first_controlnet_entry["T1_ratio"] == 0.25
assert first_controlnet_entry["T2_ratio"] == 0.1
patched_model.unet.num_upsamplers = 99
patched_model.unet.layer.info = {"hooks": [hook]}
patched_model.unet.layer.__class__ = PatchedSubmodule
raise RuntimeError("hidiffusion boom")
try:
with (
patch("invokeai.backend.hidiffusion.hidiffusion.apply_hidiffusion", side_effect=fake_apply_hidiffusion),
patch(
"invokeai.backend.hidiffusion.hidiffusion.remove_hidiffusion",
wraps=real_remove_hidiffusion,
) as mock_remove_hidiffusion,
):
with pytest.raises(RuntimeError, match="hidiffusion boom"):
with hidiffusion_patch(
model,
name_or_path="patched-model-name",
t1_ratio=0.25,
t2_ratio=0.1,
):
pass
assert mock_remove_hidiffusion.call_count == 1
assert switching_threshold_ratio_dict == original_switching
assert text_to_img_controlnet_switching_threshold_ratio_dict == original_controlnet
assert model.unet.num_upsamplers == 3
assert model.unet.layer.__class__ is DummySubmodule
assert model.unet.layer.info["hooks"] == []
hook.remove.assert_called_once()
assert model._name_or_path == "original-model-name"
assert model.config._name_or_path == "original-config-name"
finally:
switching_threshold_ratio_dict.clear()
switching_threshold_ratio_dict.update(original_switching)
text_to_img_controlnet_switching_threshold_ratio_dict.clear()
text_to_img_controlnet_switching_threshold_ratio_dict.update(original_controlnet)
def test_hidiffusion_patch_restores_state_before_propagating_remove_error():
original_switching = copy.deepcopy(switching_threshold_ratio_dict)
original_controlnet = copy.deepcopy(text_to_img_controlnet_switching_threshold_ratio_dict)
model = SimpleNamespace(
unet=DummyUNet(),
_name_or_path="original-model-name",
config=SimpleNamespace(_name_or_path="original-config-name"),
)
def fake_apply_hidiffusion(patched_model, **_kwargs):
patched_model.unet.num_upsamplers = 99
try:
with (
patch("invokeai.backend.hidiffusion.hidiffusion.apply_hidiffusion", side_effect=fake_apply_hidiffusion),
patch(
"invokeai.backend.hidiffusion.hidiffusion.remove_hidiffusion",
side_effect=RuntimeError("remove boom"),
),
):
with pytest.raises(RuntimeError, match="remove boom"):
with hidiffusion_patch(
model,
name_or_path="patched-model-name",
t1_ratio=0.25,
t2_ratio=0.1,
):
pass
assert switching_threshold_ratio_dict == original_switching
assert text_to_img_controlnet_switching_threshold_ratio_dict == original_controlnet
assert model.unet.num_upsamplers == 3
assert model._name_or_path == "original-model-name"
assert model.config._name_or_path == "original-config-name"
finally:
switching_threshold_ratio_dict.clear()
switching_threshold_ratio_dict.update(original_switching)
text_to_img_controlnet_switching_threshold_ratio_dict.clear()
text_to_img_controlnet_switching_threshold_ratio_dict.update(original_controlnet)
def test_hidiffusion_patch_removes_spoofed_name_from_config_internal_dict():
class InternalDictConfig:
def __init__(self):
self._internal_dict = {}
def __getattr__(self, name):
try:
return self._internal_dict[name]
except KeyError as error:
raise AttributeError(name) from error
config = InternalDictConfig()
model = SimpleNamespace(unet=DummyUNet(), config=config)
with (
patch("invokeai.backend.hidiffusion.hidiffusion.apply_hidiffusion"),
patch("invokeai.backend.hidiffusion.hidiffusion.remove_hidiffusion"),
):
with hidiffusion_patch(model, name_or_path="patched-model-name"):
assert config._internal_dict["_name_or_path"] == "patched-model-name"
assert "_name_or_path" not in config._internal_dict