Shard video_condition to prevent OOM in WAN 2.2 I2V

[prishajain1]

This change is needed to resolve the OOM issues faced when running the WAN 2.2 I2V pipeline for 720px for higher batch sizes.

Files changed:
wan_pipeline.py : In prepare_latents_i2v_base, added an explicit JAX sharding constraint to the video_condition tensor before it is passed to the VAE encoder.

Changes were made only to prepare_latents_i2v_base which is called only by I2V pipelines, hence tested the I2V pipelines.

Tested:

• WAN 2.2 I2V logs: https://paste.googleplex.com/4715082825203712
  • Video links:
    • https://storage.mtls.cloud.google.com/prishajain-wan/videos/wan_2_2_vid1.mp4
    • https://storage.mtls.cloud.google.com/prishajain-wan/videos/wan_2_2_vid2.mp4
• WAN 2.1 I2V logs: https://paste.googleplex.com/4505529961545728
  • Video links:
    • https://storage.mtls.cloud.google.com/prishajain-wan/videos/wan_2_1_vid1.mp4
    • https://storage.mtls.cloud.google.com/prishajain-wan/videos/wan_2_1_vid2.mp4

[github-actions]

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

[entrpn]

@prishajain1 can you do one run on multihost to ensure it works. Thanks!

[prishajain1]

@prishajain1 can you do one run on multihost to ensure it works. Thanks!

https://console.cloud.google.com/kubernetes/service/southamerica-west1/mperf-1xv6e16-multinic/default/prishajain-ubench-bk5y/details?project=tpu-prod-env-one-vm