From e1f35117187b08e1c8566adc0aa51bf4f0f0c09a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 28 Jun 2026 22:10:31 +0000 Subject: [PATCH 1/2] Fix FDG disabled-level crash and MagCache single-block step reset FrequencyDecoupledGuidance appended pred_cond_freq in the disabled-level else branch, causing NameError when the first level was disabled and stale values when later levels were disabled. Use pred_cond_pyramid[level] instead. MagCache single-block models never advanced step_index on cache skips because the head hook returned early without invoking the co-located tail hook. Advance step state from the head hook skip path for single-block setups. Co-authored-by: Simon Lynch --- .../guiders/frequency_decoupled_guidance.py | 2 +- src/diffusers/hooks/mag_cache.py | 46 +++++++++++-------- .../test_frequency_decoupled_guidance.py | 36 +++++++++++++++ tests/hooks/test_mag_cache.py | 32 +++++++++++++ 4 files changed, 96 insertions(+), 20 deletions(-) create mode 100644 tests/guiders/test_frequency_decoupled_guidance.py diff --git a/src/diffusers/guiders/frequency_decoupled_guidance.py b/src/diffusers/guiders/frequency_decoupled_guidance.py index b92ddf2..917d67a 100644 --- a/src/diffusers/guiders/frequency_decoupled_guidance.py +++ b/src/diffusers/guiders/frequency_decoupled_guidance.py @@ -275,7 +275,7 @@ def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = No pred_guided_pyramid.append(pred) else: # Add the current pred_cond_pyramid level as the "non-FDG" prediction - pred_guided_pyramid.append(pred_cond_freq) + pred_guided_pyramid.append(pred_cond_pyramid[level]) # Convert from frequency space back to data (e.g. pixel) space by applying inverse freq transform pred = build_image_from_pyramid(pred_guided_pyramid) diff --git a/src/diffusers/hooks/mag_cache.py b/src/diffusers/hooks/mag_cache.py index d28cd2d..f54d208 100644 --- a/src/diffusers/hooks/mag_cache.py +++ b/src/diffusers/hooks/mag_cache.py @@ -168,12 +168,29 @@ def reset(self): self.calibration_ratios = [] +def _advance_mag_cache_step(state: MagCacheState, config: MagCacheConfig) -> None: + state.step_index += 1 + if state.step_index >= config.num_inference_steps: + if config.calibrate: + print("\n[MagCache] Calibration Complete. Copy these values to MagCacheConfig(mag_ratios=...):") + print(f"{state.calibration_ratios}\n") + logger.info(f"MagCache Calibration Results: {state.calibration_ratios}") + + state.step_index = 0 + state.accumulated_ratio = 1.0 + state.accumulated_steps = 0 + state.accumulated_err = 0.0 + state.previous_residual = None + state.calibration_ratios = [] + + class MagCacheHeadHook(ModelHook): _is_stateful = True - def __init__(self, state_manager: StateManager, config: MagCacheConfig): + def __init__(self, state_manager: StateManager, config: MagCacheConfig, advance_on_skip: bool = False): self.state_manager = state_manager self.config = config + self.advance_on_skip = advance_on_skip self._metadata = None def initialize_hook(self, module): @@ -260,6 +277,9 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): "Cannot apply residual safely. Returning input without residual." ) + if self.advance_on_skip: + _advance_mag_cache_step(state, self.config) + if self._metadata.return_encoder_hidden_states_index is not None: original_encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs( "encoder_hidden_states", args, kwargs @@ -377,21 +397,7 @@ def _perform_calibration_step(self, state: MagCacheState, current_residual: torc state.calibration_ratios.append(ratio) def _advance_step(self, state: MagCacheState): - state.step_index += 1 - if state.step_index >= self.config.num_inference_steps: - # End of inference loop - if self.config.calibrate: - print("\n[MagCache] Calibration Complete. Copy these values to MagCacheConfig(mag_ratios=...):") - print(f"{state.calibration_ratios}\n") - logger.info(f"MagCache Calibration Results: {state.calibration_ratios}") - - # Reset state - state.step_index = 0 - state.accumulated_ratio = 1.0 - state.accumulated_steps = 0 - state.accumulated_err = 0.0 - state.previous_residual = None - state.calibration_ratios = [] + _advance_mag_cache_step(state, self.config) def apply_mag_cache(module: torch.nn.Module, config: MagCacheConfig) -> None: @@ -425,7 +431,7 @@ def apply_mag_cache(module: torch.nn.Module, config: MagCacheConfig) -> None: name, block = remaining_blocks[0] logger.info(f"MagCache: Applying Head+Tail Hooks to single block '{name}'") _apply_mag_cache_block_hook(block, state_manager, config, is_tail=True) - _apply_mag_cache_head_hook(block, state_manager, config) + _apply_mag_cache_head_hook(block, state_manager, config, advance_on_skip=True) return head_block_name, head_block = remaining_blocks.pop(0) @@ -441,14 +447,16 @@ def apply_mag_cache(module: torch.nn.Module, config: MagCacheConfig) -> None: _apply_mag_cache_block_hook(tail_block, state_manager, config, is_tail=True) -def _apply_mag_cache_head_hook(block: torch.nn.Module, state_manager: StateManager, config: MagCacheConfig) -> None: +def _apply_mag_cache_head_hook( + block: torch.nn.Module, state_manager: StateManager, config: MagCacheConfig, advance_on_skip: bool = False +) -> None: registry = HookRegistry.check_if_exists_or_initialize(block) # Automatically remove existing hook to allow re-application (e.g. switching modes) if registry.get_hook(_MAG_CACHE_LEADER_BLOCK_HOOK) is not None: registry.remove_hook(_MAG_CACHE_LEADER_BLOCK_HOOK) - hook = MagCacheHeadHook(state_manager, config) + hook = MagCacheHeadHook(state_manager, config, advance_on_skip=advance_on_skip) registry.register_hook(hook, _MAG_CACHE_LEADER_BLOCK_HOOK) diff --git a/tests/guiders/test_frequency_decoupled_guidance.py b/tests/guiders/test_frequency_decoupled_guidance.py new file mode 100644 index 0000000..5be3910 --- /dev/null +++ b/tests/guiders/test_frequency_decoupled_guidance.py @@ -0,0 +1,36 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers.guiders.frequency_decoupled_guidance import FrequencyDecoupledGuidance + + +pytest.importorskip("kornia") + + +def test_frequency_decoupled_guidance_disabled_level_uses_cond_pyramid(): + guider = FrequencyDecoupledGuidance( + guidance_scales=[0.0, 7.5], + use_original_formulation=True, + ) + guider.set_state(step=0, num_inference_steps=10, timestep=torch.tensor([500])) + pred_cond = torch.randn(1, 3, 32, 32) + pred_uncond = torch.randn(1, 3, 32, 32) + + output = guider.forward(pred_cond, pred_uncond) + + assert output.pred is not None + assert output.pred.shape == pred_cond.shape diff --git a/tests/hooks/test_mag_cache.py b/tests/hooks/test_mag_cache.py index e5a500d..7dade55 100644 --- a/tests/hooks/test_mag_cache.py +++ b/tests/hooks/test_mag_cache.py @@ -244,3 +244,35 @@ def test_mag_cache_calibration(): # Let's ensure list is empty after reset (end of step 1) ratios_after = _get_calibration_data(model) assert ratios_after == [] + + +class SingleBlockDummyTransformer(ModelMixin): + def __init__(self): + super().__init__() + self.transformer_blocks = torch.nn.ModuleList([DummyBlock()]) + + def forward(self, hidden_states, encoder_hidden_states=None): + for block in self.transformer_blocks: + hidden_states = block(hidden_states, encoder_hidden_states=encoder_hidden_states) + return hidden_states + + +def test_mag_cache_single_block_step_reset(): + """Single-block models must advance step_index on cache skips.""" + model = SingleBlockDummyTransformer() + config = MagCacheConfig( + threshold=100.0, num_inference_steps=2, retention_ratio=0.0, mag_ratios=np.array([1.0, 1.0]) + ) + apply_mag_cache(model, config) + _set_context(model, "test_context") + + out0 = model(torch.tensor([[[10.0]]])) + assert torch.allclose(out0, torch.tensor([[[20.0]]])) + + out1 = model(torch.tensor([[[11.0]]])) + assert torch.allclose(out1, torch.tensor([[[21.0]]])) + + out2 = model(torch.tensor([[[12.0]]])) + assert torch.allclose(out2, torch.tensor([[[24.0]]])), ( + f"Expected compute after reset (24.0), got {out2.item()}" + ) From ba6459e828d0bbd2e8b7fd4c88caa4637f744e49 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 28 Jun 2026 22:10:51 +0000 Subject: [PATCH 2/2] Use python -m ruff in check_copies for environments without ruff on PATH Co-authored-by: Simon Lynch --- utils/check_copies.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/utils/check_copies.py b/utils/check_copies.py index 001366c..338a25d 100644 --- a/utils/check_copies.py +++ b/utils/check_copies.py @@ -18,6 +18,7 @@ import os import re import subprocess +import sys # All paths are set with the intent you should run this script from the root of the repo with the command @@ -94,7 +95,7 @@ def get_indent(code): def run_ruff(code): - command = ["ruff", "format", "-", "--config", "pyproject.toml", "--silent"] + command = [sys.executable, "-m", "ruff", "format", "-", "--config", "pyproject.toml", "--silent"] process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE) stdout, _ = process.communicate(input=code.encode()) return stdout.decode()