Fix HiDiffusion ratio overrides and SDXL T2 behavior

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
 

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,
@@ -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()

invokeai/app/invocations/metadata_linked.py

@@ -624,7 +624,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:

invokeai/backend/hidiffusion/hidiffusion.py

@@ -136,6 +136,32 @@ 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]
+
+    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
+    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):
     class sdxl_controlnet_ppl(block_class):
         # Save for unpatching later
@@ -1559,32 +1585,29 @@ 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
-                    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 = 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,33 +1716,30 @@ 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
-                    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 = 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,32 +1850,29 @@ 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
-                    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 = 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,32 +1928,29 @@ 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
-                    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 = 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 +2059,12 @@ 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,
+    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.
@@ -2120,14 +2140,19 @@ 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},
     }
     model.info = diffusion_model.info
     hook_diffusion_model(diffusion_model)

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.