diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md index b4acfb7f528a..9b6e53e79c0f 100644 --- a/docs/source/en/api/pipelines/cosmos3.md +++ b/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. diff --git a/src/diffusers/modular_pipelines/cosmos/encoders.py b/src/diffusers/modular_pipelines/cosmos/encoders.py index 81f181d4e5a2..d85798118fbc 100644 --- a/src/diffusers/modular_pipelines/cosmos/encoders.py +++ b/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) diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index db1365210c93..589e0ed3d6b0 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/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 diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index 9c4addb69566..b5bf5fb04393 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/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() diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py index b701a35a9ebe..c61b265a2cd9 100644 --- a/tests/pipelines/cosmos/test_cosmos3.py +++ b/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