fix

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -912,15 +912,13 @@ def nearest_interp(src, target_len):
     indices = np.round(np.linspace(0, src_len - 1, target_len)).astype(np.int32)
     return src[indices]
 
-def init_magcache(num_inference_steps, retention_ratio, mag_ratios_base, split_step=None, model_type="T2V"):
+def init_magcache(num_inference_steps, retention_ratio, mag_ratios_base):
     """Initialize MagCache variables and interpolate ratios.
     
     Args:
         num_inference_steps: Number of inference steps.
         retention_ratio: Retention ratio of unchanged steps.
         mag_ratios_base: Base magnitude ratios array or list.
-        split_step: Step at which model switches (e.g. high -> low noise for 2.2).
-        model_type: Pipeline mode ("T2V" or "I2V").
     """
     import numpy as np
 
@@ -953,8 +951,6 @@ def init_magcache(num_inference_steps, retention_ratio, mag_ratios_base, split_s
         cached_residual,
         skip_warmup,
         mag_ratios,
-        split_step,
-        model_type,
     )
 
 def magcache_step(
@@ -964,10 +960,6 @@ def magcache_step(
     magcache_thresh,
     magcache_K,
     skip_warmup,
-    split_step=None,
-    model_type="T2V",
-    num_steps=None,
-    retention_ratio=0.2,
 ):
     """Update MagCache accumulated state and decide if to skip.
     
@@ -978,10 +970,6 @@ def magcache_step(
         magcache_thresh: Error threshold.
         magcache_K: Max skip steps.
         skip_warmup: Warmup steps threshold.
-        split_step: Optional step index where the model switches (e.g., from high to low noise).
-        model_type: Pipeline type ("T2V" or "I2V").
-        num_steps: Total inference steps, used to calculate post-split warmups.
-        retention_ratio: Used to calculate post-split warmups.
     """
     import numpy as np
 
@@ -998,16 +986,8 @@ def magcache_step(
     cur_mag_ratio_uncond = mag_ratios[step * 2 + 1]
 
     use_magcache = True
-    if split_step is not None:
-        if model_type == "I2V":
-            if step < int(split_step + (num_steps - split_step) * retention_ratio):
-                use_magcache = False
-        else:
-            if step < int(split_step * retention_ratio) or (step <= ((num_steps - split_step) * retention_ratio + split_step) and step >= split_step):
-                use_magcache = False
-    else:
-        if step < skip_warmup:
-            use_magcache = False
+    if step < skip_warmup:
+        use_magcache = False
 
     skip_blocks = False
     if use_magcache:

src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py

@@ -246,8 +246,6 @@ def run_inference_2_1(
         cached_residual,
         skip_warmup,
         mag_ratios,
-        split_step,
-        model_type,
     ) = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base)
 
     for step in range(num_inference_steps):
@@ -263,7 +261,7 @@ def run_inference_2_1(
           accumulated_steps_uncond,
       )
       skip_blocks, accumulated_state = magcache_step(
-          step, mag_ratios, accumulated_state, magcache_thresh, magcache_K, skip_warmup, split_step=split_step, model_type=model_type, num_steps=num_inference_steps, retention_ratio=retention_ratio
+          step, mag_ratios, accumulated_state, magcache_thresh, magcache_K, skip_warmup
       )
       (
           accumulated_ratio_cond,

src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py

@@ -456,7 +456,6 @@ def run_inference_2_2(
   if use_magcache and do_classifier_free_guidance:
     timesteps_np = np.array(scheduler_state.timesteps, dtype=np.int32)
     step_uses_high = [bool(timesteps_np[s] >= boundary) for s in range(num_inference_steps)]
-    split_step = sum(step_uses_high)
 
     (
         accumulated_ratio_cond,
@@ -468,9 +467,7 @@ def run_inference_2_2(
         cached_residual,
         skip_warmup,
         mag_ratios,
-        split_step,
-        model_type,
-    ) = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base, split_step=split_step, model_type="T2V")
+    ) = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base)
 
     prompt_embeds_combined = jnp.concatenate([prompt_embeds, negative_prompt_embeds], axis=0)
 
@@ -486,7 +483,7 @@ def run_inference_2_2(
           accumulated_steps_uncond,
       )
       skip_blocks, accumulated_state = magcache_step(
-          step, mag_ratios, accumulated_state, magcache_thresh, magcache_K, skip_warmup, split_step=split_step, model_type=model_type, num_steps=num_inference_steps, retention_ratio=retention_ratio
+          step, mag_ratios, accumulated_state, magcache_thresh, magcache_K, skip_warmup
       )
       (
           accumulated_ratio_cond,
@@ -503,6 +500,7 @@ def run_inference_2_2(
       else:
         graphdef, state, rest = low_noise_graphdef, low_noise_state, low_noise_rest
         guidance_scale = guidance_scale_low
+        skip_blocks = False  # Reference MagCache only caches high_noise_model
 
       timestep = jnp.broadcast_to(t, bsz * 2)
 

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -295,8 +295,6 @@ def run_inference_2_1_i2v(
       cached_residual,
       skip_warmup,
       mag_ratios,
-      split_step,
-      model_type,
   ) = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base)
 
   if do_classifier_free_guidance:
@@ -322,7 +320,7 @@ def run_inference_2_1_i2v(
           accumulated_steps_uncond,
       )
       skip_blocks, accumulated_state = magcache_step(
-          step, mag_ratios, accumulated_state, magcache_thresh, magcache_K, skip_warmup, split_step=split_step, model_type=model_type, num_steps=num_inference_steps, retention_ratio=retention_ratio
+          step, mag_ratios, accumulated_state, magcache_thresh, magcache_K, skip_warmup
       )
       (
           accumulated_ratio_cond,

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py

@@ -449,7 +449,6 @@ def run_inference_2_2_i2v(
   if use_magcache and do_classifier_free_guidance:
     timesteps_np = np.array(scheduler_state.timesteps, dtype=np.int32)
     step_uses_high = [bool(timesteps_np[s] >= boundary) for s in range(num_inference_steps)]
-    split_step = sum(step_uses_high)
 
     (
         accumulated_ratio_cond,
@@ -461,9 +460,7 @@ def run_inference_2_2_i2v(
         cached_residual,
         skip_warmup,
         mag_ratios,
-        split_step,
-        model_type,
-    ) = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base, split_step=split_step, model_type="I2V")
+    ) = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base)
 
 
     prompt_embeds_combined = jnp.concatenate([prompt_embeds, negative_prompt_embeds], axis=0)
@@ -486,7 +483,7 @@ def run_inference_2_2_i2v(
           accumulated_steps_uncond,
       )
       skip_blocks, accumulated_state = magcache_step(
-          step, mag_ratios, accumulated_state, magcache_thresh, magcache_K, skip_warmup, split_step=split_step, model_type=model_type, num_steps=num_inference_steps, retention_ratio=retention_ratio
+          step, mag_ratios, accumulated_state, magcache_thresh, magcache_K, skip_warmup
       )
       (
           accumulated_ratio_cond,
@@ -503,6 +500,7 @@ def run_inference_2_2_i2v(
       else:
         graphdef, state, rest = low_noise_graphdef, low_noise_state, low_noise_rest
         guidance_scale = guidance_scale_low
+        skip_blocks = False  # Reference MagCache only caches high_noise_model
 
       timestep = jnp.broadcast_to(t, bsz * 2)
       latents_doubled = jnp.concatenate([latents, latents], axis=0)