reformatted

Author: prishajain1

src/maxdiffusion/models/wan/transformers/transformer_wan.py

@@ -650,7 +650,7 @@ def scan_fn(carry, block):
 
         rematted_block_forward = self.gradient_checkpoint.apply(
             scan_fn, self.names_which_can_be_saved, self.names_which_can_be_offloaded, prevent_cse=not self.scan_layers
-         )
+        )
         initial_carry = (h, rngs)
         final_carry, _ = nnx.scan(
             rematted_block_forward,
@@ -676,7 +676,10 @@ def layer_forward(hidden_states):
             )
 
           rematted_layer_forward = self.gradient_checkpoint.apply(
-              layer_forward, self.names_which_can_be_saved, self.names_which_can_be_offloaded, prevent_cse=not self.scan_layers
+              layer_forward,
+              self.names_which_can_be_saved,
+              self.names_which_can_be_offloaded,
+              prevent_cse=not self.scan_layers,
           )
           h_out = rematted_layer_forward(h_out)
       return h_out
@@ -702,7 +705,7 @@ def layer_forward(hidden_states):
     )
     hidden_states = jnp.transpose(hidden_states, (0, 7, 1, 4, 2, 5, 3, 6))
     hidden_states = hidden_states.reshape(batch_size, -1, num_frames, height, width)
-    
+
     if return_residual:
       return hidden_states, residual_x
     return hidden_states

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -775,7 +775,7 @@ def transformer_forward_pass(
       cached_residual=cached_residual,
       return_residual=return_residual,
   )
-  
+
   if return_residual:
     noise_pred, residual_x = outputs
   else:
@@ -899,56 +899,61 @@ def transformer_forward_pass_cfg_cache(
   noise_pred_merged = noise_uncond_approx + guidance_scale * (noise_cond - noise_uncond_approx)
   return noise_pred_merged, noise_cond
 
+
 def nearest_interp(src, target_len):
-    """Nearest neighbor interpolation for ratio scaling layout."""
-    src_len = len(src)
-    if target_len == 1: 
-        import numpy as np
-        return np.array([src[-1]])
+  """Nearest neighbor interpolation for ratio scaling layout."""
+  src_len = len(src)
+  if target_len == 1:
     import numpy as np
-    indices = np.round(np.linspace(0, src_len - 1, target_len)).astype(np.int32)
-    return src[indices]
+
+    return np.array([src[-1]])
+  import numpy as np
+
+  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):
-    """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.
-    """
-    import numpy as np
-    
-    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
-    cached_residual = None
-
-    skip_warmup = int(num_inference_steps * retention_ratio)
-
-    mag_ratios_base = np.array(mag_ratios_base)
-
-    if len(mag_ratios_base) != num_inference_steps * 2:
-        mag_cond = nearest_interp(mag_ratios_base[0::2], num_inference_steps)
-        mag_uncond = nearest_interp(mag_ratios_base[1::2], num_inference_steps)
-        mag_ratios = np.concatenate([mag_cond.reshape(-1, 1), mag_uncond.reshape(-1, 1)], axis=1).reshape(-1)
-    else:
-        mag_ratios = mag_ratios_base
-
-    return (
-        accumulated_ratio_cond,
-        accumulated_ratio_uncond,
-        accumulated_err_cond,
-        accumulated_err_uncond,
-        accumulated_steps_cond,
-        accumulated_steps_uncond,
-        cached_residual,
-        skip_warmup,
-        mag_ratios,
-    )
+  """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.
+  """
+  import numpy as np
+
+  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
+  cached_residual = None
+
+  skip_warmup = int(num_inference_steps * retention_ratio)
+
+  mag_ratios_base = np.array(mag_ratios_base)
+
+  if len(mag_ratios_base) != num_inference_steps * 2:
+    mag_cond = nearest_interp(mag_ratios_base[0::2], num_inference_steps)
+    mag_uncond = nearest_interp(mag_ratios_base[1::2], num_inference_steps)
+    mag_ratios = np.concatenate([mag_cond.reshape(-1, 1), mag_uncond.reshape(-1, 1)], axis=1).reshape(-1)
+  else:
+    mag_ratios = mag_ratios_base
+
+  return (
+      accumulated_ratio_cond,
+      accumulated_ratio_uncond,
+      accumulated_err_cond,
+      accumulated_err_uncond,
+      accumulated_steps_cond,
+      accumulated_steps_uncond,
+      cached_residual,
+      skip_warmup,
+      mag_ratios,
+  )
+
 
 def magcache_step(
     step,
@@ -959,71 +964,71 @@ def magcache_step(
     skip_warmup=0,
     use_magcache=None,
 ):
-    """Update MagCache accumulated state and decide if to skip.
-    
-    Args:
-        step: Current inference step.
-        mag_ratios: Interpolated magnitude ratios array.
-        accumulated_state: Tuple containing accumulated variables.
-        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
-    
-    (
-        accumulated_ratio_cond,
-        accumulated_ratio_uncond,
-        accumulated_err_cond,
-        accumulated_err_uncond,
-        accumulated_steps_cond,
-        accumulated_steps_uncond,
-    ) = accumulated_state
-
-    cur_mag_ratio_cond = mag_ratios[step * 2]
-    cur_mag_ratio_uncond = mag_ratios[step * 2 + 1]
-
-    if use_magcache is None:
-        use_magcache = True
-        if step < skip_warmup:
-            use_magcache = False
-
-    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 = 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,
-        accumulated_ratio_uncond,
-        accumulated_err_cond,
-        accumulated_err_uncond,
-        accumulated_steps_cond,
-        accumulated_steps_uncond,
-    )
-    return skip_blocks, new_state
+  """Update MagCache accumulated state and decide if to skip.
+
+  Args:
+      step: Current inference step.
+      mag_ratios: Interpolated magnitude ratios array.
+      accumulated_state: Tuple containing accumulated variables.
+      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
+
+  (
+      accumulated_ratio_cond,
+      accumulated_ratio_uncond,
+      accumulated_err_cond,
+      accumulated_err_uncond,
+      accumulated_steps_cond,
+      accumulated_steps_uncond,
+  ) = accumulated_state
+
+  cur_mag_ratio_cond = mag_ratios[step * 2]
+  cur_mag_ratio_uncond = mag_ratios[step * 2 + 1]
+
+  if use_magcache is None:
+    use_magcache = True
+    if step < skip_warmup:
+      use_magcache = False
+
+  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 = 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,
+      accumulated_ratio_uncond,
+      accumulated_err_cond,
+      accumulated_err_uncond,
+      accumulated_steps_cond,
+      accumulated_steps_uncond,
+  )
+  return skip_blocks, new_state

src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py

@@ -12,16 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .wan_pipeline import WanPipeline, transformer_forward_pass, transformer_forward_pass_full_cfg, transformer_forward_pass_cfg_cache, nearest_interp, init_magcache, magcache_step
+from .wan_pipeline import WanPipeline, transformer_forward_pass, transformer_forward_pass_full_cfg, transformer_forward_pass_cfg_cache, init_magcache, magcache_step
 from ...models.wan.transformers.transformer_wan import WanModel
-from typing import List, Union, Optional, Any
+from typing import List, Union, Optional
 from ...pyconfig import HyperParameters
 from functools import partial
 from flax import nnx
 from flax.linen import partitioning as nn_partitioning
 import jax
 import jax.numpy as jnp
-import numpy as np
 from ...schedulers.scheduling_unipc_multistep_flax import FlaxUniPCMultistepScheduler
 
 

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -14,13 +14,12 @@
 
 from maxdiffusion import max_logging
 from maxdiffusion.image_processor import PipelineImageInput
-from .wan_pipeline import WanPipeline, transformer_forward_pass, nearest_interp, init_magcache, magcache_step
+from .wan_pipeline import WanPipeline, transformer_forward_pass, init_magcache, magcache_step
 from ...models.wan.transformers.transformer_wan import WanModel
 from typing import List, Union, Optional, Tuple
 from ...pyconfig import HyperParameters
 from functools import partial
 from flax import nnx
-import numpy as np
 from flax.linen import partitioning as nn_partitioning
 import jax
 import jax.numpy as jnp
@@ -315,7 +314,7 @@ def run_inference_2_1_i2v(
 
   for step in range(num_inference_steps):
     t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
-    
+
     skip_blocks = False
     if use_magcache and do_cfg:
       accumulated_state = (
@@ -345,7 +344,7 @@ def run_inference_2_1_i2v(
     latent_model_input = jnp.concatenate([latents_input, condition_combined], axis=-1)
     timestep = jnp.broadcast_to(t, latents_input.shape[0])
     latent_model_input = jnp.transpose(latent_model_input, (0, 4, 1, 2, 3))
-    
+
     outputs = transformer_forward_pass(
         graphdef,
         sharded_state,

src/maxdiffusion/tests/wan_magcache_test.py

@@ -152,7 +152,6 @@ class Wan21I2VMagCacheSmokeTest(unittest.TestCase):
 
   @classmethod
   def setUpClass(cls):
-
     pyconfig.initialize(
         [
             None,
@@ -224,4 +223,3 @@ def test_magcache_speedup_and_fidelity(self):
     self.assertGreaterEqual(ssim, 0.98)
     self.assertGreater(speedup, 1.0)
     self.assertGreaterEqual(psnr, 30.0)
-