Fix for multiple videos

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -176,10 +176,19 @@ def __call__(
         prompt_embeds, negative_prompt_embeds, image_embeds, last_image
     )
 
-    image_tensor = self.video_processor.preprocess(image, height=height, width=width)
-    last_image_tensor = None
-    if last_image:
-        last_image_tensor = self.video_processor.preprocess(last_image, height=height, width=width)
+    def _process_image_input(img_input, height, width, num_videos_per_prompt):
+        if img_input is None:
+            return None
+        tensor = self.video_processor.preprocess(img_input, height=height, width=width)
+        jax_array = jnp.array(tensor.cpu().numpy())
+        if jax_array.ndim == 3:
+            jax_array = jax_array[None, ...] # Add batch dimension
+        if num_videos_per_prompt > 1:
+            jax_array = jnp.repeat(jax_array, num_videos_per_prompt, axis=0)
+        return jax_array
+
+    image_tensor = _process_image_input(image, height, width, effective_batch_size)
+    last_image_tensor = _process_image_input(last_image, height, width, effective_batch_size)
 
     if rng is None:
         rng = jax.random.key(self.config.seed)