Optimize WAN pipelines: Reduce host overhead, hoist transposes, and fix Cache bugs

[Perseus14]

Overview

This PR introduces several performance optimizations and bug fixes to the Wan 2.1 and 2.2 pipelines (both T2V and I2V), significantly reducing host overhead during the denoising loop and ensuring stability under data-parallel configurations.

Key Optimizations

• Hoisted I2V Tensor Operations: Pre-transposed and concatenated the condition tensors outside of the main denoising loop in the I2V pipelines (wan_pipeline_i2v_2p1.py, wan_pipeline_i2v_2p2.py) to eliminate redundant jnp.transpose and jnp.concatenate churn on every inference step.
• Hoisted Timesteps: Hoisted timestep calculations outside of the loop where possible.

Bug Fixes & Cleanups

• CFG Cache + Ulysses Attention Crash Fix: Added dynamic checks in both Wan 2.1 and Wan 2.2 pipelines to disable use_cfg_cache if batch_size % data_shards != 0. Previously, cfg_cache evaluated only the unconditional branch (dropping the batch size to 1), which caused a ValueError in ulysses_attention's shard_map if the user configured data_shards > 1.
• Logging Standardization: Replaced raw print() statements with max_logging.log() across the pipelines for cleaner distributed logging.
• Formatting: Ran ruff and pyink to clean up linter warnings and unused imports.

[github-actions]

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

[entrpn]

are these changes also needed in the wan animate pipelines? @ninatu

[ninatu]

@entrpn

are these changes also needed in the wan animate pipelines?

not needed, since magcache is not implemented in wan animate and in wan-vace at the moment. But we it would be great to have magcache implemented in wan-vace as well!

[Perseus14]

@ninatu @entrpn PTAL