Feature: HiDiffusion integration

LICENSE-HiDiffusion.txt

@@ -0,0 +1,65 @@
+HiDiffusion - License notice
+
+Original project: https://github.com/megvii-research/HiDiffusion
+Vendored from:    https://github.com/monofy-org/HiDiffusion
+Vendored under:   invokeai/backend/hidiffusion/
+
+================================================================================
+CODE (Apache License 2.0)
+================================================================================
+
+The vendored HiDiffusion source code is licensed under the Apache License,
+Version 2.0. The Apache 2.0 license text is available in the root LICENSE file
+and at:
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+================================================================================
+MODULE KEYS (BSD 3-Clause License)
+================================================================================
+
+The SD 1.5 and SDXL module-key lists under
+invokeai/backend/hidiffusion/sd_module_key/ were sourced from dgenerate at:
+
+    https://github.com/Teriks/dgenerate/tree/d83b839033cc22c5101fb0f987bd4eb2de3d5d12/dgenerate/extras/hidiffusion/sd_module_key
+
+BSD 3-Clause License
+
+Copyright (c) 2023, Teriks
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+================================================================================
+LOCAL MODIFICATIONS
+================================================================================
+
+The following changes were applied when integrating HiDiffusion into InvokeAI:
+
+* The code was vendored under the invokeai.backend.hidiffusion package.
+* apply_hidiffusion() was extended to accept a torch.Generator, enabling
+  deterministic image generation from a seed.
+* HiDiffusion patching was integrated with InvokeAI's model loading and
+  generation lifecycle.

docs/src/content/docs/features/hidiffusion.mdx

@@ -0,0 +1,39 @@
+---
+title: HiDiffusion
+---
+
+# HiDiffusion
+
+HiDiffusion is an optional denoising enhancement that can improve detail and structure at higher resolutions for SD 1.5, SD 2.1, and SDXL. It modifies the UNet during denoising and is most noticeable at 1536px and above.
+
+The vendored implementation includes module-key definitions for SD 1.5 and SDXL. SD 2.1 intentionally reuses the SD 1.5 module-key mapping.
+
+Learn more: https://github.com/megvii-research/HiDiffusion
+
+## Where to find the switches
+
+1. Open the **Canvas** tab.
+2. Expand **Advanced Settings**.
+3. In the **Advanced** grid, enable **HiDiffusion** and optionally adjust the two sub‑toggles and ratios:
+   - **HiDiffusion: RAU‑Net**
+   - **HiDiffusion: Window Attention**
+   - **HiDiffusion: T1 Ratio**
+   - **HiDiffusion: T2 Ratio**
+
+## What the switches do
+
+- **HiDiffusion**: Enables the HiDiffusion patch for denoising. Use this for high‑resolution generations; the effect is subtle at lower sizes.
+
+- **HiDiffusion: RAU‑Net**: Enables RAU‑Net blocks. This typically improves structure and mid‑frequency detail, especially at larger resolutions.
+
+- **HiDiffusion: Window Attention**: Enables windowed attention blocks. This can boost local texture/detail, but may slightly affect global coherence in some prompts.
+
+- **HiDiffusion: T1 Ratio**: Controls when HiDiffusion switches into its mid‑stage behavior. Lower values switch earlier; higher values preserve global structure longer.
+
+- **HiDiffusion: T2 Ratio**: Controls when HiDiffusion switches into its late‑stage behavior. Higher values keep window attention active longer and can sharpen local detail.
+
+## Tips
+
+- Try **1536–2048 px** for the clearest benefits (SDXL).
+- If results look worse, disable **Window Attention** first, then RAU‑Net.
+- Effects vary by scheduler and model; compare with the same seed for a fair test.

invokeai/app/invocations/denoise_latents.py

@@ -1,7 +1,7 @@
 # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
 import inspect
 import os
-from contextlib import ExitStack
+from contextlib import ExitStack, nullcontext
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 import torch
@@ -65,6 +65,7 @@
 from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType
 from invokeai.backend.stable_diffusion.extensions.controlnet import ControlNetExt
 from invokeai.backend.stable_diffusion.extensions.freeu import FreeUExt
+from invokeai.backend.stable_diffusion.extensions.hidiffusion import HiDiffusionExt
 from invokeai.backend.stable_diffusion.extensions.inpaint import InpaintExt
 from invokeai.backend.stable_diffusion.extensions.inpaint_model import InpaintModelExt
 from invokeai.backend.stable_diffusion.extensions.lora import LoRAExt
@@ -73,6 +74,7 @@
 from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt
 from invokeai.backend.stable_diffusion.extensions.t2i_adapter import T2IAdapterExt
 from invokeai.backend.stable_diffusion.extensions_manager import ExtensionsManager
+from invokeai.backend.stable_diffusion.hidiffusion_utils import hidiffusion_patch
 from invokeai.backend.stable_diffusion.schedulers import SCHEDULER_MAP
 from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES
 from invokeai.backend.util.devices import TorchDevice
@@ -130,7 +132,7 @@ def get_scheduler(
     title="Denoise - SD1.5, SDXL",
     tags=["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"],
     category="latents",
-    version="1.5.4",
+    version="1.6.0",
 )
 class DenoiseLatentsInvocation(BaseInvocation):
     """Denoises noisy latents to decodable images"""
@@ -191,6 +193,35 @@ class DenoiseLatentsInvocation(BaseInvocation):
     cfg_rescale_multiplier: float = InputField(
         title="CFG Rescale Multiplier", default=0, ge=0, lt=1, description=FieldDescriptions.cfg_rescale_multiplier
     )
+    hidiffusion: bool = InputField(
+        default=False,
+        description=FieldDescriptions.hidiffusion,
+        title="HiDiffusion",
+    )
+    hidiffusion_raunet: bool = InputField(
+        default=True,
+        description=FieldDescriptions.hidiffusion_raunet,
+        title="HiDiffusion: RAU-Net",
+    )
+    hidiffusion_window_attn: bool = InputField(
+        default=True,
+        description=FieldDescriptions.hidiffusion_window_attn,
+        title="HiDiffusion: Window Attention",
+    )
+    hidiffusion_t1_ratio: float = InputField(
+        default=0.4,
+        ge=0,
+        le=1,
+        description=FieldDescriptions.hidiffusion_t1_ratio,
+        title="HiDiffusion: T1 Ratio",
+    )
+    hidiffusion_t2_ratio: float = InputField(
+        default=0.0,
+        ge=0,
+        le=1,
+        description=FieldDescriptions.hidiffusion_t2_ratio,
+        title="HiDiffusion: T2 Ratio",
+    )
     latents: Optional[LatentsField] = InputField(
         default=None,
         description=FieldDescriptions.latents,
@@ -486,6 +517,14 @@ def prep_control_data(
 
         return controlnet_data
 
+    @staticmethod
+    def _get_hidiffusion_name_or_path(unet_config: AnyModelConfig) -> Optional[str]:
+        return (
+            getattr(unet_config, "source", None)
+            or getattr(unet_config, "path", None)
+            or getattr(unet_config, "name", None)
+        )
+
     @staticmethod
     def parse_controlnet_field(
         exit_stack: ExitStack,
@@ -837,6 +876,7 @@ def _new_invoke(self, context: InvocationContext) -> LatentsOutput:
 
         # get the unet's config so that we can pass the base to sd_step_callback()
         unet_config = context.models.get_config(self.unet.unet.key)
+        hidiffusion_name_or_path = self._get_hidiffusion_name_or_path(unet_config)
 
         conditioning_data = self.get_conditioning_data(
             context=context,
@@ -874,6 +914,16 @@ def step_callback(state: PipelineIntermediateState) -> None:
             context.util.sd_step_callback(state, unet_config.base)
 
         ext_manager.add_extension(PreviewExt(step_callback))
+        if self.hidiffusion:
+            ext_manager.add_extension(
+                HiDiffusionExt(
+                    name_or_path=hidiffusion_name_or_path,
+                    apply_raunet=self.hidiffusion_raunet,
+                    apply_window_attn=self.hidiffusion_window_attn,
+                    t1_ratio=self.hidiffusion_t1_ratio,
+                    t2_ratio=self.hidiffusion_t2_ratio,
+                )
+            )
 
         ### cfg rescale
         if self.cfg_rescale_multiplier > 0:
@@ -940,14 +990,17 @@ def step_callback(state: PipelineIntermediateState) -> None:
             # ext: t2i/ip adapter
             ext_manager.run_callback(ExtensionCallbackType.SETUP, denoise_ctx)
 
-            with (
-                context.models.load(self.unet.unet).model_on_device() as (cached_weights, unet),
-                ModelPatcher.patch_unet_attention_processor(unet, denoise_ctx.inputs.attention_processor_cls),
+            with ExitStack() as unet_stack:
+                cached_weights, unet = unet_stack.enter_context(context.models.load(self.unet.unet).model_on_device())
+                unet._num_timesteps = timesteps.shape[0]
+                unet_stack.enter_context(
+                    ModelPatcher.patch_unet_attention_processor(unet, denoise_ctx.inputs.attention_processor_cls)
+                )
                 # ext: controlnet
-                ext_manager.patch_extensions(denoise_ctx),
-                # ext: freeu, seamless, ip adapter, lora
-                ext_manager.patch_unet(unet, cached_weights),
-            ):
+                unet_stack.enter_context(ext_manager.patch_extensions(denoise_ctx))
+                # ext: freeu, seamless, ip adapter, lora, hidiffusion
+                unet_stack.enter_context(ext_manager.patch_unet(unet, cached_weights))
+
                 sd_backend = StableDiffusionBackend(unet, scheduler)
                 denoise_ctx.unet = unet
                 result_latents = sd_backend.latents_from_embeddings(denoise_ctx, ext_manager)
@@ -997,6 +1050,7 @@ def _old_invoke(self, context: InvocationContext) -> LatentsOutput:
 
         # get the unet's config so that we can pass the base to sd_step_callback()
         unet_config = context.models.get_config(self.unet.unet.key)
+        hidiffusion_name_or_path = self._get_hidiffusion_name_or_path(unet_config)
 
         def step_callback(state: PipelineIntermediateState) -> None:
             context.util.sd_step_callback(state, unet_config.base)
@@ -1083,23 +1137,36 @@ def _lora_loader() -> Iterator[PatchSpec]:
                 denoising_end=self.denoising_end,
                 seed=seed,
             )
+            pipeline._num_timesteps = timesteps.shape[0]
 
-            result_latents = pipeline.latents_from_embeddings(
-                latents=latents,
-                timesteps=timesteps,
-                init_timestep=init_timestep,
-                noise=noise,
-                seed=seed,
-                mask=mask,
-                masked_latents=masked_latents,
-                is_gradient_mask=gradient_mask,
-                scheduler_step_kwargs=scheduler_step_kwargs,
-                conditioning_data=conditioning_data,
-                control_data=controlnet_data,
-                ip_adapter_data=ip_adapter_data,
-                t2i_adapter_data=t2i_adapter_data,
-                callback=step_callback,
-            )
+            with (
+                hidiffusion_patch(
+                    pipeline,
+                    name_or_path=hidiffusion_name_or_path,
+                    apply_raunet=self.hidiffusion_raunet,
+                    apply_window_attn=self.hidiffusion_window_attn,
+                    t1_ratio=self.hidiffusion_t1_ratio,
+                    t2_ratio=self.hidiffusion_t2_ratio,
+                )
+                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

@@ -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

@@ -622,7 +622,7 @@ class LatentsMetaOutput(LatentsOutput, MetadataOutput):
     title=f"{DenoiseLatentsInvocation.UIConfig.title} + Metadata",
     tags=["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"],
     category="metadata",
-    version="1.1.1",
+    version="1.2.0",
 )
 class DenoiseLatentsMetaInvocation(DenoiseLatentsInvocation, WithMetadata):
     def invoke(self, context: InvocationContext) -> LatentsMetaOutput:
@@ -661,6 +661,11 @@ def _loras_to_json(obj: Union[Any, list[Any]]):
         md.update({"denoising_end": self.denoising_end})
         md.update({"scheduler": self.scheduler})
         md.update({"model": self.unet.unet})
+        md.update({"hidiffusion": self.hidiffusion})
+        md.update({"hidiffusion_raunet": self.hidiffusion_raunet})
+        md.update({"hidiffusion_window_attn": self.hidiffusion_window_attn})
+        md.update({"hidiffusion_t1_ratio": self.hidiffusion_t1_ratio})
+        md.update({"hidiffusion_t2_ratio": self.hidiffusion_t2_ratio})
         if isinstance(self.control, ControlField) or (isinstance(self.control, list) and len(self.control) > 0):
             md.update({"controlnets": _to_json(self.control)})
         if isinstance(self.ip_adapter, IPAdapterField) or (

invokeai/backend/hidiffusion/__init__.py

@@ -0,0 +1,3 @@
+from invokeai.backend.hidiffusion.hidiffusion import apply_hidiffusion, remove_hidiffusion
+
+__all__ = ["apply_hidiffusion", "remove_hidiffusion"]