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 8158f4980a1..6bc9e1cf364 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,12 +29,27 @@ 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 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: 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 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**: 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 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 + +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. +- 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 2d48dd87607..7ad29273342 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 @@ -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, @@ -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], @@ -926,6 +936,9 @@ 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, + denoising_start=self.denoising_start, + denoising_end=self.denoising_end, ) ) @@ -1157,6 +1170,9 @@ 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, + 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/app/invocations/metadata_linked.py b/invokeai/app/invocations/metadata_linked.py index 5f8c18ecc98..be23ab6d85e 100644 --- a/invokeai/app/invocations/metadata_linked.py +++ b/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: diff --git a/invokeai/backend/hidiffusion/hidiffusion.py b/invokeai/backend/hidiffusion/hidiffusion.py index 33437853214..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 @@ -96,8 +95,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 +117,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,6 +135,105 @@ def _get_max_timesteps(info_dict: dict) -> int: return len(pipeline.scheduler.timesteps) +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] + + 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.") + + 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 _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"]: + return playground_is_aggressive_raunet + 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. 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) + 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 + if module.switching_threshold_ratio == "T1_ratio": + phase_start = early_ratio + else: + 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 + + +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 @@ -1337,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)) @@ -1386,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 @@ -1421,30 +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 - if generator is not None: - rand_num = torch.rand(1, generator=generator, device=generator.device) - else: - rand_num = torch.rand(1) - 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: - shift_size = (0, 0) - if rand_num > 0.25 and rand_num <= 0.5: - shift_size = (widow_size[0] // 4, widow_size[1] // 4) - if rand_num > 0.5 and rand_num <= 0.75: - shift_size = (widow_size[0] // 4 * 2, widow_size[1] // 4 * 2) - if rand_num > 0.75 and rand_num <= 1: - 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, @@ -1456,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: @@ -1535,11 +1619,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 @@ -1557,40 +1639,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 = switching_threshold_ratio_dict["sd15_1024"][self.switching_threshold_ratio] - else: - self.T1_ratio = switching_threshold_ratio_dict["sd15_2048"][self.switching_threshold_ratio] - 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 - ] - else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] - - 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 - else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_4096"][self.switching_threshold_ratio] - elif self.model == "sdxl_turbo": - self.T1_ratio = switching_threshold_ratio_dict["sdxl_turbo_1024"][self.switching_threshold_ratio] - else: - raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") - - if self.aggressive_raunet: - # 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 = () @@ -1637,13 +1687,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,) @@ -1655,9 +1707,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 @@ -1670,11 +1720,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 @@ -1691,41 +1739,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 = switching_threshold_ratio_dict["sd15_1024"][self.switching_threshold_ratio] - else: - self.T1_ratio = switching_threshold_ratio_dict["sd15_2048"][self.switching_threshold_ratio] - 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 - ] - else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] - - 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 - - else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_4096"][self.switching_threshold_ratio] - elif self.model == "sdxl_turbo": - self.T1_ratio = switching_threshold_ratio_dict["sdxl_turbo_1024"][self.switching_threshold_ratio] - else: - raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") - - if self.aggressive_raunet: - # 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) @@ -1792,11 +1807,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" ) @@ -1805,9 +1817,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 @@ -1820,50 +1830,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 = switching_threshold_ratio_dict["sd15_1024"][self.switching_threshold_ratio] - else: - self.T1_ratio = switching_threshold_ratio_dict["sd15_2048"][self.switching_threshold_ratio] - 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 - ] - else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] - - 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 - else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_4096"][self.switching_threshold_ratio] - elif self.model == "sdxl_turbo": - self.T1_ratio = switching_threshold_ratio_dict["sdxl_turbo_1024"][self.switching_threshold_ratio] - else: - raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") - - if self.aggressive_raunet: - # 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) @@ -1875,9 +1856,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( @@ -1886,9 +1865,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 @@ -1900,50 +1877,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 = switching_threshold_ratio_dict["sd15_1024"][self.switching_threshold_ratio] - else: - self.T1_ratio = switching_threshold_ratio_dict["sd15_2048"][self.switching_threshold_ratio] - 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 - ] - else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_2048"][self.switching_threshold_ratio] - - 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 - else: - self.T1_ratio = switching_threshold_ratio_dict["sdxl_4096"][self.switching_threshold_ratio] - elif self.model == "sdxl_turbo": - self.T1_ratio = switching_threshold_ratio_dict["sdxl_turbo_1024"][self.switching_threshold_ratio] - else: - raise Exception("Error model. HiDiffusion now only supports sd15, sd21, sdxl, sdxl-turbo.") - - if self.aggressive_raunet: - # 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: @@ -1977,15 +1923,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", @@ -2045,6 +1992,11 @@ 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, + 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 +2072,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, + "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/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 792c97114da..c14edf0b8a0 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 = 32, ): - """Initialize a `IPAdapterConditioningData` object.""" + """Initialize an `IPAdapterConditioningData` object. + + SD1/SD2 with both HiDiffusion RAU-Net stages can reach a 32x 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..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 = 8, + 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. + 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 = 8 + 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 13a1763f35e..befeea87672 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 @@ -22,6 +24,9 @@ def __init__( generator: torch.Generator | None = None, 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 @@ -29,10 +34,19 @@ 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._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( @@ -45,5 +59,8 @@ 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, + 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 f6e6e1681b3..03518975c69 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( @@ -26,6 +18,9 @@ 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, + 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 @@ -70,14 +65,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,26 +90,18 @@ 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, + is_inpainting_task=is_inpainting_task, + denoising_start=denoising_start, + denoising_end=denoising_end, ) yield finally: @@ -133,11 +112,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 7e9315af642..17b7a88192c 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -24700,28 +24700,42 @@ "type": "boolean" }, "hidiffusion_t1_ratio": { - "default": 0.4, - "description": "Override HiDiffusion early switch threshold (T1 ratio)", + "anyOf": [ + { + "maximum": 1, + "minimum": 0, + "type": "number" + }, + { + "type": "null" + } + ], + "default": null, + "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", - "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, - "description": "Override HiDiffusion late switch threshold (T2 ratio)", + "anyOf": [ + { + "maximum": 1, + "minimum": 0, + "type": "number" + }, + { + "type": "null" + } + ], + "default": null, + "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", - "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": [ @@ -24769,7 +24783,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" } @@ -25113,28 +25127,42 @@ "type": "boolean" }, "hidiffusion_t1_ratio": { - "default": 0.4, - "description": "Override HiDiffusion early switch threshold (T1 ratio)", + "anyOf": [ + { + "maximum": 1, + "minimum": 0, + "type": "number" + }, + { + "type": "null" + } + ], + "default": null, + "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", - "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, - "description": "Override HiDiffusion late switch threshold (T2 ratio)", + "anyOf": [ + { + "maximum": 1, + "minimum": 0, + "type": "number" + }, + { + "type": "null" + } + ], + "default": null, + "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", - "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": [ @@ -25182,7 +25210,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 96741b4f075..f3dcec72839 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -1736,6 +1736,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", @@ -2190,18 +2191,25 @@ "Can boost local detail, but may affect global coherence." ] }, + "hidiffusionRatiosAuto": { + "heading": "HiDiffusion: Automatic Ratios", + "paragraphs": [ + "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 window attention 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/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..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'; @@ -172,8 +173,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. @@ -806,6 +808,7 @@ const POST_V4_BUMP_DEFAULTED_KEYS = [ 'hiDiffusionEnabled', 'hiDiffusionRauNetEnabled', 'hiDiffusionWindowAttnEnabled', + 'hiDiffusionAutoRatios', 'hiDiffusionT1Ratio', 'hiDiffusionT2Ratio', ] as const satisfies readonly (keyof typeof zParamsState.shape)[]; @@ -893,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; @@ -900,12 +904,13 @@ 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(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.hiDiffusionAutoRatios).toBe(true); expect(result.hiDiffusionT1Ratio).toBe(0.4); expect(result.hiDiffusionT2Ratio).toBe(0.0); // Existing params should be preserved @@ -916,6 +921,31 @@ 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.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); + }); + it('merges the separate Klein / dev VAE slots into flux2VaeModel when migrating from v3', () => { expect(migrate).toBeDefined(); @@ -938,7 +968,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(v3State) as ReturnType & Record; - expect(result._version).toBe(5); + 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(); @@ -972,7 +1002,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(v3State) as ReturnType; - expect(result._version).toBe(5); + expect(result._version).toBe(7); expect(result.krea2VaeModel).toBeNull(); expect(result.krea2Qwen3VlEncoderModel).toBeNull(); expect(result.krea2SeedVarianceEnabled).toBe(false); @@ -1006,7 +1036,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(mainV4State) as ReturnType & Record; - expect(result._version).toBe(5); + 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. @@ -1034,7 +1064,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(devV4State) as ReturnType & Record; - expect(result._version).toBe(5); + 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'); @@ -1091,7 +1121,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(blob) as ReturnType; - expect(result._version).toBe(5); + expect(result._version).toBe(7); expect(result.positivePrompt).toBe('a fluffy cat'); expect(result.seed).toBe(42); expect(result.shouldRandomizeSeed).toBe(false); @@ -1139,7 +1169,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 +1287,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(blob) as ReturnType; - expect(result._version).toBe(5); + expect(result._version).toBe(7); expect(result.dimensions).toEqual(getInitialParamsState().dimensions); expect(result.positivePrompt).toBe('a fluffy cat'); expect(result.seed).toBe(7); @@ -1274,7 +1304,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(blob) as ReturnType; - expect(result._version).toBe(5); + expect(result._version).toBe(7); expect(result.positivePromptHistory).toEqual([]); expect(result.qwenImageVaeModel).toBeNull(); expect(result.wanVaeModel).toBeNull(); @@ -1284,18 +1314,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: 8, positivePrompt: 'a fluffy cat' }); const { backfilled, reset } = repairParamsState(blob); expect(backfilled).toEqual([]); expect(reset).toEqual([]); - expect(blob._version).toBe(6); + 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(); @@ -1311,7 +1341,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(blob) as ReturnType; - expect(result._version).toBe(5); + expect(result._version).toBe(7); expect(result.positivePrompt).toBe('a fluffy cat'); expect(result.seed).toBe(7); expect(result.dimensions).toBeDefined(); @@ -1328,7 +1358,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(v3State) as ReturnType; - expect(result._version).toBe(5); + expect(result._version).toBe(7); expect(result.wanTransformerLowNoise).toBeNull(); expect(result.wanComponentSource).toBeNull(); expect(result.wanVaeModel).toBeNull(); @@ -1350,7 +1380,7 @@ describe('paramsSliceConfig persisted state migration', () => { const result = migrate?.(v2State) as ReturnType; - expect(result._version).toBe(5); + expect(result._version).toBe(7); expect(result.fluxScheduler).toBe('euler'); expect(result.zImageScheduler).toBe('euler'); expect(result.colorCompensation).toBe(false); @@ -1410,7 +1440,7 @@ describe('paramsSliceConfig persisted state migration', () => { expect('hiDiffusionEnabled' in blob).toBe(false); applyParamsVersionMigrations(blob); - expect(blob._version).toBe(5); + 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 @@ -1435,6 +1465,7 @@ describe('paramsSliceConfig persisted state migration', () => { expect(result.hiDiffusionEnabled).toBe(false); expect(result.hiDiffusionRauNetEnabled).toBe(true); expect(result.hiDiffusionWindowAttnEnabled).toBe(true); + expect(result.hiDiffusionAutoRatios).toBe(true); expect(result.hiDiffusionT1Ratio).toBe(0.4); expect(result.hiDiffusionT2Ratio).toBe(0.0); expect(result.positivePrompt).toBe('a fluffy cat'); @@ -1554,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 f7c8420dade..f58a9ef3813 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts @@ -182,6 +182,9 @@ const slice = createSlice({ setHiDiffusionWindowAttnEnabled: (state, action: PayloadAction) => { state.hiDiffusionWindowAttnEnabled = action.payload; }, + setHiDiffusionAutoRatios: (state, action: PayloadAction) => { + state.hiDiffusionAutoRatios = action.payload; + }, setHiDiffusionT1Ratio: (state, action: PayloadAction) => { state.hiDiffusionT1Ratio = action.payload; }, @@ -895,6 +898,7 @@ export const { setHiDiffusionEnabled, setHiDiffusionRauNetEnabled, setHiDiffusionWindowAttnEnabled, + setHiDiffusionAutoRatios, setHiDiffusionT1Ratio, setHiDiffusionT2Ratio, setSeamlessXAxis, @@ -1162,6 +1166,27 @@ 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; + } + + 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 = { @@ -1277,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 2143e58f997..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(5), + _version: z.literal(7), maskBlur: z.number(), maskBlurMethod: zParameterMaskBlurMethod, canvasCoherenceMode: zParameterCanvasCoherenceMode, @@ -839,6 +839,7 @@ export const zParamsState = z.object({ hiDiffusionEnabled: z.boolean().default(false), hiDiffusionRauNetEnabled: z.boolean().default(true), hiDiffusionWindowAttnEnabled: z.boolean().default(true), + hiDiffusionAutoRatios: z.boolean().default(true), hiDiffusionT1Ratio: z.number().default(0.4), hiDiffusionT2Ratio: z.number().default(0.0), iterations: z.number(), @@ -965,7 +966,7 @@ export const zParamsState = z.object({ }); export type ParamsState = z.infer; export const getInitialParamsState = (): ParamsState => ({ - _version: 5, + _version: 7, maskBlur: 16, maskBlurMethod: 'box', canvasCoherenceMode: 'Gaussian Blur', @@ -983,6 +984,7 @@ export const getInitialParamsState = (): ParamsState => ({ hiDiffusionEnabled: false, hiDiffusionRauNetEnabled: true, hiDiffusionWindowAttnEnabled: true, + hiDiffusionAutoRatios: true, hiDiffusionT1Ratio: 0.4, hiDiffusionT2Ratio: 0.0, iterations: 1, diff --git a/invokeai/frontend/web/src/features/metadata/parsing.test.ts b/invokeai/frontend/web/src/features/metadata/parsing.test.ts index 6072aa1c10f..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,10 @@ import type { AppStore } from 'app/store/store'; -import { setHiDiffusionEnabled } from 'features/controlLayers/store/paramsSlice'; +import { + setHiDiffusionAutoRatios, + setHiDiffusionEnabled, + setHiDiffusionT1Ratio, + setHiDiffusionT2Ratio, +} from 'features/controlLayers/store/paramsSlice'; import { describe, expect, it, vi } from 'vitest'; import { ImageMetadataHandlers, MetadataUtils, parseMetadataHandler } from './parsing'; @@ -163,12 +168,49 @@ 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(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 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) { + hiDiffusionT2Ratio = action.payload; } return action; }), @@ -193,5 +235,8 @@ describe('HiDiffusion metadata parsing', () => { expect(store.dispatch).toHaveBeenCalledWith(setHiDiffusionEnabled(false)); expect(hiDiffusionEnabled).toBe(false); + 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 a653afc10b8..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, @@ -787,38 +788,49 @@ 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) => { - store.dispatch(setHiDiffusionT1Ratio(value)); + store.dispatch(setHiDiffusionAutoRatios(value === null)); + if (value !== null) { + store.dispatch(setHiDiffusionT1Ratio(value)); + } }, 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) => { - store.dispatch(setHiDiffusionT2Ratio(value)); + store.dispatch(setHiDiffusionAutoRatios(value === null)); + if (value !== null) { + store.dispatch(setHiDiffusionT2Ratio(value)); + } }, 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.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 11a5333c7d2..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, - hidiffusion_t2_ratio: hiDiffusionEnabled ? hiDiffusionT2Ratio : 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..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,6 +32,7 @@ const defaultParams = { cfgRescaleMultiplier: 0, hiDiffusionEnabled: false, hiDiffusionRauNetEnabled: false, + hiDiffusionAutoRatios: false, hiDiffusionT1Ratio: 0.25, hiDiffusionT2Ratio: 0.1, hiDiffusionWindowAttnEnabled: false, @@ -193,4 +194,42 @@ 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.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'); + 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(); + }); + + 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 2a8d15b1e75..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 hiDiffusionAutoRatios = useAppSelector(selectHiDiffusionAutoRatios); + const dispatch = useAppDispatch(); + const { t } = useTranslation(); + + const onChange = useCallback( + (event: ChangeEvent) => { + dispatch(setHiDiffusionAutoRatios(event.target.checked)); + }, + [dispatch] + ); + + return ( + + + + {t('parameters.hiDiffusionRatiosAuto')} + + + + + ); +}); + +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 ( - + {t('parameters.hiDiffusionT1Ratio')} @@ -160,6 +208,9 @@ ParamHiDiffusionT1Ratio.displayName = 'ParamHiDiffusionT1Ratio'; export const ParamHiDiffusionT2Ratio = 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(); @@ -167,7 +218,10 @@ export const ParamHiDiffusionT2Ratio = memo(() => { const onChange = useCallback((value: number) => dispatch(setHiDiffusionT2Ratio(value)), [dispatch]); return ( - + {t('parameters.hiDiffusionT2Ratio')} @@ -177,7 +231,7 @@ export const ParamHiDiffusionT2Ratio = memo(() => { value={hiDiffusionT2Ratio} defaultValue={RATIO_CONSTRAINTS.t2.initial} min={RATIO_CONSTRAINTS.t2.sliderMin} - max={RATIO_CONSTRAINTS.t2.sliderMax} + max={hiDiffusionT1Ratio} step={RATIO_CONSTRAINTS.t2.coarseStep} fineStep={RATIO_CONSTRAINTS.t2.fineStep} onChange={onChange} @@ -187,7 +241,7 @@ export const ParamHiDiffusionT2Ratio = memo(() => { value={hiDiffusionT2Ratio} defaultValue={RATIO_CONSTRAINTS.t2.initial} min={RATIO_CONSTRAINTS.t2.numberInputMin} - max={RATIO_CONSTRAINTS.t2.numberInputMax} + max={hiDiffusionT1Ratio} step={RATIO_CONSTRAINTS.t2.coarseStep} fineStep={RATIO_CONSTRAINTS.t2.fineStep} onChange={onChange} diff --git a/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx b/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx index 374822d4058..4635a4618eb 100644 --- a/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx +++ b/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx @@ -28,6 +28,7 @@ import ParamClipSkip from 'features/parameters/components/Advanced/ParamClipSkip import ParamFlux2DevModelSelect from 'features/parameters/components/Advanced/ParamFlux2DevModelSelect'; import ParamFlux2KleinModelSelect from 'features/parameters/components/Advanced/ParamFlux2KleinModelSelect'; import { + ParamHiDiffusionAutoRatiosToggle, ParamHiDiffusionRauNetToggle, ParamHiDiffusionT1Ratio, ParamHiDiffusionT2Ratio, @@ -170,6 +171,7 @@ export const AdvancedSettingsAccordion = memo(() => { + diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 600a97bff25..afc826469c6 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -9361,16 +9361,16 @@ export type components = { hidiffusion_window_attn?: boolean; /** * HiDiffusion: T1 Ratio - * @description Override HiDiffusion early switch threshold (T1 ratio) - * @default 0.4 + * @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; + hidiffusion_t1_ratio?: number | null; /** * HiDiffusion: T2 Ratio - * @description Override HiDiffusion late switch threshold (T2 ratio) - * @default 0 + * @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; + hidiffusion_t2_ratio?: number | null; /** * @description Latents tensor * @default null @@ -9509,16 +9509,16 @@ export type components = { hidiffusion_window_attn?: boolean; /** * HiDiffusion: T1 Ratio - * @description Override HiDiffusion early switch threshold (T1 ratio) - * @default 0.4 + * @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; + hidiffusion_t1_ratio?: number | null; /** * HiDiffusion: T2 Ratio - * @description Override HiDiffusion late switch threshold (T2 ratio) - * @default 0 + * @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; + 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..d8c01084793 100644 --- a/tests/backend/stable_diffusion/test_hidiffusion_utils.py +++ b/tests/backend/stable_diffusion/test_hidiffusion_utils.py @@ -4,15 +4,22 @@ import pytest import torch +from diffusers.models.unets.unet_2d_blocks import CrossAttnDownBlock2D +from invokeai.app.invocations.denoise_latents import DenoiseLatentsInvocation from invokeai.backend.hidiffusion.hidiffusion import ( + _get_raunet_step_range, + _get_resolution_aware_switching_threshold_ratio, _resize_controlnet_residual, + make_diffusers_cross_attn_down_block, + make_diffusers_downsampler_block, switching_threshold_ratio_dict, text_to_img_controlnet_switching_threshold_ratio_dict, ) from invokeai.backend.hidiffusion.hidiffusion import ( remove_hidiffusion as real_remove_hidiffusion, ) +from invokeai.backend.stable_diffusion.extensions.hidiffusion import HiDiffusionExt from invokeai.backend.stable_diffusion.hidiffusion_utils import hidiffusion_patch @@ -42,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) @@ -125,6 +137,70 @@ def run_with_global_seed(global_seed: int) -> 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) + + +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() @@ -184,22 +260,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 @@ -256,16 +330,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 +441,386 @@ 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_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, + 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( + ("size", "threshold", "override", "expected_ratio"), + [ + ((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_use_upstream_discrete_presets( + size: tuple[int, int], threshold: str, override: float | None, expected_ratio: float +): + module = SimpleNamespace( + model="sdxl", + switching_threshold_ratio=threshold, + info={ + "switching_threshold_overrides": {"T1_ratio": override, "T2_ratio": override}, + "text_to_img_controlnet": False, + }, + ) + + 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( + ("size", "threshold", "is_inpainting", "t2_override", "expected"), + [ + ((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_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={ + "switching_threshold_overrides": {"T1_ratio": None, "T2_ratio": t2_override}, + "text_to_img_controlnet": False, + "is_inpainting_task": is_inpainting, + "is_playground": False, + }, + ) + + 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( + 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_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( + 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) + + module.info["step_index"] = 8 + assert module(hidden_states).shape[-2:] == (4, 4) + + +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(0.4)) + + +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 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..32680b20be6 --- /dev/null +++ b/tests/backend/stable_diffusion/test_regional_ip_data.py @@ -0,0 +1,33 @@ +import pytest +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=6 * 6) + + 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 new file mode 100644 index 00000000000..1733b808c1c --- /dev/null +++ b/tests/backend/stable_diffusion/test_regional_prompt_data.py @@ -0,0 +1,38 @@ +import pytest +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=6 * 6, key_seq_len=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)