feat: add optional batched text encoder and diffusion loop

Author: Perseus14

src/maxdiffusion/configs/base_wan_14b.yml

@@ -348,6 +348,13 @@ flow_shift: 3.0
 # Skips the unconditional forward pass on ~35% of steps via residual compensation.
 # See: FasterCache (Lv et al. 2024), WAN 2.1 paper §4.4.2
 use_cfg_cache: False
+
+# Batch positive and negative prompts in text encoder to save compute.
+use_batched_text_encoder: False
+
+# Use jax.lax.scan for the diffusion loop (non-cache path only).
+# Note: Enabling this will disable per-step profiling.
+scan_diffusion_loop: False
 use_magcache: False
 magcache_thresh: 0.12
 magcache_K: 2

src/maxdiffusion/configs/base_wan_1_3b.yml

@@ -302,6 +302,13 @@ flow_shift: 3.0
 # Diffusion CFG cache (FasterCache-style, WAN 2.1 T2V only)
 use_cfg_cache: False
 
+# Batch positive and negative prompts in text encoder to save compute.
+use_batched_text_encoder: False
+
+# Use jax.lax.scan for the diffusion loop (non-cache path only).
+# Note: Enabling this will disable per-step profiling.
+scan_diffusion_loop: False
+
 # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
 guidance_rescale: 0.0
 num_inference_steps: 30

src/maxdiffusion/configs/base_wan_27b.yml

@@ -323,6 +323,13 @@ boundary_ratio: 0.875
 
 # Diffusion CFG cache (FasterCache-style)
 use_cfg_cache: False
+
+# Batch positive and negative prompts in text encoder to save compute.
+use_batched_text_encoder: False
+
+# Use jax.lax.scan for the diffusion loop (non-cache path only).
+# Note: Enabling this will disable per-step profiling.
+scan_diffusion_loop: False
 # SenCache: Sensitivity-Aware Caching (arXiv:2602.24208) — skip forward pass
 # when predicted output change (based on accumulated latent/timestep drift) is small
 use_sen_cache: False

src/maxdiffusion/configs/base_wan_i2v_14b.yml

@@ -306,6 +306,13 @@ flow_shift: 5.0
 
 # Diffusion CFG cache (FasterCache-style)
 use_cfg_cache: False
+
+# Batch positive and negative prompts in text encoder to save compute.
+use_batched_text_encoder: False
+
+# Use jax.lax.scan for the diffusion loop (non-cache path only).
+# Note: Enabling this will disable per-step profiling.
+scan_diffusion_loop: False
 # SenCache: Sensitivity-Aware Caching (arXiv:2602.24208)
 use_sen_cache: False
 use_magcache: False

src/maxdiffusion/configs/base_wan_i2v_27b.yml

@@ -318,6 +318,13 @@ boundary_ratio: 0.875
 
 # Diffusion CFG cache (FasterCache-style)
 use_cfg_cache: False
+
+# Batch positive and negative prompts in text encoder to save compute.
+use_batched_text_encoder: False
+
+# Use jax.lax.scan for the diffusion loop (non-cache path only).
+# Note: Enabling this will disable per-step profiling.
+scan_diffusion_loop: False
 # SenCache: Sensitivity-Aware Caching (arXiv:2602.24208)
 use_sen_cache: False
 

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -501,24 +501,45 @@ def encode_prompt(
       negative_prompt_embeds: jax.Array = None,
   ):
     prompt = [prompt] if isinstance(prompt, str) else prompt
-    if prompt_embeds is None:
-      prompt_embeds = self._get_t5_prompt_embeds(
-          prompt=prompt,
-          num_videos_per_prompt=num_videos_per_prompt,
-          max_sequence_length=max_sequence_length,
-      )
-      prompt_embeds = jnp.array(prompt_embeds.detach().float().numpy(), dtype=jnp.float32)
-
-    if negative_prompt_embeds is None:
-      batch_size = len(prompt_embeds)
-      negative_prompt = negative_prompt or ""
-      negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
-      negative_prompt_embeds = self._get_t5_prompt_embeds(
-          prompt=negative_prompt,
+    batch_size = len(prompt)
+
+    if negative_prompt is None:
+      negative_prompt = [""] * batch_size
+    elif isinstance(negative_prompt, str):
+      negative_prompt = [negative_prompt] * batch_size
+
+    use_batched_text_encoder = getattr(self.config, "use_batched_text_encoder", False)
+    if use_batched_text_encoder and prompt_embeds is None and negative_prompt_embeds is None:
+      # Batch both together
+      combined_prompts = prompt + negative_prompt
+      combined_embeds = self._get_t5_prompt_embeds(
+          prompt=combined_prompts,
           num_videos_per_prompt=num_videos_per_prompt,
           max_sequence_length=max_sequence_length,
       )
-      negative_prompt_embeds = jnp.array(negative_prompt_embeds.detach().float().numpy(), dtype=jnp.float32)
+      combined_embeds = jnp.array(combined_embeds.detach().float().numpy(), dtype=jnp.float32)
+
+      # Split back
+      prompt_embeds = combined_embeds[: batch_size * num_videos_per_prompt]
+      negative_prompt_embeds = combined_embeds[batch_size * num_videos_per_prompt :]
+
+    else:
+      # Fallback to separate encoding if one of them is already provided
+      if prompt_embeds is None:
+        prompt_embeds = self._get_t5_prompt_embeds(
+            prompt=prompt,
+            num_videos_per_prompt=num_videos_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+        prompt_embeds = jnp.array(prompt_embeds.detach().float().numpy(), dtype=jnp.float32)
+
+      if negative_prompt_embeds is None:
+        negative_prompt_embeds = self._get_t5_prompt_embeds(
+            prompt=negative_prompt,
+            num_videos_per_prompt=num_videos_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+        negative_prompt_embeds = jnp.array(negative_prompt_embeds.detach().float().numpy(), dtype=jnp.float32)
 
     return prompt_embeds, negative_prompt_embeds
 

src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py

@@ -471,6 +471,61 @@ def run_inference_2_2(
   profiler_steps = config.profiler_steps if config else 0
   last_profiling_step = np.clip(first_profiling_step + profiler_steps - 1, first_profiling_step, num_inference_steps - 1)
 
+  scan_diffusion_loop = getattr(config, "scan_diffusion_loop", False)
+
+  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):
+      current_latents, current_scheduler_state = carry
+
+      if do_classifier_free_guidance:
+        model_latents = jnp.concatenate([current_latents] * 2)
+      else:
+        model_latents = current_latents
+
+      timestep = jnp.broadcast_to(t, model_latents.shape[0])
+      use_high_noise = jnp.greater_equal(t, boundary)
+
+      def high_branch(_):
+        return transformer_forward_pass(
+            high_noise_graphdef,
+            high_noise_state,
+            high_noise_rest,
+            model_latents,
+            timestep,
+            prompt_embeds_combined,
+            do_classifier_free_guidance,
+            guidance_scale_high,
+        )
+
+      def low_branch(_):
+        return transformer_forward_pass(
+            low_noise_graphdef,
+            low_noise_state,
+            low_noise_rest,
+            model_latents,
+            timestep,
+            prompt_embeds_combined,
+            do_classifier_free_guidance,
+            guidance_scale_low,
+        )
+
+      noise_pred, latents_out = jax.lax.cond(use_high_noise, high_branch, low_branch, operand=None)
+
+      new_latents, new_scheduler_state = scheduler.step(current_scheduler_state, noise_pred, t, latents_out).to_tuple()
+
+      return (new_latents, new_scheduler_state), None
+
+    initial_carry = (latents, scheduler_state)
+
+    final_carry, _ = jax.lax.scan(scan_body, initial_carry, timesteps)
+
+    final_latents, _ = final_carry
+    return final_latents
+
   profiler = None
   for step in range(num_inference_steps):
     if config and max_utils.profiler_enabled(config) and step == first_profiling_step: