Feature: HiDiffusion integration

[DustyShoe]

Summary

Added HiDiffusion support across backend and UI. The denoise pipeline can now apply HiDiffusion with optional RAU‑Net and Window Attention toggles, including metadata persistence/recall and graph wiring for SD1/SDXL and the SDXL refiner. The UI exposes aligned advanced switches with informational popovers.
Added HiDiffusion documentation and a “Learn more” link.

Related Issues / Discussions

https://github.com/megvii-research/HiDiffusion
Closes #8780

QA Instructions

1. UI: Canvas → Advanced Settings → confirm HiDiffusion + RAU‑Net + Window Attention toggles visible, aligned, truncation works at narrow width.
2. Popovers: hover labels for HiDiffusion/RAU‑Net/Window Attention; popover appears and HiDiffusion shows Learn More link.
3. Metadata: generate once with HiDiffusion ON + custom RAU/Window; verify metadata fields exist and Recall restores toggles.
4. SDXL: compare OFF vs ON at 1536+ (same seed). Then toggle RAU off / Window off to confirm output changes and no errors.
5. SD1.5: repeat step 4 at highest feasible res.
6. Regression: generate with HiDiffusion OFF; no errors/warnings.

Merge Plan

Should merge without issue.

• Just to be aware of changes in lockfiles: uv.lock, pyproject.toml.
• HiDiffusion package availability: it pulls from PyPI upstream. If a version change introduces breaking behavior, we may need to pin a specific version.
• name_or_path hack: we’re forcing _name_or_path to avoid HiDiffusion errors. If diffusers changes internals, this might need adjustment.
• No DB/schema breaking changes.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[DustyShoe]

I might be spoiled by SDXL too much, but SD1.5 did not impress me much. It produces better results than default denoise. Especially when generating whole image. Inpainting is a tricky one. You might want to try different schedulers. Hands are particularly messy at 1.5k px.
SDXL on other hand, behaves much better and i'm pretty happy with it.
Also noticed ~10% improvement in speed.

[DustyShoe]

I just realised that updated version of HiDiffusion is here: https://github.com/Teriks/dgenerate/tree/master/dgenerate/extras/hidiffusion
How we can use it if it works better?

[iwr-redmond]

You may wish to consider vendoring the code inside the Invoke package, e.g. invokeai.backend.hidiffusion. Both packages are licensed with Apache 2.0.

[JPPhoto]

You need to bump versions of nodes that you've added fields to:

• invokeai/app/invocations/denoise_latents.py: bump DenoiseLatentsInvocation from 1.5.4 to a new version, likely 1.6.0 because new public inputs were added.

• invokeai/app/invocations/metadata_linked.py: bump DenoiseLatentsMetaInvocation from 1.1.1 to a new version, likely 1.2.0 for the inherited new inputs and added metadata keys.

• The generated openapi.json must also be regenerated so the frontend schema advertises the bumped versions.

Very minor, but the license file LICENSE-HiDiffusion.txt at the repository root looks like it's missing a newline at the end.

As soon as you fix these, I'll approve this to merge!

[DustyShoe]

@JPPhoto

• Bumped DenoiseLatentsInvocation to 1.6.0 and DenoiseLatentsMetaInvocation to 1.2.0.
• Regenerated openapi.json and the frontend schema; schema.ts had no resulting content changes because invocation versions are not represented in its generated types.
• Added the missing newline at the end of LICENSE-HiDiffusion.txt

[JPPhoto]

One more that I found:

• invokeai/backend/stable_diffusion/hidiffusion_utils.py:114: hidiffusion_patch() never accepts or passes invocation-seeded torch.Generator; vendored window attention therefore reaches torch.rand(1) process-global RNG at invokeai/backend/hidiffusion/hidiffusion.py:1420. Same seed can produce different images when window attention is enabled. Test by running two identical HiDiffusion window-attention denoises with the same seed while perturbing global Torch RNG between runs.

[DustyShoe]

Fixed both modular and legacy denoise paths to create a dedicated torch.Generator seeded from the invocation seed and pass it through hidiffusion_patch() into the vendored window-attention implementation. Added a regression test that runs the real window-attention patch twice with the same seed while perturbing the process-global Torch RNG; both results are identical.

[JPPhoto]

Good to merge!

[DustyShoe]

HiDiffusion runtime state could persist on cached UNet modules after the patch was removed. Re-enabling HiDiffusion or changing the bounding-box dimensions could therefore reuse stale timestep, RAU-Net, and convolution attributes, causing duplicated anatomy and other generation artifacts.

Added fix fix: stale HiDiffusion state in cached UNet resets HiDiffusion-specific runtime state before patching the relevant RAU-Net blocks again. Regular generations with HiDiffusion disabled remain unaffected. A regression test covers the ON → OFF → ON sequence using the same cached UNet instance.

[JPPhoto]

The latest changes have an issue:

• invokeai/backend/hidiffusion/hidiffusion.py:2182: remove_hidiffusion() restores classes but never calls _reset_hidiffusion_runtime_state(). If downsampler forward fails after geometry changes at lines 1860-1862 but before restoration, cached stride, padding, and dilation remain mutated; next HiDiffusion-disabled run uses corrupted convolution geometry. Re-enable reset is too late and skipped when RAU-Net is disabled. Test: force F.conv2d to raise during patched downsampler forward, exit patch context, then assert geometry is restored before an unpatched forward.

Here's a regression test you can modify and drop in that shows the fault:

from unittest.mock import patch

import pytest
import torch

from invokeai.backend.stable_diffusion.hidiffusion_utils import hidiffusion_patch


class ModelMixin(torch.nn.Module):
    """Minimal cached UNet accepted by HiDiffusion's diffusers type check."""

    def __init__(self) -> None:
        super().__init__()
        self.num_upsamplers = 3
        self._num_timesteps = 10
        self.block = torch.nn.Conv2d(1, 1, kernel_size=3, stride=2, padding=1)


def test_hidiffusion_teardown_restores_downsampler_geometry_after_forward_error() -> None:
    module_keys = {
        "down_module_key": ["block"],
        "down_module_key_extra": [],
        "up_module_key": [],
        "up_module_key_extra": [],
        "windown_attn_module_key": [],
    }
    model = ModelMixin()
    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))

    # The cached module is now unpatched. A HiDiffusion-disabled generation
    # must see its original convolution geometry.
    assert model.block.stride == original_stride
    assert model.block.padding == original_padding
    assert model.block.dilation == original_dilation

[DustyShoe]

Implemented transactional HiDiffusion teardown. Downsampler geometry is now passed to F.conv2d() through local variables and is never written to the cached module. Per-module state is snapshotted during patch application and restored during teardown, which remains protected by finally. The regression test now covers ON → failed forward → OFF and verifies that the unpatched cached UNet retains its original convolution geometry.

[JPPhoto]

Good to merge again!

[lstein]

Nice work!