anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Cosmos3 logging and prompt improvements (#6325)

Browse Source 
Signed-off-by: Maciej Bala <mbala@nvidia.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
This commit is contained in:
MaciejBalaNVandRoger Wang authored and GitHub committed 2026-08-20 12:04:49 -07:00
1 parent b2011a3edd
commit 13c2de684c
4 files changed
+335 -38
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
+139 -2
View File
@@ -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)
vllm_omni/diffusion/models/cosmos3/action.py
+2 -32
View File
@@ -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
vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
+163 -4
View File
@@ -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] = []
vllm_omni/diffusion/models/cosmos3/utils.py
+31
View File
@@ -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