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.

[mbohlool]

🔍 Code Review: PR #469 — Repeat conditioning images along batch dimension in prepare_latents

📌 Summary & Verdict

Recommendation: Approve with minor suggestions (🟡 Non-blocking improvements).

The PR cleanly resolves the TypeError batch-dimension mismatch during latent condition concatenation in both WanPipelineI2V_2_1 and WanPipelineI2V_2_2 when batch_size > 1, and includes defensive divisibility checks.

---

🌟 Strengths & Praise

• Fixes WAN 2.2 I2V as well: Great catch updating both wan_pipeline_i2v_2p1.py and wan_pipeline_i2v_2p2.py. WAN 2.2 I2V previously lacked any image repetition in prepare_latents, so this fix simultaneously prevents crashes for both batched prompts and num_videos_per_prompt > 1.
• Defensive divisibility validation: Validating that batch_size % image.shape[0] == 0 prevents cryptic shape mismatch errors downstream.

---

🔍 Findings & Suggestions

1. 🟡 Test Coverage

• Location: src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py:L109-L117 & src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py:L161-L170
• Observation: This PR fixes a real crash bug, but does not include automated unit tests.
• Suggestion: Consider adding a focused unit test in src/maxdiffusion/tests/wan/ targeting prepare_latents to verify:
  1. Correct repetition when batch_size > 1 and image.shape[0] == 1.
  2. ValueError is raised when batch_size % image.shape[0] != 0.
  3. Proper handling when last_image is present.
     (Testing prepare_latents can be done with synthetic dummy arrays without requiring checkpoint loading or full TPU inference.)

2. 💡 Alignment Check Between image and last_image

• Location: src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py:L113-L117 & src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py:L165-L169
• Observation: Both image and last_image check divisibility against batch_size independently:

  if batch_size % image.shape[0] != 0: ...
  if last_image is not None and batch_size % last_image.shape[0] != 0: ...

  If a caller passes multiple start and end images where both divide batch_size but do not match each other (e.g. batch_size = 6, image.shape[0] = 2, last_image.shape[0] = 3), image will repeat 3× ([img0, img0, img0, img1, img1, img1]) and last_image will repeat 2× ([last0, last0, last1, last1, last2, last2]). The start and end image frames will become misaligned across video generations without raising an error.
• Suggestion: Consider validating that if last_image is not None and both image.shape[0] > 1 and last_image.shape[0] > 1, their batch sizes must match (image.shape[0] == last_image.shape[0]).

3. 🟢 Unused Signature Parameter

• Location: src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py:L97 & src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py:L149
• Observation: In prepare_latents, num_videos_per_prompt: int = 1 is now unused since repetition is fully driven by batch_size (which already equals effective_batch_size = len(prompt) * num_videos_per_prompt).
• Suggestion: If kept for backwards compatibility of internal call sites, consider adding a brief docstring note or cleaning up unused references in a follow-up.

4. 🟢 Incidental Formatting Changes

• Location: src/maxdiffusion/models/wan/autoencoder_kl_wan.py:L204-L216
• Observation: The changes in autoencoder_kl_wan.py are purely whitespace/yapf formatting and contain no functional diffs. They are harmless and clean, but keeping formatting diffs separate from functional bugfix PRs helps maintain cleaner git blame history.