From 8bd6e6c06e97be1f7364eaa7aecf97bbdd76ed82 Mon Sep 17 00:00:00 2001 From: lcheng Date: Fri, 3 Jul 2026 14:09:26 +0800 Subject: [PATCH 1/2] fix(mag_cache): correct per-transformer step count for Wan 2.2 two-stage denoising --- src/diffusers/hooks/__init__.py | 2 +- src/diffusers/hooks/mag_cache.py | 14 ++++++++++++++ src/diffusers/pipelines/wan/pipeline_wan.py | 9 +++++++++ src/diffusers/pipelines/wan/pipeline_wan_i2v.py | 9 +++++++++ src/diffusers/pipelines/wan/pipeline_wan_vace.py | 9 +++++++++ 5 files changed, 42 insertions(+), 1 deletion(-) diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py index 2a9aa81608e7..f6e582d1642b 100644 --- a/src/diffusers/hooks/__init__.py +++ b/src/diffusers/hooks/__init__.py @@ -23,7 +23,7 @@ from .hooks import HookRegistry, ModelHook from .layer_skip import LayerSkipConfig, apply_layer_skip from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook - from .mag_cache import MagCacheConfig, apply_mag_cache + from .mag_cache import MagCacheConfig, apply_mag_cache, update_mag_cache_num_steps from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast from .smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig from .taylorseer_cache import TaylorSeerCacheConfig, apply_taylorseer_cache diff --git a/src/diffusers/hooks/mag_cache.py b/src/diffusers/hooks/mag_cache.py index e5f0aaebc01a..31735f4d347e 100644 --- a/src/diffusers/hooks/mag_cache.py +++ b/src/diffusers/hooks/mag_cache.py @@ -130,6 +130,7 @@ def __post_init__(self): if not torch.is_tensor(self.mag_ratios): self.mag_ratios = torch.tensor(self.mag_ratios) + self._original_mag_ratios = self.mag_ratios.clone() if len(self.mag_ratios) != self.num_inference_steps: logger.debug( f"Interpolating mag_ratios from length {len(self.mag_ratios)} to {self.num_inference_steps}" @@ -407,6 +408,7 @@ def apply_mag_cache(module: torch.nn.Module, config: MagCacheConfig) -> None: # Initialize registry on the root module so the Pipeline can set context. HookRegistry.check_if_exists_or_initialize(module) + module._mag_cache_config = config state_manager = StateManager(MagCacheState, (), {}) remaining_blocks = [] @@ -466,3 +468,15 @@ def _apply_mag_cache_block_hook( hook = MagCacheBlockHook(state_manager, is_tail, config) registry.register_hook(hook, _MAG_CACHE_BLOCK_HOOK) + + +def update_mag_cache_num_steps(module: torch.nn.Module, num_steps: int) -> None: + config: MagCacheConfig = getattr(module, "_mag_cache_config", None) + if config is None: + return + original_ratios = getattr(config, "_original_mag_ratios", config.mag_ratios) + config.num_inference_steps = num_steps + if original_ratios is not None: + config.mag_ratios = nearest_interp(original_ratios, num_steps) + + diff --git a/src/diffusers/pipelines/wan/pipeline_wan.py b/src/diffusers/pipelines/wan/pipeline_wan.py index be2d53f17932..536cd7b9e31b 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan.py +++ b/src/diffusers/pipelines/wan/pipeline_wan.py @@ -586,6 +586,15 @@ def __call__( else: boundary_timestep = None + if boundary_timestep is not None: + from ...hooks.mag_cache import update_mag_cache_num_steps + n_steps_t1 = sum(1 for t in timesteps if t >= boundary_timestep) + n_steps_t2 = len(timesteps) - n_steps_t1 + if self.transformer is not None: + update_mag_cache_num_steps(self.transformer, n_steps_t1) + if self.transformer_2 is not None: + update_mag_cache_num_steps(self.transformer_2, n_steps_t2) + with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): if self.interrupt: diff --git a/src/diffusers/pipelines/wan/pipeline_wan_i2v.py b/src/diffusers/pipelines/wan/pipeline_wan_i2v.py index 8061f67ab6b9..4989ab1b20b6 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_i2v.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_i2v.py @@ -741,6 +741,15 @@ def __call__( else: boundary_timestep = None + if boundary_timestep is not None: + from ...hooks.mag_cache import update_mag_cache_num_steps + n_steps_t1 = sum(1 for t in timesteps if t >= boundary_timestep) + n_steps_t2 = len(timesteps) - n_steps_t1 + if self.transformer is not None: + update_mag_cache_num_steps(self.transformer, n_steps_t1) + if self.transformer_2 is not None: + update_mag_cache_num_steps(self.transformer_2, n_steps_t2) + with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): if self.interrupt: diff --git a/src/diffusers/pipelines/wan/pipeline_wan_vace.py b/src/diffusers/pipelines/wan/pipeline_wan_vace.py index b0896d382d67..22ba7fdd2bc0 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_vace.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_vace.py @@ -955,6 +955,15 @@ def __call__( else: boundary_timestep = None + if boundary_timestep is not None: + from ...hooks.mag_cache import update_mag_cache_num_steps + n_steps_t1 = sum(1 for t in timesteps if t >= boundary_timestep) + n_steps_t2 = len(timesteps) - n_steps_t1 + if self.transformer is not None: + update_mag_cache_num_steps(self.transformer, n_steps_t1) + if self.transformer_2 is not None: + update_mag_cache_num_steps(self.transformer_2, n_steps_t2) + with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): if self.interrupt: From 6d1959717f4d0effac2b168371cf26a52df9e5ce Mon Sep 17 00:00:00 2001 From: lcheng Date: Wed, 8 Jul 2026 14:18:54 +0800 Subject: [PATCH 2/2] fix(mag_cache): stop mutating MagCacheConfig; broadcast per-transformer step count via plain attribute MagCacheConfig is no longer mutated at runtime. Instead, pipelines broadcast each transformer's actual step count via a plain _mag_cache_expected_steps attribute on the root module, which MagCacheHeadHook/MagCacheBlockHook read live through _get_effective_num_steps/_get_effective_mag_ratios (falling back to config.num_inference_steps when unset, so single-transformer pipelines are unaffected). Addresses both review comments on #14027: - Sharing one MagCacheConfig across transformer/transformer_2 under an asymmetric step split can no longer corrupt it, since it's never mutated. - Pipelines no longer import hook internals (update_mag_cache_num_steps is removed); they only set a plain attribute, keeping pipeline code decoupled from hook implementation details. --- src/diffusers/hooks/__init__.py | 2 +- src/diffusers/hooks/mag_cache.py | 69 ++++++++++++------- src/diffusers/pipelines/wan/pipeline_wan.py | 5 +- .../pipelines/wan/pipeline_wan_i2v.py | 5 +- .../pipelines/wan/pipeline_wan_vace.py | 5 +- 5 files changed, 51 insertions(+), 35 deletions(-) diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py index f6e582d1642b..2a9aa81608e7 100644 --- a/src/diffusers/hooks/__init__.py +++ b/src/diffusers/hooks/__init__.py @@ -23,7 +23,7 @@ from .hooks import HookRegistry, ModelHook from .layer_skip import LayerSkipConfig, apply_layer_skip from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook - from .mag_cache import MagCacheConfig, apply_mag_cache, update_mag_cache_num_steps + from .mag_cache import MagCacheConfig, apply_mag_cache from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast from .smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig from .taylorseer_cache import TaylorSeerCacheConfig, apply_taylorseer_cache diff --git a/src/diffusers/hooks/mag_cache.py b/src/diffusers/hooks/mag_cache.py index 31735f4d347e..b3b37578d6c5 100644 --- a/src/diffusers/hooks/mag_cache.py +++ b/src/diffusers/hooks/mag_cache.py @@ -169,12 +169,25 @@ def reset(self): self.calibration_ratios = [] +def _get_effective_num_steps(config: MagCacheConfig, root_module: torch.nn.Module) -> int: + expected = getattr(root_module, "_mag_cache_expected_steps", None) + return expected if expected is not None else config.num_inference_steps + + +def _get_effective_mag_ratios(config: MagCacheConfig, effective_num_steps: int) -> torch.Tensor: + if effective_num_steps == config.num_inference_steps: + return config.mag_ratios + original = getattr(config, "_original_mag_ratios", config.mag_ratios) + return nearest_interp(original, effective_num_steps) + + class MagCacheHeadHook(ModelHook): _is_stateful = True - def __init__(self, state_manager: StateManager, config: MagCacheConfig): + def __init__(self, state_manager: StateManager, config: MagCacheConfig, root_module: torch.nn.Module): self.state_manager = state_manager self.config = config + self.root_module = root_module self._metadata = None def initialize_hook(self, module): @@ -200,13 +213,16 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): should_compute = True else: # MagCache Logic + effective_num_steps = _get_effective_num_steps(self.config, self.root_module) + effective_mag_ratios = _get_effective_mag_ratios(self.config, effective_num_steps) + current_step = state.step_index - if current_step >= len(self.config.mag_ratios): + if current_step >= len(effective_mag_ratios): current_scale = 1.0 else: - current_scale = self.config.mag_ratios[current_step] + current_scale = effective_mag_ratios[current_step] - retention_step = int(self.config.retention_ratio * self.config.num_inference_steps + 0.5) + retention_step = int(self.config.retention_ratio * effective_num_steps + 0.5) if current_step >= retention_step: state.accumulated_ratio *= current_scale @@ -286,11 +302,18 @@ def reset_state(self, module): class MagCacheBlockHook(ModelHook): - def __init__(self, state_manager: StateManager, is_tail: bool = False, config: MagCacheConfig = None): + def __init__( + self, + state_manager: StateManager, + is_tail: bool = False, + config: MagCacheConfig = None, + root_module: torch.nn.Module = None, + ): super().__init__() self.state_manager = state_manager self.is_tail = is_tail self.config = config + self.root_module = root_module self._metadata = None def initialize_hook(self, module): @@ -379,7 +402,8 @@ def _perform_calibration_step(self, state: MagCacheState, current_residual: torc def _advance_step(self, state: MagCacheState): state.step_index += 1 - if state.step_index >= self.config.num_inference_steps: + effective_num_steps = _get_effective_num_steps(self.config, self.root_module) + if state.step_index >= effective_num_steps: # End of inference loop if self.config.calibrate: print("\n[MagCache] Calibration Complete. Copy these values to MagCacheConfig(mag_ratios=...):") @@ -426,31 +450,36 @@ def apply_mag_cache(module: torch.nn.Module, config: MagCacheConfig) -> None: if len(remaining_blocks) == 1: 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_block_hook(block, state_manager, config, module, is_tail=True) + _apply_mag_cache_head_hook(block, state_manager, config, module) return head_block_name, head_block = remaining_blocks.pop(0) tail_block_name, tail_block = remaining_blocks.pop(-1) logger.info(f"MagCache: Applying Head Hook to {head_block_name}") - _apply_mag_cache_head_hook(head_block, state_manager, config) + _apply_mag_cache_head_hook(head_block, state_manager, config, module) for name, block in remaining_blocks: - _apply_mag_cache_block_hook(block, state_manager, config) + _apply_mag_cache_block_hook(block, state_manager, config, module) logger.info(f"MagCache: Applying Tail Hook to {tail_block_name}") - _apply_mag_cache_block_hook(tail_block, state_manager, config, is_tail=True) + _apply_mag_cache_block_hook(tail_block, state_manager, config, module, 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, + root_module: torch.nn.Module, +) -> 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, root_module) registry.register_hook(hook, _MAG_CACHE_LEADER_BLOCK_HOOK) @@ -458,6 +487,7 @@ def _apply_mag_cache_block_hook( block: torch.nn.Module, state_manager: StateManager, config: MagCacheConfig, + root_module: torch.nn.Module, is_tail: bool = False, ) -> None: registry = HookRegistry.check_if_exists_or_initialize(block) @@ -466,17 +496,6 @@ def _apply_mag_cache_block_hook( if registry.get_hook(_MAG_CACHE_BLOCK_HOOK) is not None: registry.remove_hook(_MAG_CACHE_BLOCK_HOOK) - hook = MagCacheBlockHook(state_manager, is_tail, config) + hook = MagCacheBlockHook(state_manager, is_tail, config, root_module) registry.register_hook(hook, _MAG_CACHE_BLOCK_HOOK) - -def update_mag_cache_num_steps(module: torch.nn.Module, num_steps: int) -> None: - config: MagCacheConfig = getattr(module, "_mag_cache_config", None) - if config is None: - return - original_ratios = getattr(config, "_original_mag_ratios", config.mag_ratios) - config.num_inference_steps = num_steps - if original_ratios is not None: - config.mag_ratios = nearest_interp(original_ratios, num_steps) - - diff --git a/src/diffusers/pipelines/wan/pipeline_wan.py b/src/diffusers/pipelines/wan/pipeline_wan.py index 536cd7b9e31b..0171e99393d6 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan.py +++ b/src/diffusers/pipelines/wan/pipeline_wan.py @@ -587,13 +587,12 @@ def __call__( boundary_timestep = None if boundary_timestep is not None: - from ...hooks.mag_cache import update_mag_cache_num_steps n_steps_t1 = sum(1 for t in timesteps if t >= boundary_timestep) n_steps_t2 = len(timesteps) - n_steps_t1 if self.transformer is not None: - update_mag_cache_num_steps(self.transformer, n_steps_t1) + self.transformer._mag_cache_expected_steps = n_steps_t1 if self.transformer_2 is not None: - update_mag_cache_num_steps(self.transformer_2, n_steps_t2) + self.transformer_2._mag_cache_expected_steps = n_steps_t2 with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): diff --git a/src/diffusers/pipelines/wan/pipeline_wan_i2v.py b/src/diffusers/pipelines/wan/pipeline_wan_i2v.py index 4989ab1b20b6..eef715f0b0d3 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_i2v.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_i2v.py @@ -742,13 +742,12 @@ def __call__( boundary_timestep = None if boundary_timestep is not None: - from ...hooks.mag_cache import update_mag_cache_num_steps n_steps_t1 = sum(1 for t in timesteps if t >= boundary_timestep) n_steps_t2 = len(timesteps) - n_steps_t1 if self.transformer is not None: - update_mag_cache_num_steps(self.transformer, n_steps_t1) + self.transformer._mag_cache_expected_steps = n_steps_t1 if self.transformer_2 is not None: - update_mag_cache_num_steps(self.transformer_2, n_steps_t2) + self.transformer_2._mag_cache_expected_steps = n_steps_t2 with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): diff --git a/src/diffusers/pipelines/wan/pipeline_wan_vace.py b/src/diffusers/pipelines/wan/pipeline_wan_vace.py index 22ba7fdd2bc0..63b0f8d9c64e 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_vace.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_vace.py @@ -956,13 +956,12 @@ def __call__( boundary_timestep = None if boundary_timestep is not None: - from ...hooks.mag_cache import update_mag_cache_num_steps n_steps_t1 = sum(1 for t in timesteps if t >= boundary_timestep) n_steps_t2 = len(timesteps) - n_steps_t1 if self.transformer is not None: - update_mag_cache_num_steps(self.transformer, n_steps_t1) + self.transformer._mag_cache_expected_steps = n_steps_t1 if self.transformer_2 is not None: - update_mag_cache_num_steps(self.transformer_2, n_steps_t2) + self.transformer_2._mag_cache_expected_steps = n_steps_t2 with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps):