diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py index a3c1ff514..52ad66c46 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py @@ -97,6 +97,11 @@ def prepare_latents( num_videos_per_prompt: int = 1, trace: Optional[dict] = None, ) -> Tuple[jax.Array, jax.Array, Optional[jax.Array]]: + """Prepares latents and condition for I2V inference. + + Note: num_videos_per_prompt is kept for backwards compatibility; repetition + is driven by batch_size. + """ if hasattr(image, "detach"): image = image.detach().cpu().numpy() image = jnp.array(image) @@ -106,10 +111,22 @@ 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 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 image.shape[0] > 1 and last_image.shape[0] > 1 and image.shape[0] != last_image.shape[0]: + raise ValueError( + f"image batch size ({image.shape[0]}) and last_image batch size ({last_image.shape[0]}) must match when" + " both are greater than 1." + ) + if last_image.shape[0] < batch_size: + last_image = jnp.repeat(last_image, batch_size // last_image.shape[0], axis=0) + + if image.shape[0] < batch_size: + image = jnp.repeat(image, batch_size // image.shape[0], axis=0) num_channels_latents = self.vae.z_dim num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py index 7e7997030..91e77f1b8 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py @@ -149,6 +149,11 @@ def prepare_latents( num_videos_per_prompt: int = 1, trace: Optional[dict] = None, ) -> Tuple[jax.Array, jax.Array, Optional[jax.Array]]: + """Prepares latents and condition for I2V inference. + + Note: num_videos_per_prompt is kept for backwards compatibility; repetition + is driven by batch_size. + """ if hasattr(image, "detach"): image = image.detach().cpu().numpy() image = jnp.array(image) @@ -158,6 +163,23 @@ 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 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 image.shape[0] > 1 and last_image.shape[0] > 1 and image.shape[0] != last_image.shape[0]: + raise ValueError( + f"image batch size ({image.shape[0]}) and last_image batch size ({last_image.shape[0]}) must match when" + " both are greater than 1." + ) + if last_image.shape[0] < batch_size: + last_image = jnp.repeat(last_image, batch_size // last_image.shape[0], axis=0) + + if image.shape[0] < batch_size: + image = jnp.repeat(image, batch_size // 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 diff --git a/src/maxdiffusion/tests/wan/wan_i2v_prepare_latents_test.py b/src/maxdiffusion/tests/wan/wan_i2v_prepare_latents_test.py new file mode 100644 index 000000000..b5161d2b2 --- /dev/null +++ b/src/maxdiffusion/tests/wan/wan_i2v_prepare_latents_test.py @@ -0,0 +1,178 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import unittest +from unittest.mock import MagicMock +import jax +import jax.numpy as jnp + +from maxdiffusion.pipelines.wan.wan_pipeline_i2v_2p1 import WanPipelineI2V_2_1 +from maxdiffusion.pipelines.wan.wan_pipeline_i2v_2p2 import WanPipelineI2V_2_2 + + +class WanI2VPrepareLatentsTest(unittest.TestCase): + + def _create_mock_pipeline(self, pipeline_cls): + """Creates a mock pipeline instance with required VAE attributes.""" + pipeline = object.__new__(pipeline_cls) + pipeline.vae = MagicMock(z_dim=16) + pipeline.vae_scale_factor_temporal = 4 + pipeline.vae_scale_factor_spatial = 8 + + def mock_prepare_latents_i2v_base(image, num_frames, dtype, last_image=None, trace=None): + num_latent_frames = (num_frames - 1) // pipeline.vae_scale_factor_temporal + 1 + latent_height = 32 // pipeline.vae_scale_factor_spatial + latent_width = 32 // pipeline.vae_scale_factor_spatial + latent_condition = jnp.zeros( + (image.shape[0], num_latent_frames, latent_height, latent_width, pipeline.vae.z_dim), + dtype=dtype, + ) + return latent_condition, None + + pipeline.prepare_latents_i2v_base = MagicMock(side_effect=mock_prepare_latents_i2v_base) + return pipeline + + def test_single_image_repetition(self): + """Verifies that a single conditioning image is repeated when batch_size > 1.""" + rng = jax.random.key(0) + for pipeline_cls in (WanPipelineI2V_2_1, WanPipelineI2V_2_2): + with self.subTest(pipeline=pipeline_cls.__name__): + pipeline = self._create_mock_pipeline(pipeline_cls) + image = jnp.zeros((1, 3, 32, 32)) + batch_size = 4 + latents, condition, _ = pipeline.prepare_latents( + image=image, + batch_size=batch_size, + height=32, + width=32, + num_frames=5, + dtype=jnp.float32, + rng=rng, + ) + self.assertEqual(latents.shape[0], batch_size) + self.assertEqual(condition.shape[0], batch_size) + call_image = pipeline.prepare_latents_i2v_base.call_args[0][0] + self.assertEqual(call_image.shape[0], batch_size) + + def test_batched_image_repetition(self): + """Verifies that multiple conditioning images are repeated correctly when divisible.""" + rng = jax.random.key(0) + for pipeline_cls in (WanPipelineI2V_2_1, WanPipelineI2V_2_2): + with self.subTest(pipeline=pipeline_cls.__name__): + pipeline = self._create_mock_pipeline(pipeline_cls) + image = jnp.zeros((2, 3, 32, 32)) + batch_size = 4 + latents, condition, _ = pipeline.prepare_latents( + image=image, + batch_size=batch_size, + height=32, + width=32, + num_frames=5, + dtype=jnp.float32, + rng=rng, + ) + self.assertEqual(latents.shape[0], batch_size) + self.assertEqual(condition.shape[0], batch_size) + call_image = pipeline.prepare_latents_i2v_base.call_args[0][0] + self.assertEqual(call_image.shape[0], batch_size) + + def test_with_last_image(self): + """Verifies that both start and last images are repeated when provided.""" + rng = jax.random.key(0) + for pipeline_cls in (WanPipelineI2V_2_1, WanPipelineI2V_2_2): + with self.subTest(pipeline=pipeline_cls.__name__): + pipeline = self._create_mock_pipeline(pipeline_cls) + image = jnp.zeros((1, 3, 32, 32)) + last_image = jnp.zeros((1, 3, 32, 32)) + batch_size = 3 + latents, condition, _ = pipeline.prepare_latents( + image=image, + batch_size=batch_size, + height=32, + width=32, + num_frames=5, + dtype=jnp.float32, + rng=rng, + last_image=last_image, + ) + self.assertEqual(latents.shape[0], batch_size) + self.assertEqual(condition.shape[0], batch_size) + call_image = pipeline.prepare_latents_i2v_base.call_args[0][0] + call_last_image = pipeline.prepare_latents_i2v_base.call_args[0][3] + self.assertEqual(call_image.shape[0], batch_size) + self.assertEqual(call_last_image.shape[0], batch_size) + + def test_indivisible_image_batch_size_raises(self): + """Verifies ValueError when batch_size is not divisible by image batch size.""" + rng = jax.random.key(0) + for pipeline_cls in (WanPipelineI2V_2_1, WanPipelineI2V_2_2): + with self.subTest(pipeline=pipeline_cls.__name__): + pipeline = self._create_mock_pipeline(pipeline_cls) + image = jnp.zeros((2, 3, 32, 32)) + with self.assertRaisesRegex(ValueError, "divisible by image batch size"): + pipeline.prepare_latents( + image=image, + batch_size=3, + height=32, + width=32, + num_frames=5, + dtype=jnp.float32, + rng=rng, + ) + + def test_indivisible_last_image_batch_size_raises(self): + """Verifies ValueError when batch_size is not divisible by last_image batch size.""" + rng = jax.random.key(0) + for pipeline_cls in (WanPipelineI2V_2_1, WanPipelineI2V_2_2): + with self.subTest(pipeline=pipeline_cls.__name__): + pipeline = self._create_mock_pipeline(pipeline_cls) + image = jnp.zeros((1, 3, 32, 32)) + last_image = jnp.zeros((2, 3, 32, 32)) + with self.assertRaisesRegex(ValueError, "divisible by last_image batch size"): + pipeline.prepare_latents( + image=image, + batch_size=3, + height=32, + width=32, + num_frames=5, + dtype=jnp.float32, + rng=rng, + last_image=last_image, + ) + + def test_mismatched_image_and_last_image_batch_sizes_raises(self): + """Verifies ValueError when image and last_image have conflicting batch sizes > 1.""" + rng = jax.random.key(0) + for pipeline_cls in (WanPipelineI2V_2_1, WanPipelineI2V_2_2): + with self.subTest(pipeline=pipeline_cls.__name__): + pipeline = self._create_mock_pipeline(pipeline_cls) + image = jnp.zeros((2, 3, 32, 32)) + last_image = jnp.zeros((3, 3, 32, 32)) + with self.assertRaisesRegex(ValueError, "must match when both are greater than 1"): + pipeline.prepare_latents( + image=image, + batch_size=6, + height=32, + width=32, + num_frames=5, + dtype=jnp.float32, + rng=rng, + last_image=last_image, + ) + + +if __name__ == "__main__": + unittest.main()