fix: image preprocessing for cosmos3 (#14519)

fix: preserve Cosmos3 conditioning image aspect ratio

Author: atharvajoshi10

docs/source/en/api/pipelines/cosmos3.md

@@ -205,7 +205,7 @@ result.video[0].save("cosmos3_t2i.jpg", format="JPEG", quality=85)
 
 ## Image-to-video
 
-Pass a conditioning image via `image=`. The pipeline anchors frame 0 to the supplied image and denoises the rest. Upsample with `--mode image2video` to produce the JSON prompt.
+Pass a conditioning image via `image=`. The pipeline anchors frame 0 to the supplied image and denoises the rest. The image is resized while preserving its aspect ratio, center-cropped to the requested output size, and normalized with uint8-equivalent rounding. Upsample with `--mode image2video` to produce the JSON prompt.
 
 <hfoptions id="model">
 <hfoption id="Nano">

src/diffusers/modular_pipelines/cosmos/encoders.py

@@ -1,4 +1,9 @@
+import math
+
+import numpy as np
 import torch
+import torch.nn.functional as F
+from PIL import Image
 from transformers import AutoTokenizer
 
 from ...configuration_utils import FrozenDict
@@ -17,6 +22,49 @@
 logger = logging.get_logger(__name__)
 
 
+# Copied from diffusers.pipelines.cosmos.pipeline_cosmos3_omni._preprocess_conditioning_image
+def _preprocess_conditioning_image(
+    image: Image.Image | np.ndarray | torch.Tensor, height: int, width: int
+) -> torch.Tensor:
+    """Preprocess one Cosmos3 conditioning image to ``[1, 3, H, W]`` in ``[-1, 1]``."""
+    if isinstance(image, Image.Image):
+        image = torch.from_numpy(np.array(image.convert("RGB"), copy=True)).permute(2, 0, 1).unsqueeze(0)
+    elif isinstance(image, np.ndarray):
+        image = torch.from_numpy(image)
+        image = image.unsqueeze(0) if image.ndim == 3 else image
+        image = image.permute(0, 3, 1, 2)
+    else:
+        image = image.unsqueeze(0) if image.ndim == 3 else image
+
+    if image.ndim != 4 or image.shape[0] != 1 or image.shape[1] != 3:
+        raise ValueError(f"`image` must describe one RGB image, got shape {tuple(image.shape)}.")
+
+    is_integer_input = not image.is_floating_point()
+    image = image.to(dtype=torch.float32)
+    if not is_integer_input:
+        if image.min() < 0:
+            image = (image + 1.0) * 127.5
+        elif image.max() <= 1.0:
+            image = image * 255.0
+
+    source_height, source_width = image.shape[-2:]
+    scale = max(width / source_width, height / source_height)
+    resized_height = math.ceil(scale * source_height)
+    resized_width = math.ceil(scale * source_width)
+    image = F.interpolate(
+        image,
+        size=(resized_height, resized_width),
+        mode="bilinear",
+        align_corners=False,
+        antialias=True,
+    )
+    crop_top = round((resized_height - height) / 2)
+    crop_left = round((resized_width - width) / 2)
+    image = image[:, :, crop_top : crop_top + height, crop_left : crop_left + width]
+    image = image.round().clamp(0, 255) / 127.5 - 1.0
+    return image
+
+
 # Transfer conditions on control signals (edge/blur/depth/seg/wsm), so it uses its own system prompt instead of the
 # plain image/video ones. Defined here (not on the task pipeline) so the transfer text block is self-contained.
 _SYSTEM_PROMPT_TRANSFER = (
@@ -587,12 +635,6 @@ def description(self) -> str:
     def expected_components(self) -> list[ComponentSpec]:
         return [
             ComponentSpec("vae", AutoencoderKLWan),
-            ComponentSpec(
-                "video_processor",
-                VideoProcessor,
-                config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
-                default_creation_method="from_config",
-            ),
         ]
 
     @property
@@ -645,7 +687,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState)
                 f"`height` and `width` must be multiples of {sf}, got ({block_state.height}, {block_state.width})."
             )
 
-        conditioning_frame_2d = components.video_processor.preprocess(
+        conditioning_frame_2d = _preprocess_conditioning_image(
             block_state.image, height=block_state.height, width=block_state.width
         ).to(device=device, dtype=dtype)
 

src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py

@@ -52,6 +52,48 @@ def __init__(self, *args, **kwargs):
             )
 
 
+def _preprocess_conditioning_image(
+    image: Image.Image | np.ndarray | torch.Tensor, height: int, width: int
+) -> torch.Tensor:
+    """Preprocess one Cosmos3 conditioning image to ``[1, 3, H, W]`` in ``[-1, 1]``."""
+    if isinstance(image, Image.Image):
+        image = torch.from_numpy(np.array(image.convert("RGB"), copy=True)).permute(2, 0, 1).unsqueeze(0)
+    elif isinstance(image, np.ndarray):
+        image = torch.from_numpy(image)
+        image = image.unsqueeze(0) if image.ndim == 3 else image
+        image = image.permute(0, 3, 1, 2)
+    else:
+        image = image.unsqueeze(0) if image.ndim == 3 else image
+
+    if image.ndim != 4 or image.shape[0] != 1 or image.shape[1] != 3:
+        raise ValueError(f"`image` must describe one RGB image, got shape {tuple(image.shape)}.")
+
+    is_integer_input = not image.is_floating_point()
+    image = image.to(dtype=torch.float32)
+    if not is_integer_input:
+        if image.min() < 0:
+            image = (image + 1.0) * 127.5
+        elif image.max() <= 1.0:
+            image = image * 255.0
+
+    source_height, source_width = image.shape[-2:]
+    scale = max(width / source_width, height / source_height)
+    resized_height = math.ceil(scale * source_height)
+    resized_width = math.ceil(scale * source_width)
+    image = F.interpolate(
+        image,
+        size=(resized_height, resized_width),
+        mode="bilinear",
+        align_corners=False,
+        antialias=True,
+    )
+    crop_top = round((resized_height - height) / 2)
+    crop_left = round((resized_width - width) / 2)
+    image = image[:, :, crop_top : crop_top + height, crop_left : crop_left + width]
+    image = image.round().clamp(0, 255) / 127.5 - 1.0
+    return image
+
+
 # ============================================================================
 # Sequence layout: data structures + builders for the joint token sequence
 # ============================================================================
@@ -714,7 +756,7 @@ def _remove_action_video_padding_from_latent(
 
     def prepare_latents(
         self,
-        image: torch.Tensor | None = None,
+        image: Image.Image | np.ndarray | torch.Tensor | None = None,
         video: list[Image.Image] | torch.Tensor | np.ndarray | None = None,
         condition_frame_indexes_vision: Iterable[int] = (0, 1),
         condition_video_keep: Literal["first", "last"] = "first",
@@ -754,10 +796,9 @@ def prepare_latents(
         # Video-to-video conditioning: a top-level `video` without an action run.
         has_video_condition = video is not None and action is None
 
-        # video_processor.preprocess handles PIL/np/tensor → [1, 3, H, W] in [-1, 1], resized to (height, width).
         conditioning_frame_2d: torch.Tensor | None = None
         if image is not None:
-            conditioning_frame_2d = self.video_processor.preprocess(image, height=height, width=width).to(
+            conditioning_frame_2d = _preprocess_conditioning_image(image, height=height, width=width).to(
                 device=device, dtype=dtype
             )
 
@@ -1272,7 +1313,7 @@ def __call__(
         self,
         prompt: str | list[str],
         negative_prompt: str | list[str] | None = None,
-        image: torch.Tensor | None = None,
+        image: Image.Image | np.ndarray | torch.Tensor | None = None,
         video: list[Image.Image] | torch.Tensor | np.ndarray | None = None,
         condition_frame_indexes_vision: Iterable[int] = (0, 1),
         condition_video_keep: Literal["first", "last"] = "first",
@@ -1314,9 +1355,10 @@ def __call__(
                 per call.
             negative_prompt (`str` or `List[str]`, *optional*):
                 The negative prompt used for classifier-free guidance. When `None`, the empty string is used.
-            image (`torch.Tensor` or `PIL.Image.Image`, *optional*):
+            image (`PIL.Image.Image`, `np.ndarray`, or `torch.Tensor`, *optional*):
                 Optional conditioning frame for image-to-video. The pipeline anchors frame 0 to this image and denoises
-                the remaining frames. Ignored when `num_frames == 1`. Not used for action runs (pass `action` instead).
+                the remaining frames. The image is resized while preserving its aspect ratio, then center-cropped to
+                `height` and `width`. Ignored when `num_frames == 1`. Not used for action runs (pass `action` instead).
                 Mutually exclusive with `video`.
             video (`List[PIL.Image.Image]`, `torch.Tensor`, or `np.ndarray`, *optional*):
                 Optional conditioning clip for video-to-video. The leading frames are kept clean at the latent indexes

tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 import pytest
 import torch
 from PIL import Image
@@ -208,6 +209,35 @@ def test_vae_encoder_is_standalone_and_validates_conditioning_inputs(self):
         with pytest.raises(ValueError, match="image-to-image generation is not supported"):
             pipe(**inputs, output=self.output_name)
 
+    def test_image_encoder_uses_native_aspect_preserving_center_crop(self):
+        pipe = self.get_pipeline()
+        image_encoder = pipe.blocks.sub_blocks["vae_encoder"].sub_blocks["image_conditioning"]
+        image_pipe = image_encoder.init_pipeline(self.pretrained_model_name_or_path)
+        image_pipe.load_components(dtype=torch.float32)
+
+        image = np.zeros((32, 64, 3), dtype=np.uint8)
+        image[:, :16] = [255, 0, 0]
+        image[:, 16:48] = [0, 255, 0]
+        image[:, 48:] = [0, 0, 255]
+        center_crop = Image.fromarray(image[:, 16:48])
+
+        wide_outputs = image_pipe(
+            image=Image.fromarray(image),
+            num_frames=5,
+            height=32,
+            width=32,
+            output=["x0_tokens_vision"],
+        )
+        crop_outputs = image_pipe(
+            image=center_crop,
+            num_frames=5,
+            height=32,
+            width=32,
+            output=["x0_tokens_vision"],
+        )
+
+        torch.testing.assert_close(wide_outputs["x0_tokens_vision"], crop_outputs["x0_tokens_vision"])
+
     @pytest.mark.parametrize("prompt_name", ["prompt", "negative_prompt"])
     def test_rejects_batched_prompts(self, prompt_name):
         pipe = self.get_pipeline()

tests/pipelines/cosmos/test_cosmos3.py

@@ -15,10 +15,13 @@
 import unittest
 from unittest import mock
 
+import numpy as np
 import torch
+from PIL import Image
 from transformers import AutoTokenizer
 
 from diffusers import AutoencoderKLWan, Cosmos3OmniPipeline, Cosmos3OmniTransformer, UniPCMultistepScheduler
+from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import _preprocess_conditioning_image
 
 from ...testing_utils import enable_full_determinism, torch_device
 from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
@@ -130,6 +133,39 @@ def test_cosmos3_tokenize_prompt_uses_checkpoint_system_prompt_default(self):
 
         assert all(call.args[0][0]["role"] == "user" for call in apply_chat_template.call_args_list)
 
+    def test_i2v_image_preprocessing_preserves_aspect_ratio(self):
+        image = np.zeros((2, 4, 3), dtype=np.uint8)
+        image[:, 0] = [255, 0, 0]
+        image[:, 1] = [0, 255, 0]
+        image[:, 2] = [0, 0, 255]
+        image[:, 3] = [255, 255, 255]
+
+        actual = _preprocess_conditioning_image(Image.fromarray(image), height=2, width=2)
+        expected_pixels = torch.tensor(
+            [[[[0, 0], [0, 0]], [[255, 0], [255, 0]], [[0, 255], [0, 255]]]], dtype=torch.float32
+        )
+        expected = expected_pixels / 127.5 - 1.0
+
+        torch.testing.assert_close(actual, expected)
+
+    def test_i2v_pipeline_uses_native_preprocessing(self):
+        pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
+        pipeline.set_progress_bar_config(disable=None)
+
+        image = np.zeros((16, 32, 3), dtype=np.uint8)
+        image[:, :8] = [255, 0, 0]
+        image[:, 8:24] = [0, 255, 0]
+        image[:, 24:] = [0, 0, 255]
+        center_crop = Image.fromarray(image[:, 8:24])
+        inputs = self.get_dummy_inputs(torch_device)
+        inputs.update(image=Image.fromarray(image), num_frames=5, output_type="latent")
+
+        wide_output = pipeline(**inputs).video
+        inputs.update(image=center_crop, generator=torch.Generator(device="cpu").manual_seed(0))
+        crop_output = pipeline(**inputs).video
+
+        torch.testing.assert_close(wide_output, crop_output)
+
     @unittest.skip("Cosmos3 currently supports one prompt per pipeline call.")
     def test_inference_batch_consistent(self):
         pass