Optimize WAN pipeline: hoist timesteps, batched text encoding, SenCache to jax.lax.cond, I2V concat/transpose churn reduction, replace prints with max_logging

Author: Perseus14

src/maxdiffusion/loaders/ltx2_lora_nnx_loader.py

@@ -76,8 +76,6 @@ def translate_fn(nnx_path_str):
     # the merge_fn warns about unmatched keys in each dict, so we only warn about any leftovers
     unmatched_keys = set(h_state_dict) - set(transformer_state_dict) - set(connector_state_dict)
     if unmatched_keys:
-      max_logging.log(
-          f"{len(unmatched_keys)} key(s) in LoRA dictionary routed to no merge target: {unmatched_keys}"
-      )
+      max_logging.log(f"{len(unmatched_keys)} key(s) in LoRA dictionary routed to no merge target: {unmatched_keys}")
 
     return pipeline

src/maxdiffusion/models/quantizations.py

@@ -25,6 +25,7 @@
 import jax.numpy as jnp
 from jax.tree_util import tree_flatten_with_path, tree_unflatten
 from typing import Tuple, Sequence
+from maxdiffusion import max_logging
 
 # Params used to define mixed precision quantization configs
 DEFAULT = "__default__"  # default config
@@ -139,7 +140,7 @@ def _get_quant_config(config):
     else:
       drhs_bits = 8
       drhs_accumulator_dtype = jnp.int32
-      print(config.quantization_local_shard_count)  # -1
+      max_logging.log(config.quantization_local_shard_count)  # -1
       drhs_local_aqt = aqt_config.LocalAqt(contraction_axis_shard_count=config.quantization_local_shard_count)
     return aqt_config.config_v4(
         fwd_bits=8,

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -997,7 +997,8 @@ def _prepare_model_inputs_i2v(
 
     prompt_embeds = jax.device_put(prompt_embeds, data_sharding)
     negative_prompt_embeds = jax.device_put(negative_prompt_embeds, data_sharding)
-    image_embeds = jax.device_put(image_embeds, data_sharding)
+    if image_embeds is not None:
+      image_embeds = jax.device_put(image_embeds, data_sharding)
 
     return prompt_embeds, negative_prompt_embeds, image_embeds, effective_batch_size
 
@@ -1332,7 +1333,8 @@ def magcache_step(
       skip_warmup: Warmup steps threshold.
       use_magcache: Optional manual override boolean to enable/disable cache for this step.
   """
-  import numpy as np
+  import jax.numpy as jnp
+  import jax
 
   (
       accumulated_ratio_cond,
@@ -1347,38 +1349,97 @@ def magcache_step(
   cur_mag_ratio_uncond = mag_ratios[step * 2 + 1]
 
   if use_magcache is None:
-    use_magcache = True
-    if step < skip_warmup:
-      use_magcache = False
+    use_magcache = step >= skip_warmup
 
-  skip_blocks = False
-  if use_magcache:
+  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
     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
+    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,
+      ),
+  )
 
   new_state = (
       accumulated_ratio_cond,

src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py

@@ -101,6 +101,7 @@ def __call__(
       magcache_K: Optional[int] = None,
       retention_ratio: Optional[float] = None,
       use_kv_cache: bool = False,
+      output_type: str = "pil",
   ):
     config = getattr(self, "config", None)
     if max_sequence_length is None:
@@ -170,6 +171,9 @@ def __call__(
       latents.block_until_ready()
     trace["denoise_total"] = time.perf_counter() - t_denoise_start
 
+    if output_type == "latent":
+      return latents, trace
+
     t_decode_start = time.perf_counter()
     video = self._decode_latents_to_video(latents, trace=trace)
     if hasattr(video, "block_until_ready"):
@@ -306,10 +310,9 @@ def run_inference_2_1(
   )
 
   scan_diffusion_loop = getattr(config, "scan_diffusion_loop", False) if config else False
+  timesteps = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)
 
   if scan_diffusion_loop and not use_magcache and not use_cfg_cache:
-    timesteps = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)
-
     scheduler_state = scheduler_state.replace(last_sample=jnp.zeros_like(latents), step_index=jnp.array(0, dtype=jnp.int32))
 
     def scan_body(carry, t):
@@ -365,7 +368,7 @@ def scan_body(carry, t):
       profiler = max_utils.Profiler(config)
       profiler.start()
 
-    t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
+    t = timesteps[step]
 
     if use_magcache and do_cfg:
       timestep = jnp.broadcast_to(t, bsz * 2 if do_cfg else bsz)

src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py

@@ -118,6 +118,7 @@ def __call__(
       use_cfg_cache: bool = False,
       use_sen_cache: bool = False,
       use_kv_cache: bool = False,
+      output_type: str = "pil",
   ):
     config = getattr(self, "config", None)
     if max_sequence_length is None:
@@ -203,6 +204,9 @@ def __call__(
       latents.block_until_ready()
     trace["denoise_total"] = time.perf_counter() - t_denoise_start
 
+    if output_type == "latent":
+      return latents, trace
+
     t_decode_start = time.perf_counter()
     video = self._decode_latents_to_video(latents, trace=trace)
     if hasattr(video, "block_until_ready"):
@@ -279,6 +283,8 @@ def run_inference_2_2(
     high_transformer = nnx.merge(high_noise_graphdef, high_noise_state, high_noise_rest)
     kv_cache_high, encoder_attention_mask_high = high_transformer.compute_kv_cache(prompt_embeds_combined)
 
+  timesteps = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)
+
   # ── SenCache path (arXiv:2602.24208) ──
   if use_sen_cache and do_classifier_free_guidance:
     timesteps_np = np.array(scheduler_state.timesteps, dtype=np.int32)
@@ -303,16 +309,18 @@ def run_inference_2_2(
     num_train_timesteps = float(scheduler.config.num_train_timesteps)
 
     # SenCache state
-    ref_noise_pred = None  # y^r: cached denoiser output
-    ref_latent = None  # x^r: latent at last cache refresh
-    ref_timestep = 0.0  # t^r: timestep (normalized to [0,1]) at last cache refresh
-    accum_dx = 0.0  # accumulated ||Δx|| since last refresh
-    accum_dt = 0.0  # accumulated |Δt| since last refresh
-    reuse_count = 0  # consecutive cache reuses
-    cache_count = 0
+    ref_noise_pred = jnp.zeros(
+        (bsz * 2, latents.shape[1], latents.shape[2], latents.shape[3], latents.shape[4]), dtype=latents.dtype
+    )
+    ref_latent = jnp.zeros_like(latents)
+    ref_timestep = jnp.array(0.0, dtype=jnp.float32)
+    accum_dx = jnp.array(0.0, dtype=jnp.float32)
+    accum_dt = jnp.array(0.0, dtype=jnp.float32)
+    reuse_count = jnp.array(0, dtype=jnp.int32)
+    cache_count = jnp.array(0, dtype=jnp.int32)
 
     for step in range(num_inference_steps):
-      t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
+      t = timesteps[step]
       t_float = float(timesteps_np[step]) / num_train_timesteps  # normalize to [0, 1]
 
       # Select transformer and guidance scale
@@ -358,29 +366,30 @@ def run_inference_2_2(
         )
         ref_noise_pred = noise_pred
         ref_latent = latents
-        ref_timestep = t_float
-        accum_dx = 0.0
-        accum_dt = 0.0
-        reuse_count = 0
+        ref_timestep = jnp.array(t_float, dtype=jnp.float32)
+        accum_dx = jnp.array(0.0, dtype=jnp.float32)
+        accum_dt = jnp.array(0.0, dtype=jnp.float32)
+        reuse_count = jnp.array(0, dtype=jnp.int32)
         latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
         continue
 
       # Accumulate deltas since last full compute
-      dx_norm = float(jnp.sqrt(jnp.mean((latents - ref_latent) ** 2)))
-      dt = abs(t_float - ref_timestep)
+      dx_norm = jnp.sqrt(jnp.mean((latents - ref_latent) ** 2))
+      dt = jnp.abs(t_float - ref_timestep)
       accum_dx += dx_norm
       accum_dt += dt
 
       # Sensitivity score (Eq. 9)
       score = alpha_x * accum_dx + alpha_t * accum_dt
 
-      if score <= sen_epsilon and reuse_count < max_reuse:
-        # Cache hit: reuse previous output
-        noise_pred = ref_noise_pred
-        reuse_count += 1
-        cache_count += 1
-      else:
-        # Cache miss: full CFG forward pass
+      cond = jnp.logical_and(score <= sen_epsilon, reuse_count < max_reuse)
+
+      def hit_fn(args):
+        ref_noise_pred, ref_latent, ref_timestep, accum_dx, accum_dt, reuse_count, cache_count = args
+        return ref_noise_pred, ref_latent, ref_timestep, accum_dx, accum_dt, reuse_count + 1, cache_count + 1
+
+      def miss_fn(args):
+        _, _, _, _, _, _, cache_count = args
         latents_doubled = jnp.concatenate([latents] * 2)
         timestep = jnp.broadcast_to(t, bsz * 2)
         noise_pred, _, _ = transformer_forward_pass_full_cfg(
@@ -395,12 +404,19 @@ def run_inference_2_2(
             rotary_emb=rotary_emb,
             encoder_attention_mask=encoder_attention_mask,
         )
-        ref_noise_pred = noise_pred
-        ref_latent = latents
-        ref_timestep = t_float
-        accum_dx = 0.0
-        accum_dt = 0.0
-        reuse_count = 0
+        return (
+            noise_pred,
+            latents,
+            jnp.array(t_float, dtype=jnp.float32),
+            jnp.array(0.0, dtype=jnp.float32),
+            jnp.array(0.0, dtype=jnp.float32),
+            jnp.array(0, dtype=jnp.int32),
+            cache_count,
+        )
+
+      noise_pred, ref_latent, ref_timestep, accum_dx, accum_dt, reuse_count, cache_count = jax.lax.cond(
+          cond, hit_fn, miss_fn, (ref_noise_pred, ref_latent, ref_timestep, accum_dx, accum_dt, reuse_count, cache_count)
+      )
 
       latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
 
@@ -470,7 +486,7 @@ def run_inference_2_2(
     cached_noise_uncond = None
 
     for step in range(num_inference_steps):
-      t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
+      t = timesteps[step]
       is_cache_step = step_is_cache[step]
 
       # Select transformer and guidance scale based on precomputed schedule
@@ -607,8 +623,6 @@ def low_noise_branch(operands):
     )
 
   if scan_diffusion_loop:
-    timesteps = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)
-
     scheduler_state = scheduler_state.replace(last_sample=jnp.zeros_like(latents), step_index=jnp.array(0, dtype=jnp.int32))
 
     def scan_body(carry, t):
@@ -657,7 +671,7 @@ def scan_body(carry, t):
       profiler = max_utils.Profiler(config)
       profiler.start()
 
-    t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
+    t = timesteps[step]
 
     if step_uses_high[step]:
       graphdef, state, rest = (