Fix host slowdown regression in magcache_step by reverting to pure numpy

Author: Perseus14

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -1333,8 +1333,7 @@ def magcache_step(
       skip_warmup: Warmup steps threshold.
       use_magcache: Optional manual override boolean to enable/disable cache for this step.
   """
-  import jax.numpy as jnp
-  import jax
+  import numpy as np
 
   (
       accumulated_ratio_cond,
@@ -1349,97 +1348,38 @@ def magcache_step(
   cur_mag_ratio_uncond = mag_ratios[step * 2 + 1]
 
   if use_magcache is None:
-    use_magcache = step >= skip_warmup
+    use_magcache = True
+    if step < skip_warmup:
+      use_magcache = False
 
-  def compute_magcache_hit(args):
-    (
-        accumulated_ratio_cond,
-        accumulated_ratio_uncond,
-        accumulated_err_cond,
-        accumulated_err_uncond,
-        accumulated_steps_cond,
-        accumulated_steps_uncond,
-        cur_mag_ratio_cond,
-        cur_mag_ratio_uncond,
-    ) = args
+  skip_blocks = False
+  if use_magcache:
     new_ratio_cond = accumulated_ratio_cond * cur_mag_ratio_cond
     new_ratio_uncond = accumulated_ratio_uncond * cur_mag_ratio_uncond
 
-    err_cond = jnp.abs(1.0 - new_ratio_cond)
-    err_uncond = jnp.abs(1.0 - new_ratio_uncond)
-
-    cond = jnp.logical_and(accumulated_err_cond + err_cond < magcache_thresh, accumulated_steps_cond < magcache_K)
-    cond = jnp.logical_and(cond, accumulated_err_uncond + err_uncond < magcache_thresh)
-    cond = jnp.logical_and(cond, accumulated_steps_uncond < magcache_K)
-
-    def hit_fn(_):
-      return (
-          True,
-          new_ratio_cond,
-          new_ratio_uncond,
-          accumulated_err_cond + err_cond,
-          accumulated_err_uncond + err_uncond,
-          accumulated_steps_cond + 1,
-          accumulated_steps_uncond + 1,
-      )
-
-    def miss_fn(_):
-      return (
-          False,
-          jnp.array(1.0, dtype=jnp.float32),
-          jnp.array(1.0, dtype=jnp.float32),
-          jnp.array(0.0, dtype=jnp.float32),
-          jnp.array(0.0, dtype=jnp.float32),
-          jnp.array(0, dtype=jnp.int32),
-          jnp.array(0, dtype=jnp.int32),
-      )
-
-    return jax.lax.cond(cond, hit_fn, miss_fn, None)
-
-  def skip_magcache(args):
-    (
-        accumulated_ratio_cond,
-        accumulated_ratio_uncond,
-        accumulated_err_cond,
-        accumulated_err_uncond,
-        accumulated_steps_cond,
-        accumulated_steps_uncond,
-        _,
-        _,
-    ) = args
-    return (
-        False,
-        accumulated_ratio_cond,
-        accumulated_ratio_uncond,
-        accumulated_err_cond,
-        accumulated_err_uncond,
-        accumulated_steps_cond,
-        accumulated_steps_uncond,
-    )
-
-  (
-      skip_blocks,
-      accumulated_ratio_cond,
-      accumulated_ratio_uncond,
-      accumulated_err_cond,
-      accumulated_err_uncond,
-      accumulated_steps_cond,
-      accumulated_steps_uncond,
-  ) = jax.lax.cond(
-      use_magcache,
-      compute_magcache_hit,
-      skip_magcache,
-      (
-          accumulated_ratio_cond,
-          accumulated_ratio_uncond,
-          accumulated_err_cond,
-          accumulated_err_uncond,
-          accumulated_steps_cond,
-          accumulated_steps_uncond,
-          cur_mag_ratio_cond,
-          cur_mag_ratio_uncond,
-      ),
-  )
+    err_cond = np.abs(1.0 - new_ratio_cond)
+    err_uncond = np.abs(1.0 - new_ratio_uncond)
+
+    if (
+        accumulated_err_cond + err_cond < magcache_thresh
+        and accumulated_steps_cond < magcache_K
+        and accumulated_err_uncond + err_uncond < magcache_thresh
+        and accumulated_steps_uncond < magcache_K
+    ):
+      skip_blocks = True
+      accumulated_ratio_cond = new_ratio_cond
+      accumulated_ratio_uncond = new_ratio_uncond
+      accumulated_err_cond += err_cond
+      accumulated_err_uncond += err_uncond
+      accumulated_steps_cond += 1
+      accumulated_steps_uncond += 1
+    else:
+      accumulated_ratio_cond = 1.0
+      accumulated_ratio_uncond = 1.0
+      accumulated_err_cond = 0.0
+      accumulated_err_uncond = 0.0
+      accumulated_steps_cond = 0
+      accumulated_steps_uncond = 0
 
   new_state = (
       accumulated_ratio_cond,