From 13c2de684c3fca3882e508b7ed5a1add1e405504 Mon Sep 17 00:00:00 2001 From: MaciejBalaNV Date: Thu, 20 Aug 2026 21:04:49 +0200 Subject: [PATCH] Cosmos3 logging and prompt improvements (#6325) Signed-off-by: Maciej Bala Co-authored-by: Roger Wang --- .../models/cosmos3/test_cosmos3_pipeline.py | 141 ++++++++++++++- vllm_omni/diffusion/models/cosmos3/action.py | 34 +--- .../models/cosmos3/pipeline_cosmos3.py | 167 +++++++++++++++++- vllm_omni/diffusion/models/cosmos3/utils.py | 31 ++++ 4 files changed, 335 insertions(+), 38 deletions(-) diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index a0eab3b5e..f2527bd99 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -1095,6 +1095,7 @@ def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest ) additional = preprocess(request).prompt["additional_information"] assert (request.sampling_params.height, request.sampling_params.width) == (192, 320) + assert request.sampling_params.extra_args["_cosmos3_transfer_requested_size"] == {"width": 32, "height": 16} assert tuple(additional["preprocessed_transfer_video"].shape) == (1, 3, 4, 192, 320) assert additional["transfer_input_fps"] == 12.5 assert "preprocessed_video" not in additional @@ -1520,6 +1521,140 @@ def test_format_and_tokenize_prompts_rewrites_json_object_metadata(make_cosmos3_ ) +def test_format_and_tokenize_prompts_transfer_aspect_ratio_override(make_cosmos3_pipeline) -> None: + import json + + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + + pipeline._format_and_tokenize_prompts( + '{"caption": "A robot", "aspect_ratio": "4,3"}', + "", + num_frames=48, + frame_rate=24, + height=480, + width=832, + max_sequence_length=32, + sp=SimpleNamespace(extra_args={"aspect_ratio": "4:3"}), + is_t2i=False, + aspect_ratio_override="16,9", + ) + + assert json.loads(calls[0]["text"])["aspect_ratio"] == "16,9" + + +@pytest.mark.parametrize("rank_zero", [True, False]) +def test_format_and_tokenize_prompts_corrects_conflicting_aspect_ratio_metadata_on_every_rank( + make_cosmos3_pipeline, + monkeypatch: pytest.MonkeyPatch, + mocker, + rank_zero: bool, +) -> None: + import json + + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 + + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + monkeypatch.setattr(pipeline_cosmos3, "_is_rank_zero", lambda: rank_zero) + warning = mocker.patch.object(pipeline_cosmos3.logger, "warning") + + pipeline._format_and_tokenize_prompts( + '{"caption": "A robot", "aspect_ratio": "4,3"}', + "", + num_frames=48, + frame_rate=24, + height=720, + width=1280, + max_sequence_length=32, + sp=SimpleNamespace(extra_args={}), + is_t2i=False, + ) + + assert json.loads(calls[0]["text"])["aspect_ratio"] == "16,9" + if rank_zero: + rendered_warning = warning.call_args.args[0] % warning.call_args.args[1:] + assert "JSON prompt aspect_ratio='4,3' conflicts with the generated 1280x720 canvas" in rendered_warning + else: + warning.assert_not_called() + + +@pytest.mark.parametrize( + ("width", "height", "aspect_ratio"), + [ + (1104, 816, "4,3"), + (832, 468, "16,9"), + ], +) +def test_format_and_tokenize_prompts_keeps_nearest_canonical_aspect_ratio( + make_cosmos3_pipeline, + mocker, + width: int, + height: int, + aspect_ratio: str, +) -> None: + import json + + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 + + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + warning = mocker.patch.object(pipeline_cosmos3.logger, "warning") + + pipeline._format_and_tokenize_prompts( + json.dumps({"caption": "A robot", "aspect_ratio": aspect_ratio}), + "", + num_frames=48, + frame_rate=24, + height=height, + width=width, + max_sequence_length=32, + sp=SimpleNamespace(extra_args={}), + is_t2i=False, + ) + + assert json.loads(calls[0]["text"])["aspect_ratio"] == aspect_ratio + warning.assert_not_called() + + +def test_transfer_bucket_selection_warns_on_size_and_aspect_ratio_conflicts( + make_cosmos3_pipeline, + mocker, +) -> None: + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 + + pipeline = make_cosmos3_pipeline() + warning = mocker.patch.object(pipeline_cosmos3.logger, "warning") + sp = make_sampling_params( + width=832, + height=480, + extra_args={ + "resolution": "480", + "aspect_ratio": "4:3", + "_cosmos3_transfer_requested_size": {"width": 1280, "height": 720}, + }, + ) + + height, width, aspect_ratio = pipeline._transfer_bucket_size(sp, (480, 832)) + assert (height, width, aspect_ratio) == (480, 832, "16,9") + + pipeline._warn_transfer_bucket_conflicts( + sp, + "A transfer prompt", + source_hw=(480, 832), + height=height, + width=width, + aspect_ratio=aspect_ratio, + ) + + rendered_warnings = [call.args[0] % call.args[1:] for call in warning.call_args_list] + assert any("ignores requested size=1280x720 (WxH)" in message for message in rendered_warnings) + assert any( + "requested aspect_ratio='4:3' conflicts with the control-selected 16,9 bucket" in message + for message in rendered_warnings + ) + + def test_format_and_tokenize_prompts_removes_video_metadata_from_t2i_json(make_cosmos3_pipeline) -> None: import json @@ -1989,6 +2124,7 @@ def test_forward_transfer_uses_transfer_prompt_contract_and_source_fps_except_ws captured["negative_prompt"] = negative_prompt captured["format_frame_rate"] = frame_rate captured["format_kwargs"] = kwargs + captured["format_aspect_ratio"] = kwargs["aspect_ratio_override"] return _ids(2), _mask(), _ids(1), _mask() def fake_encode(video: torch.Tensor) -> torch.Tensor: @@ -2011,7 +2147,7 @@ def test_forward_transfer_uses_transfer_prompt_contract_and_source_fps_except_ws captured["shared_kwargs"] = kwargs["shared_kwargs"] return kwargs["latents"] - pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16) + pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16, "1,1") pipeline._format_and_tokenize_prompts = fake_format pipeline._encode_video_tensor = fake_encode pipeline._prepare_transfer_latents = fake_prepare @@ -2056,6 +2192,7 @@ def test_forward_transfer_uses_transfer_prompt_contract_and_source_fps_except_ws output = pipeline.forward(request) assert captured["format_frame_rate"] == expected_fps + assert captured["format_aspect_ratio"] == "1,1" expected_negative_prompt = "" if negative_prompt is None else negative_prompt assert captured["negative_prompt"] == expected_negative_prompt assert captured["format_kwargs"]["use_system_prompt"] is True @@ -2082,7 +2219,7 @@ def test_forward_transfer_runs_multichunk_overlap_path( pipeline = make_cosmos3_pipeline() captured: dict[str, Any] = {"targets": [], "conditional_frames": []} - pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16) + pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16, "1,1") pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (_ids(2), _mask(), _ids(1), _mask()) pipeline._set_flow_shift = lambda target, **_kwargs: captured.setdefault("flow_shifts", []).append(target) diff --git a/vllm_omni/diffusion/models/cosmos3/action.py b/vllm_omni/diffusion/models/cosmos3/action.py index c5fe41a91..dc755b39c 100644 --- a/vllm_omni/diffusion/models/cosmos3/action.py +++ b/vllm_omni/diffusion/models/cosmos3/action.py @@ -15,6 +15,8 @@ from typing import Any import numpy as np import torch +from .utils import VIDEO_RES_SIZE_INFO + ACTION_MODE_POLICY = "policy" ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics" ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics" @@ -46,38 +48,6 @@ EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = { } -VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { - "256": { - "1,1": (256, 256), - "4,3": (320, 256), - "3,4": (256, 320), - "16,9": (320, 192), - "9,16": (192, 320), - }, - "480": { - "1,1": (640, 640), - "4,3": (736, 544), - "3,4": (544, 736), - "16,9": (832, 480), - "9,16": (480, 832), - }, - "704": { - "1,1": (960, 960), - "4,3": (1088, 832), - "3,4": (832, 1088), - "16,9": (1280, 704), - "9,16": (704, 1280), - }, - "720": { - "1,1": (960, 960), - "4,3": (1104, 832), - "3,4": (832, 1104), - "16,9": (1280, 720), - "9,16": (720, 1280), - }, -} - - def normalize_action_mode(mode: Any) -> str | None: if mode is None: return None diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 2304c0212..7a924dd52 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -126,6 +126,7 @@ from .utils import ( ROBOLAB_DEFAULT_RAW_ACTION_DIM, ROBOLAB_DEFAULT_RESOLUTION, ROBOLAB_MIDTRAIN_RAW_ACTION_DIM, + VIDEO_RES_SIZE_INFO, RoboLabActionPostprocessInputs, RoboLabPolicyInputs, build_abs_pose_from_components, @@ -199,6 +200,7 @@ COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS = "FlowMatchEulerDiscreteScheduler" # are tokenized to their natural length (no padding); this only bounds the # UND pathway / GEN cross-attention cost for pathologically long prompts. COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096 +COSMOS3_TRANSFER_REQUESTED_SIZE_KEY = "_cosmos3_transfer_requested_size" def _ceil_video_num_frames(num_frames: int, temporal_compression_factor: int) -> int: @@ -243,6 +245,50 @@ def _format_json_object_prompt( return json.dumps(prompt_obj) +def _json_object_aspect_ratio(prompt: str) -> Any | None: + """Return an aspect-ratio field from a JSON-object prompt, if present.""" + try: + prompt_obj = json.loads(prompt) + except (json.JSONDecodeError, TypeError, ValueError): + return None + if not isinstance(prompt_obj, dict): + return None + return prompt_obj.get("aspect_ratio") + + +def _normalize_aspect_ratio(value: Any) -> str | None: + """Normalize supported ratio spellings for comparison and prompt metadata.""" + if value is None: + return None + parts = str(value).strip().replace(":", ",").split(",") + if len(parts) != 2: + return None + try: + width, height = (int(part.strip()) for part in parts) + except ValueError: + return None + if width <= 0 or height <= 0: + return None + divisor = math.gcd(width, height) + return f"{width // divisor},{height // divisor}" + + +def _aspect_ratio_for_dimensions(width: int, height: int) -> str: + """Return the nearest supported Cosmos aspect-ratio label.""" + if width <= 0 or height <= 0: + raise ValueError(f"Cosmos3 canvas dimensions must be positive, got {width}x{height}.") + canvas_ratio = width / height + + def ratio_value(label: str) -> float: + ratio_width, ratio_height = (int(part) for part in label.split(",")) + return ratio_width / ratio_height + + return min( + (label for sizes in VIDEO_RES_SIZE_INFO.values() for label in sizes), + key=lambda label: abs(canvas_ratio - ratio_value(label)), + ) + + def resolve_cosmos3_transformer_cls(model_config: Any) -> type[Cosmos3VFMTransformer]: """Select the Cosmos3 transformer implementation from transformer/config.json.""" backbone_type = _tf_config_get(model_config, "backbone_type", None) @@ -500,6 +546,17 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): # Resolve missing H/W. if transfer_requested: + requested_width = request.sampling_params.width + requested_height = request.sampling_params.height + if ( + requested_width is not None + and requested_height is not None + and COSMOS3_TRANSFER_REQUESTED_SIZE_KEY not in extra + ): + extra[COSMOS3_TRANSFER_REQUESTED_SIZE_KEY] = { + "width": int(requested_width), + "height": int(requested_height), + } _set_transfer_size_from_image(request, image) elif request.sampling_params.height is None or request.sampling_params.width is None: if action_mode is not None: @@ -2081,6 +2138,7 @@ class Cosmos3OmniDiffusersPipeline( use_duration_template: bool | None = None, use_resolution_template: bool | None = None, negative_metadata_mode: str | None = None, + aspect_ratio_override: str | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Format prompts with metadata templates and tokenize. @@ -2115,13 +2173,40 @@ class Cosmos3OmniDiffusersPipeline( res_tmpl = COSMOS3_IMAGE_RESOLUTION_TEMPLATE if is_t2i else COSMOS3_RESOLUTION_TEMPLATE else: res_tmpl = None + requested_aspect_ratio = self._get_sp_param(sp, "aspect_ratio", None) + prompt_aspect_ratio = _json_object_aspect_ratio(prompt) + metadata_aspect_ratio = requested_aspect_ratio if requested_aspect_ratio is not None else prompt_aspect_ratio + canvas_aspect_ratio = aspect_ratio_override or _aspect_ratio_for_dimensions(width, height) + normalized_metadata_ratio = _normalize_aspect_ratio(metadata_aspect_ratio) + has_aspect_ratio_conflict = ( + aspect_ratio_override is None + and metadata_aspect_ratio is not None + and normalized_metadata_ratio != canvas_aspect_ratio + ) + if has_aspect_ratio_conflict: + aspect_ratio_source = ( + "requested aspect_ratio" if requested_aspect_ratio is not None else "JSON prompt aspect_ratio" + ) + if _is_rank_zero(): + logger.warning( + "Cosmos3 %s=%r conflicts with the generated %dx%d canvas (WxH); using %s for prompt metadata.", + aspect_ratio_source, + metadata_aspect_ratio, + width, + height, + canvas_aspect_ratio, + ) + metadata_aspect_ratio = canvas_aspect_ratio + elif aspect_ratio_override is not None: + metadata_aspect_ratio = aspect_ratio_override + json_prompt = _format_json_object_prompt( prompt, num_frames=num_frames, frame_rate=frame_rate, height=height, width=width, - aspect_ratio=self._get_sp_param(sp, "aspect_ratio", None), + aspect_ratio=metadata_aspect_ratio, ) if json_prompt is not None: prompt = json_prompt @@ -2893,11 +2978,76 @@ class Cosmos3OmniDiffusersPipeline( self, sp: OmniDiffusionSamplingParams, source_hw: tuple[int, int] | None, - ) -> tuple[int, int]: + ) -> tuple[int, int, str]: resolution = self._get_sp_param(sp, "resolution", self._get_sp_param(sp, "image_size", 720)) source_h, source_w = source_hw or (COSMOS3_T2V_DEFAULT_HEIGHT, COSMOS3_T2V_DEFAULT_WIDTH) target_w, target_h = find_closest_target_size(int(source_h), int(source_w), resolution) - return int(target_h), int(target_w) + aspect_ratio = next( + ratio for ratio, size in VIDEO_RES_SIZE_INFO[str(resolution)].items() if size == (target_w, target_h) + ) + return int(target_h), int(target_w), aspect_ratio + + def _warn_transfer_bucket_conflicts( + self, + sp: OmniDiffusionSamplingParams, + prompt: str, + *, + source_hw: tuple[int, int] | None, + height: int, + width: int, + aspect_ratio: str, + ) -> None: + """Explain transfer's bucket selection when it overrides request metadata.""" + if not _is_rank_zero(): + return + + requested_size_data = self._get_sp_param(sp, COSMOS3_TRANSFER_REQUESTED_SIZE_KEY, None) + if isinstance(requested_size_data, Mapping): + requested_width = requested_size_data.get("width") + requested_height = requested_size_data.get("height") + else: + requested_width = getattr(sp, "width", None) + requested_height = getattr(sp, "height", None) + if requested_width is not None and requested_height is not None: + try: + requested_size = (int(requested_width), int(requested_height)) + except (TypeError, ValueError): + requested_size = None + if requested_size is not None and requested_size != (width, height): + source_description = ( + f"control input {source_hw[1]}x{source_hw[0]} (WxH)" + if source_hw is not None + else "the default transfer geometry" + ) + resolution = self._get_sp_param(sp, "resolution", self._get_sp_param(sp, "image_size", 720)) + logger.warning( + "Cosmos3 transfer ignores requested size=%dx%d (WxH); %s selected the %s resolution=%s " + "bucket, %dx%d (WxH). Transfer supports canonical Cosmos buckets, not arbitrary output sizes.", + *requested_size, + source_description, + aspect_ratio, + resolution, + width, + height, + ) + + request_aspect_ratio = self._get_sp_param(sp, "aspect_ratio", None) + aspect_ratio_source = "requested aspect_ratio" + if request_aspect_ratio is None: + request_aspect_ratio = _json_object_aspect_ratio(prompt) + aspect_ratio_source = "JSON prompt aspect_ratio" + normalized_requested_ratio = _normalize_aspect_ratio(request_aspect_ratio) + if normalized_requested_ratio is not None and normalized_requested_ratio != aspect_ratio: + logger.warning( + "Cosmos3 transfer %s=%r conflicts with the control-selected %s bucket, %dx%d (WxH); " + "using %s for prompt metadata.", + aspect_ratio_source, + request_aspect_ratio, + aspect_ratio, + width, + height, + aspect_ratio, + ) @staticmethod def _first_transfer_control_hw(transfer_config: Cosmos3TransferConfig) -> tuple[int, int] | None: @@ -3083,7 +3233,15 @@ class Cosmos3OmniDiffusersPipeline( source_hw = (int(input_frames.shape[-2]), int(input_frames.shape[-1])) else: source_hw = self._first_transfer_control_hw(transfer_config) - height, width = self._transfer_bucket_size(sp, source_hw) + height, width, transfer_aspect_ratio = self._transfer_bucket_size(sp, source_hw) + self._warn_transfer_bucket_conflicts( + sp, + prompt, + source_hw=source_hw, + height=height, + width=width, + aspect_ratio=transfer_aspect_ratio, + ) if input_frames is not None: if tuple(input_frames.shape[-2:]) != (height, width): @@ -3190,6 +3348,7 @@ class Cosmos3OmniDiffusersPipeline( use_duration_template=bool(self._get_sp_param(sp, "use_duration_template", True)), use_resolution_template=bool(self._get_sp_param(sp, "use_resolution_template", True)), negative_metadata_mode=str(self._get_sp_param(sp, "negative_metadata_mode", "same")), + aspect_ratio_override=transfer_aspect_ratio, ) output_chunks: list[torch.Tensor] = [] diff --git a/vllm_omni/diffusion/models/cosmos3/utils.py b/vllm_omni/diffusion/models/cosmos3/utils.py index bfa3fd4db..76df64b3b 100644 --- a/vllm_omni/diffusion/models/cosmos3/utils.py +++ b/vllm_omni/diffusion/models/cosmos3/utils.py @@ -27,6 +27,37 @@ COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP = "first" # pixel-frame budgets without instantiating the VAE. COSMOS3_VAE_TEMPORAL_COMPRESSION = 4 +VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { + "256": { + "1,1": (256, 256), + "4,3": (320, 256), + "3,4": (256, 320), + "16,9": (320, 192), + "9,16": (192, 320), + }, + "480": { + "1,1": (640, 640), + "4,3": (736, 544), + "3,4": (544, 736), + "16,9": (832, 480), + "9,16": (480, 832), + }, + "704": { + "1,1": (960, 960), + "4,3": (1088, 832), + "3,4": (832, 1088), + "16,9": (1280, 704), + "9,16": (704, 1280), + }, + "720": { + "1,1": (960, 960), + "4,3": (1104, 832), + "3,4": (832, 1104), + "16,9": (1280, 720), + "9,16": (720, 1280), + }, +} + ROBOLAB_DEFAULT_CONDITIONING_FPS = 15.0 ROBOLAB_DEFAULT_ACTION_CHUNK_SIZE = 32 ROBOLAB_DEFAULT_IMAGE_HEIGHT = 540