correction needed for wan 2.2

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -912,13 +912,15 @@ 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):
+def init_magcache(num_inference_steps, retention_ratio, mag_ratios_base, split_step=None, model_type="T2V"):
     """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
 
@@ -951,6 +953,8 @@ def init_magcache(num_inference_steps, retention_ratio, mag_ratios_base):
         cached_residual,
         skip_warmup,
         mag_ratios,
+        split_step,
+        model_type,
     )
 
 def magcache_step(
@@ -960,6 +964,10 @@ 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.
     
@@ -970,6 +978,10 @@ 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
 
@@ -985,8 +997,20 @@ def magcache_step(
     cur_mag_ratio_cond = mag_ratios[step * 2]
     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
+
     skip_blocks = False
-    if step >= skip_warmup:
+    if use_magcache:
         new_ratio_cond = accumulated_ratio_cond * cur_mag_ratio_cond
         new_ratio_uncond = accumulated_ratio_uncond * cur_mag_ratio_uncond
 

src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py

@@ -246,6 +246,8 @@ 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):
@@ -261,7 +263,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
+          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
       )
       (
           accumulated_ratio_cond,

src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py

@@ -454,6 +454,10 @@ def run_inference_2_2(
 
   # ── MagCache path ──
   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,
         accumulated_ratio_uncond,
@@ -464,10 +468,10 @@ def run_inference_2_2(
         cached_residual,
         skip_warmup,
         mag_ratios,
-    ) = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base)
+        split_step,
+        model_type,
+    ) = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base, split_step=split_step, model_type=self.config.model_type)
 
-    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)]
     prompt_embeds_combined = jnp.concatenate([prompt_embeds, negative_prompt_embeds], axis=0)
 
     for step in range(num_inference_steps):
@@ -482,7 +486,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
+          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
       )
       (
           accumulated_ratio_cond,

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -295,6 +295,8 @@ 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:
@@ -320,7 +322,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
+          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
       )
       (
           accumulated_ratio_cond,

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py

@@ -447,6 +447,10 @@ def run_inference_2_2_i2v(
 
   # ── MagCache path ──
   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,
         accumulated_ratio_uncond,
@@ -457,10 +461,10 @@ def run_inference_2_2_i2v(
         cached_residual,
         skip_warmup,
         mag_ratios,
-    ) = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base)
+        split_step,
+        model_type,
+    ) = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base, split_step=split_step, model_type=self.config.model_type)
 
-    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)]
 
     prompt_embeds_combined = jnp.concatenate([prompt_embeds, negative_prompt_embeds], axis=0)
     if image_embeds is not None:
@@ -482,7 +486,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
+          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
       )
       (
           accumulated_ratio_cond,