Skip to content

fix(mag_cache): correct per-transformer step count for dual-transformer pipelines#14027

Open
lcheng321 wants to merge 1 commit into
huggingface:mainfrom
lcheng321:patch-1
Open

fix(mag_cache): correct per-transformer step count for dual-transformer pipelines#14027
lcheng321 wants to merge 1 commit into
huggingface:mainfrom
lcheng321:patch-1

Conversation

@lcheng321

Copy link
Copy Markdown
Contributor

Fixes #14025

Root Cause

MagCacheConfig is applied to both transformers with num_inference_steps equal to the full pipeline step count. In dual-transformer pipelines (Wan 2.2), transformer runs only the high-noise stage and transformer_2 only the low-noise stage. Each executes a subset of the total steps, never the full count.

Three consequences:

  1. _advance_step: state.step_index >= config.num_inference_steps never fires, so state never resets and calibration output is never printed.
  2. Retention window: int(retention_ratio * num_inference_steps) is computed against the full step count, protecting more actual steps from skipping than intended.
  3. mag_ratios indexing: step_index no longer maps to the correct denoising timestep.

Changes

src/diffusers/hooks/mag_cache.py

  • Store _original_mag_ratios before interpolation so ratios can be re-interpolated to the corrected count.
  • Pin _mag_cache_config on the module after apply_mag_cache so pipelines can reach it.
  • Add update_mag_cache_num_steps helper.

src/diffusers/pipelines/wan/pipeline_wan.py, pipeline_wan_i2v.py, pipeline_wan_vace.py

After boundary_timestep is computed, count actual steps per transformer and correct each config before the denoising loop. No-op when MagCache is not applied.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

MagCache on Wan 2.2 Dual-Transformer Pipelines: Incorrect Step Accounting and Limited Effectiveness on a 4-Step Distilled Model (e.g. Wan2.2)

1 participant