feat: add pretrained orbax cache for fast WAN inference loads

[csgoogle]

Pretrained Orbax cache for fast WAN inference loads

Caches WAN weights as an Orbax checkpoint so inference can skip the slow diffusers load on repeat runs (~10× faster). The first run loads from diffusers and writes the cache; subsequent runs restore directly from it. Works with both local paths and GCS buckets.

Changes

• Added a new pretrained_orbax_dir config option to all 6 WAN configs (empty by default, so the feature is off unless set). Accepts a local path or a gs:// bucket path.
• Load priority: training checkpoint → pretrained cache → diffusers (the cache is populated automatically on a miss).
• generate_wan and generate_wan_animate now load with use_pretrained_cache=True.
• Refactored the shared Orbax logic into the WanCheckpointer base class; subclasses now only declare model_name, checkpoint_state_item_names, and pretrained_state_sources.
• Added an animate checkpointer along with tests covering the cache hit / miss / save paths.

[github-actions]

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

[Perseus14]

Can you also add this for the other WAN models?

[csgoogle]

Can you also add this for the other WAN models?

done

[github-actions]

🤖 Hi @csgoogle, 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 Pull Request introduces a pretrained Orbax cache for WAN models, providing a significant performance boost (~10×) for subsequent inference loads by skipping the slow diffusers load. The implementation is well-structured, with a clean refactoring of the shared Orbax logic into a base WanCheckpointer class, and includes comprehensive tests for the new caching strategy.

🔍 General Feedback

• Architecture: The refactoring of WanCheckpointer into an abstract base class is a great improvement, making it easy to add support for new WAN model variants with minimal boilerplate.
• Performance: The automatic population of the cache on the first run is a user-friendly feature that drastically improves the startup time for repeated tasks.
• Correctness: The restructuring of multi-transformer checkpoints (e.g., WAN 2.2) is handled correctly to maintain compatibility with existing pipeline loading logic.
• Robustness: I've suggested adding exception handling to the cache-saving logic to ensure that intermittent write failures (e.g., GCS permissions) do not disrupt the main execution flow.

[Perseus14]

Would this change have an impact on any training runs?

[csgoogle]

Would this change have an impact on any training runs?

It won't there are unit tests covering the training load, also the flag for the change only get's enabled if it's true

[csgoogle]

Would this change have an impact on any training runs?

Made it simple, now we are not touching training code.