fix and refactor

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py

@@ -96,12 +96,13 @@ def __call__(
       magcache_K: Optional[int] = None,
       retention_ratio: Optional[float] = None,
   ):
+    config = getattr(self, "config", None)
     if magcache_thresh is None:
-      magcache_thresh = getattr(self.config, "magcache_thresh", 0.12)
+      magcache_thresh = getattr(config, "magcache_thresh", 0.12)
     if magcache_K is None:
-      magcache_K = getattr(self.config, "magcache_K", 2)
+      magcache_K = getattr(config, "magcache_K", 2)
     if retention_ratio is None:
-      retention_ratio = getattr(self.config, "retention_ratio", 0.2)
+      retention_ratio = getattr(config, "retention_ratio", 0.2)
 
     if use_cfg_cache and guidance_scale <= 1.0:
       raise ValueError(
@@ -138,7 +139,7 @@ def __call__(
         magcache_K=magcache_K,
         retention_ratio=retention_ratio,
         height=height,
-        mag_ratios_base=self.config.mag_ratios_base if hasattr(self.config, "mag_ratios_base") else None,
+        mag_ratios_base=getattr(config, "mag_ratios_base", None),
     )
 
     with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
@@ -244,43 +245,23 @@ def run_inference_2_1(
   cached_noise_uncond = None
 
   if use_magcache and do_cfg:
-    (
-        accumulated_ratio_cond,
-        accumulated_ratio_uncond,
-        accumulated_err_cond,
-        accumulated_err_uncond,
-        accumulated_steps_cond,
-        accumulated_steps_uncond,
-        cached_residual,
-        skip_warmup,
-        mag_ratios,
-    ) = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base)
-
-    for step in range(num_inference_steps):
-      t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
+    magcache_init = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base)
+    accumulated_state = magcache_init[:6]
+    cached_residual = magcache_init[6]
+    skip_warmup = magcache_init[7]
+    mag_ratios = magcache_init[8]
+
+  for step in range(num_inference_steps):
+    t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
+
+    if use_magcache and do_cfg:
       timestep = jnp.broadcast_to(t, bsz * 2 if do_cfg else bsz)
 
-      accumulated_state = (
-          accumulated_ratio_cond,
-          accumulated_ratio_uncond,
-          accumulated_err_cond,
-          accumulated_err_uncond,
-          accumulated_steps_cond,
-          accumulated_steps_uncond,
-      )
       skip_blocks, accumulated_state = magcache_step(
           step, mag_ratios, accumulated_state, magcache_thresh, magcache_K, skip_warmup
       )
-      (
-          accumulated_ratio_cond,
-          accumulated_ratio_uncond,
-          accumulated_err_cond,
-          accumulated_err_uncond,
-          accumulated_steps_cond,
-          accumulated_steps_uncond,
-      ) = accumulated_state
-
-      outputs = transformer_forward_pass(
+
+      noise_pred, latents, residual_x_cur = transformer_forward_pass(
           graphdef,
           sharded_state,
           rest_of_state,
@@ -294,18 +275,10 @@ def run_inference_2_1(
           return_residual=True,
       )
 
-      noise_pred, latents_returned, residual_x_cur = outputs
-
       if not skip_blocks:
         cached_residual = residual_x_cur
 
-      latents = latents_returned
-      latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
-    return latents
-
-  else:
-    for step in range(num_inference_steps):
-      t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
+    else:
       is_cache_step = step_is_cache[step]
 
       if is_cache_step:
@@ -351,5 +324,6 @@ def run_inference_2_1(
             guidance_scale=guidance_scale,
         )
 
-      latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
-    return latents
+    latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
+
+  return latents

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -154,12 +154,13 @@ def __call__(
       magcache_K: Optional[int] = None,
       retention_ratio: Optional[float] = None,
   ):
+    config = getattr(self, "config", None)
     if magcache_thresh is None:
-      magcache_thresh = getattr(self.config, "magcache_thresh", 0.04)
+      magcache_thresh = getattr(config, "magcache_thresh", 0.04)
     if magcache_K is None:
-      magcache_K = getattr(self.config, "magcache_K", 2)
+      magcache_K = getattr(config, "magcache_K", 2)
     if retention_ratio is None:
-      retention_ratio = getattr(self.config, "retention_ratio", 0.2)
+      retention_ratio = getattr(config, "retention_ratio", 0.2)
 
     height = height or self.config.height
     width = width or self.config.width
@@ -291,17 +292,11 @@ def run_inference_2_1_i2v(
   do_cfg = guidance_scale > 1.0
 
   if use_magcache and do_cfg:
-    (
-        accumulated_ratio_cond,
-        accumulated_ratio_uncond,
-        accumulated_err_cond,
-        accumulated_err_uncond,
-        accumulated_steps_cond,
-        accumulated_steps_uncond,
-        cached_residual,
-        skip_warmup,
-        mag_ratios,
-    ) = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base)
+    magcache_init = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base)
+    accumulated_state = magcache_init[:6]
+    cached_residual = magcache_init[6]
+    skip_warmup = magcache_init[7]
+    mag_ratios = magcache_init[8]
 
   if do_cfg:
     prompt_embeds_combined = jnp.concatenate([prompt_embeds, negative_prompt_embeds], axis=0)
@@ -317,25 +312,9 @@ def run_inference_2_1_i2v(
 
     skip_blocks = False
     if use_magcache and do_cfg:
-      accumulated_state = (
-          accumulated_ratio_cond,
-          accumulated_ratio_uncond,
-          accumulated_err_cond,
-          accumulated_err_uncond,
-          accumulated_steps_cond,
-          accumulated_steps_uncond,
-      )
       skip_blocks, accumulated_state = magcache_step(
           step, mag_ratios, accumulated_state, magcache_thresh, magcache_K, skip_warmup
       )
-      (
-          accumulated_ratio_cond,
-          accumulated_ratio_uncond,
-          accumulated_err_cond,
-          accumulated_err_uncond,
-          accumulated_steps_cond,
-          accumulated_steps_uncond,
-      ) = accumulated_state
 
     latents_input = latents
     if do_cfg: