fix(wan_i2v): repeat conditioning images along batch dimension in prepare_latents

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -201,13 +201,19 @@ def __call__(self, x: jax.Array) -> jax.Array:
     n, h, w, c = in_shape
     target_h = int(h * self.scale_factor[0])
     target_w = int(w * self.scale_factor[1])
-    if self.method == "nearest" and self.scale_factor[0] == int(self.scale_factor[0]) and self.scale_factor[1] == int(self.scale_factor[1]):
+    if (
+        self.method == "nearest"
+        and self.scale_factor[0] == int(self.scale_factor[0])
+        and self.scale_factor[1] == int(self.scale_factor[1])
+    ):
       scale_h = int(self.scale_factor[0])
       scale_w = int(self.scale_factor[1])
       out = jnp.repeat(jnp.repeat(x, scale_h, axis=1), scale_w, axis=2)
     else:
       if self.method == "nearest":
-        max_logging.log(f"Warning: WanUpsample2D nearest method requested but scale_factor {self.scale_factor} is not integer. Falling back to jax.image.resize.")
+        max_logging.log(
+            f"Warning: WanUpsample2D nearest method requested but scale_factor {self.scale_factor} is not integer. Falling back to jax.image.resize."
+        )
       out = jax.image.resize(x.astype(jnp.float32), (n, target_h, target_w, c), method=self.method)
       out = out.astype(input_dtype)
     return out
@@ -1234,7 +1240,10 @@ def scan_fn(carry, chunk):
       if spatial_sharding is not None:
         out_chunk = jax.lax.with_sharding_constraint(out_chunk, spatial_sharding)
       next_feat_map = jax.tree_util.tree_map(
-          lambda x: jax.lax.with_sharding_constraint(x, spatial_sharding) if spatial_sharding is not None and hasattr(x, "shape") and x.ndim == len(spatial_sharding.spec) else x, next_feat_map
+          lambda x: jax.lax.with_sharding_constraint(x, spatial_sharding)
+          if spatial_sharding is not None and hasattr(x, "shape") and x.ndim == len(spatial_sharding.spec)
+          else x,
+          next_feat_map,
       )
       return next_feat_map, out_chunk
 
@@ -1333,7 +1342,9 @@ def scan_fn(carry, chunk_in):
           if spatial_sharding is not None:
             out_chunk = jax.lax.with_sharding_constraint(out_chunk, spatial_sharding)
           next_feat_map = jax.tree_util.tree_map(
-              lambda x: jax.lax.with_sharding_constraint(x, spatial_sharding) if spatial_sharding is not None and hasattr(x, "shape") and x.ndim == len(spatial_sharding.spec) else x,
+              lambda x: jax.lax.with_sharding_constraint(x, spatial_sharding)
+              if spatial_sharding is not None and hasattr(x, "shape") and x.ndim == len(spatial_sharding.spec)
+              else x,
               next_feat_map,
           )
           return next_feat_map, out_chunk

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -106,10 +106,15 @@ def prepare_latents(
         last_image = last_image.detach().cpu().numpy()
       last_image = jnp.array(last_image)
 
-    if num_videos_per_prompt > 1:
-      image = jnp.repeat(image, num_videos_per_prompt, axis=0)
-      if last_image is not None:
-        last_image = jnp.repeat(last_image, num_videos_per_prompt, axis=0)
+    if batch_size % image.shape[0] != 0:
+      raise ValueError(f"Batch size ({batch_size}) must be divisible by image batch size ({image.shape[0]}).")
+    if image.shape[0] < batch_size:
+      image = jnp.repeat(image, batch_size // image.shape[0], axis=0)
+    if last_image is not None:
+      if batch_size % last_image.shape[0] != 0:
+        raise ValueError(f"Batch size ({batch_size}) must be divisible by last_image batch size ({last_image.shape[0]}).")
+      if last_image.shape[0] < batch_size:
+        last_image = jnp.repeat(last_image, batch_size // last_image.shape[0], axis=0)
 
     num_channels_latents = self.vae.z_dim
     num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py

@@ -158,6 +158,16 @@ def prepare_latents(
         last_image = last_image.detach().cpu().numpy()
       last_image = jnp.array(last_image)
 
+    if batch_size % image.shape[0] != 0:
+      raise ValueError(f"Batch size ({batch_size}) must be divisible by image batch size ({image.shape[0]}).")
+    if image.shape[0] < batch_size:
+      image = jnp.repeat(image, batch_size // image.shape[0], axis=0)
+    if last_image is not None:
+      if batch_size % last_image.shape[0] != 0:
+        raise ValueError(f"Batch size ({batch_size}) must be divisible by last_image batch size ({last_image.shape[0]}).")
+      if last_image.shape[0] < batch_size:
+        last_image = jnp.repeat(last_image, batch_size // last_image.shape[0], axis=0)
+
     num_channels_latents = self.vae.z_dim
     num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
     latent_height = height // self.vae_scale_factor_spatial