wan 2.2 fixes

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -959,7 +959,8 @@ def magcache_step(
     accumulated_state,
     magcache_thresh,
     magcache_K,
-    skip_warmup,
+    skip_warmup=0,
+    use_magcache=None,
 ):
     """Update MagCache accumulated state and decide if to skip.
     
@@ -970,6 +971,7 @@ def magcache_step(
         magcache_thresh: Error threshold.
         magcache_K: Max skip steps.
         skip_warmup: Warmup steps threshold.
+        use_magcache: Optional manual override boolean to enable/disable cache for this step.
     """
     import numpy as np
 
@@ -985,9 +987,10 @@ def magcache_step(
     cur_mag_ratio_cond = mag_ratios[step * 2]
     cur_mag_ratio_uncond = mag_ratios[step * 2 + 1]
 
-    use_magcache = True
-    if step < skip_warmup:
-        use_magcache = False
+    if use_magcache is None:
+        use_magcache = True
+        if step < skip_warmup:
+            use_magcache = False
 
     skip_blocks = False
     if use_magcache:

src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py

@@ -471,9 +471,17 @@ def run_inference_2_2(
 
     prompt_embeds_combined = jnp.concatenate([prompt_embeds, negative_prompt_embeds], axis=0)
 
+    high_noise_steps = sum(step_uses_high)
+
     for step in range(num_inference_steps):
       t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
 
+      use_magcache = True
+      if step < int(high_noise_steps * retention_ratio):
+        use_magcache = False
+      elif step >= high_noise_steps and step <= int(high_noise_steps + (num_inference_steps - high_noise_steps) * retention_ratio):
+        use_magcache = False
+
       accumulated_state = (
           accumulated_ratio_cond,
           accumulated_ratio_uncond,
@@ -483,7 +491,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, use_magcache=use_magcache
       )
       (
           accumulated_ratio_cond,
@@ -500,7 +508,6 @@ 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_2p2.py

@@ -471,9 +471,15 @@ def run_inference_2_2_i2v(
 
     condition_combined = jnp.concatenate([condition] * 2)
 
+    high_noise_steps = sum(step_uses_high)
+
     for step in range(num_inference_steps):
       t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
 
+      use_magcache = True
+      if step < int(high_noise_steps + (num_inference_steps - high_noise_steps) * retention_ratio):
+        use_magcache = False
+
       accumulated_state = (
           accumulated_ratio_cond,
           accumulated_ratio_uncond,
@@ -483,7 +489,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, use_magcache=use_magcache
       )
       (
           accumulated_ratio_cond,
@@ -500,7 +506,6 @@ 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)