feat(wan): Add text encoder batching and optional scan loop for diffusion

[Perseus14]

This PR introduces several key optimizations for the WAN pipelines (T2V 2.1/2.2 and I2V 2.1/2.2) to improve performance, CPU resource utilization, and TPU execution efficiency:

1. T5 Text Encoder CPU Optimizations:
   • Dynamic bfloat16 loading: Dynamically maps JAX config.weights_dtype to torch_dtype (enabling bfloat16 CPU execution) to cut memory bandwidth consumption in half on the CPU host.
   • JIT Compilation (torch.compile): Compiles the T5 model for CPU using PyTorch's compiler to fuse kernels and maximize weight reuse in the CPU cache, achieving ultra-fast CPU inference.
2. Optional Text Encoder Batching: Allows batching positive and negative prompts together when calling the heavy T5 text encoder. This reduces the number of calls from 2 to 1, saving compute on the CPU.
3. Optional Scan Loop for Diffusion: Implements a hybrid scan loop using jax.lax.scan for the non-cache path of the diffusion process in all four main WAN pipelines. This avoids Python loop overhead while remaining compatible with scan_layers: true at the layer level. For WAN 2.2 pipelines, it uses jax.lax.cond to switch between the dual transformers at each step.
4. Timing Instrumentation: Added timing measurement (trace dictionary) to all pipelines to support the TIMING SUMMARY printout in generate_wan.py, providing visibility into Conditioning, Denoise Total, and VAE Decode times.

Changes

maxdiffusion/pipelines/wan

[MODIFY] wan_pipeline.py

• Dynamically map JAX config.weights_dtype to PyTorch torch_dtype using getattr.
• Enabled torch.compile(text_encoder) inside load_text_encoder for CPU optimization.
• Refactored encode_prompt to batch positive and negative prompts when use_batched_text_encoder is enabled in the config.

[MODIFY] wan_pipeline_2_2.py, wan_pipeline_i2v_2p2.py

• Implemented a hybrid scan loop in run_inference methods using jax.lax.scan and jax.lax.cond.

[MODIFY] wan_pipeline_2_1.py, wan_pipeline_i2v_2p1.py

• Implemented a similar scan loop using jax.lax.scan (without needing lax.cond as they use a single transformer).
• Added trace dictionary return from __call__ to support timing summary.

maxdiffusion/configs

[MODIFY] All 5 WAN config files (base_wan_*.yml)

• Added use_batched_text_encoder: False by default.
• Added scan_diffusion_loop: False by default, with a warning that enabling it will disable per-step profiling.

Generation Time

Environment & Configuration:

• Config: 720p, 81 frames
• Model: WAN2.2 - T2V
• Hardware: TPU v7x-8
• JAX Version: 0.10.0

Command: https://paste.googleplex.com/6221970925551616

==================================================
  TIMING SUMMARY
==================================================
  Load (checkpoint):     132.1s
  Compile:               164.1s
  Inference:             132.5s
  ────────────────────────────────────────
  Conditioning:            1.6s
  Denoise Total:         127.4s
  VAE Decode:              3.6s
==================================================

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

[github-actions]

🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details.

[github-actions]

🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This PR introduces two valuable performance optimizations for the WAN pipeline: batched text encoding and a jax.lax.scan-based diffusion loop. These changes improve compute efficiency and reduce Python loop overhead during inference. The implementation is clean and integrates well with the existing architecture.

🔍 General Feedback

• Optimization Consistency: The batched text encoder logic correctly handles the partitioning of embeddings back into positive and negative sets, ensuring compatibility with the existing API.
• Robustness: I've identified one potential unsafe access to the config object in the scan loop path which could lead to a crash if config is None. A simple fix has been suggested.
• Performance: The use of jax.lax.scan for the non-cache path is a great addition for performance-sensitive workloads on TPU/GPU.

[Perseus14]

Done! PTAL @mbohlool

[eltsai]

This is great @Perseus14! Do we know how much speed gain we get from (1) text encoding batching and (2) scan diffusion look respectively?