Add conditional frame control to Cosmos2 Video2World

src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py

@@ -507,6 +507,7 @@ def __call__(
         callback_on_step_end_tensor_inputs: list[str] = ["latents"],
         max_sequence_length: int = 512,
         sigma_conditioning: float = 0.0001,
+        num_latent_conditional_frames: int = 1,
     ):
         r"""
         The call function to the pipeline for generation.
@@ -573,6 +574,9 @@ def __call__(
             sigma_conditioning (`float`, defaults to `0.0001`):
                 The sigma value used for scaling conditioning latents. Ideally, it should not be changed or should be
                 set to a small value close to zero.
+            num_latent_conditional_frames (`int`, defaults to `1`):
+                Number of latent conditional frames to use for Video2World conditioning. Must be 1 or 2.
+                With the default VAE, this conditions on the last 1 or 5 input video frames, respectively.
 
         Examples:
 
@@ -652,7 +656,21 @@ def __call__(
         if image is not None:
             video = self.video_processor.preprocess(image, height, width).unsqueeze(2)
         else:
+            if num_latent_conditional_frames not in [1, 2]:
+                raise ValueError(
+                    f"`num_latent_conditional_frames` must be 1 or 2, but got {num_latent_conditional_frames}"
+                )
+
             video = self.video_processor.preprocess_video(video, height, width)
+
+            frames_to_extract = 4 * (num_latent_conditional_frames - 1) + 1
+            if video.shape[2] < frames_to_extract:
+                raise ValueError(
+                    f"Input video has only {video.shape[2]} frames but Video2World requires at least "
+                    f"{frames_to_extract} frames for conditioning."
+                )
+
+            video = video[:, :, -frames_to_extract:]
         video = video.to(device=device, dtype=vae_dtype)
 
         num_channels_latents = self.transformer.config.in_channels - 1

tests/pipelines/cosmos/test_cosmos2_video2world.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import PIL.Image
+import pytest
 import torch
 from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
 
@@ -48,7 +49,15 @@ class Cosmos2VideoToWorldPipelineTesterConfig(BasePipelineTesterConfig):
     output_shape = (9, 3, 32, 32)
     # Cosmos2 video-to-world is a video pipeline: it exposes `num_videos_per_prompt`, not `num_images_per_prompt`.
     optional_input_params = frozenset(
-        ["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"]
+        [
+            "num_inference_steps",
+            "num_videos_per_prompt",
+            "generator",
+            "latents",
+            "output_type",
+            "return_dict",
+            "num_latent_conditional_frames",
+        ]
     )
 
     def get_dummy_components(self):
@@ -135,6 +144,42 @@ def test_inference(self):
         generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]])
         assert_tensors_close(generated_slice, expected_slice, atol=1e-3)
 
+    @pytest.mark.parametrize(
+        ("num_latent_conditional_frames", "num_input_frames"),
+        [(None, 1), (2, 5)],
+    )
+    def test_video_conditioning_uses_last_frames(self, num_latent_conditional_frames, num_input_frames):
+        pipe = self.get_pipeline()
+
+        video = [PIL.Image.new("RGB", (32, 32), color=(i * 20, 0, 0)) for i in range(9)]
+
+        inputs = self.get_dummy_inputs()
+        inputs.pop("image")
+        inputs["video"] = video
+
+        if num_latent_conditional_frames is not None:
+            inputs["num_latent_conditional_frames"] = num_latent_conditional_frames
+
+        inputs["generator"] = self.get_generator(0)
+        output_full_video = pipe(**inputs).frames
+
+        inputs["video"] = video[-num_input_frames:]
+        inputs["generator"] = self.get_generator(0)
+        output_conditioning_frames = pipe(**inputs).frames
+
+        assert_tensors_close(output_full_video, output_conditioning_frames, atol=1e-4)
+
+    def test_video_conditioning_rejects_invalid_num_latent_frames(self):
+        pipe = self.get_pipeline()
+
+        inputs = self.get_dummy_inputs()
+        inputs.pop("image")
+        inputs["video"] = [PIL.Image.new("RGB", (32, 32)) for _ in range(9)]
+        inputs["num_latent_conditional_frames"] = 3
+
+        with pytest.raises(ValueError, match="must be 1 or 2"):
+            pipe(**inputs)
+
     def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-2):
         super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)