From 81838942998f691d1ebfb2331d078dfd6c7d9afb Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Sun, 16 Aug 2026 23:50:05 +0300 Subject: [PATCH 1/7] Fix HiDiffusion ratio handling --- .../src/content/docs/features/hidiffusion.mdx | 7 +- invokeai/app/invocations/denoise_latents.py | 10 +- invokeai/app/invocations/metadata_linked.py | 2 +- invokeai/backend/hidiffusion/hidiffusion.py | 91 +++++++++++-------- .../stable_diffusion/hidiffusion_utils.py | 36 +------- invokeai/frontend/web/openapi.json | 80 ++++++++++------ invokeai/frontend/web/public/locales/en.json | 10 +- .../InformationalPopover/constants.ts | 2 + .../controlLayers/store/paramsSlice.test.ts | 68 +++++++++----- .../controlLayers/store/paramsSlice.ts | 13 ++- .../src/features/controlLayers/store/types.ts | 12 +-- .../web/src/features/metadata/parsing.test.ts | 27 +++++- .../web/src/features/metadata/parsing.tsx | 17 ++-- .../util/graph/generation/addSDXLRefiner.ts | 4 +- .../generation/buildHiDiffusionGraph.test.ts | 21 ++++- .../util/graph/generation/buildSD1Graph.ts | 4 +- .../util/graph/generation/buildSDXLGraph.ts | 4 +- .../Advanced/ParamHiDiffusionToggle.tsx | 66 ++++++++++++-- .../AdvancedSettingsAccordion.tsx | 2 + .../frontend/web/src/services/api/schema.ts | 16 ++-- .../test_hidiffusion_utils.py | 81 +++++++++++++++-- 21 files changed, 402 insertions(+), 171 deletions(-) diff --git a/docs/src/content/docs/features/hidiffusion.mdx b/docs/src/content/docs/features/hidiffusion.mdx index 8158f4980a1..1a8a35b276b 100644 --- a/docs/src/content/docs/features/hidiffusion.mdx +++ b/docs/src/content/docs/features/hidiffusion.mdx @@ -17,6 +17,7 @@ Learn more: https://github.com/megvii-research/HiDiffusion 3. In the **Advanced** grid, enable **HiDiffusion** and optionally adjust the two sub‑toggles and ratios: - **HiDiffusion: RAU‑Net** - **HiDiffusion: Window Attention** + - **HiDiffusion: Automatic Ratios** - **HiDiffusion: T1 Ratio** - **HiDiffusion: T2 Ratio** @@ -28,9 +29,11 @@ Learn more: https://github.com/megvii-research/HiDiffusion - **HiDiffusion: Window Attention**: Enables windowed attention blocks. This can boost local texture/detail, but may slightly affect global coherence in some prompts. -- **HiDiffusion: T1 Ratio**: Controls when HiDiffusion switches into its mid‑stage behavior. Lower values switch earlier; higher values preserve global structure longer. +- **HiDiffusion: Automatic Ratios**: Uses HiDiffusion's model- and resolution-specific T1/T2 presets. Disable it to enter manual overrides. -- **HiDiffusion: T2 Ratio**: Controls when HiDiffusion switches into its late‑stage behavior. Higher values keep window attention active longer and can sharpen local detail. +- **HiDiffusion: T1 Ratio**: Controls the first RAU-Net switching threshold. Lower values switch earlier; higher values keep the first resolution-aware stage active longer. + +- **HiDiffusion: T2 Ratio**: Controls the second RAU-Net switching threshold used for extreme-resolution generation. It does not control window attention. Higher values keep the second resolution-aware stage active longer. ## Tips diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index 2d48dd87607..3682d80145a 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -133,7 +133,7 @@ def get_scheduler( title="Denoise - SD1.5, SDXL", tags=["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"], category="latents", - version="1.6.0", + version="1.7.0", ) class DenoiseLatentsInvocation(BaseInvocation): """Denoises noisy latents to decodable images""" @@ -209,15 +209,15 @@ class DenoiseLatentsInvocation(BaseInvocation): description=FieldDescriptions.hidiffusion_window_attn, title="HiDiffusion: Window Attention", ) - hidiffusion_t1_ratio: float = InputField( - default=0.4, + hidiffusion_t1_ratio: Optional[float] = InputField( + default=None, ge=0, le=1, description=FieldDescriptions.hidiffusion_t1_ratio, title="HiDiffusion: T1 Ratio", ) - hidiffusion_t2_ratio: float = InputField( - default=0.0, + hidiffusion_t2_ratio: Optional[float] = InputField( + default=None, ge=0, le=1, description=FieldDescriptions.hidiffusion_t2_ratio, diff --git a/invokeai/app/invocations/metadata_linked.py b/invokeai/app/invocations/metadata_linked.py index 3ee70440436..d76695d4413 100644 --- a/invokeai/app/invocations/metadata_linked.py +++ b/invokeai/app/invocations/metadata_linked.py @@ -622,7 +622,7 @@ class LatentsMetaOutput(LatentsOutput, MetadataOutput): title=f"{DenoiseLatentsInvocation.UIConfig.title} + Metadata", tags=["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"], category="metadata", - version="1.2.0", + version="1.3.0", ) class DenoiseLatentsMetaInvocation(DenoiseLatentsInvocation, WithMetadata): def invoke(self, context: InvocationContext) -> LatentsMetaOutput: diff --git a/invokeai/backend/hidiffusion/hidiffusion.py b/invokeai/backend/hidiffusion/hidiffusion.py index 33437853214..5c7d83b84ee 100644 --- a/invokeai/backend/hidiffusion/hidiffusion.py +++ b/invokeai/backend/hidiffusion/hidiffusion.py @@ -136,6 +136,14 @@ def _get_max_timesteps(info_dict: dict) -> int: return len(pipeline.scheduler.timesteps) +def _get_switching_threshold_ratio(module: torch.nn.Module, presets: dict, preset_key: str) -> float: + """Resolve a threshold ratio without mutating process-global presets.""" + override = module.info["switching_threshold_overrides"].get(module.switching_threshold_ratio) + if override is not None: + return override + return presets[preset_key][module.switching_threshold_ratio] + + def make_diffusers_sdxl_controlnet_ppl(block_class): class sdxl_controlnet_ppl(block_class): # Save for unpatching later @@ -1559,17 +1567,19 @@ def forward( ori_H, ori_W = self.info["size"] if self.model == "sd15": if ori_H < 256 or ori_W < 256: - self.T1_ratio = switching_threshold_ratio_dict["sd15_1024"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_1024") else: - self.T1_ratio = switching_threshold_ratio_dict["sd15_2048"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_2048") elif self.model == "sdxl": if ori_H < 512 or ori_W < 512: if self.info["text_to_img_controlnet"]: - self.T1_ratio = text_to_img_controlnet_switching_threshold_ratio_dict["sdxl_2048"][ - self.switching_threshold_ratio - ] + self.T1_ratio = _get_switching_threshold_ratio( + self, text_to_img_controlnet_switching_threshold_ratio_dict, "sdxl_2048" + ) else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio( + self, switching_threshold_ratio_dict, "sdxl_2048" + ) if self.info["is_inpainting_task"]: self.aggressive_raunet = inpainting_is_aggressive_raunet @@ -1578,13 +1588,13 @@ def forward( else: self.aggressive_raunet = is_aggressive_raunet else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_4096"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_4096") elif self.model == "sdxl_turbo": - self.T1_ratio = switching_threshold_ratio_dict["sdxl_turbo_1024"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_turbo_1024") else: raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") - if self.aggressive_raunet: + if self.aggressive_raunet and self.switching_threshold_ratio == "T1_ratio": # self.T1_start = min(int(self.max_timestep * self.T1_ratio * 0.4), int(8/50 * self.max_timestep)) self.T1_start = int(aggressive_step / 50 * self.max_timestep) self.T1_end = int(self.max_timestep * self.T1_ratio) @@ -1693,17 +1703,19 @@ def forward( ori_H, ori_W = self.info["size"] if self.model == "sd15": if ori_H < 256 or ori_W < 256: - self.T1_ratio = switching_threshold_ratio_dict["sd15_1024"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_1024") else: - self.T1_ratio = switching_threshold_ratio_dict["sd15_2048"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_2048") elif self.model == "sdxl": if ori_H < 512 or ori_W < 512: if self.info["text_to_img_controlnet"]: - self.T1_ratio = text_to_img_controlnet_switching_threshold_ratio_dict["sdxl_2048"][ - self.switching_threshold_ratio - ] + self.T1_ratio = _get_switching_threshold_ratio( + self, text_to_img_controlnet_switching_threshold_ratio_dict, "sdxl_2048" + ) else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio( + self, switching_threshold_ratio_dict, "sdxl_2048" + ) if self.info["is_inpainting_task"]: self.aggressive_raunet = inpainting_is_aggressive_raunet @@ -1713,13 +1725,13 @@ def forward( self.aggressive_raunet = is_aggressive_raunet else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_4096"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_4096") elif self.model == "sdxl_turbo": - self.T1_ratio = switching_threshold_ratio_dict["sdxl_turbo_1024"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_turbo_1024") else: raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") - if self.aggressive_raunet: + if self.aggressive_raunet and self.switching_threshold_ratio == "T1_ratio": # self.T1_start = min(int(self.max_timestep * self.T1_ratio * 0.4), int(8/50 * self.max_timestep)) self.T1_start = int(aggressive_step / 50 * self.max_timestep) self.T1_end = int(self.max_timestep * self.T1_ratio) @@ -1830,17 +1842,19 @@ def forward(self, hidden_states: torch.Tensor, scale=1.0) -> torch.Tensor: ori_H, ori_W = self.info["size"] if self.model == "sd15": if ori_H < 256 or ori_W < 256: - self.T1_ratio = switching_threshold_ratio_dict["sd15_1024"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_1024") else: - self.T1_ratio = switching_threshold_ratio_dict["sd15_2048"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_2048") elif self.model == "sdxl": if ori_H < 512 or ori_W < 512: if self.info["text_to_img_controlnet"]: - self.T1_ratio = text_to_img_controlnet_switching_threshold_ratio_dict["sdxl_2048"][ - self.switching_threshold_ratio - ] + self.T1_ratio = _get_switching_threshold_ratio( + self, text_to_img_controlnet_switching_threshold_ratio_dict, "sdxl_2048" + ) else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio( + self, switching_threshold_ratio_dict, "sdxl_2048" + ) if self.info["is_inpainting_task"]: self.aggressive_raunet = inpainting_is_aggressive_raunet @@ -1849,13 +1863,13 @@ def forward(self, hidden_states: torch.Tensor, scale=1.0) -> torch.Tensor: else: self.aggressive_raunet = is_aggressive_raunet else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_4096"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_4096") elif self.model == "sdxl_turbo": - self.T1_ratio = switching_threshold_ratio_dict["sdxl_turbo_1024"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_turbo_1024") else: raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") - if self.aggressive_raunet: + if self.aggressive_raunet and self.switching_threshold_ratio == "T1_ratio": # self.T1 = min(int(self.max_timestep * self.T1_ratio), int(8/50 * self.max_timestep)) self.T1 = int(aggressive_step / 50 * self.max_timestep) else: @@ -1911,17 +1925,19 @@ def forward(self, hidden_states: torch.Tensor, scale=1.0) -> torch.Tensor: ori_H, ori_W = self.info["size"] if self.model == "sd15": if ori_H < 256 or ori_W < 256: - self.T1_ratio = switching_threshold_ratio_dict["sd15_1024"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_1024") else: - self.T1_ratio = switching_threshold_ratio_dict["sd15_2048"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_2048") elif self.model == "sdxl": if ori_H < 512 or ori_W < 512: if self.info["text_to_img_controlnet"]: - self.T1_ratio = text_to_img_controlnet_switching_threshold_ratio_dict["sdxl_2048"][ - self.switching_threshold_ratio - ] + self.T1_ratio = _get_switching_threshold_ratio( + self, text_to_img_controlnet_switching_threshold_ratio_dict, "sdxl_2048" + ) else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio( + self, switching_threshold_ratio_dict, "sdxl_2048" + ) if self.info["is_inpainting_task"]: self.aggressive_raunet = inpainting_is_aggressive_raunet @@ -1930,13 +1946,13 @@ def forward(self, hidden_states: torch.Tensor, scale=1.0) -> torch.Tensor: else: self.aggressive_raunet = is_aggressive_raunet else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_4096"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_4096") elif self.model == "sdxl_turbo": - self.T1_ratio = switching_threshold_ratio_dict["sdxl_turbo_1024"][self.switching_threshold_ratio] + self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_turbo_1024") else: raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") - if self.aggressive_raunet: + if self.aggressive_raunet and self.switching_threshold_ratio == "T1_ratio": # self.T1 = min(int(self.max_timestep * self.T1_ratio), int(8/50 * self.max_timestep)) self.T1 = int(aggressive_step / 50 * self.max_timestep) else: @@ -2045,6 +2061,8 @@ def apply_hidiffusion( generator: torch.Generator | None = None, has_controlnet: bool = False, is_controlnet_text_to_image: bool = False, + t1_ratio: float | None = None, + t2_ratio: float | None = None, ): """ model: diffusers model. We support SD 1.5, 2.1, XL, XL Turbo. @@ -2128,6 +2146,7 @@ def apply_hidiffusion( "is_inpainting_task": model.__class__ in auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING.values(), "is_playground": is_playground, "pipeline": model, + "switching_threshold_overrides": {"T1_ratio": t1_ratio, "T2_ratio": t2_ratio}, } model.info = diffusion_model.info hook_diffusion_model(diffusion_model) diff --git a/invokeai/backend/stable_diffusion/hidiffusion_utils.py b/invokeai/backend/stable_diffusion/hidiffusion_utils.py index f6e6e1681b3..fa83cf8bc67 100644 --- a/invokeai/backend/stable_diffusion/hidiffusion_utils.py +++ b/invokeai/backend/stable_diffusion/hidiffusion_utils.py @@ -1,19 +1,11 @@ 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( @@ -70,14 +62,6 @@ def _set_name_or_path_on_config(cfg, value: str) -> bool: original_num_upsamplers = getattr(target, "num_upsamplers", None) - ratio_overrides = None - ratio_dicts = None - if t1_ratio is not None or t2_ratio is not None: - 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: @@ -103,23 +87,12 @@ def _set_name_or_path_on_config(cfg, value: str) -> bool: 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, + t1_ratio=t1_ratio, + t2_ratio=t2_ratio, has_controlnet=has_controlnet, is_controlnet_text_to_image=is_controlnet_text_to_image, generator=generator, @@ -133,11 +106,6 @@ def _apply_ratio_overrides(ratio_dict: dict) -> None: 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: diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 63051f90f00..5d227fb87d6 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -23986,28 +23986,42 @@ "type": "boolean" }, "hidiffusion_t1_ratio": { - "default": 0.4, + "anyOf": [ + { + "maximum": 1, + "minimum": 0, + "type": "number" + }, + { + "type": "null" + } + ], + "default": null, "description": "Override HiDiffusion early switch threshold (T1 ratio)", "field_kind": "input", "input": "any", - "maximum": 1, - "minimum": 0, - "orig_default": 0.4, + "orig_default": null, "orig_required": false, - "title": "HiDiffusion: T1 Ratio", - "type": "number" + "title": "HiDiffusion: T1 Ratio" }, "hidiffusion_t2_ratio": { - "default": 0.0, + "anyOf": [ + { + "maximum": 1, + "minimum": 0, + "type": "number" + }, + { + "type": "null" + } + ], + "default": null, "description": "Override HiDiffusion late switch threshold (T2 ratio)", "field_kind": "input", "input": "any", - "maximum": 1, - "minimum": 0, - "orig_default": 0.0, + "orig_default": null, "orig_required": false, - "title": "HiDiffusion: T2 Ratio", - "type": "number" + "title": "HiDiffusion: T2 Ratio" }, "latents": { "anyOf": [ @@ -24055,7 +24069,7 @@ "tags": ["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"], "title": "Denoise - SD1.5, SDXL", "type": "object", - "version": "1.6.0", + "version": "1.7.0", "output": { "$ref": "#/components/schemas/LatentsOutput" } @@ -24399,28 +24413,42 @@ "type": "boolean" }, "hidiffusion_t1_ratio": { - "default": 0.4, + "anyOf": [ + { + "maximum": 1, + "minimum": 0, + "type": "number" + }, + { + "type": "null" + } + ], + "default": null, "description": "Override HiDiffusion early switch threshold (T1 ratio)", "field_kind": "input", "input": "any", - "maximum": 1, - "minimum": 0, - "orig_default": 0.4, + "orig_default": null, "orig_required": false, - "title": "HiDiffusion: T1 Ratio", - "type": "number" + "title": "HiDiffusion: T1 Ratio" }, "hidiffusion_t2_ratio": { - "default": 0.0, + "anyOf": [ + { + "maximum": 1, + "minimum": 0, + "type": "number" + }, + { + "type": "null" + } + ], + "default": null, "description": "Override HiDiffusion late switch threshold (T2 ratio)", "field_kind": "input", "input": "any", - "maximum": 1, - "minimum": 0, - "orig_default": 0.0, + "orig_default": null, "orig_required": false, - "title": "HiDiffusion: T2 Ratio", - "type": "number" + "title": "HiDiffusion: T2 Ratio" }, "latents": { "anyOf": [ @@ -24468,7 +24496,7 @@ "tags": ["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"], "title": "Denoise - SD1.5, SDXL + Metadata", "type": "object", - "version": "1.2.0", + "version": "1.3.0", "output": { "$ref": "#/components/schemas/LatentsMetaOutput" } diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json index b7dbb0a82d7..2c04b120a33 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -1071,6 +1071,7 @@ "hiDiffusion": "HiDiffusion", "hiDiffusionRauNet": "HiDiffusion: RAU-Net", "hiDiffusionWindowAttn": "HiDiffusion: Window Attention", + "hiDiffusionRatiosAuto": "HiDiffusion: Automatic Ratios", "hiDiffusionT1Ratio": "HiDiffusion: T1 Ratio", "hiDiffusionT2Ratio": "HiDiffusion: T2 Ratio", "imageDetails": "Image Details", @@ -2173,6 +2174,13 @@ "Can boost local detail, but may affect global coherence." ] }, + "hidiffusionRatiosAuto": { + "heading": "HiDiffusion: Automatic Ratios", + "paragraphs": [ + "Uses the model- and resolution-specific T1 and T2 ratios recommended by HiDiffusion.", + "Disable this to set explicit ratio overrides below." + ] + }, "hidiffusionT1Ratio": { "heading": "HiDiffusion: T1 Ratio", "paragraphs": [ @@ -2184,7 +2192,7 @@ "heading": "HiDiffusion: T2 Ratio", "paragraphs": [ "Controls the late switch point for HiDiffusion (T2).", - "Higher values keep window attention active longer." + "Higher values keep the second RAU-Net resolution-aware stage active longer." ] }, "clipSkip": { diff --git a/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts b/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts index d91077f6607..edbe78ddab2 100644 --- a/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts +++ b/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts @@ -78,6 +78,7 @@ export type Feature = | 'hidiffusion' | 'hidiffusionRauNet' | 'hidiffusionWindowAttn' + | 'hidiffusionRatiosAuto' | 'hidiffusionT1Ratio' | 'hidiffusionT2Ratio' | 'colorCompensation' @@ -256,6 +257,7 @@ export const POPOVER_DATA: { [key in Feature]?: PopoverData } = { }, hidiffusionRauNet: {}, hidiffusionWindowAttn: {}, + hidiffusionRatiosAuto: {}, hidiffusionT1Ratio: {}, hidiffusionT2Ratio: {}, } as const; diff --git a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.test.ts b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.test.ts index e2a2713aa41..30c2690c72a 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.test.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.test.ts @@ -172,8 +172,9 @@ describe('paramsSlice selectors for external models', () => { * `animaT5EncoderModel`, since removed from the schema. * - v4: the narrowest v4 blob is not a release at all — it is the one written by the build that did * the bump, `1aeb05bbf0` (97 keys). Releases writing v4 start at v6.14.0-rc1. - * - v5: the current version, reached by the FLUX.2 [dev] merge `f10d2a4f5a`, which is also the - * build that wrote the narrowest v5 blob. Pinning the fixture at the bump commit is what keeps + * - v5: reached by the FLUX.2 [dev] merge `f10d2a4f5a`, which is also the build that wrote the + * narrowest v5 blob. The current v6 migration converts legacy HiDiffusion defaults to Auto. + * Pinning the fixture at the v5 bump commit is what keeps * the invariant below meaningful for the current tier: the version steps can never cover it (a * v5 blob matches no branch), so every key added since the bump has to carry a zod default, and * this entry is what proves it does. @@ -900,14 +901,14 @@ describe('paramsSliceConfig persisted state migration', () => { // v2 migrates all the way through the current chain (v2 -> v3 adds Qwen fields, // v3 -> v4 adds Krea-2 and PiD fields). - expect(result._version).toBe(5); + expect(result._version).toBe(6); 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); + expect(result.hiDiffusionT1Ratio).toBeNull(); + expect(result.hiDiffusionT2Ratio).toBeNull(); // Existing params should be preserved expect(result.positivePrompt).toBe('a fluffy cat'); expect(result.seed).toBe(42); @@ -916,6 +917,29 @@ describe('paramsSliceConfig persisted state migration', () => { expect(result.dimensions.height).toBe(768); }); + it('migrates old HiDiffusion defaults to automatic ratios and preserves custom overrides', () => { + expect(migrate).toBeDefined(); + const initial = getInitialParamsState(); + + const oldDefaults = migrate?.({ + ...initial, + _version: 5, + hiDiffusionT1Ratio: 0.4, + hiDiffusionT2Ratio: 0.0, + }) as ReturnType; + const customOverrides = migrate?.({ + ...initial, + _version: 5, + hiDiffusionT1Ratio: 0.65, + hiDiffusionT2Ratio: 0.25, + }) as ReturnType; + + expect(oldDefaults.hiDiffusionT1Ratio).toBeNull(); + expect(oldDefaults.hiDiffusionT2Ratio).toBeNull(); + expect(customOverrides.hiDiffusionT1Ratio).toBe(0.65); + expect(customOverrides.hiDiffusionT2Ratio).toBe(0.25); + }); + it('merges the separate Klein / dev VAE slots into flux2VaeModel when migrating from v3', () => { expect(migrate).toBeDefined(); @@ -938,7 +962,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(v3State) as ReturnType & Record; - expect(result._version).toBe(5); + expect(result._version).toBe(6); expect((result.flux2VaeModel as { key: string } | null)?.key).toBe('klein-vae'); // The new standalone dev Mistral encoder slot must be seeded, not left undefined. expect(result.flux2DevMistralEncoderModel).toBeNull(); @@ -972,7 +996,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(v3State) as ReturnType; - expect(result._version).toBe(5); + expect(result._version).toBe(6); expect(result.krea2VaeModel).toBeNull(); expect(result.krea2Qwen3VlEncoderModel).toBeNull(); expect(result.krea2SeedVarianceEnabled).toBe(false); @@ -1006,7 +1030,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(mainV4State) as ReturnType & Record; - expect(result._version).toBe(5); + expect(result._version).toBe(6); expect((result.flux2VaeModel as { key: string } | null)?.key).toBe('klein-vae'); expect(result.flux2DevMistralEncoderModel).toBeNull(); // main's own v4 values must survive untouched. @@ -1034,7 +1058,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(devV4State) as ReturnType & Record; - expect(result._version).toBe(5); + expect(result._version).toBe(6); // The branch's own v4 values must survive untouched. expect((result.flux2VaeModel as { key: string } | null)?.key).toBe('flux2-vae'); expect(result.pidMode).toBe('off'); @@ -1091,7 +1115,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(blob) as ReturnType; - expect(result._version).toBe(5); + expect(result._version).toBe(6); expect(result.positivePrompt).toBe('a fluffy cat'); expect(result.seed).toBe(42); expect(result.shouldRandomizeSeed).toBe(false); @@ -1139,7 +1163,7 @@ describe('paramsSliceConfig persisted state migration', () => { expect( backfilled, - version === getInitialParamsState()._version + Number(version) === getInitialParamsState()._version ? `Keys missing from a blob written at ${release}, the commit that bumped _version to ${version}. ` + `A blob already at the current version matches no branch in the migration chain, so no step can ` + `seed these — each needs a zod default, or upgrading throws in zParamsState.parse() and wipes ` + @@ -1257,7 +1281,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(blob) as ReturnType; - expect(result._version).toBe(5); + expect(result._version).toBe(6); expect(result.dimensions).toEqual(getInitialParamsState().dimensions); expect(result.positivePrompt).toBe('a fluffy cat'); expect(result.seed).toBe(7); @@ -1274,7 +1298,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(blob) as ReturnType; - expect(result._version).toBe(5); + expect(result._version).toBe(6); expect(result.positivePromptHistory).toEqual([]); expect(result.qwenImageVaeModel).toBeNull(); expect(result.wanVaeModel).toBeNull(); @@ -1284,18 +1308,18 @@ describe('paramsSliceConfig persisted state migration', () => { it('never repairs _version, so version detection cannot be bypassed', () => { // `_version` is the input to the version steps, so the net must leave it alone. If it repaired // it, any blob whose version is not the current literal — including one written by a *newer* - // build — would be silently stamped v5 having run no step, and its stale field values would be + // build — would be silently stamped with the current version having run no step, and its stale field values would be // accepted as current. Deliberately not routed through migrate(): the version steps normalise // `_version` before the net ever sees it, so only calling the net directly tests the guard. // The blob is otherwise complete (the current tier's key set), so `_version` is the only thing // the parse below can object to. - const blob = buildReleaseBlob('f10d2a4f5a', { _version: 6, positivePrompt: 'a fluffy cat' }); + const blob = buildReleaseBlob('f10d2a4f5a', { _version: 7, positivePrompt: 'a fluffy cat' }); const { backfilled, reset } = repairParamsState(blob); expect(backfilled).toEqual([]); expect(reset).toEqual([]); - expect(blob._version).toBe(6); + expect(blob._version).toBe(7); // Still fatal, which is the correct outcome for a downgrade: that slice really was written by a // schema this build does not know. expect(() => zParamsState.parse(blob)).toThrow(); @@ -1311,7 +1335,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(blob) as ReturnType; - expect(result._version).toBe(5); + expect(result._version).toBe(6); expect(result.positivePrompt).toBe('a fluffy cat'); expect(result.seed).toBe(7); expect(result.dimensions).toBeDefined(); @@ -1328,7 +1352,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(v3State) as ReturnType; - expect(result._version).toBe(5); + expect(result._version).toBe(6); expect(result.wanTransformerLowNoise).toBeNull(); expect(result.wanComponentSource).toBeNull(); expect(result.wanVaeModel).toBeNull(); @@ -1350,7 +1374,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(v2State) as ReturnType; - expect(result._version).toBe(5); + expect(result._version).toBe(6); expect(result.fluxScheduler).toBe('euler'); expect(result.zImageScheduler).toBe('euler'); expect(result.colorCompensation).toBe(false); @@ -1410,7 +1434,7 @@ describe('paramsSliceConfig persisted state migration', () => { expect('hiDiffusionEnabled' in blob).toBe(false); applyParamsVersionMigrations(blob); - expect(blob._version).toBe(5); + expect(blob._version).toBe(6); // The value assertions below cannot, on their own, prove the defaults exist: three mechanisms // produce the identical values, so any two can hide the third being reverted. Parsing directly @@ -1435,8 +1459,8 @@ describe('paramsSliceConfig persisted state migration', () => { 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); + expect(result.hiDiffusionT1Ratio).toBeNull(); + expect(result.hiDiffusionT2Ratio).toBeNull(); expect(result.positivePrompt).toBe('a fluffy cat'); expect(result.seed).toBe(42); }); diff --git a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts index f7c8420dade..3db8e6b356d 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts @@ -182,10 +182,10 @@ const slice = createSlice({ setHiDiffusionWindowAttnEnabled: (state, action: PayloadAction) => { state.hiDiffusionWindowAttnEnabled = action.payload; }, - setHiDiffusionT1Ratio: (state, action: PayloadAction) => { + setHiDiffusionT1Ratio: (state, action: PayloadAction) => { state.hiDiffusionT1Ratio = action.payload; }, - setHiDiffusionT2Ratio: (state, action: PayloadAction) => { + setHiDiffusionT2Ratio: (state, action: PayloadAction) => { state.hiDiffusionT2Ratio = action.payload; }, setSeamlessXAxis: (state, action: PayloadAction) => { @@ -1162,6 +1162,15 @@ export const applyParamsVersionMigrations = (state: any): void => { state.gemma2EncoderModel = state.gemma2EncoderModel ?? null; state.pidSteps = state.pidSteps ?? 4; } + + if (state._version === 5) { + // v5 -> v6: numeric HiDiffusion defaults unintentionally overrode the library's + // resolution-aware presets. Treat the old default values as automatic thresholds while + // preserving values that users explicitly changed. + state._version = 6; + state.hiDiffusionT1Ratio = state.hiDiffusionT1Ratio === 0.4 ? null : state.hiDiffusionT1Ratio; + state.hiDiffusionT2Ratio = state.hiDiffusionT2Ratio === 0.0 ? null : state.hiDiffusionT2Ratio; + } }; export const paramsSliceConfig: SliceConfig = { diff --git a/invokeai/frontend/web/src/features/controlLayers/store/types.ts b/invokeai/frontend/web/src/features/controlLayers/store/types.ts index 2143e58f997..b6fff1580f7 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/types.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/types.ts @@ -817,7 +817,7 @@ const zPidMode = z.enum(['off', 'fit', 'native']); export type PidMode = z.infer; export const zParamsState = z.object({ - _version: z.literal(5), + _version: z.literal(6), maskBlur: z.number(), maskBlurMethod: zParameterMaskBlurMethod, canvasCoherenceMode: zParameterCanvasCoherenceMode, @@ -839,8 +839,8 @@ export const zParamsState = z.object({ hiDiffusionEnabled: z.boolean().default(false), hiDiffusionRauNetEnabled: z.boolean().default(true), hiDiffusionWindowAttnEnabled: z.boolean().default(true), - hiDiffusionT1Ratio: z.number().default(0.4), - hiDiffusionT2Ratio: z.number().default(0.0), + hiDiffusionT1Ratio: z.number().nullable().default(null), + hiDiffusionT2Ratio: z.number().nullable().default(null), iterations: z.number(), scheduler: zParameterScheduler, fluxScheduler: zParameterFluxScheduler, @@ -965,7 +965,7 @@ export const zParamsState = z.object({ }); export type ParamsState = z.infer; export const getInitialParamsState = (): ParamsState => ({ - _version: 5, + _version: 6, maskBlur: 16, maskBlurMethod: 'box', canvasCoherenceMode: 'Gaussian Blur', @@ -983,8 +983,8 @@ export const getInitialParamsState = (): ParamsState => ({ hiDiffusionEnabled: false, hiDiffusionRauNetEnabled: true, hiDiffusionWindowAttnEnabled: true, - hiDiffusionT1Ratio: 0.4, - hiDiffusionT2Ratio: 0.0, + hiDiffusionT1Ratio: null, + hiDiffusionT2Ratio: null, iterations: 1, scheduler: 'dpmpp_3m_k', fluxScheduler: 'euler', diff --git a/invokeai/frontend/web/src/features/metadata/parsing.test.ts b/invokeai/frontend/web/src/features/metadata/parsing.test.ts index 6072aa1c10f..12c8a96fa9e 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.test.ts +++ b/invokeai/frontend/web/src/features/metadata/parsing.test.ts @@ -1,5 +1,9 @@ import type { AppStore } from 'app/store/store'; -import { setHiDiffusionEnabled } from 'features/controlLayers/store/paramsSlice'; +import { + setHiDiffusionEnabled, + setHiDiffusionT1Ratio, + setHiDiffusionT2Ratio, +} from 'features/controlLayers/store/paramsSlice'; import { describe, expect, it, vi } from 'vitest'; import { ImageMetadataHandlers, MetadataUtils, parseMetadataHandler } from './parsing'; @@ -163,12 +167,31 @@ describe('Qwen metadata parsing', () => { }); describe('HiDiffusion metadata parsing', () => { + it('recalls null ratios as automatic thresholds', async () => { + const store = createStore(); + const metadata = { hidiffusion_t1_ratio: null, hidiffusion_t2_ratio: null }; + + const t1 = await parseMetadataHandler(metadata, ImageMetadataHandlers.HiDiffusionT1Ratio, store); + const t2 = await parseMetadataHandler(metadata, ImageMetadataHandlers.HiDiffusionT2Ratio, store); + ImageMetadataHandlers.HiDiffusionT1Ratio.recall(t1, store); + ImageMetadataHandlers.HiDiffusionT2Ratio.recall(t2, store); + + expect(store.dispatch).toHaveBeenCalledWith(setHiDiffusionT1Ratio(null)); + expect(store.dispatch).toHaveBeenCalledWith(setHiDiffusionT2Ratio(null)); + }); + it('disables HiDiffusion when recalling all metadata from an older image', async () => { let hiDiffusionEnabled = true; + let hiDiffusionT1Ratio: number | null = 0.8; + let hiDiffusionT2Ratio: number | null = 0.6; const store = { dispatch: vi.fn((action) => { if (action.type === setHiDiffusionEnabled.type) { hiDiffusionEnabled = action.payload; + } else if (action.type === setHiDiffusionT1Ratio.type) { + hiDiffusionT1Ratio = action.payload; + } else if (action.type === setHiDiffusionT2Ratio.type) { + hiDiffusionT2Ratio = action.payload; } return action; }), @@ -193,5 +216,7 @@ describe('HiDiffusion metadata parsing', () => { expect(store.dispatch).toHaveBeenCalledWith(setHiDiffusionEnabled(false)); expect(hiDiffusionEnabled).toBe(false); + expect(hiDiffusionT1Ratio).toBeNull(); + expect(hiDiffusionT2Ratio).toBeNull(); }); }); diff --git a/invokeai/frontend/web/src/features/metadata/parsing.tsx b/invokeai/frontend/web/src/features/metadata/parsing.tsx index 86ca976a192..95ce69a51b9 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.tsx +++ b/invokeai/frontend/web/src/features/metadata/parsing.tsx @@ -773,12 +773,17 @@ const HiDiffusionWindowAttn: SingleMetadataHandler = { //#endregion HiDiffusionWindowAttn //#region HiDiffusionT1Ratio -const HiDiffusionT1Ratio: SingleMetadataHandler = { +const HiDiffusionRatioValue = ({ value }: SingleMetadataValueProps) => { + const { t } = useTranslation(); + return ; +}; + +const HiDiffusionT1Ratio: SingleMetadataHandler = { [SingleMetadataKey]: true, type: 'HiDiffusionT1Ratio', parse: (metadata, _store) => { const raw = getProperty(metadata, 'hidiffusion_t1_ratio'); - const parsed = z.number().parse(raw); + const parsed = raw === undefined ? null : z.number().nullable().parse(raw); return Promise.resolve(parsed); }, recall: (value, store) => { @@ -786,17 +791,17 @@ const HiDiffusionT1Ratio: SingleMetadataHandler = { }, i18nKey: 'metadata.hiDiffusionT1Ratio', LabelComponent: MetadataLabel, - ValueComponent: ({ value }: SingleMetadataValueProps) => , + ValueComponent: HiDiffusionRatioValue, }; //#endregion HiDiffusionT1Ratio //#region HiDiffusionT2Ratio -const HiDiffusionT2Ratio: SingleMetadataHandler = { +const HiDiffusionT2Ratio: SingleMetadataHandler = { [SingleMetadataKey]: true, type: 'HiDiffusionT2Ratio', parse: (metadata, _store) => { const raw = getProperty(metadata, 'hidiffusion_t2_ratio'); - const parsed = z.number().parse(raw); + const parsed = raw === undefined ? null : z.number().nullable().parse(raw); return Promise.resolve(parsed); }, recall: (value, store) => { @@ -804,7 +809,7 @@ const HiDiffusionT2Ratio: SingleMetadataHandler = { }, i18nKey: 'metadata.hiDiffusionT2Ratio', LabelComponent: MetadataLabel, - ValueComponent: ({ value }: SingleMetadataValueProps) => , + ValueComponent: HiDiffusionRatioValue, }; //#endregion HiDiffusionT2Ratio diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.ts index 11a5333c7d2..2f755377df9 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.ts @@ -65,8 +65,8 @@ export const addSDXLRefiner = async ( hidiffusion: hiDiffusionEnabled, hidiffusion_raunet: hiDiffusionRauNetEnabled, hidiffusion_window_attn: hiDiffusionWindowAttnEnabled, - hidiffusion_t1_ratio: hiDiffusionEnabled ? hiDiffusionT1Ratio : undefined, - hidiffusion_t2_ratio: hiDiffusionEnabled ? hiDiffusionT2Ratio : undefined, + hidiffusion_t1_ratio: hiDiffusionEnabled ? (hiDiffusionT1Ratio ?? undefined) : undefined, + hidiffusion_t2_ratio: hiDiffusionEnabled ? (hiDiffusionT2Ratio ?? undefined) : undefined, denoising_start: refinerStart, denoising_end: 1, }); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildHiDiffusionGraph.test.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildHiDiffusionGraph.test.ts index 8853af97baa..6ed9b077a82 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildHiDiffusionGraph.test.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildHiDiffusionGraph.test.ts @@ -32,8 +32,8 @@ const defaultParams = { cfgRescaleMultiplier: 0, hiDiffusionEnabled: false, hiDiffusionRauNetEnabled: false, - hiDiffusionT1Ratio: 0.25, - hiDiffusionT2Ratio: 0.1, + hiDiffusionT1Ratio: 0.25 as number | null, + hiDiffusionT2Ratio: 0.1 as number | null, hiDiffusionWindowAttnEnabled: false, scheduler: 'euler', steps: 20, @@ -193,4 +193,21 @@ describe('HiDiffusion graph metadata', () => { expect(metadata.hidiffusion_t1_ratio).toBe(0.25); expect(metadata.hidiffusion_t2_ratio).toBe(0.1); }); + + it('omits automatic ratio overrides from the SDXL denoise and metadata nodes', async () => { + currentModel = sdxlModel; + params.hiDiffusionEnabled = true; + params.hiDiffusionRauNetEnabled = true; + params.hiDiffusionT1Ratio = null; + params.hiDiffusionT2Ratio = null; + + const { g } = await buildSDXLGraph(buildGraphArg()); + const denoise = g.getNodes().find((node) => node.type === 'denoise_latents'); + const metadata = getMetadata(g); + + expect(denoise?.hidiffusion_t1_ratio).toBeUndefined(); + expect(denoise?.hidiffusion_t2_ratio).toBeUndefined(); + expect(metadata.hidiffusion_t1_ratio).toBeNull(); + expect(metadata.hidiffusion_t2_ratio).toBeNull(); + }); }); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts index 2a8d15b1e75..8ac62a08885 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts @@ -109,8 +109,8 @@ export const buildSD1Graph = async (arg: GraphBuilderArg): Promise { + const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled); + const hiDiffusionRauNetEnabled = useAppSelector(selectHiDiffusionRauNetEnabled); + const hiDiffusionT1Ratio = useAppSelector(selectHiDiffusionT1Ratio); + const hiDiffusionT2Ratio = useAppSelector(selectHiDiffusionT2Ratio); + const dispatch = useAppDispatch(); + const { t } = useTranslation(); + const isAutomatic = hiDiffusionT1Ratio === null && hiDiffusionT2Ratio === null; + + const onChange = useCallback( + (event: ChangeEvent) => { + if (event.target.checked) { + dispatch(setHiDiffusionT1Ratio(null)); + dispatch(setHiDiffusionT2Ratio(null)); + } else { + dispatch(setHiDiffusionT1Ratio(RATIO_CONSTRAINTS.t1.initial)); + dispatch(setHiDiffusionT2Ratio(RATIO_CONSTRAINTS.t2.initial)); + } + }, + [dispatch] + ); + + return ( + + + + {t('parameters.hiDiffusionRatiosAuto')} + + + + + ); +}); + +ParamHiDiffusionAutoRatiosToggle.displayName = 'ParamHiDiffusionAutoRatiosToggle'; + export const ParamHiDiffusionT1Ratio = memo(() => { const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled); + const hiDiffusionRauNetEnabled = useAppSelector(selectHiDiffusionRauNetEnabled); const hiDiffusionT1Ratio = useAppSelector(selectHiDiffusionT1Ratio); + const hiDiffusionT2Ratio = useAppSelector(selectHiDiffusionT2Ratio); const dispatch = useAppDispatch(); const { t } = useTranslation(); const onChange = useCallback((value: number) => dispatch(setHiDiffusionT1Ratio(value)), [dispatch]); return ( - + {t('parameters.hiDiffusionT1Ratio')} { marks /> { const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled); + const hiDiffusionRauNetEnabled = useAppSelector(selectHiDiffusionRauNetEnabled); + const hiDiffusionT1Ratio = useAppSelector(selectHiDiffusionT1Ratio); const hiDiffusionT2Ratio = useAppSelector(selectHiDiffusionT2Ratio); const dispatch = useAppDispatch(); const { t } = useTranslation(); @@ -167,14 +216,19 @@ export const ParamHiDiffusionT2Ratio = memo(() => { const onChange = useCallback((value: number) => dispatch(setHiDiffusionT2Ratio(value)), [dispatch]); return ( - + {t('parameters.hiDiffusionT2Ratio')} { marks /> { + diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index f05af125fe6..a94da660a5e 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -9117,15 +9117,15 @@ export type components = { /** * HiDiffusion: T1 Ratio * @description Override HiDiffusion early switch threshold (T1 ratio) - * @default 0.4 + * @default null */ - hidiffusion_t1_ratio?: number; + hidiffusion_t1_ratio?: number | null; /** * HiDiffusion: T2 Ratio * @description Override HiDiffusion late switch threshold (T2 ratio) - * @default 0 + * @default null */ - hidiffusion_t2_ratio?: number; + hidiffusion_t2_ratio?: number | null; /** * @description Latents tensor * @default null @@ -9265,15 +9265,15 @@ export type components = { /** * HiDiffusion: T1 Ratio * @description Override HiDiffusion early switch threshold (T1 ratio) - * @default 0.4 + * @default null */ - hidiffusion_t1_ratio?: number; + hidiffusion_t1_ratio?: number | null; /** * HiDiffusion: T2 Ratio * @description Override HiDiffusion late switch threshold (T2 ratio) - * @default 0 + * @default null */ - hidiffusion_t2_ratio?: number; + hidiffusion_t2_ratio?: number | null; /** * @description Latents tensor * @default null diff --git a/tests/backend/stable_diffusion/test_hidiffusion_utils.py b/tests/backend/stable_diffusion/test_hidiffusion_utils.py index 5f8619a4882..b46930a9d90 100644 --- a/tests/backend/stable_diffusion/test_hidiffusion_utils.py +++ b/tests/backend/stable_diffusion/test_hidiffusion_utils.py @@ -7,6 +7,7 @@ from invokeai.backend.hidiffusion.hidiffusion import ( _resize_controlnet_residual, + make_diffusers_downsampler_block, switching_threshold_ratio_dict, text_to_img_controlnet_switching_threshold_ratio_dict, ) @@ -256,16 +257,14 @@ def test_hidiffusion_patch_restores_state_when_apply_hidiffusion_raises(): ) hook = MagicMock() - def fake_apply_hidiffusion(patched_model, **_kwargs): + 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 + assert kwargs["t1_ratio"] == 0.25 + assert kwargs["t2_ratio"] == 0.1 + assert switching_threshold_ratio_dict == original_switching + assert text_to_img_controlnet_switching_threshold_ratio_dict == original_controlnet patched_model.unet.num_upsamplers = 99 patched_model.unet.layer.info = {"hooks": [hook]} @@ -369,3 +368,71 @@ def __getattr__(self, name): assert config._internal_dict["_name_or_path"] == "patched-model-name" assert "_name_or_path" not in config._internal_dict + + +def test_hidiffusion_ratio_overrides_are_isolated_between_overlapping_patches(): + original_switching = copy.deepcopy(switching_threshold_ratio_dict) + original_controlnet = copy.deepcopy(text_to_img_controlnet_switching_threshold_ratio_dict) + first_model = SimpleNamespace(unet=DummyUNet()) + second_model = SimpleNamespace(unet=DummyUNet()) + applied_overrides: list[tuple[object, float | None, float | None]] = [] + + def fake_apply_hidiffusion(model, **kwargs): + applied_overrides.append((model, kwargs["t1_ratio"], kwargs["t2_ratio"])) + + with ( + patch("invokeai.backend.hidiffusion.hidiffusion.apply_hidiffusion", side_effect=fake_apply_hidiffusion), + patch("invokeai.backend.hidiffusion.hidiffusion.remove_hidiffusion"), + ): + first_patch = hidiffusion_patch(first_model, name_or_path="first", t1_ratio=0.2, t2_ratio=0.1) + second_patch = hidiffusion_patch(second_model, name_or_path="second", t1_ratio=0.8, t2_ratio=0.9) + first_patch.__enter__() + second_patch.__enter__() + first_patch.__exit__(None, None, None) + second_patch.__exit__(None, None, None) + + assert applied_overrides == [(first_model, 0.2, 0.1), (second_model, 0.8, 0.9)] + assert switching_threshold_ratio_dict == original_switching + assert text_to_img_controlnet_switching_threshold_ratio_dict == original_controlnet + + +def test_sdxl_t2_override_controls_downsampler_at_2048_resolution(): + patched_conv = make_diffusers_downsampler_block(torch.nn.Conv2d) + hidden_states = torch.arange(64, dtype=torch.float32).reshape(1, 1, 8, 8) + + def run(t2_ratio: float) -> torch.Tensor: + module = patched_conv(1, 1, kernel_size=3, stride=2, padding=1, bias=False) + torch.nn.init.constant_(module.weight, 1.0) + module.info = { + "size": (256, 256), + "pipeline": SimpleNamespace(_num_timesteps=30), + "text_to_img_controlnet": False, + "is_inpainting_task": False, + "is_playground": False, + "switching_threshold_overrides": {"T1_ratio": None, "T2_ratio": t2_ratio}, + } + module.model = "sdxl" + module.switching_threshold_ratio = "T2_ratio" + return module(hidden_states) + + assert not torch.equal(run(0.0), run(1.0)) + + +def test_sdxl_automatic_ratios_preserve_extreme_resolution_preset(): + patched_conv = make_diffusers_downsampler_block(torch.nn.Conv2d) + module = patched_conv(1, 1, kernel_size=3, stride=2, padding=1, bias=False) + module.info = { + "size": (512, 512), + "pipeline": SimpleNamespace(_num_timesteps=30), + "text_to_img_controlnet": False, + "is_inpainting_task": False, + "is_playground": False, + "switching_threshold_overrides": {"T1_ratio": None, "T2_ratio": None}, + } + module.model = "sdxl" + module.switching_threshold_ratio = "T2_ratio" + + module(torch.ones(1, 1, 8, 8)) + + assert module.T1_ratio == 0.3 + assert module.T1 == 9 From c0c99493d139f5d5804068d6a96fa6eaa936cb07 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Mon, 17 Aug 2026 00:08:16 +0300 Subject: [PATCH 2/7] Fix HiDiffusion automatic ratio label --- invokeai/frontend/web/public/locales/en.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json index 2c04b120a33..2bc7831128f 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -1071,7 +1071,6 @@ "hiDiffusion": "HiDiffusion", "hiDiffusionRauNet": "HiDiffusion: RAU-Net", "hiDiffusionWindowAttn": "HiDiffusion: Window Attention", - "hiDiffusionRatiosAuto": "HiDiffusion: Automatic Ratios", "hiDiffusionT1Ratio": "HiDiffusion: T1 Ratio", "hiDiffusionT2Ratio": "HiDiffusion: T2 Ratio", "imageDetails": "Image Details", @@ -1734,6 +1733,7 @@ "hiDiffusion": "HiDiffusion", "hiDiffusionRauNet": "HiDiffusion: RAU-Net", "hiDiffusionWindowAttn": "HiDiffusion: Window Attention", + "hiDiffusionRatiosAuto": "HiDiffusion: Automatic Ratios", "hiDiffusionT1Ratio": "HiDiffusion: T1 Ratio", "hiDiffusionT2Ratio": "HiDiffusion: T2 Ratio", "coherenceMode": "Mode", From cc7e29482b5ba46eaad0377c5df911ee2d58f9c0 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Mon, 17 Aug 2026 00:58:04 +0300 Subject: [PATCH 3/7] Fix HiDiffusion regional guidance masks --- .../diffusion/regional_ip_data.py | 7 +++++-- .../diffusion/regional_prompt_data.py | 6 +++--- .../stable_diffusion/test_regional_ip_data.py | 18 ++++++++++++++++ .../test_regional_prompt_data.py | 21 +++++++++++++++++++ 4 files changed, 47 insertions(+), 5 deletions(-) create mode 100644 tests/backend/stable_diffusion/test_regional_ip_data.py create mode 100644 tests/backend/stable_diffusion/test_regional_prompt_data.py diff --git a/invokeai/backend/stable_diffusion/diffusion/regional_ip_data.py b/invokeai/backend/stable_diffusion/diffusion/regional_ip_data.py index 792c97114da..eb2a9eea56b 100644 --- a/invokeai/backend/stable_diffusion/diffusion/regional_ip_data.py +++ b/invokeai/backend/stable_diffusion/diffusion/regional_ip_data.py @@ -11,9 +11,12 @@ def __init__( masks: list[torch.Tensor], dtype: torch.dtype, device: torch.device, - max_downscale_factor: int = 8, + max_downscale_factor: int = 16, ): - """Initialize a `IPAdapterConditioningData` object.""" + """Initialize an `IPAdapterConditioningData` object. + + HiDiffusion's RAU-Net requires one mask level beyond the standard UNet's 8x downscale. + """ assert len(image_prompt_embeds) == len(scales) == len(masks) # The image prompt embeddings. diff --git a/invokeai/backend/stable_diffusion/diffusion/regional_prompt_data.py b/invokeai/backend/stable_diffusion/diffusion/regional_prompt_data.py index eddd31f0c42..bebe00b35b1 100644 --- a/invokeai/backend/stable_diffusion/diffusion/regional_prompt_data.py +++ b/invokeai/backend/stable_diffusion/diffusion/regional_prompt_data.py @@ -19,7 +19,7 @@ def __init__( regions: list[TextConditioningRegions], device: torch.device, dtype: torch.dtype, - max_downscale_factor: int = 8, + max_downscale_factor: int = 16, ): """Initialize a `RegionalPromptData` object. Args: @@ -28,7 +28,7 @@ def __init__( device (torch.device): The device to use for the attention masks. dtype (torch.dtype): The data type to use for the attention masks. max_downscale_factor: Spatial masks will be prepared for downscale factors from 1 to max_downscale_factor - in steps of 2x. + in steps of 2x. HiDiffusion's RAU-Net requires one level beyond the standard UNet's 8x downscale. """ self._regions = regions self._device = device @@ -41,7 +41,7 @@ def __init__( self._negative_cross_attn_mask_score = -10000.0 def _prepare_spatial_masks( - self, regions: list[TextConditioningRegions], max_downscale_factor: int = 8 + self, regions: list[TextConditioningRegions], max_downscale_factor: int = 16 ) -> list[dict[int, torch.Tensor]]: """Prepare the spatial masks for all downscaling factors.""" # batch_masks_by_seq_len[b][s] contains the spatial masks for the b'th batch sample with a query sequence length diff --git a/tests/backend/stable_diffusion/test_regional_ip_data.py b/tests/backend/stable_diffusion/test_regional_ip_data.py new file mode 100644 index 00000000000..734fac305e3 --- /dev/null +++ b/tests/backend/stable_diffusion/test_regional_ip_data.py @@ -0,0 +1,18 @@ +import torch + +from invokeai.backend.stable_diffusion.diffusion.regional_ip_data import RegionalIPData + + +def test_regional_ip_data_supports_hidiffusion_raunet_downscale() -> None: + regional_ip_data = RegionalIPData( + image_prompt_embeds=[torch.zeros((1, 1, 4, 8))], + scales=[1.0], + masks=[torch.ones((1, 1, 192, 192))], + dtype=torch.float32, + device=torch.device("cpu"), + ) + + masks = regional_ip_data.get_masks(query_seq_len=12 * 12) + + assert masks.shape == (1, 1, 12 * 12, 1) + assert torch.count_nonzero(masks) == masks.numel() diff --git a/tests/backend/stable_diffusion/test_regional_prompt_data.py b/tests/backend/stable_diffusion/test_regional_prompt_data.py new file mode 100644 index 00000000000..e1990838c2c --- /dev/null +++ b/tests/backend/stable_diffusion/test_regional_prompt_data.py @@ -0,0 +1,21 @@ +import torch + +from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Range, TextConditioningRegions +from invokeai.backend.stable_diffusion.diffusion.regional_prompt_data import RegionalPromptData + + +def test_regional_prompt_data_supports_hidiffusion_raunet_downscale() -> None: + regions = TextConditioningRegions( + masks=torch.ones((1, 1, 192, 192), dtype=torch.bool), + ranges=[Range(start=0, end=4)], + ) + regional_prompt_data = RegionalPromptData( + regions=[regions], + device=torch.device("cpu"), + dtype=torch.float32, + ) + + attention_mask = regional_prompt_data.get_cross_attn_mask(query_seq_len=12 * 12, key_seq_len=4) + + assert attention_mask.shape == (1, 12 * 12, 4) + assert torch.count_nonzero(attention_mask) == 0 From cc121badbefa39570753de797976faa2611b707c Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Mon, 24 Aug 2026 11:29:32 +0300 Subject: [PATCH 4/7] Stabilize HiDiffusion denoising transitions --- invokeai/app/invocations/denoise_latents.py | 8 ++ invokeai/backend/hidiffusion/hidiffusion.py | 60 ++++++++------- .../extensions/hidiffusion.py | 12 +++ .../stable_diffusion/hidiffusion_utils.py | 8 ++ .../test_hidiffusion_utils.py | 75 +++++++++++++++++++ 5 files changed, 136 insertions(+), 27 deletions(-) diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index 3682d80145a..e00e48cbbb4 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -926,6 +926,10 @@ def step_callback(state: PipelineIntermediateState) -> None: t1_ratio=self.hidiffusion_t1_ratio, t2_ratio=self.hidiffusion_t2_ratio, generator=torch.Generator(device="cpu").manual_seed(seed), + is_inpainting_task=self.denoise_mask is not None, + use_aggressive_raunet=False, + denoising_start=self.denoising_start, + denoising_end=self.denoising_end, ) ) @@ -1157,6 +1161,10 @@ def _lora_loader() -> Iterator[PatchSpec]: t1_ratio=self.hidiffusion_t1_ratio, t2_ratio=self.hidiffusion_t2_ratio, generator=torch.Generator(device="cpu").manual_seed(seed), + is_inpainting_task=self.denoise_mask is not None, + use_aggressive_raunet=False, + denoising_start=self.denoising_start, + denoising_end=self.denoising_end, ) if self.hidiffusion else nullcontext() diff --git a/invokeai/backend/hidiffusion/hidiffusion.py b/invokeai/backend/hidiffusion/hidiffusion.py index 5c7d83b84ee..6bd8fec67b0 100644 --- a/invokeai/backend/hidiffusion/hidiffusion.py +++ b/invokeai/backend/hidiffusion/hidiffusion.py @@ -137,11 +137,29 @@ def _get_max_timesteps(info_dict: dict) -> int: def _get_switching_threshold_ratio(module: torch.nn.Module, presets: dict, preset_key: str) -> float: - """Resolve a threshold ratio without mutating process-global presets.""" + """Resolve a threshold ratio for the executed part of the denoising schedule.""" override = module.info["switching_threshold_overrides"].get(module.switching_threshold_ratio) + full_schedule_ratio = override if override is not None else presets[preset_key][module.switching_threshold_ratio] + + denoising_start = module.info.get("denoising_start", 0.0) + denoising_end = module.info.get("denoising_end", 1.0) + if denoising_end <= denoising_start: + return 0.0 + + executed_schedule_ratio = (full_schedule_ratio - denoising_start) / (denoising_end - denoising_start) + return max(0.0, min(1.0, executed_schedule_ratio)) + + +def _should_use_aggressive_raunet(module: torch.nn.Module) -> bool: + """Resolve whether RAU-Net should be activated after denoising has already started.""" + override = module.info.get("use_aggressive_raunet") if override is not None: return override - return presets[preset_key][module.switching_threshold_ratio] + if module.info["is_inpainting_task"]: + return inpainting_is_aggressive_raunet + if module.info["is_playground"]: + return playground_is_aggressive_raunet + return is_aggressive_raunet def make_diffusers_sdxl_controlnet_ppl(block_class): @@ -1581,12 +1599,7 @@ def forward( self, switching_threshold_ratio_dict, "sdxl_2048" ) - if self.info["is_inpainting_task"]: - self.aggressive_raunet = inpainting_is_aggressive_raunet - elif self.info["is_playground"]: - self.aggressive_raunet = playground_is_aggressive_raunet - else: - self.aggressive_raunet = is_aggressive_raunet + self.aggressive_raunet = _should_use_aggressive_raunet(self) else: self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_4096") elif self.model == "sdxl_turbo": @@ -1717,12 +1730,7 @@ def forward( self, switching_threshold_ratio_dict, "sdxl_2048" ) - if self.info["is_inpainting_task"]: - self.aggressive_raunet = inpainting_is_aggressive_raunet - elif self.info["is_playground"]: - self.aggressive_raunet = playground_is_aggressive_raunet - else: - self.aggressive_raunet = is_aggressive_raunet + self.aggressive_raunet = _should_use_aggressive_raunet(self) else: self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_4096") @@ -1856,12 +1864,7 @@ def forward(self, hidden_states: torch.Tensor, scale=1.0) -> torch.Tensor: self, switching_threshold_ratio_dict, "sdxl_2048" ) - if self.info["is_inpainting_task"]: - self.aggressive_raunet = inpainting_is_aggressive_raunet - elif self.info["is_playground"]: - self.aggressive_raunet = playground_is_aggressive_raunet - else: - self.aggressive_raunet = is_aggressive_raunet + self.aggressive_raunet = _should_use_aggressive_raunet(self) else: self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_4096") elif self.model == "sdxl_turbo": @@ -1939,12 +1942,7 @@ def forward(self, hidden_states: torch.Tensor, scale=1.0) -> torch.Tensor: self, switching_threshold_ratio_dict, "sdxl_2048" ) - if self.info["is_inpainting_task"]: - self.aggressive_raunet = inpainting_is_aggressive_raunet - elif self.info["is_playground"]: - self.aggressive_raunet = playground_is_aggressive_raunet - else: - self.aggressive_raunet = is_aggressive_raunet + self.aggressive_raunet = _should_use_aggressive_raunet(self) else: self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_4096") elif self.model == "sdxl_turbo": @@ -2063,6 +2061,10 @@ def apply_hidiffusion( is_controlnet_text_to_image: bool = False, t1_ratio: float | None = None, t2_ratio: float | None = None, + is_inpainting_task: bool | None = None, + use_aggressive_raunet: bool | None = None, + denoising_start: float = 0.0, + denoising_end: float = 1.0, ): """ model: diffusers model. We support SD 1.5, 2.1, XL, XL Turbo. @@ -2138,13 +2140,17 @@ def apply_hidiffusion( elif set(sdxl_module_key) < set(diffusion_model_module_key): name_or_path = "stabilityai/stable-diffusion-xl-base-1.0" + detected_inpainting_task = model.__class__ in auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING.values() diffusion_model.info = { "size": None, "upsample_size": None, "hooks": [], "text_to_img_controlnet": has_controlnet and is_controlnet_text_to_image, - "is_inpainting_task": model.__class__ in auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING.values(), + "is_inpainting_task": detected_inpainting_task if is_inpainting_task is None else is_inpainting_task, "is_playground": is_playground, + "use_aggressive_raunet": use_aggressive_raunet, + "denoising_start": denoising_start, + "denoising_end": denoising_end, "pipeline": model, "switching_threshold_overrides": {"T1_ratio": t1_ratio, "T2_ratio": t2_ratio}, } diff --git a/invokeai/backend/stable_diffusion/extensions/hidiffusion.py b/invokeai/backend/stable_diffusion/extensions/hidiffusion.py index 13a1763f35e..464e6a1d197 100644 --- a/invokeai/backend/stable_diffusion/extensions/hidiffusion.py +++ b/invokeai/backend/stable_diffusion/extensions/hidiffusion.py @@ -22,6 +22,10 @@ def __init__( generator: torch.Generator | None = None, has_controlnet: bool = False, is_controlnet_text_to_image: bool = False, + is_inpainting_task: bool | None = None, + use_aggressive_raunet: bool | None = None, + denoising_start: float = 0.0, + denoising_end: float = 1.0, ): super().__init__() self._name_or_path = name_or_path @@ -29,6 +33,10 @@ def __init__( self._apply_window_attn = apply_window_attn self._has_controlnet = has_controlnet self._is_controlnet_text_to_image = is_controlnet_text_to_image + self._is_inpainting_task = is_inpainting_task + self._use_aggressive_raunet = use_aggressive_raunet + self._denoising_start = denoising_start + self._denoising_end = denoising_end self._t1_ratio = t1_ratio self._t2_ratio = t2_ratio self._generator = generator @@ -45,5 +53,9 @@ def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeigh t1_ratio=self._t1_ratio, t2_ratio=self._t2_ratio, generator=self._generator, + is_inpainting_task=self._is_inpainting_task, + use_aggressive_raunet=self._use_aggressive_raunet, + denoising_start=self._denoising_start, + denoising_end=self._denoising_end, ): yield None diff --git a/invokeai/backend/stable_diffusion/hidiffusion_utils.py b/invokeai/backend/stable_diffusion/hidiffusion_utils.py index fa83cf8bc67..371da5ca6bd 100644 --- a/invokeai/backend/stable_diffusion/hidiffusion_utils.py +++ b/invokeai/backend/stable_diffusion/hidiffusion_utils.py @@ -18,6 +18,10 @@ def hidiffusion_patch( generator: torch.Generator | None = None, has_controlnet: bool = False, is_controlnet_text_to_image: bool = False, + is_inpainting_task: bool | None = None, + use_aggressive_raunet: bool | None = None, + denoising_start: float = 0.0, + denoising_end: float = 1.0, ): """Context manager that applies HiDiffusion and restores the model on exit.""" from invokeai.backend.hidiffusion.hidiffusion import apply_hidiffusion, remove_hidiffusion @@ -96,6 +100,10 @@ def _set_name_or_path_on_config(cfg, value: str) -> bool: has_controlnet=has_controlnet, is_controlnet_text_to_image=is_controlnet_text_to_image, generator=generator, + is_inpainting_task=is_inpainting_task, + use_aggressive_raunet=use_aggressive_raunet, + denoising_start=denoising_start, + denoising_end=denoising_end, ) yield finally: diff --git a/tests/backend/stable_diffusion/test_hidiffusion_utils.py b/tests/backend/stable_diffusion/test_hidiffusion_utils.py index b46930a9d90..b7e9bad9130 100644 --- a/tests/backend/stable_diffusion/test_hidiffusion_utils.py +++ b/tests/backend/stable_diffusion/test_hidiffusion_utils.py @@ -6,7 +6,9 @@ import torch from invokeai.backend.hidiffusion.hidiffusion import ( + _get_switching_threshold_ratio, _resize_controlnet_residual, + _should_use_aggressive_raunet, make_diffusers_downsampler_block, switching_threshold_ratio_dict, text_to_img_controlnet_switching_threshold_ratio_dict, @@ -396,6 +398,79 @@ def fake_apply_hidiffusion(model, **kwargs): assert text_to_img_controlnet_switching_threshold_ratio_dict == original_controlnet +def test_hidiffusion_patch_forwards_generation_context(): + model = SimpleNamespace(unet=DummyUNet()) + + with ( + patch("invokeai.backend.hidiffusion.hidiffusion.apply_hidiffusion") as mock_apply_hidiffusion, + patch("invokeai.backend.hidiffusion.hidiffusion.remove_hidiffusion"), + ): + with hidiffusion_patch( + model, + name_or_path="stabilityai/stable-diffusion-xl-base-1.0", + is_inpainting_task=True, + use_aggressive_raunet=False, + denoising_start=0.6, + denoising_end=1.0, + ): + pass + + kwargs = mock_apply_hidiffusion.call_args.kwargs + assert kwargs["is_inpainting_task"] is True + assert kwargs["use_aggressive_raunet"] is False + assert kwargs["denoising_start"] == 0.6 + assert kwargs["denoising_end"] == 1.0 + + +@pytest.mark.parametrize( + ("denoising_start", "denoising_end", "expected_ratio"), + [ + (0.0, 1.0, 0.4), + (0.6, 1.0, 0.0), + (0.0, 0.2, 1.0), + (0.2, 0.6, 0.5), + ], +) +def test_hidiffusion_ratios_are_mapped_to_the_executed_denoising_range( + denoising_start: float, denoising_end: float, expected_ratio: float +): + module = SimpleNamespace( + switching_threshold_ratio="T1_ratio", + info={ + "switching_threshold_overrides": {"T1_ratio": None, "T2_ratio": None}, + "denoising_start": denoising_start, + "denoising_end": denoising_end, + }, + ) + + ratio = _get_switching_threshold_ratio(module, switching_threshold_ratio_dict, "sdxl_2048") + + assert ratio == pytest.approx(expected_ratio) + + +@pytest.mark.parametrize( + ("is_inpainting_task", "override", "expected"), + [ + (False, False, False), + (True, True, True), + (True, None, False), + (False, None, True), + ], +) +def test_explicit_aggressive_raunet_setting_takes_precedence( + is_inpainting_task: bool, override: bool | None, expected: bool +): + module = SimpleNamespace( + info={ + "is_inpainting_task": is_inpainting_task, + "is_playground": False, + "use_aggressive_raunet": override, + } + ) + + assert _should_use_aggressive_raunet(module) is expected + + def test_sdxl_t2_override_controls_downsampler_at_2048_resolution(): patched_conv = make_diffusers_downsampler_block(torch.nn.Conv2d) hidden_states = torch.arange(64, dtype=torch.float32).reshape(1, 1, 8, 8) From b5b842408733e5de1819fa6ffe4e23d785416219 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Thu, 3 Sep 2026 10:35:42 +0300 Subject: [PATCH 5/7] fix(hidiffusion): correct scheduling and ratio behavior Restore staged logical-step scheduling and guidance residual compatibility. Preserve manual ratios behind Auto, validate T2 <= T1, and exclude the unsupported refiner path. --- LICENSE-HiDiffusion.txt | 22 +- .../src/content/docs/features/hidiffusion.mdx | 18 +- invokeai/app/invocations/denoise_latents.py | 18 +- invokeai/app/invocations/fields.py | 11 +- invokeai/backend/hidiffusion/hidiffusion.py | 307 +++++++----------- .../stable_diffusion/diffusers_pipeline.py | 5 + .../diffusion/regional_ip_data.py | 4 +- .../diffusion/regional_prompt_data.py | 6 +- .../extensions/hidiffusion.py | 19 +- .../stable_diffusion/hidiffusion_utils.py | 6 - invokeai/frontend/web/public/locales/en.json | 12 +- .../controlLayers/store/paramsSlice.test.ts | 66 ++-- .../controlLayers/store/paramsSlice.ts | 21 +- .../src/features/controlLayers/store/types.ts | 14 +- .../web/src/features/metadata/parsing.test.ts | 32 +- .../web/src/features/metadata/parsing.tsx | 11 +- .../graph/generation/addSDXLRefiner.test.ts | 61 ++++ .../util/graph/generation/addSDXLRefiner.ts | 10 - .../generation/buildHiDiffusionGraph.test.ts | 30 +- .../util/graph/generation/buildSD1Graph.ts | 9 +- .../util/graph/generation/buildSDXLGraph.ts | 9 +- .../Advanced/ParamHiDiffusionToggle.tsx | 48 +-- .../test_hidiffusion_utils.py | 293 +++++++++++++++-- .../stable_diffusion/test_regional_ip_data.py | 19 +- .../test_regional_prompt_data.py | 21 +- 25 files changed, 721 insertions(+), 351 deletions(-) create mode 100644 invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.test.ts diff --git a/LICENSE-HiDiffusion.txt b/LICENSE-HiDiffusion.txt index 73095184ff3..fc31ebfac3e 100644 --- a/LICENSE-HiDiffusion.txt +++ b/LICENSE-HiDiffusion.txt @@ -1,7 +1,9 @@ HiDiffusion - License notice Original project: https://github.com/megvii-research/HiDiffusion -Vendored from: https://github.com/monofy-org/HiDiffusion +Intermediate fork: https://github.com/monofy-org/HiDiffusion +Vendored implementation baseline: + https://github.com/Teriks/dgenerate/tree/d83b839033cc22c5101fb0f987bd4eb2de3d5d12/dgenerate/extras/hidiffusion Vendored under: invokeai/backend/hidiffusion/ ================================================================================ @@ -14,6 +16,18 @@ and at: https://www.apache.org/licenses/LICENSE-2.0 +The dgenerate distribution includes the following NOTICE for its HiDiffusion +fork: + + This code is vendored from: https://github.com/monofy-org/HiDiffusion + + This repository is a fork which implements some fixes to HiDiffusion. + + It further modifies HiDiffusion so that apply_hidiffusion accepts a + torch.Generator object, enabling deterministic images based on a seed. + + HiDiffusion is here: https://github.com/megvii-research/HiDiffusion + ================================================================================ MODULE KEYS (BSD 3-Clause License) ================================================================================ @@ -59,7 +73,9 @@ 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. +* Automatic model- and resolution-specific ratios can be overridden per + invocation without mutating the vendored global preset tables. +* Inpainting, guidance, adapter, and cached-model lifecycle behavior was + adapted to InvokeAI's denoising pipelines. diff --git a/docs/src/content/docs/features/hidiffusion.mdx b/docs/src/content/docs/features/hidiffusion.mdx index 1a8a35b276b..cf1ed6e0c2f 100644 --- a/docs/src/content/docs/features/hidiffusion.mdx +++ b/docs/src/content/docs/features/hidiffusion.mdx @@ -29,14 +29,24 @@ Learn more: https://github.com/megvii-research/HiDiffusion - **HiDiffusion: Window Attention**: Enables windowed attention blocks. This can boost local texture/detail, but may slightly affect global coherence in some prompts. -- **HiDiffusion: Automatic Ratios**: Uses HiDiffusion's model- and resolution-specific T1/T2 presets. Disable it to enter manual overrides. +- **HiDiffusion: Automatic Ratios**: Uses HiDiffusion's original discrete model- and resolution-specific T1/T2 presets behind the scenes. SDXL uses the 2048 preset while either latent dimension is below the 4096 reference threshold, and the 4096 preset only when both dimensions reach it. The manual sliders retain their values while automatic ratios are enabled, so disabling automatic ratios restores the previous manual settings. -- **HiDiffusion: T1 Ratio**: Controls the first RAU-Net switching threshold. Lower values switch earlier; higher values keep the first resolution-aware stage active longer. +- **HiDiffusion: T1 Ratio**: Controls how long the primary RAU-Net stage remains active. At extreme resolutions, this is the later of the two RAU-Net cutoffs. Lower values switch earlier. -- **HiDiffusion: T2 Ratio**: Controls the second RAU-Net switching threshold used for extreme-resolution generation. It does not control window attention. Higher values keep the second resolution-aware stage active longer. +- **HiDiffusion: T2 Ratio**: Manually controls the end of the additional early RAU-Net stage. For ordinary SDXL generation, this also controls when the primary stage begins. When automatic ratios are enabled, the original implementation uses a fixed `8 / 50` boundary for this early stage at the 2048 preset. Excessive manual values can reduce pose and background diversity or introduce artifacts. T2 cannot exceed T1 and does not control window attention. + +The T1/T2 input names are retained for API compatibility with the upstream implementation's code keys. T2 is the earlier boundary and T1 is the later boundary; do not interpret the names as chronological order. With automatic ratios at the SDXL 2048 preset, the additional stage runs for the first 8 of every 50 denoising steps, the primary stage then runs until T1, and the ordinary UNet is used afterward. Inpainting follows the upstream non-aggressive schedule: its primary stage starts at the first denoising step. + +### Manual ratio constraint + +When setting the ratios manually, **T2 must be less than or equal to T1**. T2 controls an additional resolution-reduction stage that depends on the primary T1 stage. Allowing T2 to remain active after T1 has ended would produce an invalid RAU-Net stage order and can cause severe structural artifacts. + +The UI limits the maximum T2 value to the current T1 value. If T1 is reduced below the current T2 value, the UI also reduces T2 to match it. Workflows or API requests that explicitly provide `T2 > T1` are rejected by the backend instead of being silently modified. ## Tips -- Try **1536–2048 px** for the clearest benefits (SDXL). +- Start around **1536–2048 px** for SDXL. Automatic ratios do not interpolate: the upstream 4096 preset is selected only when both image dimensions reach 4096 px (512 latent pixels). Test intermediate and non-square sizes explicitly before increasing either ratio manually. +- Ordinary SDXL generation uses the upstream staged schedule: the additional RAU-Net path is active first, the primary path takes over at the early boundary, and the ordinary UNet takes over at T1. Inpainting does not use the initial aggressive stage. +- HiDiffusion is applied only to the SDXL base denoise stage, not the optional refiner stage. - If results look worse, disable **Window Attention** first, then RAU‑Net. - Effects vary by scheduler and model; compare with the same seed for a fair test. diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index e00e48cbbb4..83a1d8cb846 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -16,7 +16,7 @@ from diffusers.schedulers.scheduling_tcd import TCDScheduler from diffusers.schedulers.scheduling_utils import SchedulerMixin as Scheduler from PIL import Image -from pydantic import field_validator +from pydantic import field_validator, model_validator from torchvision.transforms.functional import resize as tv_resize from transformers import CLIPVisionModelWithProjection @@ -248,6 +248,16 @@ def ge_one(cls, v: Union[List[float], float]) -> Union[List[float], float]: raise ValueError("cfg_scale must be greater than 1") return v + @model_validator(mode="after") + def validate_hidiffusion_ratio_order(self): + if ( + self.hidiffusion_t1_ratio is not None + and self.hidiffusion_t2_ratio is not None + and self.hidiffusion_t2_ratio > self.hidiffusion_t1_ratio + ): + raise ValueError("HiDiffusion T2 ratio must be less than or equal to the T1 ratio") + return self + @staticmethod def _get_text_embeddings_and_masks( cond_list: list[ConditioningField], @@ -927,9 +937,6 @@ def step_callback(state: PipelineIntermediateState) -> None: t2_ratio=self.hidiffusion_t2_ratio, generator=torch.Generator(device="cpu").manual_seed(seed), is_inpainting_task=self.denoise_mask is not None, - use_aggressive_raunet=False, - denoising_start=self.denoising_start, - denoising_end=self.denoising_end, ) ) @@ -1162,9 +1169,6 @@ def _lora_loader() -> Iterator[PatchSpec]: t2_ratio=self.hidiffusion_t2_ratio, generator=torch.Generator(device="cpu").manual_seed(seed), is_inpainting_task=self.denoise_mask is not None, - use_aggressive_raunet=False, - denoising_start=self.denoising_start, - denoising_end=self.denoising_end, ) if self.hidiffusion else nullcontext() diff --git a/invokeai/app/invocations/fields.py b/invokeai/app/invocations/fields.py index 3826a340012..3326f4dd257 100644 --- a/invokeai/app/invocations/fields.py +++ b/invokeai/app/invocations/fields.py @@ -152,8 +152,15 @@ class FieldDescriptions: hidiffusion = "Apply HiDiffusion (RAU-Net + MSW-MSA) for higher-resolution denoising" hidiffusion_raunet = "Apply HiDiffusion RAU-Net blocks" hidiffusion_window_attn = "Apply HiDiffusion window attention blocks" - hidiffusion_t1_ratio = "Override HiDiffusion early switch threshold (T1 ratio)" - hidiffusion_t2_ratio = "Override HiDiffusion late switch threshold (T2 ratio)" + hidiffusion_t1_ratio = ( + "Override the duration of HiDiffusion's primary RAU-Net stage (upstream code key T1_ratio). " + "At extreme resolutions this is the later of the two RAU-Net cutoffs." + ) + hidiffusion_t2_ratio = ( + "Override the duration of HiDiffusion's additional extreme-resolution RAU-Net stage (upstream code key " + "T2_ratio). This is the earlier cutoff when both stages are active and cannot exceed T1; excessive values " + "can reduce composition diversity or introduce artifacts." + ) scheduler = "Scheduler to use during inference" positive_cond = "Positive conditioning tensor" negative_cond = "Negative conditioning tensor" diff --git a/invokeai/backend/hidiffusion/hidiffusion.py b/invokeai/backend/hidiffusion/hidiffusion.py index 6bd8fec67b0..5daad401011 100644 --- a/invokeai/backend/hidiffusion/hidiffusion.py +++ b/invokeai/backend/hidiffusion/hidiffusion.py @@ -96,8 +96,9 @@ def sdxl_turbo_hidiffusion_key(): ] -# T1_ratio: see T1 introduced in the main paper. T1 = number_inference_step * T1_ratio. A higher T1_ratio can better mitigate object duplication. We set T1_ratio=0.4 by default. You'd better adjust it to fit your prompt. Only active when apply_raunet=True. -# T2_ratio: see T2 introduced in the appendix, used in extreme resolution image generation. T2 = number_inference_step * T2_ratio. A higher T2_ratio can better mitigate object duplication. Only active when apply_raunet=True +# These names are retained from the upstream implementation. T1_ratio controls the primary RAU-Net +# stage, while T2_ratio controls the additional stage used at extreme resolutions. When both are +# nonzero, the T2_ratio cutoff occurs first even though the paper labels the transitions chronologically. switching_threshold_ratio_dict = { "sd15_1024": {"T1_ratio": 0.4, "T2_ratio": 0.0}, "sd15_2048": {"T1_ratio": 0.7, "T2_ratio": 0.3}, @@ -117,7 +118,6 @@ def sdxl_turbo_hidiffusion_key(): inpainting_is_aggressive_raunet = False playground_is_aggressive_raunet = False - with importlib.resources.open_text(f"{__package__}.sd_module_key", "sd15_module_key.txt", encoding="utf-8") as f: sd15_module_key = f.read().splitlines() @@ -136,25 +136,49 @@ def _get_max_timesteps(info_dict: dict) -> int: return len(pipeline.scheduler.timesteps) -def _get_switching_threshold_ratio(module: torch.nn.Module, presets: dict, preset_key: str) -> float: - """Resolve a threshold ratio for the executed part of the denoising schedule.""" - override = module.info["switching_threshold_overrides"].get(module.switching_threshold_ratio) - full_schedule_ratio = override if override is not None else presets[preset_key][module.switching_threshold_ratio] +def _get_automatic_switching_threshold_ratio(module: torch.nn.Module, height: int, width: int, threshold: str) -> float: + """Select the discrete model- and resolution-specific preset used by upstream HiDiffusion.""" + if module.model == "sdxl_turbo": + return switching_threshold_ratio_dict["sdxl_turbo_1024"][threshold] - denoising_start = module.info.get("denoising_start", 0.0) - denoising_end = module.info.get("denoising_end", 1.0) - if denoising_end <= denoising_start: - return 0.0 + if module.model == "sd15": + preset_key = "sd15_1024" if height < 256 or width < 256 else "sd15_2048" + elif module.model == "sdxl": + preset_key = "sdxl_2048" if height < 512 or width < 512 else "sdxl_4096" + else: + raise ValueError("HiDiffusion only supports sd15, sd21, sdxl, and sdxl-turbo.") - executed_schedule_ratio = (full_schedule_ratio - denoising_start) / (denoising_end - denoising_start) - return max(0.0, min(1.0, executed_schedule_ratio)) + if module.model == "sdxl" and preset_key == "sdxl_2048" and module.info["text_to_img_controlnet"]: + return text_to_img_controlnet_switching_threshold_ratio_dict[preset_key][threshold] + return switching_threshold_ratio_dict[preset_key][threshold] -def _should_use_aggressive_raunet(module: torch.nn.Module) -> bool: - """Resolve whether RAU-Net should be activated after denoising has already started.""" - override = module.info.get("use_aggressive_raunet") - if override is not None: - return override +def _get_resolution_aware_switching_threshold_ratio(module: torch.nn.Module, height: int, width: int) -> float: + """Resolve a manual override or the upstream discrete preset for the current latent size.""" + override = module.info["switching_threshold_overrides"].get(module.switching_threshold_ratio) + ratio = ( + override + if override is not None + else _get_automatic_switching_threshold_ratio(module, height, width, module.switching_threshold_ratio) + ) + + if module.switching_threshold_ratio == "T2_ratio": + t1_override = module.info["switching_threshold_overrides"].get("T1_ratio") + t1_ratio = ( + t1_override + if t1_override is not None + else _get_automatic_switching_threshold_ratio(module, height, width, "T1_ratio") + ) + if ratio > t1_ratio: + raise ValueError("HiDiffusion T2 ratio must be less than or equal to the T1 ratio.") + + return ratio + + +def _uses_aggressive_raunet(module: torch.nn.Module, height: int, width: int) -> bool: + """Return whether upstream's staged SDXL schedule applies to this generation.""" + if module.model != "sdxl" or (height >= 512 and width >= 512): + return False if module.info["is_inpainting_task"]: return inpainting_is_aggressive_raunet if module.info["is_playground"]: @@ -162,6 +186,41 @@ def _should_use_aggressive_raunet(module: torch.nn.Module) -> bool: return is_aggressive_raunet +def _get_raunet_step_range(module: torch.nn.Module, height: int, width: int) -> tuple[float, int, int]: + """Resolve the active half-open step range for one patched RAU-Net module. + + At ordinary SDXL resolutions, upstream uses the extra (T2-position) modules first, then the + primary (T1-position) modules. A manual T2 override replaces upstream's fixed 8/50 boundary so + that InvokeAI's explicit T2 control remains effective. + """ + ratio = _get_resolution_aware_switching_threshold_ratio(module, height, width) + start = 0 + end = int(module.max_timestep * ratio) + + if _uses_aggressive_raunet(module, height, width): + t2_override = module.info["switching_threshold_overrides"].get("T2_ratio") + early_ratio = aggressive_step / 50 if t2_override is None else t2_override + early_end = int(module.max_timestep * early_ratio) + if module.switching_threshold_ratio == "T1_ratio": + start = early_end + else: + end = early_end + + return ratio, start, end + + +def _get_current_step(module: torch.nn.Module) -> int: + """Return the logical denoising step when managed by InvokeAI, or the upstream per-forward fallback.""" + step_index = module.info.get("step_index") + return module.timestep if step_index is None else step_index + + +def _advance_fallback_step(module: torch.nn.Module) -> None: + """Preserve upstream behavior for callers that do not provide a logical denoising step.""" + if module.info.get("step_index") is None: + module.timestep = (module.timestep + 1) % module.max_timestep + + def make_diffusers_sdxl_controlnet_ppl(block_class): class sdxl_controlnet_ppl(block_class): # Save for unpatching later @@ -1448,23 +1507,29 @@ def window_reverse(windows, window_size, H, W, shift_size): norm_hidden_states = self.pos_embed(norm_hidden_states) # MSW-MSA - if generator is not None: - rand_num = torch.rand(1, generator=generator, device=generator.device) - else: - rand_num = torch.rand(1) + logical_step = self.info.get("step_index") + if logical_step is None or self.__dict__.get("_hidiffusion_window_shift_step") != logical_step: + if generator is not None: + rand_num = torch.rand(1, generator=generator, device=generator.device) + else: + rand_num = torch.rand(1) + self._hidiffusion_window_shift_step = logical_step + self._hidiffusion_window_shift_bucket = min(int(rand_num.item() * 4), 3) + + shift_bucket = self._hidiffusion_window_shift_bucket B, N, C = hidden_states.shape ori_H, ori_W = self.info["size"] downsample_ratio = round(((ori_H * ori_W) / N) ** 0.5) H, W = (math.ceil(ori_H / downsample_ratio), math.ceil(ori_W / downsample_ratio)) widow_size = (math.ceil(H / 2), math.ceil(W / 2)) - if rand_num <= 0.25: + if shift_bucket == 0: shift_size = (0, 0) - if rand_num > 0.25 and rand_num <= 0.5: + elif shift_bucket == 1: shift_size = (widow_size[0] // 4, widow_size[1] // 4) - if rand_num > 0.5 and rand_num <= 0.75: + elif shift_bucket == 2: shift_size = (widow_size[0] // 4 * 2, widow_size[1] // 4 * 2) - if rand_num > 0.75 and rand_num <= 1: + else: shift_size = (widow_size[0] // 4 * 3, widow_size[1] // 4 * 3) norm_hidden_states = window_partition(norm_hidden_states, widow_size, shift_size, H, W) # 2. Prepare GLIGEN inputs @@ -1561,11 +1626,9 @@ class cross_attn_down_block(block_class): # Save for unpatching later _parent = block_class timestep = 0 - aggressive_raunet = False T1_ratio = 0 T1_start = 0 T1_end = 0 - aggressive_raunet = False T1 = 0 # to avoid confict with sdxl-turbo max_timestep = 50 info: dict = None @@ -1583,37 +1646,8 @@ def forward( ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]: self.max_timestep = _get_max_timesteps(self.info) ori_H, ori_W = self.info["size"] - if self.model == "sd15": - if ori_H < 256 or ori_W < 256: - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_1024") - else: - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_2048") - elif self.model == "sdxl": - if ori_H < 512 or ori_W < 512: - if self.info["text_to_img_controlnet"]: - self.T1_ratio = _get_switching_threshold_ratio( - self, text_to_img_controlnet_switching_threshold_ratio_dict, "sdxl_2048" - ) - else: - self.T1_ratio = _get_switching_threshold_ratio( - self, switching_threshold_ratio_dict, "sdxl_2048" - ) - - self.aggressive_raunet = _should_use_aggressive_raunet(self) - else: - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_4096") - elif self.model == "sdxl_turbo": - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_turbo_1024") - else: - raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") - - if self.aggressive_raunet and self.switching_threshold_ratio == "T1_ratio": - # self.T1_start = min(int(self.max_timestep * self.T1_ratio * 0.4), int(8/50 * self.max_timestep)) - self.T1_start = int(aggressive_step / 50 * self.max_timestep) - self.T1_end = int(self.max_timestep * self.T1_ratio) - self.T1 = 0 # to avoid confict with sdxl-turbo - else: - self.T1 = int(self.max_timestep * self.T1_ratio) + self.T1_ratio, self.T1_start, self.T1_end = _get_raunet_step_range(self, ori_H, ori_W) + self.T1 = self.T1_end output_states = () @@ -1660,13 +1694,15 @@ def custom_forward(*inputs): # apply additional residuals to the output of the last pair of resnet and attention blocks if i == len(blocks) - 1 and additional_residuals is not None: + if additional_residuals.shape[-2:] != hidden_states.shape[-2:]: + additional_residuals = F.adaptive_avg_pool2d( + additional_residuals, output_size=hidden_states.shape[-2:] + ) hidden_states = hidden_states + additional_residuals if i == 0: - if self.aggressive_raunet and self.timestep >= self.T1_start and self.timestep < self.T1_end: - self.info["upsample_size"] = (hidden_states.shape[2], hidden_states.shape[3]) - hidden_states = F.avg_pool2d(hidden_states, kernel_size=(2, 2), ceil_mode=True) - elif self.timestep < self.T1: + current_step = _get_current_step(self) + if self.T1_start <= current_step < self.T1_end: self.info["upsample_size"] = (hidden_states.shape[2], hidden_states.shape[3]) hidden_states = F.avg_pool2d(hidden_states, kernel_size=(2, 2), ceil_mode=True) output_states = output_states + (hidden_states,) @@ -1678,9 +1714,7 @@ def custom_forward(*inputs): output_states = output_states + (hidden_states,) - self.timestep += 1 - if self.timestep == self.max_timestep: - self.timestep = 0 + _advance_fallback_step(self) return hidden_states, output_states @@ -1693,11 +1727,9 @@ class cross_attn_up_block(block_class): # Save for unpatching later _parent = block_class timestep = 0 - aggressive_raunet = False T1_ratio = 0 T1_start = 0 T1_end = 0 - aggressive_raunet = False T1 = 0 # to avoid confict with sdxl-turbo max_timestep = 50 @@ -1714,38 +1746,8 @@ def forward( ) -> torch.FloatTensor: self.max_timestep = _get_max_timesteps(self.info) ori_H, ori_W = self.info["size"] - if self.model == "sd15": - if ori_H < 256 or ori_W < 256: - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_1024") - else: - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_2048") - elif self.model == "sdxl": - if ori_H < 512 or ori_W < 512: - if self.info["text_to_img_controlnet"]: - self.T1_ratio = _get_switching_threshold_ratio( - self, text_to_img_controlnet_switching_threshold_ratio_dict, "sdxl_2048" - ) - else: - self.T1_ratio = _get_switching_threshold_ratio( - self, switching_threshold_ratio_dict, "sdxl_2048" - ) - - self.aggressive_raunet = _should_use_aggressive_raunet(self) - - else: - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_4096") - elif self.model == "sdxl_turbo": - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_turbo_1024") - else: - raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") - - if self.aggressive_raunet and self.switching_threshold_ratio == "T1_ratio": - # self.T1_start = min(int(self.max_timestep * self.T1_ratio * 0.4), int(8/50 * self.max_timestep)) - self.T1_start = int(aggressive_step / 50 * self.max_timestep) - self.T1_end = int(self.max_timestep * self.T1_ratio) - self.T1 = 0 # to avoid confict with sdxl-turbo - else: - self.T1 = int(self.max_timestep * self.T1_ratio) + self.T1_ratio, self.T1_start, self.T1_end = _get_raunet_step_range(self, ori_H, ori_W) + self.T1 = self.T1_end is_freeu_enabled = ( getattr(self, "s1", None) @@ -1812,11 +1814,8 @@ def custom_forward(*inputs): )[0] if i == 1: - if self.aggressive_raunet and self.timestep >= self.T1_start and self.timestep < self.T1_end: - hidden_states = F.interpolate( - hidden_states, size=self.info["upsample_size"], mode="bicubic" - ) - elif self.timestep < self.T1: + current_step = _get_current_step(self) + if self.T1_start <= current_step < self.T1_end: hidden_states = F.interpolate( hidden_states, size=self.info["upsample_size"], mode="bicubic" ) @@ -1825,9 +1824,7 @@ def custom_forward(*inputs): hidden_states = upsampler(hidden_states, upsample_size) # hidden_states = upsampler(hidden_states, upsample_size, scale=lora_scale) - self.timestep += 1 - if self.timestep == self.max_timestep: - self.timestep = 0 + _advance_fallback_step(self) return hidden_states @@ -1840,47 +1837,21 @@ class downsampler_block(block_class): # Save for unpatching later _parent = block_class T1_ratio = 0 + T1_start = 0 + T1_end = 0 T1 = 0 timestep = 0 - aggressive_raunet = False max_timestep = 50 def forward(self, hidden_states: torch.Tensor, scale=1.0) -> torch.Tensor: self.max_timestep = _get_max_timesteps(self.info) ori_H, ori_W = self.info["size"] - if self.model == "sd15": - if ori_H < 256 or ori_W < 256: - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_1024") - else: - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_2048") - elif self.model == "sdxl": - if ori_H < 512 or ori_W < 512: - if self.info["text_to_img_controlnet"]: - self.T1_ratio = _get_switching_threshold_ratio( - self, text_to_img_controlnet_switching_threshold_ratio_dict, "sdxl_2048" - ) - else: - self.T1_ratio = _get_switching_threshold_ratio( - self, switching_threshold_ratio_dict, "sdxl_2048" - ) - - self.aggressive_raunet = _should_use_aggressive_raunet(self) - else: - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_4096") - elif self.model == "sdxl_turbo": - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_turbo_1024") - else: - raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") - - if self.aggressive_raunet and self.switching_threshold_ratio == "T1_ratio": - # self.T1 = min(int(self.max_timestep * self.T1_ratio), int(8/50 * self.max_timestep)) - self.T1 = int(aggressive_step / 50 * self.max_timestep) - else: - self.T1 = int(self.max_timestep * self.T1_ratio) + self.T1_ratio, self.T1_start, self.T1_end = _get_raunet_step_range(self, ori_H, ori_W) + self.T1 = self.T1_end stride = self.stride padding = self.padding dilation = self.dilation - if self.timestep < self.T1: + if self.T1_start <= _get_current_step(self) < self.T1_end: stride = (4, 4) padding = (2, 2) dilation = (2, 2) @@ -1892,9 +1863,7 @@ def forward(self, hidden_states: torch.Tensor, scale=1.0) -> torch.Tensor: hidden_states = F.conv2d( hidden_states, self.weight, self.bias, stride, padding, dilation, self.groups ) - self.timestep += 1 - if self.timestep == self.max_timestep: - self.timestep = 0 + _advance_fallback_step(self) return hidden_states else: original_outputs = F.conv2d( @@ -1903,9 +1872,7 @@ def forward(self, hidden_states: torch.Tensor, scale=1.0) -> torch.Tensor: return original_outputs + (scale * self.lora_layer(hidden_states)) else: hidden_states = F.conv2d(hidden_states, self.weight, self.bias, stride, padding, dilation, self.groups) - self.timestep += 1 - if self.timestep == self.max_timestep: - self.timestep = 0 + _advance_fallback_step(self) return hidden_states return downsampler_block @@ -1917,47 +1884,19 @@ class upsampler_block(block_class): # Save for unpatching later _parent = block_class T1_ratio = 0 + T1_start = 0 + T1_end = 0 T1 = 0 timestep = 0 - aggressive_raunet = False max_timestep = 50 info: dict = None def forward(self, hidden_states: torch.Tensor, scale=1.0) -> torch.Tensor: self.max_timestep = _get_max_timesteps(self.info) ori_H, ori_W = self.info["size"] - if self.model == "sd15": - if ori_H < 256 or ori_W < 256: - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_1024") - else: - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sd15_2048") - elif self.model == "sdxl": - if ori_H < 512 or ori_W < 512: - if self.info["text_to_img_controlnet"]: - self.T1_ratio = _get_switching_threshold_ratio( - self, text_to_img_controlnet_switching_threshold_ratio_dict, "sdxl_2048" - ) - else: - self.T1_ratio = _get_switching_threshold_ratio( - self, switching_threshold_ratio_dict, "sdxl_2048" - ) - - self.aggressive_raunet = _should_use_aggressive_raunet(self) - else: - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_4096") - elif self.model == "sdxl_turbo": - self.T1_ratio = _get_switching_threshold_ratio(self, switching_threshold_ratio_dict, "sdxl_turbo_1024") - else: - raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") - - if self.aggressive_raunet and self.switching_threshold_ratio == "T1_ratio": - # self.T1 = min(int(self.max_timestep * self.T1_ratio), int(8/50 * self.max_timestep)) - self.T1 = int(aggressive_step / 50 * self.max_timestep) - else: - self.T1 = int(self.max_timestep * self.T1_ratio) - self.timestep += 1 - if self.timestep == self.max_timestep: - self.timestep = 0 + self.T1_ratio, self.T1_start, self.T1_end = _get_raunet_step_range(self, ori_H, ori_W) + self.T1 = self.T1_end + _advance_fallback_step(self) if old_diffusers: if self.lora_layer is None: @@ -1991,15 +1930,16 @@ def hook(module, args): _HIDIFFUSION_RUNTIME_ATTRIBUTES = ( "timestep", - "aggressive_raunet", "T1_ratio", - "T1", "T1_start", "T1_end", + "T1", "max_timestep", "ori_stride", "ori_padding", "ori_dilation", + "_hidiffusion_window_shift_step", + "_hidiffusion_window_shift_bucket", ) _HIDIFFUSION_STATE_ATTRIBUTES = ( "stride", @@ -2062,9 +2002,6 @@ def apply_hidiffusion( t1_ratio: float | None = None, t2_ratio: float | None = None, is_inpainting_task: bool | None = None, - use_aggressive_raunet: bool | None = None, - denoising_start: float = 0.0, - denoising_end: float = 1.0, ): """ model: diffusers model. We support SD 1.5, 2.1, XL, XL Turbo. @@ -2148,9 +2085,7 @@ def apply_hidiffusion( "text_to_img_controlnet": has_controlnet and is_controlnet_text_to_image, "is_inpainting_task": detected_inpainting_task if is_inpainting_task is None else is_inpainting_task, "is_playground": is_playground, - "use_aggressive_raunet": use_aggressive_raunet, - "denoising_start": denoising_start, - "denoising_end": denoising_end, + "step_index": None, "pipeline": model, "switching_threshold_overrides": {"T1_ratio": t1_ratio, "T2_ratio": t2_ratio}, } diff --git a/invokeai/backend/stable_diffusion/diffusers_pipeline.py b/invokeai/backend/stable_diffusion/diffusers_pipeline.py index 6324e451bb6..06c773b2448 100644 --- a/invokeai/backend/stable_diffusion/diffusers_pipeline.py +++ b/invokeai/backend/stable_diffusion/diffusers_pipeline.py @@ -453,6 +453,11 @@ def step( # invokeai_diffuser has batched timesteps, but diffusers schedulers expect a single value timestep = t[0] + # HiDiffusion modules share this dictionary. Setting the logical step here prevents separate CFG forwards from + # advancing the RAU-Net schedule independently and lets window attention reuse one shift for the whole step. + if hasattr(self.unet, "info"): + self.unet.info["step_index"] = step_index + # Handle masked image-to-image (a.k.a inpainting). if mask_guidance is not None: # NOTE: This is intentionally done *before* self.scheduler.scale_model_input(...). diff --git a/invokeai/backend/stable_diffusion/diffusion/regional_ip_data.py b/invokeai/backend/stable_diffusion/diffusion/regional_ip_data.py index eb2a9eea56b..c14edf0b8a0 100644 --- a/invokeai/backend/stable_diffusion/diffusion/regional_ip_data.py +++ b/invokeai/backend/stable_diffusion/diffusion/regional_ip_data.py @@ -11,11 +11,11 @@ def __init__( masks: list[torch.Tensor], dtype: torch.dtype, device: torch.device, - max_downscale_factor: int = 16, + max_downscale_factor: int = 32, ): """Initialize an `IPAdapterConditioningData` object. - HiDiffusion's RAU-Net requires one mask level beyond the standard UNet's 8x downscale. + SD1/SD2 with both HiDiffusion RAU-Net stages can reach a 32x downscale. """ assert len(image_prompt_embeds) == len(scales) == len(masks) diff --git a/invokeai/backend/stable_diffusion/diffusion/regional_prompt_data.py b/invokeai/backend/stable_diffusion/diffusion/regional_prompt_data.py index bebe00b35b1..d0da01d724b 100644 --- a/invokeai/backend/stable_diffusion/diffusion/regional_prompt_data.py +++ b/invokeai/backend/stable_diffusion/diffusion/regional_prompt_data.py @@ -19,7 +19,7 @@ def __init__( regions: list[TextConditioningRegions], device: torch.device, dtype: torch.dtype, - max_downscale_factor: int = 16, + max_downscale_factor: int = 32, ): """Initialize a `RegionalPromptData` object. Args: @@ -28,7 +28,7 @@ def __init__( device (torch.device): The device to use for the attention masks. dtype (torch.dtype): The data type to use for the attention masks. max_downscale_factor: Spatial masks will be prepared for downscale factors from 1 to max_downscale_factor - in steps of 2x. HiDiffusion's RAU-Net requires one level beyond the standard UNet's 8x downscale. + in steps of 2x. SD1/SD2 with both HiDiffusion RAU-Net stages can reach a 32x downscale. """ self._regions = regions self._device = device @@ -41,7 +41,7 @@ def __init__( self._negative_cross_attn_mask_score = -10000.0 def _prepare_spatial_masks( - self, regions: list[TextConditioningRegions], max_downscale_factor: int = 16 + self, regions: list[TextConditioningRegions], max_downscale_factor: int = 32 ) -> list[dict[int, torch.Tensor]]: """Prepare the spatial masks for all downscaling factors.""" # batch_masks_by_seq_len[b][s] contains the spatial masks for the b'th batch sample with a query sequence length diff --git a/invokeai/backend/stable_diffusion/extensions/hidiffusion.py b/invokeai/backend/stable_diffusion/extensions/hidiffusion.py index 464e6a1d197..129d4c68a4b 100644 --- a/invokeai/backend/stable_diffusion/extensions/hidiffusion.py +++ b/invokeai/backend/stable_diffusion/extensions/hidiffusion.py @@ -6,7 +6,9 @@ import torch from diffusers import UNet2DConditionModel -from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase +from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext +from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType +from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase, callback from invokeai.backend.stable_diffusion.hidiffusion_utils import hidiffusion_patch from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage @@ -23,9 +25,6 @@ def __init__( has_controlnet: bool = False, is_controlnet_text_to_image: bool = False, is_inpainting_task: bool | None = None, - use_aggressive_raunet: bool | None = None, - denoising_start: float = 0.0, - denoising_end: float = 1.0, ): super().__init__() self._name_or_path = name_or_path @@ -34,13 +33,16 @@ def __init__( self._has_controlnet = has_controlnet self._is_controlnet_text_to_image = is_controlnet_text_to_image self._is_inpainting_task = is_inpainting_task - self._use_aggressive_raunet = use_aggressive_raunet - self._denoising_start = denoising_start - self._denoising_end = denoising_end self._t1_ratio = t1_ratio self._t2_ratio = t2_ratio self._generator = generator + @callback(ExtensionCallbackType.PRE_STEP, order=-1000) + def set_step_index(self, ctx: DenoiseContext) -> None: + """Keep HiDiffusion scheduling stable across all UNet forwards in one denoising step.""" + if ctx.unet is not None and hasattr(ctx.unet, "info"): + ctx.unet.info["step_index"] = ctx.step_index + @contextmanager def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeightsStorage): with hidiffusion_patch( @@ -54,8 +56,5 @@ def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeigh t2_ratio=self._t2_ratio, generator=self._generator, is_inpainting_task=self._is_inpainting_task, - use_aggressive_raunet=self._use_aggressive_raunet, - denoising_start=self._denoising_start, - denoising_end=self._denoising_end, ): yield None diff --git a/invokeai/backend/stable_diffusion/hidiffusion_utils.py b/invokeai/backend/stable_diffusion/hidiffusion_utils.py index 371da5ca6bd..29d6cdc37ce 100644 --- a/invokeai/backend/stable_diffusion/hidiffusion_utils.py +++ b/invokeai/backend/stable_diffusion/hidiffusion_utils.py @@ -19,9 +19,6 @@ def hidiffusion_patch( has_controlnet: bool = False, is_controlnet_text_to_image: bool = False, is_inpainting_task: bool | None = None, - use_aggressive_raunet: bool | None = None, - denoising_start: float = 0.0, - denoising_end: float = 1.0, ): """Context manager that applies HiDiffusion and restores the model on exit.""" from invokeai.backend.hidiffusion.hidiffusion import apply_hidiffusion, remove_hidiffusion @@ -101,9 +98,6 @@ def _set_name_or_path_on_config(cfg, value: str) -> bool: is_controlnet_text_to_image=is_controlnet_text_to_image, generator=generator, is_inpainting_task=is_inpainting_task, - use_aggressive_raunet=use_aggressive_raunet, - denoising_start=denoising_start, - denoising_end=denoising_end, ) yield finally: diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json index a63e37c9b89..68fc1034360 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -2187,22 +2187,22 @@ "hidiffusionRatiosAuto": { "heading": "HiDiffusion: Automatic Ratios", "paragraphs": [ - "Uses the model- and resolution-specific T1 and T2 ratios recommended by HiDiffusion.", - "Disable this to set explicit ratio overrides below." + "Uses HiDiffusion's original model- and resolution-specific presets behind the scenes.", + "Your manual T1 and T2 slider values are preserved when this is enabled, so you can switch back without re-entering them." ] }, "hidiffusionT1Ratio": { "heading": "HiDiffusion: T1 Ratio", "paragraphs": [ - "Controls the early switch point for HiDiffusion (T1).", - "Lower values switch earlier; higher values preserve global structure longer." + "Controls how long the primary RAU-Net stage remains active. At extreme resolutions, this is the later of the two RAU-Net cutoffs.", + "The T1 name follows the upstream code key; lower values switch earlier." ] }, "hidiffusionT2Ratio": { "heading": "HiDiffusion: T2 Ratio", "paragraphs": [ - "Controls the late switch point for HiDiffusion (T2).", - "Higher values keep the second RAU-Net resolution-aware stage active longer." + "Controls the additional extreme-resolution RAU-Net stage. When both stages are active, T2 is the earlier cutoff despite its name.", + "The T2 name follows the upstream code key. T2 cannot exceed T1; excessive values can reduce pose and background diversity or introduce artifacts." ] }, "clipSkip": { diff --git a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.test.ts b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.test.ts index 30c2690c72a..453451d0098 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.test.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.test.ts @@ -23,6 +23,7 @@ import { selectModelSupportsRefImages, selectModelSupportsSeed, selectModelSupportsSteps, + setHiDiffusionAutoRatios, setIdeogram4Steps, } from './paramsSlice'; import { getInitialParamsState, zParamsState } from './types'; @@ -807,6 +808,7 @@ const POST_V4_BUMP_DEFAULTED_KEYS = [ 'hiDiffusionEnabled', 'hiDiffusionRauNetEnabled', 'hiDiffusionWindowAttnEnabled', + 'hiDiffusionAutoRatios', 'hiDiffusionT1Ratio', 'hiDiffusionT2Ratio', ] as const satisfies readonly (keyof typeof zParamsState.shape)[]; @@ -894,6 +896,7 @@ describe('paramsSliceConfig persisted state migration', () => { delete v2State.hiDiffusionEnabled; delete v2State.hiDiffusionRauNetEnabled; delete v2State.hiDiffusionWindowAttnEnabled; + delete v2State.hiDiffusionAutoRatios; delete v2State.hiDiffusionT1Ratio; delete v2State.hiDiffusionT2Ratio; @@ -901,14 +904,15 @@ describe('paramsSliceConfig persisted state migration', () => { // v2 migrates all the way through the current chain (v2 -> v3 adds Qwen fields, // v3 -> v4 adds Krea-2 and PiD fields). - expect(result._version).toBe(6); + expect(result._version).toBe(7); 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).toBeNull(); - expect(result.hiDiffusionT2Ratio).toBeNull(); + expect(result.hiDiffusionAutoRatios).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); @@ -934,8 +938,10 @@ describe('paramsSliceConfig persisted state migration', () => { hiDiffusionT2Ratio: 0.25, }) as ReturnType; - expect(oldDefaults.hiDiffusionT1Ratio).toBeNull(); - expect(oldDefaults.hiDiffusionT2Ratio).toBeNull(); + expect(oldDefaults.hiDiffusionAutoRatios).toBe(true); + expect(oldDefaults.hiDiffusionT1Ratio).toBe(0.4); + expect(oldDefaults.hiDiffusionT2Ratio).toBe(0.0); + expect(customOverrides.hiDiffusionAutoRatios).toBe(false); expect(customOverrides.hiDiffusionT1Ratio).toBe(0.65); expect(customOverrides.hiDiffusionT2Ratio).toBe(0.25); }); @@ -962,7 +968,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(v3State) as ReturnType & Record; - expect(result._version).toBe(6); + expect(result._version).toBe(7); expect((result.flux2VaeModel as { key: string } | null)?.key).toBe('klein-vae'); // The new standalone dev Mistral encoder slot must be seeded, not left undefined. expect(result.flux2DevMistralEncoderModel).toBeNull(); @@ -996,7 +1002,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(v3State) as ReturnType; - expect(result._version).toBe(6); + expect(result._version).toBe(7); expect(result.krea2VaeModel).toBeNull(); expect(result.krea2Qwen3VlEncoderModel).toBeNull(); expect(result.krea2SeedVarianceEnabled).toBe(false); @@ -1030,7 +1036,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(mainV4State) as ReturnType & Record; - expect(result._version).toBe(6); + expect(result._version).toBe(7); expect((result.flux2VaeModel as { key: string } | null)?.key).toBe('klein-vae'); expect(result.flux2DevMistralEncoderModel).toBeNull(); // main's own v4 values must survive untouched. @@ -1058,7 +1064,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(devV4State) as ReturnType & Record; - expect(result._version).toBe(6); + expect(result._version).toBe(7); // The branch's own v4 values must survive untouched. expect((result.flux2VaeModel as { key: string } | null)?.key).toBe('flux2-vae'); expect(result.pidMode).toBe('off'); @@ -1115,7 +1121,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(blob) as ReturnType; - expect(result._version).toBe(6); + expect(result._version).toBe(7); expect(result.positivePrompt).toBe('a fluffy cat'); expect(result.seed).toBe(42); expect(result.shouldRandomizeSeed).toBe(false); @@ -1281,7 +1287,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(blob) as ReturnType; - expect(result._version).toBe(6); + expect(result._version).toBe(7); expect(result.dimensions).toEqual(getInitialParamsState().dimensions); expect(result.positivePrompt).toBe('a fluffy cat'); expect(result.seed).toBe(7); @@ -1298,7 +1304,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(blob) as ReturnType; - expect(result._version).toBe(6); + expect(result._version).toBe(7); expect(result.positivePromptHistory).toEqual([]); expect(result.qwenImageVaeModel).toBeNull(); expect(result.wanVaeModel).toBeNull(); @@ -1313,13 +1319,13 @@ describe('paramsSliceConfig persisted state migration', () => { // `_version` before the net ever sees it, so only calling the net directly tests the guard. // The blob is otherwise complete (the current tier's key set), so `_version` is the only thing // the parse below can object to. - const blob = buildReleaseBlob('f10d2a4f5a', { _version: 7, positivePrompt: 'a fluffy cat' }); + const blob = buildReleaseBlob('f10d2a4f5a', { _version: 8, positivePrompt: 'a fluffy cat' }); const { backfilled, reset } = repairParamsState(blob); expect(backfilled).toEqual([]); expect(reset).toEqual([]); - expect(blob._version).toBe(7); + expect(blob._version).toBe(8); // Still fatal, which is the correct outcome for a downgrade: that slice really was written by a // schema this build does not know. expect(() => zParamsState.parse(blob)).toThrow(); @@ -1335,7 +1341,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(blob) as ReturnType; - expect(result._version).toBe(6); + expect(result._version).toBe(7); expect(result.positivePrompt).toBe('a fluffy cat'); expect(result.seed).toBe(7); expect(result.dimensions).toBeDefined(); @@ -1352,7 +1358,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(v3State) as ReturnType; - expect(result._version).toBe(6); + expect(result._version).toBe(7); expect(result.wanTransformerLowNoise).toBeNull(); expect(result.wanComponentSource).toBeNull(); expect(result.wanVaeModel).toBeNull(); @@ -1374,7 +1380,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(v2State) as ReturnType; - expect(result._version).toBe(6); + expect(result._version).toBe(7); expect(result.fluxScheduler).toBe('euler'); expect(result.zImageScheduler).toBe('euler'); expect(result.colorCompensation).toBe(false); @@ -1434,7 +1440,7 @@ describe('paramsSliceConfig persisted state migration', () => { expect('hiDiffusionEnabled' in blob).toBe(false); applyParamsVersionMigrations(blob); - expect(blob._version).toBe(6); + expect(blob._version).toBe(7); // The value assertions below cannot, on their own, prove the defaults exist: three mechanisms // produce the identical values, so any two can hide the third being reverted. Parsing directly @@ -1459,8 +1465,9 @@ describe('paramsSliceConfig persisted state migration', () => { expect(result.hiDiffusionEnabled).toBe(false); expect(result.hiDiffusionRauNetEnabled).toBe(true); expect(result.hiDiffusionWindowAttnEnabled).toBe(true); - expect(result.hiDiffusionT1Ratio).toBeNull(); - expect(result.hiDiffusionT2Ratio).toBeNull(); + expect(result.hiDiffusionAutoRatios).toBe(true); + expect(result.hiDiffusionT1Ratio).toBe(0.4); + expect(result.hiDiffusionT2Ratio).toBe(0.0); expect(result.positivePrompt).toBe('a fluffy cat'); expect(result.seed).toBe(42); }); @@ -1578,6 +1585,25 @@ describe('paramsSlice prompt history', () => { }); }); +describe('paramsSlice HiDiffusion automatic ratios', () => { + it('changes only the automatic mode and preserves manual slider values', () => { + const initial = { + ...getInitialParamsState(), + hiDiffusionAutoRatios: false, + hiDiffusionT1Ratio: 0.65, + hiDiffusionT2Ratio: 0.2, + }; + + const automatic = paramsSliceConfig.slice.reducer(initial, setHiDiffusionAutoRatios(true)); + const manual = paramsSliceConfig.slice.reducer(automatic, setHiDiffusionAutoRatios(false)); + + expect(automatic.hiDiffusionT1Ratio).toBe(0.65); + expect(automatic.hiDiffusionT2Ratio).toBe(0.2); + expect(manual.hiDiffusionT1Ratio).toBe(0.65); + expect(manual.hiDiffusionT2Ratio).toBe(0.2); + }); +}); + describe('paramsSlice ideogram4Steps normalization (backend requires >= 2)', () => { it('keeps a valid override step count', () => { const state = paramsSliceConfig.slice.reducer(getInitialParamsState(), setIdeogram4Steps(20)); diff --git a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts index 3db8e6b356d..f58a9ef3813 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts @@ -182,10 +182,13 @@ const slice = createSlice({ setHiDiffusionWindowAttnEnabled: (state, action: PayloadAction) => { state.hiDiffusionWindowAttnEnabled = action.payload; }, - setHiDiffusionT1Ratio: (state, action: PayloadAction) => { + setHiDiffusionAutoRatios: (state, action: PayloadAction) => { + state.hiDiffusionAutoRatios = action.payload; + }, + setHiDiffusionT1Ratio: (state, action: PayloadAction) => { state.hiDiffusionT1Ratio = action.payload; }, - setHiDiffusionT2Ratio: (state, action: PayloadAction) => { + setHiDiffusionT2Ratio: (state, action: PayloadAction) => { state.hiDiffusionT2Ratio = action.payload; }, setSeamlessXAxis: (state, action: PayloadAction) => { @@ -895,6 +898,7 @@ export const { setHiDiffusionEnabled, setHiDiffusionRauNetEnabled, setHiDiffusionWindowAttnEnabled, + setHiDiffusionAutoRatios, setHiDiffusionT1Ratio, setHiDiffusionT2Ratio, setSeamlessXAxis, @@ -1171,6 +1175,18 @@ export const applyParamsVersionMigrations = (state: any): void => { state.hiDiffusionT1Ratio = state.hiDiffusionT1Ratio === 0.4 ? null : state.hiDiffusionT1Ratio; state.hiDiffusionT2Ratio = state.hiDiffusionT2Ratio === 0.0 ? null : state.hiDiffusionT2Ratio; } + + if (state._version === 6) { + // v6 -> v7: keep the automatic-ratio mode separate from the manual slider values. Older + // states used two null ratios to represent Auto, which erased the user's manual values every + // time the mode was enabled. + state._version = 7; + state.hiDiffusionAutoRatios = + (state.hiDiffusionT1Ratio === null || state.hiDiffusionT1Ratio === undefined) && + (state.hiDiffusionT2Ratio === null || state.hiDiffusionT2Ratio === undefined); + state.hiDiffusionT1Ratio = state.hiDiffusionT1Ratio ?? 0.4; + state.hiDiffusionT2Ratio = state.hiDiffusionT2Ratio ?? 0.0; + } }; export const paramsSliceConfig: SliceConfig = { @@ -1286,6 +1302,7 @@ export const selectOptimizedDenoisingEnabled = createParamsSelector((params) => export const selectHiDiffusionEnabled = createParamsSelector((params) => params.hiDiffusionEnabled); export const selectHiDiffusionRauNetEnabled = createParamsSelector((params) => params.hiDiffusionRauNetEnabled); export const selectHiDiffusionWindowAttnEnabled = createParamsSelector((params) => params.hiDiffusionWindowAttnEnabled); +export const selectHiDiffusionAutoRatios = createParamsSelector((params) => params.hiDiffusionAutoRatios); export const selectHiDiffusionT1Ratio = createParamsSelector((params) => params.hiDiffusionT1Ratio); export const selectHiDiffusionT2Ratio = createParamsSelector((params) => params.hiDiffusionT2Ratio); export const selectPositivePrompt = createParamsSelector((params) => params.positivePrompt); diff --git a/invokeai/frontend/web/src/features/controlLayers/store/types.ts b/invokeai/frontend/web/src/features/controlLayers/store/types.ts index b6fff1580f7..4913df01c86 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/types.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/types.ts @@ -817,7 +817,7 @@ const zPidMode = z.enum(['off', 'fit', 'native']); export type PidMode = z.infer; export const zParamsState = z.object({ - _version: z.literal(6), + _version: z.literal(7), maskBlur: z.number(), maskBlurMethod: zParameterMaskBlurMethod, canvasCoherenceMode: zParameterCanvasCoherenceMode, @@ -839,8 +839,9 @@ export const zParamsState = z.object({ hiDiffusionEnabled: z.boolean().default(false), hiDiffusionRauNetEnabled: z.boolean().default(true), hiDiffusionWindowAttnEnabled: z.boolean().default(true), - hiDiffusionT1Ratio: z.number().nullable().default(null), - hiDiffusionT2Ratio: z.number().nullable().default(null), + hiDiffusionAutoRatios: z.boolean().default(true), + hiDiffusionT1Ratio: z.number().default(0.4), + hiDiffusionT2Ratio: z.number().default(0.0), iterations: z.number(), scheduler: zParameterScheduler, fluxScheduler: zParameterFluxScheduler, @@ -965,7 +966,7 @@ export const zParamsState = z.object({ }); export type ParamsState = z.infer; export const getInitialParamsState = (): ParamsState => ({ - _version: 6, + _version: 7, maskBlur: 16, maskBlurMethod: 'box', canvasCoherenceMode: 'Gaussian Blur', @@ -983,8 +984,9 @@ export const getInitialParamsState = (): ParamsState => ({ hiDiffusionEnabled: false, hiDiffusionRauNetEnabled: true, hiDiffusionWindowAttnEnabled: true, - hiDiffusionT1Ratio: null, - hiDiffusionT2Ratio: null, + hiDiffusionAutoRatios: true, + hiDiffusionT1Ratio: 0.4, + hiDiffusionT2Ratio: 0.0, iterations: 1, scheduler: 'dpmpp_3m_k', fluxScheduler: 'euler', diff --git a/invokeai/frontend/web/src/features/metadata/parsing.test.ts b/invokeai/frontend/web/src/features/metadata/parsing.test.ts index 12c8a96fa9e..68cc6a3d727 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.test.ts +++ b/invokeai/frontend/web/src/features/metadata/parsing.test.ts @@ -1,5 +1,6 @@ import type { AppStore } from 'app/store/store'; import { + setHiDiffusionAutoRatios, setHiDiffusionEnabled, setHiDiffusionT1Ratio, setHiDiffusionT2Ratio, @@ -176,18 +177,36 @@ describe('HiDiffusion metadata parsing', () => { ImageMetadataHandlers.HiDiffusionT1Ratio.recall(t1, store); ImageMetadataHandlers.HiDiffusionT2Ratio.recall(t2, store); - expect(store.dispatch).toHaveBeenCalledWith(setHiDiffusionT1Ratio(null)); - expect(store.dispatch).toHaveBeenCalledWith(setHiDiffusionT2Ratio(null)); + expect(store.dispatch).toHaveBeenCalledWith(setHiDiffusionAutoRatios(true)); + expect(store.dispatch).not.toHaveBeenCalledWith(setHiDiffusionT1Ratio(expect.anything())); + expect(store.dispatch).not.toHaveBeenCalledWith(setHiDiffusionT2Ratio(expect.anything())); + }); + + it('recalls numeric ratios as manual thresholds', async () => { + const store = createStore(); + const metadata = { hidiffusion_t1_ratio: 0.65, hidiffusion_t2_ratio: 0.2 }; + + const t1 = await parseMetadataHandler(metadata, ImageMetadataHandlers.HiDiffusionT1Ratio, store); + const t2 = await parseMetadataHandler(metadata, ImageMetadataHandlers.HiDiffusionT2Ratio, store); + ImageMetadataHandlers.HiDiffusionT1Ratio.recall(t1, store); + ImageMetadataHandlers.HiDiffusionT2Ratio.recall(t2, store); + + expect(store.dispatch).toHaveBeenCalledWith(setHiDiffusionAutoRatios(false)); + expect(store.dispatch).toHaveBeenCalledWith(setHiDiffusionT1Ratio(0.65)); + expect(store.dispatch).toHaveBeenCalledWith(setHiDiffusionT2Ratio(0.2)); }); it('disables HiDiffusion when recalling all metadata from an older image', async () => { let hiDiffusionEnabled = true; - let hiDiffusionT1Ratio: number | null = 0.8; - let hiDiffusionT2Ratio: number | null = 0.6; + let hiDiffusionAutoRatios = false; + let hiDiffusionT1Ratio = 0.8; + let hiDiffusionT2Ratio = 0.6; const store = { dispatch: vi.fn((action) => { if (action.type === setHiDiffusionEnabled.type) { hiDiffusionEnabled = action.payload; + } else if (action.type === setHiDiffusionAutoRatios.type) { + hiDiffusionAutoRatios = action.payload; } else if (action.type === setHiDiffusionT1Ratio.type) { hiDiffusionT1Ratio = action.payload; } else if (action.type === setHiDiffusionT2Ratio.type) { @@ -216,7 +235,8 @@ describe('HiDiffusion metadata parsing', () => { expect(store.dispatch).toHaveBeenCalledWith(setHiDiffusionEnabled(false)); expect(hiDiffusionEnabled).toBe(false); - expect(hiDiffusionT1Ratio).toBeNull(); - expect(hiDiffusionT2Ratio).toBeNull(); + expect(hiDiffusionAutoRatios).toBe(true); + expect(hiDiffusionT1Ratio).toBe(0.8); + expect(hiDiffusionT2Ratio).toBe(0.6); }); }); diff --git a/invokeai/frontend/web/src/features/metadata/parsing.tsx b/invokeai/frontend/web/src/features/metadata/parsing.tsx index 0c57a4058d9..5219be9f57d 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.tsx +++ b/invokeai/frontend/web/src/features/metadata/parsing.tsx @@ -44,6 +44,7 @@ import { setFluxDypeScale, setFluxScheduler, setGuidance, + setHiDiffusionAutoRatios, setHiDiffusionEnabled, setHiDiffusionRauNetEnabled, setHiDiffusionT1Ratio, @@ -801,7 +802,10 @@ const HiDiffusionT1Ratio: SingleMetadataHandler = { return Promise.resolve(parsed); }, recall: (value, store) => { - store.dispatch(setHiDiffusionT1Ratio(value)); + store.dispatch(setHiDiffusionAutoRatios(value === null)); + if (value !== null) { + store.dispatch(setHiDiffusionT1Ratio(value)); + } }, i18nKey: 'metadata.hiDiffusionT1Ratio', LabelComponent: MetadataLabel, @@ -819,7 +823,10 @@ const HiDiffusionT2Ratio: SingleMetadataHandler = { return Promise.resolve(parsed); }, recall: (value, store) => { - store.dispatch(setHiDiffusionT2Ratio(value)); + store.dispatch(setHiDiffusionAutoRatios(value === null)); + if (value !== null) { + store.dispatch(setHiDiffusionT2Ratio(value)); + } }, i18nKey: 'metadata.hiDiffusionT2Ratio', LabelComponent: MetadataLabel, diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.test.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.test.ts new file mode 100644 index 00000000000..ac28804fc52 --- /dev/null +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.test.ts @@ -0,0 +1,61 @@ +import type { Invocation } from 'services/api/types'; +import { describe, expect, it, vi } from 'vitest'; + +let nextId = 0; +vi.mock('features/controlLayers/konva/util', () => ({ + getPrefixedId: (prefix: string) => `${prefix}:${nextId++}`, +})); + +const refinerModel = { + key: 'refiner-model', + hash: 'refiner-hash', + name: 'SDXL Refiner', + base: 'sdxl', + type: 'main', +}; + +vi.mock('features/metadata/util/modelFetchingHelpers', () => ({ + fetchModelConfigWithTypeGuard: vi.fn(() => Promise.resolve(refinerModel)), +})); + +import { addSDXLRefiner } from './addSDXLRefiner'; +import { Graph } from './Graph'; + +describe('addSDXLRefiner', () => { + it('does not apply HiDiffusion to the unsupported refiner stage', async () => { + const g = new Graph('test'); + const denoise = g.addNode({ type: 'denoise_latents', id: 'base-denoise' } as Invocation<'denoise_latents'>); + const posCond = g.addNode({ type: 'sdxl_compel_prompt', id: 'pos' } as Invocation<'sdxl_compel_prompt'>); + const negCond = g.addNode({ type: 'sdxl_compel_prompt', id: 'neg' } as Invocation<'sdxl_compel_prompt'>); + const l2i = g.addNode({ type: 'l2i', id: 'l2i' } as Invocation<'l2i'>); + + const state = { + params: { + refinerModel, + refinerPositiveAestheticScore: 6, + refinerNegativeAestheticScore: 2.5, + refinerSteps: 20, + refinerScheduler: 'euler', + refinerCFGScale: 7.5, + refinerStart: 0.8, + hiDiffusionEnabled: true, + hiDiffusionRauNetEnabled: true, + hiDiffusionWindowAttnEnabled: true, + hiDiffusionT1Ratio: 0.4, + hiDiffusionT2Ratio: 0.3, + }, + } as never; + + await addSDXLRefiner(state, g, denoise, null, posCond, negCond, l2i); + + const refinerDenoise = Object.values(g.getGraph().nodes).find( + (node) => node.type === 'denoise_latents' && node.id !== denoise.id + ) as Invocation<'denoise_latents'> | undefined; + expect(refinerDenoise).toBeDefined(); + expect(refinerDenoise?.hidiffusion).toBeUndefined(); + expect(refinerDenoise?.hidiffusion_raunet).toBeUndefined(); + expect(refinerDenoise?.hidiffusion_window_attn).toBeUndefined(); + expect(refinerDenoise?.hidiffusion_t1_ratio).toBeUndefined(); + expect(refinerDenoise?.hidiffusion_t2_ratio).toBeUndefined(); + }); +}); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.ts index 2f755377df9..5485834db13 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addSDXLRefiner.ts @@ -23,11 +23,6 @@ export const addSDXLRefiner = async ( refinerScheduler, refinerCFGScale, refinerStart, - hiDiffusionEnabled, - hiDiffusionRauNetEnabled, - hiDiffusionT1Ratio, - hiDiffusionT2Ratio, - hiDiffusionWindowAttnEnabled, } = state.params; assert(refinerModel, 'No refiner model found in state'); @@ -62,11 +57,6 @@ 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) : undefined, - hidiffusion_t2_ratio: hiDiffusionEnabled ? (hiDiffusionT2Ratio ?? undefined) : undefined, denoising_start: refinerStart, denoising_end: 1, }); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildHiDiffusionGraph.test.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildHiDiffusionGraph.test.ts index 6ed9b077a82..748fc7dc7db 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildHiDiffusionGraph.test.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildHiDiffusionGraph.test.ts @@ -32,8 +32,9 @@ const defaultParams = { cfgRescaleMultiplier: 0, hiDiffusionEnabled: false, hiDiffusionRauNetEnabled: false, - hiDiffusionT1Ratio: 0.25 as number | null, - hiDiffusionT2Ratio: 0.1 as number | null, + hiDiffusionAutoRatios: false, + hiDiffusionT1Ratio: 0.25, + hiDiffusionT2Ratio: 0.1, hiDiffusionWindowAttnEnabled: false, scheduler: 'euler', steps: 20, @@ -198,8 +199,9 @@ describe('HiDiffusion graph metadata', () => { currentModel = sdxlModel; params.hiDiffusionEnabled = true; params.hiDiffusionRauNetEnabled = true; - params.hiDiffusionT1Ratio = null; - params.hiDiffusionT2Ratio = null; + params.hiDiffusionAutoRatios = true; + params.hiDiffusionT1Ratio = 0.65; + params.hiDiffusionT2Ratio = 0.2; const { g } = await buildSDXLGraph(buildGraphArg()); const denoise = g.getNodes().find((node) => node.type === 'denoise_latents'); @@ -210,4 +212,24 @@ describe('HiDiffusion graph metadata', () => { expect(metadata.hidiffusion_t1_ratio).toBeNull(); expect(metadata.hidiffusion_t2_ratio).toBeNull(); }); + + it('omits automatic ratio overrides from the SD1 denoise while retaining manual slider values', async () => { + currentModel = sd1Model; + params.hiDiffusionEnabled = true; + params.hiDiffusionRauNetEnabled = true; + params.hiDiffusionAutoRatios = true; + params.hiDiffusionT1Ratio = 0.65; + params.hiDiffusionT2Ratio = 0.2; + + const { g } = await buildSD1Graph(buildGraphArg()); + const denoise = g.getNodes().find((node) => node.type === 'denoise_latents'); + const metadata = getMetadata(g); + + expect(denoise?.hidiffusion_t1_ratio).toBeUndefined(); + expect(denoise?.hidiffusion_t2_ratio).toBeUndefined(); + expect(metadata.hidiffusion_t1_ratio).toBeNull(); + expect(metadata.hidiffusion_t2_ratio).toBeNull(); + expect(params.hiDiffusionT1Ratio).toBe(0.65); + expect(params.hiDiffusionT2Ratio).toBe(0.2); + }); }); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts index 8ac62a08885..c7dce80a307 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts @@ -45,6 +45,7 @@ export const buildSD1Graph = async (arg: GraphBuilderArg): Promise { const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled); const hiDiffusionRauNetEnabled = useAppSelector(selectHiDiffusionRauNetEnabled); - const hiDiffusionT1Ratio = useAppSelector(selectHiDiffusionT1Ratio); - const hiDiffusionT2Ratio = useAppSelector(selectHiDiffusionT2Ratio); + const hiDiffusionAutoRatios = useAppSelector(selectHiDiffusionAutoRatios); const dispatch = useAppDispatch(); const { t } = useTranslation(); - const isAutomatic = hiDiffusionT1Ratio === null && hiDiffusionT2Ratio === null; const onChange = useCallback( (event: ChangeEvent) => { - if (event.target.checked) { - dispatch(setHiDiffusionT1Ratio(null)); - dispatch(setHiDiffusionT2Ratio(null)); - } else { - dispatch(setHiDiffusionT1Ratio(RATIO_CONSTRAINTS.t1.initial)); - dispatch(setHiDiffusionT2Ratio(RATIO_CONSTRAINTS.t2.initial)); - } + dispatch(setHiDiffusionAutoRatios(event.target.checked)); }, [dispatch] ); @@ -148,7 +142,7 @@ export const ParamHiDiffusionAutoRatiosToggle = memo(() => { @@ -161,18 +155,25 @@ ParamHiDiffusionAutoRatiosToggle.displayName = 'ParamHiDiffusionAutoRatiosToggle export const ParamHiDiffusionT1Ratio = memo(() => { const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled); const hiDiffusionRauNetEnabled = useAppSelector(selectHiDiffusionRauNetEnabled); + const hiDiffusionAutoRatios = useAppSelector(selectHiDiffusionAutoRatios); const hiDiffusionT1Ratio = useAppSelector(selectHiDiffusionT1Ratio); const hiDiffusionT2Ratio = useAppSelector(selectHiDiffusionT2Ratio); const dispatch = useAppDispatch(); const { t } = useTranslation(); - const onChange = useCallback((value: number) => dispatch(setHiDiffusionT1Ratio(value)), [dispatch]); + const onChange = useCallback( + (value: number) => { + dispatch(setHiDiffusionT1Ratio(value)); + if (hiDiffusionT2Ratio > value) { + dispatch(setHiDiffusionT2Ratio(value)); + } + }, + [dispatch, hiDiffusionT2Ratio] + ); return ( @@ -181,7 +182,7 @@ export const ParamHiDiffusionT1Ratio = memo(() => { { marks /> { const hiDiffusionEnabled = useAppSelector(selectHiDiffusionEnabled); const hiDiffusionRauNetEnabled = useAppSelector(selectHiDiffusionRauNetEnabled); + const hiDiffusionAutoRatios = useAppSelector(selectHiDiffusionAutoRatios); const hiDiffusionT1Ratio = useAppSelector(selectHiDiffusionT1Ratio); const hiDiffusionT2Ratio = useAppSelector(selectHiDiffusionT2Ratio); const dispatch = useAppDispatch(); @@ -217,9 +219,7 @@ export const ParamHiDiffusionT2Ratio = memo(() => { return ( @@ -228,20 +228,20 @@ export const ParamHiDiffusionT2Ratio = memo(() => { torch.Tensor: torch.testing.assert_close(first, second) +def test_hidiffusion_window_attention_reuses_shift_within_logical_step(): + module_keys = { + "down_module_key": [], + "down_module_key_extra": [], + "up_module_key": [], + "up_module_key_extra": [], + "windown_attn_module_key": ["transformer"], + } + model = WindowAttentionModelMixin() + generator = torch.Generator(device="cpu").manual_seed(1234) + hidden_states = torch.arange(64, dtype=torch.float32).reshape(1, 64, 1) + + 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) + model.info["step_index"] = 0 + first = model.transformer(hidden_states).clone() + generator_state_after_first_forward = generator.get_state().clone() + second = model.transformer(hidden_states).clone() + + torch.testing.assert_close(first, second) + torch.testing.assert_close(generator.get_state(), generator_state_after_first_forward) + + model.info["step_index"] = 1 + model.transformer(hidden_states) + assert not torch.equal(generator.get_state(), generator_state_after_first_forward) + + @pytest.mark.parametrize("is_text_to_image", [False, True]) def test_hidiffusion_patch_uses_controlnet_aware_forward_for_bare_unet(is_text_to_image: bool): model = ModelMixin() @@ -187,22 +227,20 @@ def test_hidiffusion_patch_resets_cached_runtime_state_when_reenabled(): 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.T1 = 9 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.T1 == 0 assert model.block.max_timestep == 50 @@ -409,66 +447,249 @@ def test_hidiffusion_patch_forwards_generation_context(): model, name_or_path="stabilityai/stable-diffusion-xl-base-1.0", is_inpainting_task=True, - use_aggressive_raunet=False, - denoising_start=0.6, - denoising_end=1.0, ): pass kwargs = mock_apply_hidiffusion.call_args.kwargs assert kwargs["is_inpainting_task"] is True - assert kwargs["use_aggressive_raunet"] is False - assert kwargs["denoising_start"] == 0.6 - assert kwargs["denoising_end"] == 1.0 @pytest.mark.parametrize( - ("denoising_start", "denoising_end", "expected_ratio"), + ("size", "threshold", "override", "expected_ratio"), [ - (0.0, 1.0, 0.4), - (0.6, 1.0, 0.0), - (0.0, 0.2, 1.0), - (0.2, 0.6, 0.5), + ((256, 256), "T1_ratio", None, 0.4), + ((384, 384), "T1_ratio", None, 0.4), + ((512, 512), "T1_ratio", None, 0.7), + ((384, 384), "T2_ratio", None, 0.0), + ((512, 256), "T1_ratio", None, 0.4), + ((384, 384), "T1_ratio", 0.25, 0.25), ], ) -def test_hidiffusion_ratios_are_mapped_to_the_executed_denoising_range( - denoising_start: float, denoising_end: float, expected_ratio: float +def test_hidiffusion_ratios_use_upstream_discrete_presets( + size: tuple[int, int], threshold: str, override: float | None, expected_ratio: float ): module = SimpleNamespace( - switching_threshold_ratio="T1_ratio", + model="sdxl", + switching_threshold_ratio=threshold, info={ - "switching_threshold_overrides": {"T1_ratio": None, "T2_ratio": None}, - "denoising_start": denoising_start, - "denoising_end": denoising_end, + "switching_threshold_overrides": {"T1_ratio": override, "T2_ratio": override}, + "text_to_img_controlnet": False, }, ) - ratio = _get_switching_threshold_ratio(module, switching_threshold_ratio_dict, "sdxl_2048") + ratio = _get_resolution_aware_switching_threshold_ratio(module, *size) assert ratio == pytest.approx(expected_ratio) +def test_hidiffusion_controlnet_uses_its_normal_resolution_preset(): + module = SimpleNamespace( + model="sdxl", + switching_threshold_ratio="T1_ratio", + info={ + "switching_threshold_overrides": {"T1_ratio": None, "T2_ratio": None}, + "text_to_img_controlnet": True, + }, + ) + + assert _get_resolution_aware_switching_threshold_ratio(module, 256, 256) == pytest.approx(0.5) + assert _get_resolution_aware_switching_threshold_ratio(module, 512, 512) == pytest.approx(0.7) + + @pytest.mark.parametrize( - ("is_inpainting_task", "override", "expected"), + ("size", "threshold", "is_inpainting", "t2_override", "expected"), [ - (False, False, False), - (True, True, True), - (True, None, False), - (False, None, True), + ((256, 256), "T2_ratio", False, None, (0.0, 0, 8)), + ((256, 256), "T1_ratio", False, None, (0.4, 8, 20)), + ((256, 256), "T2_ratio", True, None, (0.0, 0, 0)), + ((256, 256), "T1_ratio", True, None, (0.4, 0, 20)), + ((512, 512), "T2_ratio", False, None, (0.3, 0, 15)), + ((512, 512), "T1_ratio", False, None, (0.7, 0, 35)), + ((256, 256), "T2_ratio", False, 0.1, (0.1, 0, 5)), + ((256, 256), "T1_ratio", False, 0.1, (0.4, 5, 20)), ], ) -def test_explicit_aggressive_raunet_setting_takes_precedence( - is_inpainting_task: bool, override: bool | None, expected: bool +def test_hidiffusion_raunet_schedule_matches_upstream_stages( + size: tuple[int, int], + threshold: str, + is_inpainting: bool, + t2_override: float | None, + expected: tuple[float, int, int], ): module = SimpleNamespace( + model="sdxl", + max_timestep=50, + switching_threshold_ratio=threshold, info={ - "is_inpainting_task": is_inpainting_task, + "switching_threshold_overrides": {"T1_ratio": None, "T2_ratio": t2_override}, + "text_to_img_controlnet": False, + "is_inpainting_task": is_inpainting, "is_playground": False, - "use_aggressive_raunet": override, - } + }, + ) + + assert _get_raunet_step_range(module, *size) == expected + + +@pytest.mark.parametrize("t1_override", [None, 0.4]) +def test_hidiffusion_rejects_t2_above_the_resolved_t1(t1_override: float | None): + module = SimpleNamespace( + model="sdxl", + switching_threshold_ratio="T2_ratio", + info={ + "switching_threshold_overrides": {"T1_ratio": t1_override, "T2_ratio": 0.5}, + "text_to_img_controlnet": False, + }, + ) + + with pytest.raises(ValueError, match="T2 ratio must be less than or equal to the T1 ratio"): + _get_resolution_aware_switching_threshold_ratio(module, 256, 256) + + +def test_denoise_invocation_rejects_explicit_t2_above_t1(): + invocation = DenoiseLatentsInvocation.model_construct(hidiffusion_t1_ratio=0.4, hidiffusion_t2_ratio=0.5) + + with pytest.raises(ValueError, match="T2 ratio must be less than or equal to the T1 ratio"): + invocation.validate_hidiffusion_ratio_order() + + +def test_logical_step_prevents_sequential_guidance_from_advancing_t2_twice(): + patched_conv = make_diffusers_downsampler_block(torch.nn.Conv2d) + module = patched_conv(1, 1, kernel_size=3, stride=2, padding=1, bias=False) + module.info = { + "size": (256, 256), + "pipeline": SimpleNamespace(_num_timesteps=10), + "text_to_img_controlnet": False, + "is_inpainting_task": False, + "is_playground": False, + "step_index": 0, + "switching_threshold_overrides": {"T1_ratio": None, "T2_ratio": 0.4}, + } + module.model = "sdxl" + module.switching_threshold_ratio = "T2_ratio" + hidden_states = torch.ones(1, 1, 8, 8) + + negative = module(hidden_states) + positive = module(hidden_states) + + assert negative.shape[-2:] == (2, 2) + assert positive.shape[-2:] == (2, 2) + assert module.timestep == 0 + + module.info["step_index"] = 4 + after_t2 = module(hidden_states) + assert after_t2.shape[-2:] == (4, 4) + + +def test_hidiffusion_extension_sets_logical_step_on_patched_unet(): + unet = SimpleNamespace(info={"step_index": None}) + ctx = SimpleNamespace(unet=unet, step_index=3) + extension = HiDiffusionExt(name_or_path="runwayml/stable-diffusion-v1-5") + + extension.set_step_index(ctx) + + assert unet.info["step_index"] == 3 + + +def test_t2i_adapter_residual_is_resized_for_active_raunet(): + patched_block = make_diffusers_cross_attn_down_block(CrossAttnDownBlock2D) + module = patched_block( + in_channels=4, + out_channels=4, + temb_channels=4, + num_layers=2, + resnet_groups=1, + num_attention_heads=1, + cross_attention_dim=4, + add_downsample=False, + ) + module.info = { + "size": (64, 64), + "pipeline": SimpleNamespace(_num_timesteps=10), + "text_to_img_controlnet": False, + "is_inpainting_task": False, + "is_playground": False, + "step_index": 0, + "switching_threshold_overrides": {"T1_ratio": 1.0, "T2_ratio": 1.0}, + } + module.model = "sd15" + module.switching_threshold_ratio = "T2_ratio" + + hidden_states, output_states = module( + hidden_states=torch.randn(1, 4, 8, 8), + temb=torch.randn(1, 4), + encoder_hidden_states=torch.randn(1, 2, 4), + additional_residuals=torch.randn(1, 4, 8, 8), + ) + + assert hidden_states.shape[-2:] == (4, 4) + assert output_states[-1].shape[-2:] == (4, 4) + + +def test_sdxl_primary_raunet_is_active_after_aggressive_stage_until_t1(): + patched_block = make_diffusers_cross_attn_down_block(CrossAttnDownBlock2D) + module = patched_block( + in_channels=4, + out_channels=4, + temb_channels=4, + num_layers=2, + resnet_groups=1, + num_attention_heads=1, + cross_attention_dim=4, + add_downsample=False, ) + module.info = { + "size": (256, 256), + "pipeline": SimpleNamespace(_num_timesteps=50), + "text_to_img_controlnet": False, + "is_inpainting_task": False, + "is_playground": False, + "step_index": 0, + "switching_threshold_overrides": {"T1_ratio": None, "T2_ratio": None}, + } + module.model = "sdxl" + module.switching_threshold_ratio = "T1_ratio" + inputs = { + "hidden_states": torch.randn(1, 4, 8, 8), + "temb": torch.randn(1, 4), + "encoder_hidden_states": torch.randn(1, 2, 4), + } + + hidden_states, _ = module(**inputs) + assert hidden_states.shape[-2:] == (8, 8) + + module.info["step_index"] = 8 + hidden_states, _ = module(**inputs) + assert hidden_states.shape[-2:] == (4, 4) + + module.info["step_index"] = 20 + hidden_states, _ = module(**inputs) + assert hidden_states.shape[-2:] == (8, 8) + + +def test_sdxl_additional_raunet_is_active_before_aggressive_boundary(): + patched_conv = make_diffusers_downsampler_block(torch.nn.Conv2d) + module = patched_conv(1, 1, kernel_size=3, stride=2, padding=1, bias=False) + module.info = { + "size": (256, 256), + "pipeline": SimpleNamespace(_num_timesteps=50), + "text_to_img_controlnet": False, + "is_inpainting_task": False, + "is_playground": False, + "step_index": 0, + "switching_threshold_overrides": {"T1_ratio": None, "T2_ratio": None}, + } + module.model = "sdxl" + module.switching_threshold_ratio = "T2_ratio" + hidden_states = torch.ones(1, 1, 8, 8) + + assert module(hidden_states).shape[-2:] == (2, 2) + + module.info["step_index"] = 7 + assert module(hidden_states).shape[-2:] == (2, 2) - assert _should_use_aggressive_raunet(module) is expected + module.info["step_index"] = 8 + assert module(hidden_states).shape[-2:] == (4, 4) def test_sdxl_t2_override_controls_downsampler_at_2048_resolution(): @@ -490,7 +711,7 @@ def run(t2_ratio: float) -> torch.Tensor: module.switching_threshold_ratio = "T2_ratio" return module(hidden_states) - assert not torch.equal(run(0.0), run(1.0)) + assert not torch.equal(run(0.0), run(0.4)) def test_sdxl_automatic_ratios_preserve_extreme_resolution_preset(): diff --git a/tests/backend/stable_diffusion/test_regional_ip_data.py b/tests/backend/stable_diffusion/test_regional_ip_data.py index 734fac305e3..32680b20be6 100644 --- a/tests/backend/stable_diffusion/test_regional_ip_data.py +++ b/tests/backend/stable_diffusion/test_regional_ip_data.py @@ -1,3 +1,4 @@ +import pytest import torch from invokeai.backend.stable_diffusion.diffusion.regional_ip_data import RegionalIPData @@ -12,7 +13,21 @@ def test_regional_ip_data_supports_hidiffusion_raunet_downscale() -> None: device=torch.device("cpu"), ) - masks = regional_ip_data.get_masks(query_seq_len=12 * 12) + masks = regional_ip_data.get_masks(query_seq_len=6 * 6) - assert masks.shape == (1, 1, 12 * 12, 1) + assert masks.shape == (1, 1, 6 * 6, 1) assert torch.count_nonzero(masks) == masks.numel() + + +def test_regional_ip_data_rejects_unprepared_hidiffusion_scale() -> None: + regional_ip_data = RegionalIPData( + image_prompt_embeds=[torch.zeros((1, 1, 4, 8))], + scales=[1.0], + masks=[torch.ones((1, 1, 192, 192))], + dtype=torch.float32, + device=torch.device("cpu"), + max_downscale_factor=16, + ) + + with pytest.raises(KeyError): + regional_ip_data.get_masks(query_seq_len=6 * 6) diff --git a/tests/backend/stable_diffusion/test_regional_prompt_data.py b/tests/backend/stable_diffusion/test_regional_prompt_data.py index e1990838c2c..1733b808c1c 100644 --- a/tests/backend/stable_diffusion/test_regional_prompt_data.py +++ b/tests/backend/stable_diffusion/test_regional_prompt_data.py @@ -1,3 +1,4 @@ +import pytest import torch from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Range, TextConditioningRegions @@ -15,7 +16,23 @@ def test_regional_prompt_data_supports_hidiffusion_raunet_downscale() -> None: dtype=torch.float32, ) - attention_mask = regional_prompt_data.get_cross_attn_mask(query_seq_len=12 * 12, key_seq_len=4) + attention_mask = regional_prompt_data.get_cross_attn_mask(query_seq_len=6 * 6, key_seq_len=4) - assert attention_mask.shape == (1, 12 * 12, 4) + assert attention_mask.shape == (1, 6 * 6, 4) assert torch.count_nonzero(attention_mask) == 0 + + +def test_regional_prompt_data_rejects_unprepared_hidiffusion_scale() -> None: + regions = TextConditioningRegions( + masks=torch.ones((1, 1, 192, 192), dtype=torch.bool), + ranges=[Range(start=0, end=4)], + ) + regional_prompt_data = RegionalPromptData( + regions=[regions], + device=torch.device("cpu"), + dtype=torch.float32, + max_downscale_factor=16, + ) + + with pytest.raises(KeyError): + regional_prompt_data.get_cross_attn_mask(query_seq_len=6 * 6, key_seq_len=4) From 5fcc221b012d5c73b5eccffe5cb821e567909ee1 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Sun, 6 Sep 2026 10:02:37 +0300 Subject: [PATCH 6/7] chore: update HiDiffusion generated API schemas --- invokeai/frontend/web/openapi.json | 8 ++++---- invokeai/frontend/web/src/services/api/schema.ts | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 5254c4cd5f3..1caa298c824 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -24644,7 +24644,7 @@ } ], "default": null, - "description": "Override HiDiffusion early switch threshold (T1 ratio)", + "description": "Override the duration of HiDiffusion's primary RAU-Net stage (upstream code key T1_ratio). At extreme resolutions this is the later of the two RAU-Net cutoffs.", "field_kind": "input", "input": "any", "orig_default": null, @@ -24663,7 +24663,7 @@ } ], "default": null, - "description": "Override HiDiffusion late switch threshold (T2 ratio)", + "description": "Override the duration of HiDiffusion's additional extreme-resolution RAU-Net stage (upstream code key T2_ratio). This is the earlier cutoff when both stages are active and cannot exceed T1; excessive values can reduce composition diversity or introduce artifacts.", "field_kind": "input", "input": "any", "orig_default": null, @@ -25071,7 +25071,7 @@ } ], "default": null, - "description": "Override HiDiffusion early switch threshold (T1 ratio)", + "description": "Override the duration of HiDiffusion's primary RAU-Net stage (upstream code key T1_ratio). At extreme resolutions this is the later of the two RAU-Net cutoffs.", "field_kind": "input", "input": "any", "orig_default": null, @@ -25090,7 +25090,7 @@ } ], "default": null, - "description": "Override HiDiffusion late switch threshold (T2 ratio)", + "description": "Override the duration of HiDiffusion's additional extreme-resolution RAU-Net stage (upstream code key T2_ratio). This is the earlier cutoff when both stages are active and cannot exceed T1; excessive values can reduce composition diversity or introduce artifacts.", "field_kind": "input", "input": "any", "orig_default": null, diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 531189ddc94..db5e0b7ee20 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -9327,13 +9327,13 @@ export type components = { hidiffusion_window_attn?: boolean; /** * HiDiffusion: T1 Ratio - * @description Override HiDiffusion early switch threshold (T1 ratio) + * @description Override the duration of HiDiffusion's primary RAU-Net stage (upstream code key T1_ratio). At extreme resolutions this is the later of the two RAU-Net cutoffs. * @default null */ hidiffusion_t1_ratio?: number | null; /** * HiDiffusion: T2 Ratio - * @description Override HiDiffusion late switch threshold (T2 ratio) + * @description Override the duration of HiDiffusion's additional extreme-resolution RAU-Net stage (upstream code key T2_ratio). This is the earlier cutoff when both stages are active and cannot exceed T1; excessive values can reduce composition diversity or introduce artifacts. * @default null */ hidiffusion_t2_ratio?: number | null; @@ -9475,13 +9475,13 @@ export type components = { hidiffusion_window_attn?: boolean; /** * HiDiffusion: T1 Ratio - * @description Override HiDiffusion early switch threshold (T1 ratio) + * @description Override the duration of HiDiffusion's primary RAU-Net stage (upstream code key T1_ratio). At extreme resolutions this is the later of the two RAU-Net cutoffs. * @default null */ hidiffusion_t1_ratio?: number | null; /** * HiDiffusion: T2 Ratio - * @description Override HiDiffusion late switch threshold (T2 ratio) + * @description Override the duration of HiDiffusion's additional extreme-resolution RAU-Net stage (upstream code key T2_ratio). This is the earlier cutoff when both stages are active and cannot exceed T1; excessive values can reduce composition diversity or introduce artifacts. * @default null */ hidiffusion_t2_ratio?: number | null; From 5d7e2348f1f7e83e64746db36a7f4e8d46fabad3 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Sun, 13 Sep 2026 16:18:56 +0300 Subject: [PATCH 7/7] fix(hidiffusion): preserve partial denoise detail --- .../src/content/docs/features/hidiffusion.mdx | 5 +- invokeai/app/invocations/denoise_latents.py | 4 + invokeai/backend/hidiffusion/hidiffusion.py | 101 +++++++++--------- .../extensions/hidiffusion.py | 6 ++ .../stable_diffusion/hidiffusion_utils.py | 4 + .../test_hidiffusion_utils.py | 92 ++++++++++++++++ 6 files changed, 159 insertions(+), 53 deletions(-) diff --git a/docs/src/content/docs/features/hidiffusion.mdx b/docs/src/content/docs/features/hidiffusion.mdx index cf1ed6e0c2f..6bc9e1cf364 100644 --- a/docs/src/content/docs/features/hidiffusion.mdx +++ b/docs/src/content/docs/features/hidiffusion.mdx @@ -35,7 +35,9 @@ Learn more: https://github.com/megvii-research/HiDiffusion - **HiDiffusion: T2 Ratio**: Manually controls the end of the additional early RAU-Net stage. For ordinary SDXL generation, this also controls when the primary stage begins. When automatic ratios are enabled, the original implementation uses a fixed `8 / 50` boundary for this early stage at the 2048 preset. Excessive manual values can reduce pose and background diversity or introduce artifacts. T2 cannot exceed T1 and does not control window attention. -The T1/T2 input names are retained for API compatibility with the upstream implementation's code keys. T2 is the earlier boundary and T1 is the later boundary; do not interpret the names as chronological order. With automatic ratios at the SDXL 2048 preset, the additional stage runs for the first 8 of every 50 denoising steps, the primary stage then runs until T1, and the ordinary UNet is used afterward. Inpainting follows the upstream non-aggressive schedule: its primary stage starts at the first denoising step. +The T1/T2 input names are retained for API compatibility with the upstream implementation's code keys. T2 is the earlier boundary and T1 is the later boundary; do not interpret the names as chronological order. With automatic ratios at the SDXL 2048 preset, the additional stage runs for the first 8 of every 50 denoising steps, the primary stage then runs until T1, and the ordinary UNet is used afterward. Inpainting follows the upstream non-aggressive schedule and does not use the additional early stage. + +For img2img and inpainting, these phases retain their positions in the full denoising schedule. InvokeAI clips them to the portion selected by denoising strength instead of restarting HiDiffusion at the first remaining step. If denoising begins after T1, RAU-Net is not applied. ### Manual ratio constraint @@ -48,5 +50,6 @@ The UI limits the maximum T2 value to the current T1 value. If T1 is reduced bel - Start around **1536–2048 px** for SDXL. Automatic ratios do not interpolate: the upstream 4096 preset is selected only when both image dimensions reach 4096 px (512 latent pixels). Test intermediate and non-square sizes explicitly before increasing either ratio manually. - Ordinary SDXL generation uses the upstream staged schedule: the additional RAU-Net path is active first, the primary path takes over at the early boundary, and the ordinary UNet takes over at T1. Inpainting does not use the initial aggressive stage. - HiDiffusion is applied only to the SDXL base denoise stage, not the optional refiner stage. +- Window Attention falls back to ordinary global attention when a feature-map dimension cannot be divided into its 2x2 window layout. This avoids detail loss from resizing the feature map solely for window partitioning. - If results look worse, disable **Window Attention** first, then RAU‑Net. - Effects vary by scheduler and model; compare with the same seed for a fair test. diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index 83a1d8cb846..7ad29273342 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -937,6 +937,8 @@ def step_callback(state: PipelineIntermediateState) -> None: t2_ratio=self.hidiffusion_t2_ratio, generator=torch.Generator(device="cpu").manual_seed(seed), is_inpainting_task=self.denoise_mask is not None, + denoising_start=self.denoising_start, + denoising_end=self.denoising_end, ) ) @@ -1169,6 +1171,8 @@ def _lora_loader() -> Iterator[PatchSpec]: t2_ratio=self.hidiffusion_t2_ratio, generator=torch.Generator(device="cpu").manual_seed(seed), is_inpainting_task=self.denoise_mask is not None, + denoising_start=self.denoising_start, + denoising_end=self.denoising_end, ) if self.hidiffusion else nullcontext() diff --git a/invokeai/backend/hidiffusion/hidiffusion.py b/invokeai/backend/hidiffusion/hidiffusion.py index 5daad401011..6130928efef 100644 --- a/invokeai/backend/hidiffusion/hidiffusion.py +++ b/invokeai/backend/hidiffusion/hidiffusion.py @@ -1,6 +1,5 @@ import importlib.resources import math -import warnings from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union import diffusers @@ -191,20 +190,34 @@ def _get_raunet_step_range(module: torch.nn.Module, height: int, width: int) -> At ordinary SDXL resolutions, upstream uses the extra (T2-position) modules first, then the primary (T1-position) modules. A manual T2 override replaces upstream's fixed 8/50 boundary so - that InvokeAI's explicit T2 control remains effective. + that InvokeAI's explicit T2 control remains effective. Phase boundaries are defined against the + full denoising schedule and clipped to the partial range executed by img2img or inpainting. """ ratio = _get_resolution_aware_switching_threshold_ratio(module, height, width) - start = 0 - end = int(module.max_timestep * ratio) + phase_start = 0.0 + phase_end = ratio if _uses_aggressive_raunet(module, height, width): t2_override = module.info["switching_threshold_overrides"].get("T2_ratio") early_ratio = aggressive_step / 50 if t2_override is None else t2_override - early_end = int(module.max_timestep * early_ratio) if module.switching_threshold_ratio == "T1_ratio": - start = early_end + phase_start = early_ratio else: - end = early_end + phase_end = early_ratio + + denoising_start = module.info.get("denoising_start", 0.0) + denoising_end = module.info.get("denoising_end", 1.0) + denoising_span = denoising_end - denoising_start + if denoising_span <= 0: + return ratio, 0, 0 + + # T1/T2 are positions in the full denoising schedule, while img2img and inpainting execute only + # [denoising_start, denoising_end]. Intersect the global RAU-Net phase with that executed range, + # then convert the result to indices in the shortened local timestep list. + local_start_ratio = max(0.0, min(1.0, (phase_start - denoising_start) / denoising_span)) + local_end_ratio = max(0.0, min(1.0, (phase_end - denoising_start) / denoising_span)) + start = int(module.max_timestep * local_start_ratio) + end = int(module.max_timestep * local_end_ratio) return ratio, start, end @@ -1422,22 +1435,6 @@ def window_partition(x, window_size, shift_size, H, W): """ B, N, C = x.shape x = x.view(B, H, W, C) - if H % 2 != 0 or W % 2 != 0: - warnings.warn( - f"HiDiffusion Warning: The feature size is {(H, W)} and cannot be directly partitioned into windows. We interpolate the size to {(window_size[0] * 2, window_size[1] * 2)} " - f"to enable the window partition. Even though the generation is OK, the image quality would be largely decreased. " - f"We suggest removing window attention by setting apply_hidiffusion(pipe, apply_window_attn=False) for better image quality.", - stacklevel=2, - ) - x = ( - F.interpolate( - x.permute(0, 3, 1, 2).contiguous(), - size=(window_size[0] * 2, window_size[1] * 2), - mode="bicubic", - ) - .permute(0, 2, 3, 1) - .contiguous() - ) if type(shift_size) is list or type(shift_size) is tuple: if shift_size[0] > 0: x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1]), dims=(1, 2)) @@ -1471,12 +1468,6 @@ def window_reverse(windows, window_size, H, W, shift_size): else: if shift_size > 0: x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2)) - if H % 2 != 0 or W % 2 != 0: - x = ( - F.interpolate(x.permute(0, 3, 1, 2).contiguous(), size=(H, W), mode="bicubic") - .permute(0, 2, 3, 1) - .contiguous() - ) x = x.view(B, H * W, C) return x @@ -1506,36 +1497,37 @@ def window_reverse(windows, window_size, H, W, shift_size): if self.pos_embed is not None: norm_hidden_states = self.pos_embed(norm_hidden_states) - # MSW-MSA - logical_step = self.info.get("step_index") - if logical_step is None or self.__dict__.get("_hidiffusion_window_shift_step") != logical_step: - if generator is not None: - rand_num = torch.rand(1, generator=generator, device=generator.device) - else: - rand_num = torch.rand(1) - self._hidiffusion_window_shift_step = logical_step - self._hidiffusion_window_shift_bucket = min(int(rand_num.item() * 4), 3) - - shift_bucket = self._hidiffusion_window_shift_bucket - B, N, C = hidden_states.shape ori_H, ori_W = self.info["size"] downsample_ratio = round(((ori_H * ori_W) / N) ** 0.5) H, W = (math.ceil(ori_H / downsample_ratio), math.ceil(ori_W / downsample_ratio)) - widow_size = (math.ceil(H / 2), math.ceil(W / 2)) - if shift_bucket == 0: - shift_size = (0, 0) - elif shift_bucket == 1: - shift_size = (widow_size[0] // 4, widow_size[1] // 4) - elif shift_bucket == 2: - shift_size = (widow_size[0] // 4 * 2, widow_size[1] // 4 * 2) - else: - shift_size = (widow_size[0] // 4 * 3, widow_size[1] // 4 * 3) - norm_hidden_states = window_partition(norm_hidden_states, widow_size, shift_size, H, W) # 2. Prepare GLIGEN inputs cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} gligen_kwargs = cross_attention_kwargs.pop("gligen", None) + use_window_attention = H % 2 == 0 and W % 2 == 0 + if use_window_attention: + logical_step = self.info.get("step_index") + if logical_step is None or self.__dict__.get("_hidiffusion_window_shift_step") != logical_step: + if generator is not None: + rand_num = torch.rand(1, generator=generator, device=generator.device) + else: + rand_num = torch.rand(1) + self._hidiffusion_window_shift_step = logical_step + self._hidiffusion_window_shift_bucket = min(int(rand_num.item() * 4), 3) + + shift_bucket = self._hidiffusion_window_shift_bucket + window_size = (H // 2, W // 2) + if shift_bucket == 0: + shift_size = (0, 0) + elif shift_bucket == 1: + shift_size = (window_size[0] // 4, window_size[1] // 4) + elif shift_bucket == 2: + shift_size = (window_size[0] // 4 * 2, window_size[1] // 4 * 2) + else: + shift_size = (window_size[0] // 4 * 3, window_size[1] // 4 * 3) + norm_hidden_states = window_partition(norm_hidden_states, window_size, shift_size, H, W) + attn_output = self.attn1( norm_hidden_states, encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, @@ -1547,7 +1539,8 @@ def window_reverse(windows, window_size, H, W, shift_size): elif self.use_ada_layer_norm_single: attn_output = gate_msa * attn_output - attn_output = window_reverse(attn_output, widow_size, H, W, shift_size) + if use_window_attention: + attn_output = window_reverse(attn_output, window_size, H, W, shift_size) hidden_states = attn_output + hidden_states if hidden_states.ndim == 4: @@ -2002,6 +1995,8 @@ def apply_hidiffusion( t1_ratio: float | None = None, t2_ratio: float | None = None, is_inpainting_task: bool | None = None, + denoising_start: float = 0.0, + denoising_end: float = 1.0, ): """ model: diffusers model. We support SD 1.5, 2.1, XL, XL Turbo. @@ -2088,6 +2083,8 @@ def apply_hidiffusion( "step_index": None, "pipeline": model, "switching_threshold_overrides": {"T1_ratio": t1_ratio, "T2_ratio": t2_ratio}, + "denoising_start": denoising_start, + "denoising_end": denoising_end, } model.info = diffusion_model.info hook_diffusion_model(diffusion_model) diff --git a/invokeai/backend/stable_diffusion/extensions/hidiffusion.py b/invokeai/backend/stable_diffusion/extensions/hidiffusion.py index 129d4c68a4b..befeea87672 100644 --- a/invokeai/backend/stable_diffusion/extensions/hidiffusion.py +++ b/invokeai/backend/stable_diffusion/extensions/hidiffusion.py @@ -25,6 +25,8 @@ def __init__( has_controlnet: bool = False, is_controlnet_text_to_image: bool = False, is_inpainting_task: bool | None = None, + denoising_start: float = 0.0, + denoising_end: float = 1.0, ): super().__init__() self._name_or_path = name_or_path @@ -33,6 +35,8 @@ def __init__( self._has_controlnet = has_controlnet self._is_controlnet_text_to_image = is_controlnet_text_to_image self._is_inpainting_task = is_inpainting_task + self._denoising_start = denoising_start + self._denoising_end = denoising_end self._t1_ratio = t1_ratio self._t2_ratio = t2_ratio self._generator = generator @@ -56,5 +60,7 @@ def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeigh t2_ratio=self._t2_ratio, generator=self._generator, is_inpainting_task=self._is_inpainting_task, + denoising_start=self._denoising_start, + denoising_end=self._denoising_end, ): yield None diff --git a/invokeai/backend/stable_diffusion/hidiffusion_utils.py b/invokeai/backend/stable_diffusion/hidiffusion_utils.py index 29d6cdc37ce..03518975c69 100644 --- a/invokeai/backend/stable_diffusion/hidiffusion_utils.py +++ b/invokeai/backend/stable_diffusion/hidiffusion_utils.py @@ -19,6 +19,8 @@ def hidiffusion_patch( has_controlnet: bool = False, is_controlnet_text_to_image: bool = False, is_inpainting_task: bool | None = None, + denoising_start: float = 0.0, + denoising_end: float = 1.0, ): """Context manager that applies HiDiffusion and restores the model on exit.""" from invokeai.backend.hidiffusion.hidiffusion import apply_hidiffusion, remove_hidiffusion @@ -98,6 +100,8 @@ def _set_name_or_path_on_config(cfg, value: str) -> bool: is_controlnet_text_to_image=is_controlnet_text_to_image, generator=generator, is_inpainting_task=is_inpainting_task, + denoising_start=denoising_start, + denoising_end=denoising_end, ) yield finally: diff --git a/tests/backend/stable_diffusion/test_hidiffusion_utils.py b/tests/backend/stable_diffusion/test_hidiffusion_utils.py index e61cbf3bc6a..d8c01084793 100644 --- a/tests/backend/stable_diffusion/test_hidiffusion_utils.py +++ b/tests/backend/stable_diffusion/test_hidiffusion_utils.py @@ -49,7 +49,12 @@ def __init__(self): class WindowMeanAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.last_sequence_length: int | None = None + def forward(self, hidden_states: torch.Tensor, **_kwargs): + self.last_sequence_length = hidden_states.shape[1] return hidden_states.mean(dim=1, keepdim=True).expand_as(hidden_states) @@ -168,6 +173,34 @@ def test_hidiffusion_window_attention_reuses_shift_within_logical_step(): assert not torch.equal(generator.get_state(), generator_state_after_first_forward) +def test_hidiffusion_window_attention_falls_back_to_global_attention_for_odd_feature_maps(): + module_keys = { + "down_module_key": [], + "down_module_key_extra": [], + "up_module_key": [], + "up_module_key_extra": [], + "windown_attn_module_key": ["transformer"], + } + model = WindowAttentionModelMixin() + hidden_states = torch.arange(15, dtype=torch.float32).reshape(1, 15, 1) + + 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=torch.Generator(device="cpu").manual_seed(1234), + ), + ): + model.info["size"] = (5, 3) + output = model.transformer(hidden_states) + + assert output.shape == hidden_states.shape + assert model.transformer.attn1.last_sequence_length == 15 + + @pytest.mark.parametrize("is_text_to_image", [False, True]) def test_hidiffusion_patch_uses_controlnet_aware_forward_for_bare_unet(is_text_to_image: bool): model = ModelMixin() @@ -447,11 +480,15 @@ def test_hidiffusion_patch_forwards_generation_context(): model, name_or_path="stabilityai/stable-diffusion-xl-base-1.0", is_inpainting_task=True, + denoising_start=0.25, + denoising_end=0.75, ): pass kwargs = mock_apply_hidiffusion.call_args.kwargs assert kwargs["is_inpainting_task"] is True + assert kwargs["denoising_start"] == pytest.approx(0.25) + assert kwargs["denoising_end"] == pytest.approx(0.75) @pytest.mark.parametrize( @@ -531,6 +568,45 @@ def test_hidiffusion_raunet_schedule_matches_upstream_stages( assert _get_raunet_step_range(module, *size) == expected +@pytest.mark.parametrize( + ("threshold", "is_inpainting", "denoising_start", "denoising_end", "max_timestep", "expected"), + [ + ("T1_ratio", True, 0.5, 1.0, 25, (0.4, 0, 0)), + ("T1_ratio", False, 0.5, 1.0, 25, (0.4, 0, 0)), + ("T2_ratio", False, 0.5, 1.0, 25, (0.0, 0, 0)), + ("T1_ratio", True, 0.2, 1.0, 40, (0.4, 0, 10)), + ("T1_ratio", False, 0.2, 1.0, 40, (0.4, 0, 10)), + ("T2_ratio", False, 0.2, 1.0, 40, (0.0, 0, 0)), + ("T1_ratio", True, 0.1, 0.2, 5, (0.4, 0, 5)), + ("T1_ratio", False, 0.0, 0.1, 5, (0.4, 5, 5)), + ("T2_ratio", False, 0.0, 0.1, 5, (0.0, 0, 5)), + ], +) +def test_hidiffusion_raunet_schedule_is_clipped_to_partial_denoising_range( + threshold: str, + is_inpainting: bool, + denoising_start: float, + denoising_end: float, + max_timestep: int, + expected: tuple[float, int, int], +): + module = SimpleNamespace( + model="sdxl", + max_timestep=max_timestep, + switching_threshold_ratio=threshold, + info={ + "switching_threshold_overrides": {"T1_ratio": None, "T2_ratio": None}, + "text_to_img_controlnet": False, + "is_inpainting_task": is_inpainting, + "is_playground": False, + "denoising_start": denoising_start, + "denoising_end": denoising_end, + }, + ) + + assert _get_raunet_step_range(module, 256, 256) == expected + + @pytest.mark.parametrize("t1_override", [None, 0.4]) def test_hidiffusion_rejects_t2_above_the_resolved_t1(t1_override: float | None): module = SimpleNamespace( @@ -591,6 +667,22 @@ def test_hidiffusion_extension_sets_logical_step_on_patched_unet(): assert unet.info["step_index"] == 3 +def test_hidiffusion_extension_forwards_partial_denoising_range(): + extension = HiDiffusionExt( + name_or_path="runwayml/stable-diffusion-v1-5", + denoising_start=0.5, + denoising_end=0.9, + ) + + with patch("invokeai.backend.stable_diffusion.extensions.hidiffusion.hidiffusion_patch") as mock_patch: + with extension.patch_unet(MagicMock(), MagicMock()): + pass + + kwargs = mock_patch.call_args.kwargs + assert kwargs["denoising_start"] == pytest.approx(0.5) + assert kwargs["denoising_end"] == pytest.approx(0.9) + + def test_t2i_adapter_residual_is_resized_for_active_raunet(): patched_block = make_diffusers_cross_attn_down_block(CrossAttnDownBlock2D) module = patched_block(