Repeat conditioning images along batch dimension in prepare_latents

[prishajain1]

Overview

Fixes an issue in WAN Image-to-Video pipelines (WanPipelineI2V_2_1 and WanPipelineI2V_2_2) where running with a batch size > 1 (such as passing multiple prompts) crashes with:
TypeError: Cannot concatenate arrays with shapes that differ in dimensions other than the one being concatenated: concatenating along dimension 4 for shapes (2, 21, 90, 160, 4), (1, 21, 90, 160, 16).

Root Cause

• mask_lat_size is constructed using batch_size (e.g., shape (2, 21, 90, 160, 4)).
• latent_condition is encoded from image, which had shape[0] == 1 because image repetition only occurred when num_videos_per_prompt > 1, ignoring cases where batch_size > 1 due to multiple prompts.
• When concatenating along axis=-1, JAX raised a TypeError due to mismatched batch dimensions (2 vs 1).

Changes

• Updated prepare_latents in wan_pipeline_i2v_2p1.py and wan_pipeline_i2v_2p2.pyto repeat image (and last_image) along axis 0 whenever image.shape[0] < batch_size.

[github-actions]

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

[gemini-code-assist]

Code Review

This pull request updates the latent preparation logic in both wan_pipeline_i2v_2p1.py and wan_pipeline_i2v_2p2.py to dynamically repeat the input image and last_image arrays to match the target batch_size. The reviewer correctly identified a potential issue where a non-divisible batch size or an image batch size larger than the target could cause integer division to yield incorrect repeat factors or zero, leading to downstream shape mismatches. Implementing the suggested defensive checks to validate divisibility and raise clear errors will significantly improve the robustness of both pipelines.