diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index 95b5dbe77..72d79c304 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -160,7 +160,12 @@ JSON references currently support `image_url`/`video_url`; `file_id` references are not implemented yet. Models may expose additional V2V controls through `extra_params`. For example, Cosmos3 supports `condition_frame_indexes_vision` and `condition_video_keep` to select which -decoded reference frames are used as clean conditioning. +decoded reference frames are used as clean conditioning. Cosmos3 transfer mode +also accepts `edge`, `blur`, `depth`, `seg`, or `wsm` control hints plus +transfer options such as `control_path`, `control_guidance`, +`control_guidance_interval`, `num_video_frames_per_chunk`, +`num_conditional_frames`, `show_control_condition`, and `show_input`; see the +Cosmos3 recipe for complete examples. ### Speech-to-Video diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index 062009fe6..e617f6127 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -6,7 +6,7 @@ - Vendor: NVIDIA - Model: `nvidia/Cosmos3-Nano` -- Task: Text-to-image (T2I), text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) generation, with optional synchronized audio (video + sound), action policy +- Task: Text-to-image (T2I), text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) generation, with optional transfer controls, synchronized audio (video + sound), action policy - Mode: Online serving with the OpenAI-compatible image/video APIs, plus offline generation via the `Omni` API - Maintainer: Community @@ -25,6 +25,9 @@ mode is selected per request: reference-video latent frames; use `extra_params.condition_frame_indexes_vision` and `extra_params.condition_video_keep` to choose which prefix/tail frames guide generation. +- **Transfer V2V** — pass one or more transfer hints in `extra_params` + (`edge`, `blur`, `depth`, `seg`, `wsm`) to guide generation with control + frames. Transfer mode is video-only and cannot be combined with sound or action. - **T2VS / I2VS** — add `generate_sound=true` (and optional `sound_duration`) to a T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples. @@ -190,6 +193,18 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \ -F 'video_reference={"video_url":"https://example.com/reference.mp4"}' \ -o cosmos3_v2v_from_url.mp4 +# Transfer V2V with a precomputed depth control video. `control_path` can point +# to a local image/video; edge and blur can also be computed from `input_reference` +# by passing `"edge":true` or `"blur":true`. +curl -sS -X POST http://localhost:8000/v1/videos/sync \ + -H "Accept: video/mp4" \ + -F "model=nvidia/Cosmos3-Nano" \ + -F "prompt=Generate a realistic scene following the provided control video." \ + -F "size=1280x720" -F "num_frames=121" \ + -F "num_inference_steps=50" -F "seed=125" \ + -F 'extra_params={"depth":{"control_path":"/path/to/depth_control.mp4"},"max_frames":121,"resolution":"720","num_video_frames_per_chunk":121}' \ + -o cosmos3_transfer_depth.mp4 + # Text-to-video-with-sound curl -sS -X POST http://localhost:8000/v1/videos/sync \ -H "Accept: video/mp4" \ @@ -312,6 +327,17 @@ vllm serve nvidia/Cosmos3-Nano-Policy-DROID \ For V2V, `condition_frame_indexes_vision` selects the clean conditioned latent frame indexes (default `[0, 1]`), and `condition_video_keep` selects whether the API decodes the first or last needed reference frames (`"first"` by default). +- **Transfer controls:** `extra_params` may include `edge`, `blur`, `depth`, + `seg`, or `wsm`. Each hint accepts `true`, a path string, or an object such as + `{"control_path": "/path/to/control.mp4"}`; `edge` also accepts + `preset_edge_threshold` and `blur` accepts `preset_blur_strength`. + Transfer-level options include `control_guidance`, + `control_guidance_interval`, `num_video_frames_per_chunk` (default `93`, + `101` for WSM), `num_conditional_frames` (default `1`), + `num_first_chunk_conditional_frames`, `max_frames`, + `show_control_condition`, `show_input`, and + `share_vision_temporal_positions`. Non-WSM transfer preserves the input video + fps when available; WSM defaults to 10 fps unless `fps` is supplied. - **DROID OpenPI observations:** include a string `prompt`, either `observation/image` or the three-view DROID camera keys (`observation/wrist_image_left`, `observation/exterior_image_1_left`, diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index f77f63d50..4a5fd43d4 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -9,6 +9,7 @@ from dataclasses import dataclass from types import SimpleNamespace from typing import Any +import numpy as np import pytest import torch from PIL import Image @@ -131,9 +132,12 @@ class StubCosmos3Transformer(nn.Module): ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: token = int(text_ids.reshape(-1)[0].item()) if text_ids.numel() else 0 sound_latents = kwargs.get("sound_latents") + control_latents = kwargs.get("control_latents") + control_bonus = 100 if control_latents is not None else 0 self.calls.append( { "token": token, + "has_control": control_latents is not None, "timestep": timestep.clone(), "text_mask": text_mask.clone(), "cache_before": self.cached_kv, @@ -145,7 +149,7 @@ class StubCosmos3Transformer(nn.Module): self.cached_kv = [(marker, marker + 100)] self.cached_freqs_gen = (marker + 200, marker + 300) action_latents = kwargs.get("action_latents") - outputs: list[torch.Tensor] = [torch.full_like(hidden_states, float(token))] + outputs: list[torch.Tensor] = [torch.full_like(hidden_states, float(token + control_bonus))] if action_latents is not None: outputs.append(torch.full_like(action_latents, float(token + 20))) if sound_latents is not None: @@ -190,6 +194,7 @@ def make_cosmos3_pipeline(): pipeline._current_flow_shift = 1.0 pipeline._guidance_scale = None pipeline._num_timesteps = None + pipeline._cosmos3_branch_caches = None pipeline._cache_dit_requires_paired_cfg = False pipeline._sound_tokenizer = None pipeline.progress_bar = passthrough_progress_bar @@ -198,6 +203,13 @@ def make_cosmos3_pipeline(): return _make +@pytest.fixture +def sequential_cfg_parallel(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.distributed import cfg_parallel + + monkeypatch.setattr(cfg_parallel, "get_classifier_free_guidance_world_size", lambda: 1) + + def make_sampling_params(**overrides: Any) -> SimpleNamespace: values = { "height": None, @@ -395,6 +407,160 @@ def test_preprocess_i2v_image_and_action_video_inputs() -> None: assert additional["condition_frame_indexes_vision"] == [0, 1] +def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transfer + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + get_cosmos3_pre_process_func, + ) + + cfg = transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": True})) + assert cfg is not None + assert list(cfg.hints) == ["edge"] + assert cfg.guidance_scale == 3.0 + assert cfg.control_guidance == 1.5 + assert cfg.flow_shift == 10.0 + assert cfg.num_video_frames_per_chunk == 93 + assert cfg.share_vision_temporal_positions is True + # fps omitted (no fps/frame_rate on the sampling params) -> wsm preset default (10) applies. + defaulted_fps_cfg = transfer.resolve_transfer_config(make_sampling_params(extra_args={"wsm": True})) + assert defaulted_fps_cfg is not None + assert defaulted_fps_cfg.fps == 10 + # fps provided (frame_rate set) -> the user value wins over the preset default. + explicit_fps_cfg = transfer.resolve_transfer_config(make_sampling_params(frame_rate=24.0, extra_args={"wsm": True})) + assert explicit_fps_cfg is not None + assert explicit_fps_cfg.fps == 24.0 + assert ( + Cosmos3OmniDiffusersPipeline.reference_video_decode_spec(extra_args={"edge": True, "max_frames": 4}).max_frames + == 4 + ) + frames_for_pad = torch.arange(3 * 3, dtype=torch.uint8).reshape(1, 3, 1, 3) + assert transfer.pad_temporal_frames(frames_for_pad, 5)[0, :, 0, 0].tolist() == [0, 3, 6, 6, 3] + + real_import_module = transfer.importlib.import_module + + def raise_missing_cv2(name: str, *args: Any, **kwargs: Any): + if name == "cv2": + raise ImportError("missing cv2") + return real_import_module(name, *args, **kwargs) + + monkeypatch.setattr(transfer.importlib, "import_module", raise_missing_cv2) + with pytest.raises(ImportError, match="opencv-python"): + transfer.load_or_compute_control_frames( + cfg.hints["edge"], + height=8, + width=8, + max_frames=1, + input_frames=torch.zeros(3, 1, 8, 8, dtype=torch.uint8), + ) + + precomputed = torch.zeros(3, 2, 8, 8, dtype=torch.uint8) + precomputed_cfg = transfer.resolve_transfer_config( + make_sampling_params(extra_args={"edge": {"control": precomputed}}) + ) + assert precomputed_cfg is not None + loaded = transfer.load_or_compute_control_frames( + precomputed_cfg.hints["edge"], + height=8, + width=8, + max_frames=2, + input_frames=None, + ) + assert tuple(loaded.shape) == (3, 2, 8, 8) + + preprocess = get_cosmos3_pre_process_func(SimpleNamespace()) + + class FramesWithFps(list): + fps = 12.5 + + frames = FramesWithFps(Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue", "yellow", "black")) + request = SimpleNamespace( + prompts=[{"prompt": "transfer", "multi_modal_data": {"video": frames}}], + sampling_params=SimpleNamespace( + height=16, + width=32, + extra_args={"edge": True, "max_frames": 4, "resolution": "256"}, + ), + ) + additional = preprocess(request).prompts[0]["additional_information"] + assert (request.sampling_params.height, request.sampling_params.width) == (192, 320) + 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 + + +def test_transfer_fps_matches_resolved_frame_rate_precedence() -> None: + """When fps and frame_rate differ, transfer must pick frame_rate -- the same + precedence as OmniDiffusionSamplingParams.resolved_frame_rate. Uses the real + sampling-params dataclass so the resolved_frame_rate property is exercised.""" + from vllm_omni.diffusion.models.cosmos3 import transfer + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + sp = OmniDiffusionSamplingParams(fps=24, frame_rate=12.0, extra_args={"edge": True}) + assert sp.resolved_frame_rate == 12.0 + cfg = transfer.resolve_transfer_config(sp) + assert cfg is not None + # edge has no preset fps default, so cfg.fps comes straight from fps resolution. + assert cfg.fps == sp.resolved_frame_rate == 12.0 + + +def test_transfer_edge_uses_rgb_canny(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transfer + + class FakeCv2: + def __init__(self) -> None: + self.canny_inputs: list[np.ndarray] = [] + + def Canny(self, image, lower, upper): + assert (lower, upper) == (100, 200) + self.canny_inputs.append(image.copy()) + return np.zeros(image.shape[:2], dtype=np.uint8) + + fake_cv2 = FakeCv2() + monkeypatch.setattr(transfer, "_import_cv2", lambda _hint_key: fake_cv2) + + frames = torch.zeros(3, 1, 4, 5, dtype=torch.uint8) + frames[0] = 255 + edge = transfer.make_edge_control(frames, "medium") + + assert tuple(edge.shape) == (3, 1, 4, 5) + assert len(fake_cv2.canny_inputs) == 1 + assert fake_cv2.canny_inputs[0].shape == (4, 5, 3) + + +def test_transfer_blur_uses_scaled_bilateral(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transfer + + class FakeCv2: + INTER_AREA = 1 + INTER_LINEAR = 2 + INTER_CUBIC = 3 + + def __init__(self) -> None: + self.bilateral_calls: list[tuple[tuple[int, int, int], int, float, float]] = [] + + def resize(self, image, size, interpolation): + del interpolation + width, height = size + return np.zeros((height, width, image.shape[2]), dtype=image.dtype) + + def bilateralFilter(self, image, diameter, sigma_color, sigma_space): + self.bilateral_calls.append((image.shape, diameter, sigma_color, sigma_space)) + return image + + def GaussianBlur(self, *args, **kwargs): + raise AssertionError("Cosmos3 transfer blur should use bilateralFilter, not GaussianBlur.") + + fake_cv2 = FakeCv2() + monkeypatch.setattr(transfer, "_import_cv2", lambda _hint_key: fake_cv2) + + frames = torch.zeros(3, 1, 72, 72, dtype=torch.uint8) + blurred = transfer.make_blur_control(frames, "high") + + assert tuple(blurred.shape) == (3, 1, 72, 72) + assert fake_cv2.bilateral_calls == [((72, 72, 3), 3, 15.0, 10.0)] + + def test_postprocess_handles_image_video_audio_and_validation() -> None: from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_post_process_func @@ -643,6 +809,267 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result)) +def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_cfg_parallel) -> None: + pipeline = make_cosmos3_pipeline() + latents = torch.zeros(1, 2, 1, 1, 1) + velocity_mask = torch.ones(1, 1, 1, 1, 1) + + result = pipeline.diffuse_transfer( + latents=latents, + timesteps=torch.tensor([7]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=3.0, + control_guidance=1.5, + control_guidance_interval=None, + control_latents=[torch.zeros_like(latents)], + shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask}, + velocity_mask=velocity_mask, + condition_latents=torch.zeros_like(latents), + ) + + assert [(call["token"], call["has_control"]) for call in pipeline.transformer.calls] == [ + (2, True), + (2, False), + (1, True), + ] + torch.testing.assert_close(result, torch.full_like(latents, 254.0)) + + +def test_diffuse_transfer_skips_idle_cfg_branches(make_cosmos3_pipeline, sequential_cfg_parallel) -> None: + latents = torch.zeros(1, 2, 1, 1, 1) + velocity_mask = torch.ones(1, 1, 1, 1, 1) + + control_only = make_cosmos3_pipeline() + control_result = control_only.diffuse_transfer( + latents=latents, + timesteps=torch.tensor([7]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=1.0, + control_guidance=1.5, + control_guidance_interval=None, + control_latents=[torch.zeros_like(latents)], + shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask}, + velocity_mask=velocity_mask, + condition_latents=torch.zeros_like(latents), + ) + assert [(call["token"], call["has_control"]) for call in control_only.transformer.calls] == [ + (2, True), + (2, False), + ] + torch.testing.assert_close(control_result, torch.full_like(latents, 152.0)) + + text_only = make_cosmos3_pipeline() + text_result = text_only.diffuse_transfer( + latents=latents, + timesteps=torch.tensor([7]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=3.0, + control_guidance=1.0, + control_guidance_interval=None, + control_latents=[torch.zeros_like(latents)], + shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask}, + velocity_mask=velocity_mask, + condition_latents=torch.zeros_like(latents), + ) + assert [(call["token"], call["has_control"]) for call in text_only.transformer.calls] == [ + (2, True), + (1, True), + ] + torch.testing.assert_close(text_result, torch.full_like(latents, 104.0)) + + +def test_diffuse_transfer_interval_switches_branch_counts(make_cosmos3_pipeline, sequential_cfg_parallel) -> None: + pipeline = make_cosmos3_pipeline() + latents = torch.zeros(1, 2, 1, 1, 1) + velocity_mask = torch.ones(1, 1, 1, 1, 1) + + result = pipeline.diffuse_transfer( + latents=latents, + timesteps=torch.tensor([900, 500, 100]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=3.0, + control_guidance=1.5, + control_guidance_interval=(400.0, 1000.0), + control_latents=[torch.zeros_like(latents)], + shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask}, + velocity_mask=velocity_mask, + condition_latents=torch.zeros_like(latents), + guidance_interval=(800.0, 1000.0), + ) + + assert [(call["token"], call["has_control"]) for call in pipeline.transformer.calls] == [ + (2, True), + (2, False), + (1, True), + (2, True), + (2, False), + (2, True), + ] + torch.testing.assert_close(result, torch.full_like(latents, 508.0)) + + +@pytest.mark.parametrize(("hint_key", "expected_fps"), [("edge", 8.0), ("wsm", 10.0)]) +def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint_key: str, expected_fps: float) -> None: + pipeline = make_cosmos3_pipeline() + captured: dict[str, Any] = {} + + def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, *args, **kwargs): + del prompt, negative_prompt, num_frames, height, width, args, kwargs + captured["format_frame_rate"] = frame_rate + return _ids(2), _mask(), _ids(1), _mask() + + def fake_encode(video: torch.Tensor) -> torch.Tensor: + latent_frames = (video.shape[2] - 1) // pipeline.vae_scale_factor_temporal + 1 + return torch.ones( + 1, + 2, + latent_frames, + max(1, video.shape[-2] // pipeline.vae_scale_factor_spatial), + max(1, video.shape[-1] // pipeline.vae_scale_factor_spatial), + ) + + def fake_prepare(target_norm, current_conditional_frames, generator): + del current_conditional_frames, generator + latent = fake_encode(target_norm) + velocity_mask = torch.ones(1, 1, latent.shape[2], 1, 1) + return torch.zeros_like(latent), velocity_mask, torch.zeros_like(latent) + + def fake_diffuse_transfer(**kwargs): + captured["shared_kwargs"] = kwargs["shared_kwargs"] + return kwargs["latents"] + + pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16) + pipeline._format_and_tokenize_prompts = fake_format + pipeline._encode_video_tensor = fake_encode + pipeline._prepare_transfer_latents = fake_prepare + pipeline.diffuse_transfer = fake_diffuse_transfer + pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target) + pipeline._decode_latents = lambda latents: torch.zeros(1, 3, 5, 16, 16, device="meta") + + control = torch.zeros(3, 5, 16, 16, dtype=torch.uint8) + request = SimpleNamespace( + prompts=[ + { + "prompt": "transfer", + "modalities": ["video"], + "additional_information": { + "preprocessed_transfer_video": torch.zeros(1, 3, 5, 16, 16), + "transfer_input_fps": 8.0, + }, + } + ], + sampling_params=make_sampling_params( + height=16, + width=16, + # fps omitted (no frame_rate) -> non-wsm uses the source video fps (8), wsm uses its preset (10). + extra_args={ + hint_key: {"control": control}, + "max_frames": 5, + "num_video_frames_per_chunk": 5, + "show_control_condition": True, + }, + ), + ) + + output = pipeline.forward(request) + + assert captured["format_frame_rate"] == expected_fps + assert captured["shared_kwargs"]["fps"] == expected_fps + assert captured["flow_shifts"] == [10.0] + assert output.custom_output["fps"] == expected_fps + assert output.output["video"].device.type == "meta" + + +def test_forward_transfer_runs_multichunk_overlap_path( + make_cosmos3_pipeline, + sequential_cfg_parallel, +) -> None: + pipeline = make_cosmos3_pipeline() + captured: dict[str, Any] = {"targets": [], "conditional_frames": []} + + pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16) + pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (_ids(2), _mask(), _ids(1), _mask()) + pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target) + + original_prepare = pipeline._prepare_transfer_latents + + def recording_prepare(target_norm, current_conditional_frames, generator): + captured["targets"].append(target_norm.detach().clone()) + captured["conditional_frames"].append(current_conditional_frames) + return original_prepare(target_norm, current_conditional_frames, generator) + + pipeline._prepare_transfer_latents = recording_prepare + + decoded_chunks = [ + torch.tensor([-0.6, -0.5, -0.4, -0.3, -0.2], dtype=torch.float32), + torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5], dtype=torch.float32), + ] + + def fake_decode(latents): + chunk_values = decoded_chunks[len(captured.setdefault("decode_calls", []))] + captured["decode_calls"].append(latents.detach().clone()) + return chunk_values.view(1, 1, 5, 1, 1).expand(1, 3, 5, 16, 16).clone() + + pipeline._decode_latents = fake_decode + + transfer_video = torch.zeros(1, 3, 8, 16, 16) + transfer_video[:, :, 1] = 1.0 + control = torch.zeros(3, 8, 16, 16, dtype=torch.uint8) + request = SimpleNamespace( + prompts=[ + { + "prompt": "transfer", + "modalities": ["video"], + "additional_information": { + "preprocessed_transfer_video": transfer_video, + "transfer_input_fps": 8.0, + }, + } + ], + sampling_params=make_sampling_params( + height=16, + width=16, + num_inference_steps=1, + guidance_scale=1.0, + extra_args={ + "edge": {"control": control}, + "control_guidance": 1.0, + "max_frames": 8, + "num_video_frames_per_chunk": 5, + "num_conditional_frames": 1, + "num_first_chunk_conditional_frames": 2, + }, + ), + ) + + output = pipeline.forward(request) + + assert captured["conditional_frames"] == [2, 1] + assert len(captured["decode_calls"]) == 2 + assert output.output["video"].shape == (1, 3, 8, 16, 16) + torch.testing.assert_close( + output.output["video"][0, 0, :, 0, 0], + torch.tensor([-0.6, -0.5, -0.4, -0.3, -0.2, 0.2, 0.3, 0.4]), + ) + assert output.custom_output["transfer_controls"]["edge"].shape == (1, 3, 8, 16, 16) + torch.testing.assert_close(captured["targets"][0][:, :, 0], torch.full((1, 3, 16, 16), -1.0)) + torch.testing.assert_close(captured["targets"][0][:, :, 1], torch.full((1, 3, 16, 16), 1.0)) + torch.testing.assert_close(captured["targets"][0][:, :, 2:], torch.full((1, 3, 3, 16, 16), 1.0)) + torch.testing.assert_close(captured["targets"][1][:, :, 0], torch.full((1, 3, 16, 16), -0.2)) + + def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -> None: """With cache-dit active the uncond pass is kept even outside the guidance interval (so cache-dit's has_separate_cfg parity stays in phase), and the @@ -936,6 +1363,26 @@ class TestForwardRouting: make_sampling_params(), "only for video", ), + ( + [{"prompt": "x", "modalities": ["image"]}], + make_sampling_params(extra_args={"edge": {"control_path": "/tmp/control.mp4"}}), + "transfer inference is supported only for video outputs", + ), + ( + [{"prompt": "x", "modalities": ["video"], "generate_sound": True}], + make_sampling_params(extra_args={"edge": {"control_path": "/tmp/control.mp4"}}), + "cannot be combined with sound generation", + ), + ( + [{"prompt": "x", "modalities": ["video"]}], + make_sampling_params( + extra_args={ + "edge": {"control_path": "/tmp/control.mp4"}, + "action_mode": "policy", + } + ), + "cannot be combined with action generation", + ), ], ) def test_forward_rejects_invalid_public_requests( diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 62411fc67..cb4e28b51 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -131,6 +131,38 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No assert tuple(output.shape) == (1, 2, 1, 2, 2) +def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None)) + + model = transformer_cosmos3.Cosmos3VFMTransformer( + SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32) + ) + hidden_states = torch.zeros(1, 2, 1, 2, 2) + output = model( + hidden_states=hidden_states, + timestep=torch.tensor([1.0]), + text_ids=torch.tensor([[1, 2]], dtype=torch.long), + text_mask=torch.ones(1, 2, dtype=torch.long), + video_shape=(1, 2, 2), + fps=24.0, + control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)], + ) + + assert tuple(output.shape) == tuple(hidden_states.shape) + with pytest.raises(ValueError, match="control latent shape"): + model( + hidden_states=hidden_states, + timestep=torch.tensor([1.0]), + text_ids=torch.tensor([[1, 2]], dtype=torch.long), + text_mask=torch.ones(1, 2, dtype=torch.long), + video_shape=(1, 2, 2), + fps=24.0, + control_latents=[torch.zeros(1, 2, 2, 2, 2)], + ) + + def test_forward_gathers_gen_tokens_before_unpatchify(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 @@ -354,3 +386,33 @@ def test_compute_rope_freqs_places_text_video_action_and_sound_positions() -> No _, gen_pos = rotary.position_ids assert gen_pos.shape == (3, 1, 5) assert gen_pos[0, 0].tolist() == [102, 103, 103, 104, 102] + + rotary.position_ids.clear() + model._compute_rope_freqs( + text_mask=torch.tensor([[1, 1]], dtype=torch.long), + t=2, + hp=1, + wp=1, + fps=24.0, + device=torch.device("cpu"), + dtype=torch.float32, + num_vision_items=3, + share_vision_temporal_positions=True, + ) + _, shared_gen_pos = rotary.position_ids + assert shared_gen_pos[0, 0].tolist() == [102, 103, 102, 103, 102, 103] + + rotary.position_ids.clear() + model._compute_rope_freqs( + text_mask=torch.tensor([[1, 1]], dtype=torch.long), + t=2, + hp=1, + wp=1, + fps=24.0, + device=torch.device("cpu"), + dtype=torch.float32, + num_vision_items=3, + share_vision_temporal_positions=False, + ) + _, offset_gen_pos = rotary.position_ids + assert offset_gen_pos[0, 0].tolist() == [102, 103, 104, 105, 106, 107] diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py index 9f38f097a..cb39c0dc8 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py @@ -466,11 +466,18 @@ class TestCFGParallelForwardPath: """Test the LTX-2.3 CFG-parallel denoising path without loading model weights.""" @pytest.mark.parametrize(("cfg_rank", "expected_prompt_value"), [(0, 1.0), (1, 0.0)]) + @pytest.mark.parametrize( + ("frame_rate_input", "audio_sampling_rate", "expected_frame_rate"), + [(1.0, 1, 1.0), (None, 24, 24.0)], + ) def test_forward_cfg_parallel_steps_video_and_audio_scheduler( self, monkeypatch, cfg_rank, expected_prompt_value, + frame_rate_input, + audio_sampling_rate, + expected_frame_rate, ): from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23 from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -484,7 +491,7 @@ class TestCFGParallelForwardPath: pipe.vae_temporal_compression_ratio = 1 pipe.transformer_spatial_patch_size = 1 pipe.transformer_temporal_patch_size = 1 - pipe.audio_sampling_rate = 1 + pipe.audio_sampling_rate = audio_sampling_rate pipe.audio_hop_length = 1 pipe.audio_vae_temporal_compression_ratio = 1 pipe.audio_vae_mel_compression_ratio = 1 @@ -553,8 +560,11 @@ class TestCFGParallelForwardPath: assert prompt_embeds.shape[0] == 2 return prompt_embeds, prompt_embeds, prompt_attention_mask + rope_video_fps: list[float] = [] + class FakeRope: def prepare_video_coords(self, batch_size, num_frames, height, width, device, fps): + rope_video_fps.append(fps) return torch.zeros(batch_size, num_frames * height * width, 3, device=device) def prepare_audio_coords(self, batch_size, num_frames, device): @@ -611,7 +621,7 @@ class TestCFGParallelForwardPath: height=32, width=32, num_frames=1, - frame_rate=1.0, + frame_rate=frame_rate_input, num_inference_steps=2, guidance_scale=4.0, latents=video_latents, @@ -651,6 +661,11 @@ class TestCFGParallelForwardPath: torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1)) torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2)) + # fps regression guard: an omitted request fps (frame_rate_input=None) must resolve + # to the model's own 24.0 default, not crash on None; a provided rate is passed through. + assert rope_video_fps + assert all(fps == expected_frame_rate for fps in rope_video_fps) + class TestRegistryIntegration: """Verify all LTX-2.3 pipeline variants are registered.""" diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index dbf766391..7e6fa2a73 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -423,6 +423,7 @@ def test_decode_video_bytes_can_keep_last_frames(): ) assert len(frames) == 2 + assert frames.fps == pytest.approx(8.0) red_means = [np.asarray(frame)[:, :, 0].mean() for frame in frames] assert red_means[0] > 100 assert red_means[1] > red_means[0] @@ -520,11 +521,49 @@ def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture): engine = test_client.app.state.openai_serving_video._engine_client captured = engine.captured_sampling_params_list[0] assert captured.num_frames == 72 - assert captured.fps == 24 - assert captured.frame_rate == 24.0 + # fps omitted -> sampling params carry None (the "not provided" signal); the 24 + # default is applied only at output encoding. + assert captured.fps is None + assert captured.frame_rate is None assert fps_values == [24] +def test_model_reported_fps_wins_when_request_fps_omitted(test_client, mocker: MockerFixture): + fps_values = [] + + def _fake_encode(video, fps, audio=None, audio_sample_rate=None, **kwargs): + del video, audio, audio_sample_rate, kwargs + fps_values.append(fps) + return b"fake-video" + + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + side_effect=_fake_encode, + ) + + engine = test_client.app.state.openai_serving_video._engine_client + + async def _generate(prompt, request_id, sampling_params_list): + engine.captured_prompt = prompt + engine.captured_sampling_params_list = sampling_params_list + result = MockVideoResult([object()]) + result.multimodal_output["fps"] = 8 + yield result + + engine.generate = _generate + + response = test_client.post("/v1/videos", data={"prompt": "source fps"}) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + captured = engine.captured_sampling_params_list[0] + # fps omitted -> None on the sampling params; the model-reported fps (8) wins for output. + assert captured.fps is None + assert captured.frame_rate is None + assert fps_values == [8] + + def test_size_param_sets_width_height(test_client, mocker: MockerFixture): mocker.patch( "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index b0b1f7559..78df8f70f 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -20,7 +20,7 @@ from __future__ import annotations import math import os import time -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from dataclasses import fields from typing import Any, ClassVar @@ -50,6 +50,7 @@ from vllm_omni.diffusion.models.interface import ( from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.entrypoints.openai.video_api_utils import positive_float from vllm_omni.inputs.data import OmniDiffusionSamplingParams from .action import ( @@ -65,6 +66,19 @@ from .action import ( pad_action_to_dim, resolve_domain_id, ) +from .transfer import ( + Cosmos3TransferConfig, + has_transfer_hints, + load_or_compute_control_frames, + media_hw, + media_to_uint8_cthw, + normalized_video_to_uint8_cthw, + pad_temporal_frames, + resize_center_crop_uint8_cthw, + resolve_transfer_config, + transfer_max_frames_from_extra_args, + uint8_cthw_to_normalized_5d, +) from .transformer_cosmos3 import Cosmos3VFMTransformer, resolve_sound_gen from .utils import ( COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION, @@ -166,6 +180,14 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): def _request_action_mode(request: OmniDiffusionRequest) -> str | None: return normalize_action_mode(_extra_args(request).get("action_mode")) + def _set_transfer_size_from_image(request: OmniDiffusionRequest, image: PIL.Image.Image) -> tuple[int, int]: + extra = _extra_args(request) + resolution = extra.get("resolution", extra.get("image_size", 720)) + target_w, target_h = find_closest_target_size(image.height, image.width, resolution) + request.sampling_params.height = target_h + request.sampling_params.width = target_w + return int(target_h), int(target_w) + def _set_action_size_from_image(request: OmniDiffusionRequest, image: PIL.Image.Image) -> tuple[int, int]: sp = request.sampling_params if sp.height is not None and sp.width is not None: @@ -248,7 +270,43 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): image = image.crop((left, top, left + target_w, top + target_h)) return video_processor.preprocess(image, height=target_h, width=target_w) + def _video_payload_value(video: Any, key: str) -> Any: + if isinstance(video, Mapping): + return video.get(key) + return getattr(video, key, None) + + def _video_payload_fps(video: Any) -> float | None: + for key in ("fps", "frame_rate", "source_fps", "input_fps", "avg_fps", "average_fps"): + fps = positive_float(_video_payload_value(video, key)) + if fps is not None: + return fps + for key in ("metadata", "info"): + metadata = _video_payload_value(video, key) + if metadata is None or metadata is video: + continue + fps = _video_payload_fps(metadata) + if fps is not None: + return fps + if isinstance(video, Mapping): + for key in ("frames", "data", "video"): + nested = video.get(key) + if nested is None or nested is video: + continue + fps = _video_payload_fps(nested) + if fps is not None: + return fps + return None + + def _unwrap_video_payload(video: Any) -> Any: + if isinstance(video, Mapping): + for key in ("frames", "data", "video"): + nested = video.get(key) + if nested is not None: + return nested + return video + def _video_payload_to_frames(video: Any) -> list[Any]: + video = _unwrap_video_payload(video) if isinstance(video, list): return video if isinstance(video, torch.Tensor): @@ -315,7 +373,9 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): prompt["additional_information"] = {} raw_video_frames: list[Any] | None = None + transfer_input_fps: float | None = None if raw_video is not None: + transfer_input_fps = _video_payload_fps(raw_video) raw_video_frames = _video_payload_to_frames(raw_video) if not raw_video_frames: raise TypeError("Cosmos3 video input must be a non-empty list of PIL images or image paths.") @@ -325,9 +385,13 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): image = _pil_to_rgb(raw_video_frames[0]) else: image = _pil_to_rgb(raw_image) + extra = _extra_args(request) + transfer_requested = action_mode is None and has_transfer_hints(extra) # Auto-calculate H/W from aspect ratio (720p max area) - if request.sampling_params.height is None or request.sampling_params.width is None: + if transfer_requested: + _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: _set_action_size_from_image(request, image) else: @@ -357,27 +421,40 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): ) else: assert raw_video_frames is not None - extra = _extra_args(request) - condition_frame_indexes_vision = normalize_condition_frame_indexes_vision( - extra.get( - "condition_frame_indexes_vision", - prompt.get("condition_frame_indexes_vision"), + if transfer_requested: + if transfer_input_fps is not None: + prompt["additional_information"]["transfer_input_fps"] = transfer_input_fps + transfer_frames = media_to_uint8_cthw( + raw_video_frames, + height=int(target_h), + width=int(target_w), + max_frames=transfer_max_frames_from_extra_args(extra), + ) + prompt["additional_information"]["preprocessed_transfer_video"] = uint8_cthw_to_normalized_5d( + transfer_frames, + dtype=torch.float32, + ) + else: + condition_frame_indexes_vision = normalize_condition_frame_indexes_vision( + extra.get( + "condition_frame_indexes_vision", + prompt.get("condition_frame_indexes_vision"), + ) + ) + keep = normalize_condition_video_keep( + extra.get("condition_video_keep", prompt.get("condition_video_keep")) + ) + max_frames = condition_pixel_frame_count(condition_frame_indexes_vision) + prompt["additional_information"]["preprocessed_video"] = _preprocess_condition_video( + raw_video_frames, + int(target_h), + int(target_w), + max_frames, + keep, + ) + prompt["additional_information"]["condition_frame_indexes_vision"] = list( + condition_frame_indexes_vision ) - ) - keep = normalize_condition_video_keep( - extra.get("condition_video_keep", prompt.get("condition_video_keep")) - ) - max_frames = condition_pixel_frame_count(condition_frame_indexes_vision) - prompt["additional_information"]["preprocessed_video"] = _preprocess_condition_video( - raw_video_frames, - int(target_h), - int(target_w), - max_frames, - keep, - ) - prompt["additional_information"]["condition_frame_indexes_vision"] = list( - condition_frame_indexes_vision - ) if action_mode is not None and raw_video_frames is not None: prompt["additional_information"]["preprocessed_video"] = _preprocess_action_video( raw_video_frames, @@ -558,6 +635,12 @@ class Cosmos3OmniDiffusersPipeline( extra_args: dict[str, Any] | None = None, ) -> ReferenceVideoDecodeSpec: extra_args = extra_args if isinstance(extra_args, dict) else {} + if has_transfer_hints(extra_args): + max_frames = transfer_max_frames_from_extra_args(extra_args) + if num_frames is not None: + max_frames = min(max_frames, int(num_frames)) + return ReferenceVideoDecodeSpec(max_frames=max_frames, keep="first") + action_mode = normalize_action_mode(extra_args.get("action_mode")) if action_mode is not None: if num_frames is not None: @@ -677,6 +760,7 @@ class Cosmos3OmniDiffusersPipeline( self._guidance_scale = None self._num_timesteps = None + self._cosmos3_branch_caches: dict[str, tuple[Any, Any]] | None = None self._robolab_transform = None # Set True by ``enable_cache_for_cosmos3`` when cache-dit is enabled on @@ -850,7 +934,63 @@ class Cosmos3OmniDiffusersPipeline( The transformer returns the raw prediction: video-only as a tensor, or a tuple in video, action, sound order for multimodal generation. """ - return self.transformer(**kwargs) + cache_key = kwargs.pop("_cosmos3_cache_key", None) + if cache_key is None: + return self.transformer(**kwargs) + + branch_caches = self._cosmos3_branch_caches + if branch_caches is None: + return self.transformer(**kwargs) + + cache_key = str(cache_key) + self.transformer.cached_kv, self.transformer.cached_freqs_gen = branch_caches.get(cache_key, (None, None)) + prediction = self.transformer(**kwargs) + branch_caches[cache_key] = (self.transformer.cached_kv, self.transformer.cached_freqs_gen) + return prediction + + def combine_multi_branch_cfg_noise( + self, + predictions: list[torch.Tensor | tuple[torch.Tensor, ...]], + true_cfg_scale: float | dict[str, float], + cfg_normalize: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: + if not isinstance(true_cfg_scale, dict) or true_cfg_scale.get("mode") != "cosmos3_transfer": + return super().combine_multi_branch_cfg_noise(predictions, true_cfg_scale, cfg_normalize) + mode = str(true_cfg_scale.get("branch_mode", "control_and_text")) + guidance_scale = float(true_cfg_scale.get("guidance_scale", 1.0)) + control_guidance = float(true_cfg_scale.get("control_guidance", 1.0)) + + if mode == "control_only": + if len(predictions) != 2: + raise ValueError(f"Cosmos3 transfer control-only CFG expects 2 branches, got {len(predictions)}.") + cond_full, cond_no_control = predictions + if isinstance(cond_full, tuple) or isinstance(cond_no_control, tuple): + raise ValueError("Cosmos3 transfer control-only CFG expects video-only tensor predictions.") + cfg_reference = cond_full + combined = cond_no_control + control_guidance * (cond_full - cond_no_control) + elif mode == "text_only": + if len(predictions) != 2: + raise ValueError(f"Cosmos3 transfer text-only CFG expects 2 branches, got {len(predictions)}.") + cond_full, uncond_full = predictions + if isinstance(cond_full, tuple) or isinstance(uncond_full, tuple): + raise ValueError("Cosmos3 transfer text-only CFG expects video-only tensor predictions.") + cfg_reference = cond_full + combined = uncond_full + guidance_scale * (cond_full - uncond_full) + elif mode == "control_and_text": + if len(predictions) != 3: + raise ValueError(f"Cosmos3 transfer control+text CFG expects 3 branches, got {len(predictions)}.") + cond_full, cond_no_control, uncond_full = predictions + if isinstance(cond_full, tuple) or isinstance(cond_no_control, tuple) or isinstance(uncond_full, tuple): + raise ValueError("Cosmos3 transfer control+text CFG expects video-only tensor predictions.") + cfg_reference = cond_full + control_cond = cond_no_control + control_guidance * (cond_full - cond_no_control) + combined = uncond_full + guidance_scale * (control_cond - uncond_full) + else: + raise ValueError(f"Unknown Cosmos3 transfer CFG branch_mode={mode!r}.") + + if cfg_normalize: + combined = self.cfg_normalize_function(cfg_reference, combined) + return combined @staticmethod def _cfg_parallel_active() -> bool: @@ -2176,6 +2316,444 @@ class Cosmos3OmniDiffusersPipeline( outputs.append(sound_latents) return outputs[0] if len(outputs) == 1 else tuple(outputs) + @staticmethod + def _get_transfer_num_chunks( + total_frames: int, + frames_per_chunk: int, + conditional_frames: int, + ) -> tuple[int, int]: + if frames_per_chunk <= 0: + raise ValueError("Cosmos3 transfer frames_per_chunk must be positive.") + if total_frames <= frames_per_chunk: + return 1, frames_per_chunk + stride = frames_per_chunk - conditional_frames + if stride <= 0: + raise ValueError("Cosmos3 transfer num_conditional_frames must be smaller than num_video_frames_per_chunk.") + remaining = total_frames - frames_per_chunk + extra_chunks = remaining // stride + (1 if remaining % stride else 0) + return 1 + extra_chunks, stride + + def _prepare_transfer_latents( + self, + target_video: torch.Tensor, + current_conditional_frames: int, + generator: torch.Generator, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + condition_latents = self._encode_video_tensor(target_video) + noise = randn_tensor( + condition_latents.shape, + generator=generator, + device=self.device, + dtype=self.dtype, + ) + condition_mask = torch.zeros( + 1, + 1, + condition_latents.shape[2], + 1, + 1, + device=self.device, + dtype=self.dtype, + ) + if current_conditional_frames > 0: + latent_frames = (current_conditional_frames - 1) // self.vae_scale_factor_temporal + 1 + condition_mask[:, :, :latent_frames] = 1.0 + latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise + velocity_mask = 1.0 - condition_mask + return latents, velocity_mask, condition_mask * condition_latents + + def _transfer_bucket_size( + self, + sp: OmniDiffusionSamplingParams, + source_hw: tuple[int, int] | None, + ) -> tuple[int, int]: + 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) + + @staticmethod + def _first_transfer_control_hw(transfer_config: Cosmos3TransferConfig) -> tuple[int, int] | None: + for hint in transfer_config.ordered_hints: + if hint.control is not None: + detected = media_hw(hint.control) + if detected is not None: + return detected + if hint.control_path is not None: + detected = media_hw(hint.control_path) + if detected is not None: + return detected + return None + + def diffuse_transfer( + self, + latents: torch.Tensor, + timesteps: torch.Tensor, + cond_ids: torch.Tensor, + cond_mask: torch.Tensor, + uncond_ids: torch.Tensor, + uncond_mask: torch.Tensor, + guidance_scale: float, + control_guidance: float, + control_guidance_interval: tuple[float, float] | None, + control_latents: list[torch.Tensor], + shared_kwargs: dict[str, Any], + *, + velocity_mask: torch.Tensor, + condition_latents: torch.Tensor, + guidance_interval: tuple[float, float] | None = None, + ) -> torch.Tensor: + def _active_at(t: torch.Tensor, interval: tuple[float, float] | None) -> bool: + if interval is None: + return True + t_scalar = float(t.item()) if torch.is_tensor(t) else float(t) + lo, hi = interval + return lo <= t_scalar <= hi + + self.transformer.reset_cache() + self._cosmos3_branch_caches = {} + try: + for t in self.progress_bar(timesteps): + timestep = t.unsqueeze(0) + step_guidance = guidance_scale if _active_at(t, guidance_interval) else 1.0 + step_control = control_guidance if _active_at(t, control_guidance_interval) else 1.0 + needs_text_cfg = step_guidance > 1.0 + needs_control_cfg = step_control != 1.0 + + cond_full_kwargs = dict( + hidden_states=latents, + timestep=timestep, + text_ids=cond_ids, + text_mask=cond_mask, + control_latents=control_latents, + _cosmos3_cache_key="transfer_cond_full", + **shared_kwargs, + ) + if needs_control_cfg and needs_text_cfg: + branches_kwargs = [ + cond_full_kwargs, + dict( + hidden_states=latents, + timestep=timestep, + text_ids=cond_ids, + text_mask=cond_mask, + control_latents=None, + _cosmos3_cache_key="transfer_cond_no_control", + **shared_kwargs, + ), + dict( + hidden_states=latents, + timestep=timestep, + text_ids=uncond_ids, + text_mask=uncond_mask, + control_latents=control_latents, + _cosmos3_cache_key="transfer_uncond_full", + **shared_kwargs, + ), + ] + noise_pred = self.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale={ + "mode": "cosmos3_transfer", + "branch_mode": "control_and_text", + "guidance_scale": step_guidance, + "control_guidance": step_control, + }, + branches_kwargs=branches_kwargs, + cfg_normalize=False, + ) + elif needs_control_cfg: + branches_kwargs = [ + cond_full_kwargs, + dict( + hidden_states=latents, + timestep=timestep, + text_ids=cond_ids, + text_mask=cond_mask, + control_latents=None, + _cosmos3_cache_key="transfer_cond_no_control", + **shared_kwargs, + ), + ] + noise_pred = self.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale={ + "mode": "cosmos3_transfer", + "branch_mode": "control_only", + "control_guidance": step_control, + }, + branches_kwargs=branches_kwargs, + cfg_normalize=False, + ) + elif needs_text_cfg: + branches_kwargs = [ + cond_full_kwargs, + dict( + hidden_states=latents, + timestep=timestep, + text_ids=uncond_ids, + text_mask=uncond_mask, + control_latents=control_latents, + _cosmos3_cache_key="transfer_uncond_full", + **shared_kwargs, + ), + ] + noise_pred = self.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale={ + "mode": "cosmos3_transfer", + "branch_mode": "text_only", + "guidance_scale": step_guidance, + }, + branches_kwargs=branches_kwargs, + cfg_normalize=False, + ) + else: + noise_pred = self.predict_noise(**cond_full_kwargs) + if isinstance(noise_pred, tuple): + raise ValueError("Cosmos3 transfer diffusion expects video-only tensor predictions.") + noise_pred = noise_pred * velocity_mask + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents + finally: + self._cosmos3_branch_caches = None + self.transformer.reset_cache() + return latents + + def _forward_transfer( + self, + *, + prompt: str, + negative_prompt: str, + sp: OmniDiffusionSamplingParams, + transfer_config: Cosmos3TransferConfig, + transfer_video_tensor: torch.Tensor | None, + transfer_input_fps: float | None, + ) -> DiffusionOutput: + input_frames = None + if transfer_video_tensor is not None: + input_frames = normalized_video_to_uint8_cthw(transfer_video_tensor) + 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) + + if input_frames is not None: + if tuple(input_frames.shape[-2:]) != (height, width): + input_frames = resize_center_crop_uint8_cthw(input_frames, height, width) + input_frames = input_frames[:, : transfer_config.max_frames] + + per_hint_frames: dict[str, torch.Tensor] = {} + for hint in transfer_config.ordered_hints: + frames = load_or_compute_control_frames( + hint, + height=height, + width=width, + max_frames=transfer_config.max_frames, + input_frames=input_frames, + ) + if frames.shape[1] < 1: + raise ValueError(f"Cosmos3 transfer hint '{hint.key}' produced no frames.") + per_hint_frames[hint.key] = frames + if not per_hint_frames: + raise ValueError("Cosmos3 transfer requires at least one control hint.") + + total_frames = next(iter(per_hint_frames.values())).shape[1] + if transfer_config.num_frames is not None: + total_frames = min(total_frames, int(transfer_config.num_frames)) + total_frames = max(1, total_frames) + per_hint_frames = {key: pad_temporal_frames(frames, total_frames) for key, frames in per_hint_frames.items()} + if input_frames is not None: + input_frames = pad_temporal_frames(input_frames, total_frames) + + temporal_compression = self.vae_scale_factor_temporal + chunk_frames = 1 if total_frames == 1 else transfer_config.num_video_frames_per_chunk + chunk_frames = math.ceil((chunk_frames - 1) / temporal_compression) * temporal_compression + 1 + num_chunks, stride = self._get_transfer_num_chunks( + total_frames, + chunk_frames, + transfer_config.num_conditional_frames, + ) + padded_frames = max(total_frames, chunk_frames) + per_hint_frames = {key: pad_temporal_frames(frames, padded_frames) for key, frames in per_hint_frames.items()} + if input_frames is not None: + input_frames = pad_temporal_frames(input_frames, padded_frames) + + configured_frame_rate = positive_float(transfer_config.fps) + input_frame_rate = positive_float(transfer_input_fps) + sampling_frame_rate = ( + positive_float(self._get_sp_param(sp, "resolved_frame_rate")) + or positive_float(self._get_sp_param(sp, "frame_rate")) + or positive_float(self._get_sp_param(sp, "fps")) + ) + is_wsm_only = len(transfer_config.hints) == 1 and "wsm" in transfer_config.hints + if is_wsm_only: + frame_rate = configured_frame_rate or input_frame_rate or sampling_frame_rate or 24.0 + else: + frame_rate = input_frame_rate or configured_frame_rate or sampling_frame_rate or 24.0 + num_inference_steps = sp.num_inference_steps or COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS + guidance_scale = ( + float(transfer_config.guidance_scale) + if transfer_config.guidance_scale is not None + else float(sp.guidance_scale or COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE) + ) + flow_shift_target = float( + transfer_config.flow_shift + if transfer_config.flow_shift is not None + else self._get_sp_param(sp, "flow_shift", COSMOS3_V2V_DEFAULT_FLOW_SHIFT) + ) + max_sequence_length = ( + self._get_sp_param(sp, "max_sequence_length", COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH) + or COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH + ) + use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", False)) + + self._guidance_scale = guidance_scale + self._num_timesteps = num_inference_steps + self._set_flow_shift(flow_shift_target) + + generator = sp.generator + if generator is None: + seed = sp.seed if sp.seed is not None else 42 + generator = torch.Generator(device=self.device).manual_seed(seed) + + cond_ids, cond_mask, uncond_ids, uncond_mask = self._format_and_tokenize_prompts( + prompt, + negative_prompt, + chunk_frames, + frame_rate, + height, + width, + max_sequence_length, + sp, + use_system_prompt, + is_t2i=False, + ) + + output_chunks: list[torch.Tensor] = [] + control_chunks_per_hint: dict[str, list[torch.Tensor]] = {key: [] for key in per_hint_frames} + previous_output: torch.Tensor | None = None + + for chunk_id in range(num_chunks): + start_frame = chunk_id * stride + end_frame = min(start_frame + chunk_frames, total_frames) + control_norms = { + key: uint8_cthw_to_normalized_5d( + pad_temporal_frames(frames[:, start_frame:end_frame], chunk_frames), + dtype=self.dtype, + ) + for key, frames in per_hint_frames.items() + } + target_norm = torch.zeros_like(next(iter(control_norms.values()))) + current_conditional_frames = 0 + + if chunk_id == 0 and transfer_config.num_first_chunk_conditional_frames > 0: + if input_frames is None: + raise ValueError("Cosmos3 transfer num_first_chunk_conditional_frames > 0 requires a video input.") + current_conditional_frames = min( + transfer_config.num_first_chunk_conditional_frames, + input_frames.shape[1], + chunk_frames, + ) + if current_conditional_frames > 0: + input_cond = uint8_cthw_to_normalized_5d( + input_frames[:, :current_conditional_frames], + dtype=self.dtype, + ) + target_norm[:, :, :current_conditional_frames] = input_cond + if current_conditional_frames < chunk_frames: + fill = target_norm[:, :, current_conditional_frames - 1 : current_conditional_frames] + target_norm[:, :, current_conditional_frames:] = fill.expand( + -1, + -1, + chunk_frames - current_conditional_frames, + -1, + -1, + ) + elif chunk_id > 0 and previous_output is not None: + current_conditional_frames = min( + transfer_config.num_conditional_frames, + previous_output.shape[2], + chunk_frames, + ) + if current_conditional_frames > 0: + target_norm[:, :, :current_conditional_frames] = previous_output[ + :, :, -current_conditional_frames: + ].to(target_norm) + if current_conditional_frames < chunk_frames: + fill = target_norm[:, :, current_conditional_frames - 1 : current_conditional_frames] + target_norm[:, :, current_conditional_frames:] = fill.expand( + -1, + -1, + chunk_frames - current_conditional_frames, + -1, + -1, + ) + + control_latents = [self._encode_video_tensor(video) for video in control_norms.values()] + latents, velocity_mask, condition_latents = self._prepare_transfer_latents( + target_norm, + current_conditional_frames, + generator, + ) + video_shape = (latents.shape[2], latents.shape[3], latents.shape[4]) + shared_kwargs = dict( + video_shape=video_shape, + fps=frame_rate, + noisy_frame_mask=velocity_mask, + transfer_share_vision_temporal_positions=transfer_config.share_vision_temporal_positions, + ) + + self.scheduler.set_timesteps(num_inference_steps, device=self.device) + latents = self.diffuse_transfer( + latents=latents, + timesteps=self.scheduler.timesteps, + cond_ids=cond_ids, + cond_mask=cond_mask, + uncond_ids=uncond_ids, + uncond_mask=uncond_mask, + guidance_scale=guidance_scale, + control_guidance=transfer_config.control_guidance, + control_guidance_interval=transfer_config.control_guidance_interval, + control_latents=control_latents, + shared_kwargs=shared_kwargs, + velocity_mask=velocity_mask, + condition_latents=condition_latents, + ) + output_video = self._decode_latents(latents).clamp(-1, 1) + previous_output = output_video + + if chunk_id == 0: + output_chunks.append(output_video) + for key, control in control_norms.items(): + control_chunks_per_hint[key].append(control) + else: + output_chunks.append(output_video[:, :, current_conditional_frames:]) + for key, control in control_norms.items(): + control_chunks_per_hint[key].append(control[:, :, current_conditional_frames:]) + + full_output = torch.cat(output_chunks, dim=2)[:, :, :total_frames] + full_controls = { + key: torch.cat(chunks, dim=2)[:, :, :total_frames] for key, chunks in control_chunks_per_hint.items() + } + + if transfer_config.show_control_condition: + all_controls = torch.cat([full_controls[key] for key in per_hint_frames], dim=-1) + all_controls = all_controls.to(full_output) + full_output = torch.cat([all_controls, full_output], dim=-1) + if transfer_config.show_input and input_frames is not None: + normalized_input = uint8_cthw_to_normalized_5d(input_frames[:, :total_frames], dtype=torch.float32) + full_output = torch.cat([normalized_input.to(full_output), full_output], dim=-1) + + return DiffusionOutput( + output={"video": full_output}, + custom_output={ + "transfer_controls": full_controls, + "transfer_hints": list(per_hint_frames), + "fps": frame_rate, + }, + ) + # -- Forward (main generation entry point) ------------------------------- def forward( @@ -2199,18 +2777,30 @@ class Cosmos3OmniDiffusersPipeline( negative_prompt = None image_tensor = None video_tensor = None + transfer_video_tensor = None + transfer_input_fps = None else: prompt = prompt_data.get("prompt", "") negative_prompt = prompt_data.get("negative_prompt") additional_info = prompt_data.get("additional_information", {}) or {} image_tensor = additional_info.get("preprocessed_image") video_tensor = additional_info.get("preprocessed_video") + transfer_video_tensor = additional_info.get("preprocessed_transfer_video") + transfer_input_fps = positive_float(additional_info.get("transfer_input_fps")) is_t2i = self._is_t2i_request(req) sound_enabled = self._is_sound_request(prompt_data, sp) action_mode = self._get_action_mode(prompt_data, sp) action_enabled = action_mode is not None + transfer_config = resolve_transfer_config(sp, prompt_data) action_video_tensor = video_tensor if action_enabled else None + if transfer_config is not None: + if is_t2i: + raise ValueError("Cosmos3 transfer inference is supported only for video outputs.") + if action_enabled: + raise ValueError("Cosmos3 transfer inference cannot be combined with action generation.") + if sound_enabled: + raise ValueError("Cosmos3 transfer inference cannot be combined with sound generation.") if action_enabled and is_t2i: raise ValueError("Cosmos3 action generation is supported only for video outputs.") if action_enabled and sound_enabled: @@ -2234,6 +2824,19 @@ class Cosmos3OmniDiffusersPipeline( ) if negative_prompt is None: negative_prompt = "" + if transfer_config is not None: + if image_tensor is not None: + raise ValueError("Cosmos3 transfer inference accepts video inputs or control_path values, not images.") + if transfer_video_tensor is None and video_tensor is not None: + transfer_video_tensor = video_tensor + return self._forward_transfer( + prompt=prompt, + negative_prompt=negative_prompt, + sp=sp, + transfer_config=transfer_config, + transfer_video_tensor=transfer_video_tensor, + transfer_input_fps=transfer_input_fps, + ) if image_tensor is not None and video_tensor is not None and not action_enabled: raise ValueError("Cosmos3 non-action generation accepts either image or video input, not both.") if video_tensor is not None and is_t2i: @@ -2525,14 +3128,19 @@ class Cosmos3OmniDiffusersPipeline( video = self._decode_latents(latents) if _is_rank_zero(): logger.info("Video decoded in %.2fs", time.time() - decode_start) - logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start) + if not sound_enabled: + logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start) if sound_enabled: if sound_latents is None or target_audio_samples is None or sound_sample_rate is None: raise ValueError("Cosmos3 sound generation finished without sound latents.") if _is_rank_zero(): logger.info("Decoding sound...") + sound_decode_start = time.time() audio = self._decode_sound_latents(sound_latents, target_audio_samples) + if _is_rank_zero(): + logger.info("Sound tokenizer decoded in %.2fs", time.time() - sound_decode_start) + logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start) return DiffusionOutput(output={"video": video, "audio": audio, "audio_sample_rate": sound_sample_rate}) if action_enabled: diff --git a/vllm_omni/diffusion/models/cosmos3/transfer.py b/vllm_omni/diffusion/models/cosmos3/transfer.py new file mode 100644 index 000000000..a7930a905 --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/transfer.py @@ -0,0 +1,625 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cosmos3 transfer inference helpers. + +The reference Cosmos Framework transfer path accepts one or more control hints +and packs them before the noisy target video. This module keeps the vLLM Omni +pipeline-specific pieces small: request parsing, media normalization, and +optional on-the-fly edge/blur control generation. +""" + +from __future__ import annotations + +import importlib +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np +import PIL.Image +import torch +import torch.nn.functional as F + +TRANSFER_HINT_KEYS: tuple[str, ...] = ("edge", "blur", "depth", "seg", "wsm") +TRANSFER_SAMPLE_DEFAULTS: dict[str, Any] = { + "num_video_frames_per_chunk": 93, + "num_conditional_frames": 1, + "max_frames": 5000, + "show_control_condition": False, + "show_input": False, + "num_first_chunk_conditional_frames": 0, + "share_vision_temporal_positions": True, +} +TRANSFER_DEFAULTS: dict[str, dict[str, Any]] = { + "edge": {"guidance_scale": 3.0, "control_guidance": 1.5, "flow_shift": 10.0}, + "blur": {"guidance_scale": 3.0, "control_guidance": 1.5, "flow_shift": 10.0}, + "depth": {"guidance_scale": 3.0, "control_guidance": 1.5, "flow_shift": 10.0}, + "seg": {"guidance_scale": 3.0, "control_guidance": 2.0, "flow_shift": 10.0}, + "wsm": { + "guidance_scale": 1.0, + "control_guidance": 3.0, + "flow_shift": 10.0, + "num_frames": 101, + "fps": 10, + "num_video_frames_per_chunk": 101, + }, +} +EDGE_PRESETS: dict[str, tuple[int, int]] = { + "none": (20, 50), + "very_low": (20, 50), + "low": (50, 100), + "medium": (100, 200), + "high": (200, 300), + "very_high": (300, 400), +} +BLUR_DOWNUP_PRESETS: dict[str, int] = { + "none": 1, + "very_low": 4, + "low": 4, + "medium": 10, + "high": 16, + "very_high": 16, +} +BLUR_PRE_BLUR_DOWNSCALE_PRESETS: dict[str, int] = { + "none": 1, + "very_low": 1, + "low": 4, + "medium": 2, + "high": 1, + "very_high": 4, +} +BILATERAL_REFERENCE_RESOLUTION = 720 +BILATERAL_D = 30 +BILATERAL_SIGMA_COLOR = 150 +BILATERAL_SIGMA_SPACE = 100 +BILATERAL_ITERATIONS = 1 +IMAGE_EXTENSIONS = {".bmp", ".gif", ".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp"} + + +@dataclass +class Cosmos3TransferHint: + key: str + control_path: str | None = None + control: Any | None = None + preset_edge_threshold: str = "medium" + preset_blur_strength: str = "medium" + + +@dataclass +class Cosmos3TransferConfig: + hints: dict[str, Cosmos3TransferHint] = field(default_factory=dict) + guidance_scale: float | None = None + control_guidance: float = 1.0 + control_guidance_interval: tuple[float, float] | None = None + flow_shift: float | None = None + num_video_frames_per_chunk: int = 93 + num_conditional_frames: int = 1 + max_frames: int = 5000 + show_control_condition: bool = False + show_input: bool = False + num_first_chunk_conditional_frames: int = 0 + share_vision_temporal_positions: bool = True + num_frames: int | None = None + fps: float | None = None + + @property + def ordered_hints(self) -> list[Cosmos3TransferHint]: + return [self.hints[key] for key in TRANSFER_HINT_KEYS if key in self.hints] + + +def _extra_args(sp: Any) -> Mapping[str, Any]: + extra = getattr(sp, "extra_args", None) + return extra if isinstance(extra, Mapping) else {} + + +def _prompt_value(prompt_data: Any, key: str) -> Any: + if not isinstance(prompt_data, Mapping): + return None + if prompt_data.get(key) is not None: + return prompt_data[key] + additional = prompt_data.get("additional_information") + if isinstance(additional, Mapping): + return additional.get(key) + return None + + +def _param(extra: Mapping[str, Any], sp: Any, prompt_data: Any, key: str, default: Any = None) -> Any: + if extra.get(key) is not None: + return extra[key] + value = getattr(sp, key, None) + if value is not None: + return value + value = _prompt_value(prompt_data, key) + return default if value is None else value + + +def _param_any(extra: Mapping[str, Any], sp: Any, prompt_data: Any, keys: tuple[str, ...], default: Any = None) -> Any: + for key in keys: + value = _param(extra, sp, prompt_data, key, None) + if value is not None: + return value + return default + + +def _is_user_field(extra: Mapping[str, Any], sp: Any, prompt_data: Any, *keys: str) -> bool: + for key in keys: + if key in extra and extra[key] is not None: + return True + if getattr(sp, key, None) is not None: + return True + if _prompt_value(prompt_data, key) is not None: + return True + return False + + +# fps, frame_rate, and resolved_frame_rate are aliases for the same request-level +# control. The serving layer leaves all of them ``None`` on the sampling params when +# the user did not provide fps, so a non-``None`` value here unambiguously means the +# user set it. +_FPS_KEYS: tuple[str, ...] = ("resolved_frame_rate", "frame_rate", "fps") + + +def _transfer_fps_param(extra: Mapping[str, Any], sp: Any, prompt_data: Any) -> Any: + # extra_args is a model-specific override and wins over the generic sampling params. + for key in _FPS_KEYS: + if extra.get(key) is not None: + return extra[key] + return _param_any(extra, sp, prompt_data, _FPS_KEYS) + + +def _as_bool(value: Any, default: bool = False) -> bool: + if value is None: + return default + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on"} + return bool(value) + + +def _as_interval(value: Any) -> tuple[float, float] | None: + if value is None: + return None + if isinstance(value, str): + value = [item.strip() for item in value.split(",") if item.strip()] + if not isinstance(value, (list, tuple)) or len(value) != 2: + raise ValueError("Cosmos3 transfer control_guidance_interval must contain exactly two values.") + lo, hi = float(value[0]), float(value[1]) + if lo > hi: + raise ValueError(f"Cosmos3 transfer control_guidance_interval must be ordered as [lo, hi], got {(lo, hi)}.") + return lo, hi + + +def has_transfer_hints(extra_args: Mapping[str, Any] | None) -> bool: + if not isinstance(extra_args, Mapping): + return False + return any(extra_args.get(key) is not None for key in TRANSFER_HINT_KEYS) + + +def transfer_max_frames_from_extra_args(extra_args: Mapping[str, Any] | None) -> int: + if not isinstance(extra_args, Mapping): + return int(TRANSFER_SAMPLE_DEFAULTS["max_frames"]) + value = extra_args.get("max_frames", TRANSFER_SAMPLE_DEFAULTS["max_frames"]) + return max(1, int(value)) + + +def resolve_transfer_config(sp: Any, prompt_data: Any = None) -> Cosmos3TransferConfig | None: + extra = _extra_args(sp) + hints: dict[str, Cosmos3TransferHint] = {} + for key in TRANSFER_HINT_KEYS: + raw = _param(extra, sp, prompt_data, key, None) + if raw is None: + continue + if raw is True: + raw = {} + elif isinstance(raw, str | Path): + raw = {"control_path": str(raw)} + if not isinstance(raw, Mapping): + raise TypeError( + f"Cosmos3 transfer hint '{key}' must be an object, path string, or true; got {type(raw)!r}." + ) + hints[key] = Cosmos3TransferHint( + key=key, + control_path=str(raw["control_path"]) if raw.get("control_path") is not None else None, + control=raw.get("control"), + preset_edge_threshold=str(raw.get("preset_edge_threshold") or "medium").lower(), + preset_blur_strength=str(raw.get("preset_blur_strength") or "medium").lower(), + ) + + if not hints: + transfer_only = ( + "control_guidance", + "control_guidance_interval", + "num_video_frames_per_chunk", + "num_conditional_frames", + "max_frames", + "show_control_condition", + "show_input", + "num_first_chunk_conditional_frames", + "share_vision_temporal_positions", + ) + if any(_is_user_field(extra, sp, prompt_data, key) for key in transfer_only): + raise ValueError("Cosmos3 transfer options were provided, but no transfer hint was selected.") + return None + + fps_value = _transfer_fps_param(extra, sp, prompt_data) + config = Cosmos3TransferConfig( + hints=hints, + guidance_scale=( + float(_param_any(extra, sp, prompt_data, ("guidance_scale", "guidance"))) + if _param_any(extra, sp, prompt_data, ("guidance_scale", "guidance")) is not None + else None + ), + control_guidance=float(_param(extra, sp, prompt_data, "control_guidance", 1.0)), + control_guidance_interval=_as_interval(_param(extra, sp, prompt_data, "control_guidance_interval", None)), + flow_shift=( + float(_param_any(extra, sp, prompt_data, ("flow_shift", "shift"))) + if _param_any(extra, sp, prompt_data, ("flow_shift", "shift")) is not None + else None + ), + num_video_frames_per_chunk=int( + _param( + extra, + sp, + prompt_data, + "num_video_frames_per_chunk", + TRANSFER_SAMPLE_DEFAULTS["num_video_frames_per_chunk"], + ) + ), + num_conditional_frames=int( + _param(extra, sp, prompt_data, "num_conditional_frames", TRANSFER_SAMPLE_DEFAULTS["num_conditional_frames"]) + ), + max_frames=int(_param(extra, sp, prompt_data, "max_frames", TRANSFER_SAMPLE_DEFAULTS["max_frames"])), + show_control_condition=_as_bool( + _param(extra, sp, prompt_data, "show_control_condition", None), + bool(TRANSFER_SAMPLE_DEFAULTS["show_control_condition"]), + ), + show_input=_as_bool( + _param(extra, sp, prompt_data, "show_input", None), + bool(TRANSFER_SAMPLE_DEFAULTS["show_input"]), + ), + num_first_chunk_conditional_frames=int( + _param( + extra, + sp, + prompt_data, + "num_first_chunk_conditional_frames", + TRANSFER_SAMPLE_DEFAULTS["num_first_chunk_conditional_frames"], + ) + ), + share_vision_temporal_positions=_as_bool( + _param(extra, sp, prompt_data, "share_vision_temporal_positions", None), + bool(TRANSFER_SAMPLE_DEFAULTS["share_vision_temporal_positions"]), + ), + num_frames=( + int(_param(extra, sp, prompt_data, "num_frames")) + if _param(extra, sp, prompt_data, "num_frames", None) is not None + else None + ), + fps=(float(fps_value) if fps_value is not None else None), + ) + + if len(hints) == 1: + hint_key = next(iter(hints)) + for field_name, default_value in TRANSFER_DEFAULTS[hint_key].items(): + if field_name == "guidance_scale": + user_set = _is_user_field(extra, sp, prompt_data, "guidance_scale", "guidance") + elif field_name == "flow_shift": + user_set = _is_user_field(extra, sp, prompt_data, "flow_shift", "shift") + elif field_name == "fps": + user_set = _is_user_field(extra, sp, prompt_data, *_FPS_KEYS) + else: + user_set = _is_user_field(extra, sp, prompt_data, field_name) + if not user_set: + setattr(config, field_name, default_value) + + if config.num_video_frames_per_chunk <= 0: + raise ValueError("Cosmos3 transfer num_video_frames_per_chunk must be positive.") + if config.num_conditional_frames < 0: + raise ValueError("Cosmos3 transfer num_conditional_frames must be non-negative.") + if config.max_frames <= 0: + raise ValueError("Cosmos3 transfer max_frames must be positive.") + if config.num_first_chunk_conditional_frames < 0: + raise ValueError("Cosmos3 transfer num_first_chunk_conditional_frames must be non-negative.") + for hint in hints.values(): + if hint.key == "edge" and hint.preset_edge_threshold not in EDGE_PRESETS: + raise ValueError(f"Unsupported Cosmos3 edge preset: {hint.preset_edge_threshold!r}.") + if hint.key == "blur" and hint.preset_blur_strength not in BLUR_DOWNUP_PRESETS: + raise ValueError(f"Unsupported Cosmos3 blur preset: {hint.preset_blur_strength!r}.") + return config + + +def _import_cv2(hint_key: str): + try: + return importlib.import_module("cv2") + except ImportError as exc: + raise ImportError( + f"Cosmos3 transfer hint '{hint_key}' requires opencv-python for on-the-fly control generation. " + "Install opencv-python in the environment, or provide a precomputed control_path for this hint." + ) from exc + + +def pad_temporal_frames(frames: torch.Tensor, target_frames: int) -> torch.Tensor: + if frames.ndim != 4: + raise ValueError(f"Cosmos3 transfer frames must have shape [C, T, H, W], got {tuple(frames.shape)}.") + target_frames = int(target_frames) + if target_frames <= 0: + raise ValueError("Cosmos3 transfer target frame count must be positive.") + if frames.shape[1] >= target_frames: + return frames + if frames.shape[1] == 0: + raise ValueError("Cannot pad an empty Cosmos3 transfer frame tensor.") + padded = frames + while padded.shape[1] < target_frames: + pad_len = min(padded.shape[1] - 1, target_frames - padded.shape[1]) + if pad_len <= 0: + pad_frame = padded[:, -1:].repeat(1, target_frames - padded.shape[1], 1, 1) + padded = torch.cat([padded, pad_frame], dim=1) + break + padded = torch.cat([padded, padded.flip(dims=[1])[:, :pad_len]], dim=1) + return padded + + +def normalized_video_to_uint8_cthw(video: torch.Tensor) -> torch.Tensor: + tensor = video.detach().cpu() + if tensor.ndim == 5: + if tensor.shape[0] != 1: + raise ValueError("Cosmos3 transfer supports only batch size 1 video inputs.") + tensor = tensor[0] + if tensor.ndim != 4: + raise ValueError( + f"Cosmos3 transfer video must have shape [1, 3, T, H, W] or [3, T, H, W], got {tuple(video.shape)}." + ) + if tensor.shape[0] != 3: + if tensor.shape[-1] in (3, 4): + tensor = tensor[..., :3].permute(3, 0, 1, 2) + else: + raise ValueError(f"Cosmos3 transfer video must have 3 RGB channels, got {tuple(video.shape)}.") + if tensor.is_floating_point(): + if tensor.numel() and (tensor.min().item() < 0.0 or tensor.max().item() > 1.0): + tensor = tensor.clamp(-1.0, 1.0).mul(0.5).add(0.5) + tensor = tensor.clamp(0.0, 1.0).mul(255.0).round().to(torch.uint8) + else: + tensor = tensor.clamp(0, 255).to(torch.uint8) + return tensor.contiguous() + + +def uint8_cthw_to_normalized_5d(frames: torch.Tensor, *, dtype: torch.dtype) -> torch.Tensor: + if frames.ndim != 4 or frames.shape[0] != 3: + raise ValueError(f"Cosmos3 transfer frames must have shape [3, T, H, W], got {tuple(frames.shape)}.") + return frames.to(dtype=dtype).div(127.5).sub(1.0).unsqueeze(0).contiguous() + + +def _pil_to_uint8_rgb(value: Any) -> np.ndarray: + if isinstance(value, PIL.Image.Image): + return np.asarray(value.convert("RGB"), dtype=np.uint8) + if isinstance(value, str | Path): + return np.asarray(PIL.Image.open(value).convert("RGB"), dtype=np.uint8) + if isinstance(value, torch.Tensor): + tensor = value.detach().cpu() + if tensor.ndim == 3 and tensor.shape[0] in (3, 4): + tensor = tensor[:3].permute(1, 2, 0) + if tensor.is_floating_point(): + if tensor.numel() and (tensor.min().item() < 0.0 or tensor.max().item() > 1.0): + tensor = tensor.clamp(-1.0, 1.0).mul(0.5).add(0.5) + tensor = tensor.clamp(0.0, 1.0).mul(255.0).round().to(torch.uint8) + return tensor[..., :3].numpy().astype(np.uint8) + if isinstance(value, np.ndarray): + array = value + if array.ndim == 3 and array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4): + array = np.transpose(array[:3], (1, 2, 0)) + if np.issubdtype(array.dtype, np.floating): + if array.size and (array.min() < 0.0 or array.max() > 1.0): + array = np.clip(array, -1.0, 1.0) * 0.5 + 0.5 + array = (np.clip(array, 0.0, 1.0) * 255.0).round().astype(np.uint8) + return array[..., :3].astype(np.uint8) + raise TypeError(f"Cosmos3 transfer expected an RGB frame, got {type(value)!r}.") + + +def media_hw(value: Any) -> tuple[int, int] | None: + if isinstance(value, PIL.Image.Image): + return value.height, value.width + if isinstance(value, str | Path): + media_path = Path(value) + if media_path.suffix.lower() in IMAGE_EXTENSIONS: + with PIL.Image.open(media_path) as image: + return image.height, image.width + try: + import imageio.v3 as iio + except ImportError: + return None + for frame in iio.imiter(media_path): + array = np.asarray(frame) + if array.ndim >= 2: + return int(array.shape[0]), int(array.shape[1]) + return None + if isinstance(value, torch.Tensor): + tensor = value.detach() + if tensor.ndim == 5: + tensor = tensor[0] + if tensor.ndim == 4: + if tensor.shape[0] in (3, 4) and tensor.shape[-1] not in (3, 4): + return int(tensor.shape[-2]), int(tensor.shape[-1]) + if tensor.shape[-1] in (3, 4): + return int(tensor.shape[1]), int(tensor.shape[2]) + if tensor.ndim == 3: + if tensor.shape[0] in (3, 4) and tensor.shape[-1] not in (3, 4): + return int(tensor.shape[-2]), int(tensor.shape[-1]) + if tensor.shape[-1] in (3, 4): + return int(tensor.shape[0]), int(tensor.shape[1]) + return None + if isinstance(value, np.ndarray): + array = value + if array.ndim == 5: + array = array[0] + if array.ndim == 4: + if array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4): + return int(array.shape[-2]), int(array.shape[-1]) + if array.shape[-1] in (3, 4): + return int(array.shape[1]), int(array.shape[2]) + if array.ndim == 3: + if array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4): + return int(array.shape[-2]), int(array.shape[-1]) + if array.shape[-1] in (3, 4): + return int(array.shape[0]), int(array.shape[1]) + return None + if isinstance(value, list | tuple) and value: + return media_hw(value[0]) + return None + + +def resize_center_crop_uint8_cthw(frames: torch.Tensor, height: int, width: int) -> torch.Tensor: + if frames.ndim != 4 or frames.shape[0] != 3: + raise ValueError(f"Cosmos3 transfer frames must have shape [3, T, H, W], got {tuple(frames.shape)}.") + orig_h, orig_w = int(frames.shape[2]), int(frames.shape[3]) + scale = max(width / orig_w, height / orig_h) + resize_h = int(np.ceil(scale * orig_h)) + resize_w = int(np.ceil(scale * orig_w)) + frames_tchw = frames.permute(1, 0, 2, 3).to(dtype=torch.float32) + resized = F.interpolate(frames_tchw, size=(resize_h, resize_w), mode="bilinear", align_corners=False) + top = (resize_h - height) // 2 + left = (resize_w - width) // 2 + cropped = resized[:, :, top : top + height, left : left + width] + return cropped.round().clamp(0, 255).to(torch.uint8).permute(1, 0, 2, 3).contiguous() + + +def _path_media_to_uint8_cthw(path: str | Path, max_frames: int | None) -> torch.Tensor: + media_path = Path(path) + if not media_path.exists(): + raise FileNotFoundError(f"Missing Cosmos3 transfer control_path: {media_path}") + if media_path.suffix.lower() in IMAGE_EXTENSIONS: + array = _pil_to_uint8_rgb(media_path) + return torch.from_numpy(array).permute(2, 0, 1).unsqueeze(1).contiguous() + + try: + import imageio.v3 as iio + except ImportError as exc: + raise ImportError( + "Cosmos3 transfer video control_path loading requires imageio. " + "Install imageio[ffmpeg] or provide decoded control frames." + ) from exc + + frames: list[torch.Tensor] = [] + limit = max_frames if max_frames is not None else None + for frame in iio.imiter(media_path): + frames.append(torch.from_numpy(_pil_to_uint8_rgb(frame)).permute(2, 0, 1)) + if limit is not None and len(frames) >= int(limit): + break + if not frames: + raise ValueError(f"Cosmos3 transfer control_path produced no frames: {media_path}") + return torch.stack(frames, dim=1).contiguous() + + +def media_to_uint8_cthw(value: Any, *, height: int, width: int, max_frames: int | None = None) -> torch.Tensor: + if isinstance(value, str | Path): + frames = _path_media_to_uint8_cthw(value, max_frames=max_frames) + elif isinstance(value, torch.Tensor): + frames = normalized_video_to_uint8_cthw(value) + elif isinstance(value, np.ndarray): + frames = normalized_video_to_uint8_cthw(torch.from_numpy(value)) + elif isinstance(value, list | tuple): + if not value: + raise ValueError("Cosmos3 transfer control frames cannot be empty.") + selected = value[:max_frames] if max_frames is not None else value + tensors = [torch.from_numpy(_pil_to_uint8_rgb(frame)).permute(2, 0, 1) for frame in selected] + frames = torch.stack(tensors, dim=1).contiguous() + else: + raise TypeError(f"Unsupported Cosmos3 transfer control payload type: {type(value)!r}.") + if max_frames is not None: + frames = frames[:, : int(max_frames)] + return resize_center_crop_uint8_cthw(frames, int(height), int(width)) + + +def make_edge_control(frames: torch.Tensor, preset: str) -> torch.Tensor: + cv2 = _import_cv2("edge") + try: + lower, upper = EDGE_PRESETS[preset] + except KeyError as exc: + raise ValueError(f"Unsupported Cosmos3 edge preset: {preset!r}.") from exc + frames_np = frames.detach().cpu().numpy().astype(np.uint8) + edge_maps = [] + for idx in range(frames_np.shape[1]): + frame = np.ascontiguousarray(np.transpose(frames_np[:, idx], (1, 2, 0))) + edge_maps.append(cv2.Canny(frame, lower, upper)) + edge = np.stack(edge_maps, axis=0)[None] + return torch.from_numpy(edge).expand(3, -1, -1, -1).contiguous() + + +def _scale_for_bilateral_resolution(value: float, longest_side: int) -> float: + if longest_side <= 0: + return value + return value * (longest_side / BILATERAL_REFERENCE_RESOLUTION) + + +def _scaled_bilateral_params(height: int, width: int) -> tuple[int, float, float]: + longest_side = int(max(height, width)) + diameter = max(1, int(round(_scale_for_bilateral_resolution(float(BILATERAL_D), longest_side)))) + if diameter % 2 == 0: + diameter += 1 + sigma_color = max(1.0, _scale_for_bilateral_resolution(float(BILATERAL_SIGMA_COLOR), longest_side)) + sigma_space = max(1.0, _scale_for_bilateral_resolution(float(BILATERAL_SIGMA_SPACE), longest_side)) + return diameter, sigma_color, sigma_space + + +def make_blur_control(frames: torch.Tensor, preset: str) -> torch.Tensor: + cv2 = _import_cv2("blur") + preset = preset.lower() + if preset not in BLUR_DOWNUP_PRESETS: + raise ValueError(f"Unsupported Cosmos3 blur preset: {preset!r}.") + if preset == "none": + return frames.clone() + + frames_np = frames.detach().cpu().numpy().astype(np.uint8) + _, t, h, w = frames_np.shape + pre_blur_factor = max(1, int(BLUR_PRE_BLUR_DOWNSCALE_PRESETS[preset])) + downup_factor = max(1, int(BLUR_DOWNUP_PRESETS[preset])) + result = np.empty_like(frames_np) + for idx in range(t): + frame = np.ascontiguousarray(np.transpose(frames_np[:, idx], (1, 2, 0))) + if pre_blur_factor > 1: + small_w = max(1, w // pre_blur_factor) + small_h = max(1, h // pre_blur_factor) + frame = cv2.resize(frame, (small_w, small_h), interpolation=cv2.INTER_AREA) + diameter, sigma_color, sigma_space = _scaled_bilateral_params(frame.shape[0], frame.shape[1]) + for _ in range(BILATERAL_ITERATIONS): + frame = cv2.bilateralFilter(frame, diameter, sigma_color, sigma_space) + if pre_blur_factor > 1: + frame = cv2.resize(frame, (w, h), interpolation=cv2.INTER_LINEAR) + if downup_factor > 1: + small_w = max(1, w // downup_factor) + small_h = max(1, h // downup_factor) + frame = cv2.resize(frame, (small_w, small_h), interpolation=cv2.INTER_CUBIC) + frame = cv2.resize(frame, (w, h), interpolation=cv2.INTER_CUBIC) + result[:, idx] = np.transpose(frame, (2, 0, 1)) + return torch.from_numpy(result).contiguous() + + +def load_or_compute_control_frames( + hint: Cosmos3TransferHint, + *, + height: int, + width: int, + max_frames: int, + input_frames: torch.Tensor | None, +) -> torch.Tensor: + if hint.control is not None: + return media_to_uint8_cthw(hint.control, height=height, width=width, max_frames=max_frames) + if hint.control_path is not None: + return media_to_uint8_cthw(hint.control_path, height=height, width=width, max_frames=max_frames) + if hint.key == "edge": + if input_frames is None: + raise ValueError( + "Cosmos3 transfer hint 'edge' requires either a video input for on-the-fly control generation " + "or a precomputed control_path." + ) + return make_edge_control(input_frames[:, :max_frames], hint.preset_edge_threshold) + if hint.key == "blur": + if input_frames is None: + raise ValueError( + "Cosmos3 transfer hint 'blur' requires either a video input for on-the-fly control generation " + "or a precomputed control_path." + ) + return make_blur_control(input_frames[:, :max_frames], hint.preset_blur_strength) + raise FileNotFoundError( + f"Cosmos3 transfer hint '{hint.key}' requires a precomputed control_path; " + "on-the-fly generation is supported only for edge and blur." + ) diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 4543d6922..3eef62621 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -1234,8 +1234,12 @@ class Cosmos3VFMTransformer(nn.Module): action_start_frame_offset: int = 1, action_fps: float | None = None, t_sound: int | None = None, + num_vision_items: int = 1, + share_vision_temporal_positions: bool = False, ) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]: """Compute mRoPE cos/sin for UND text and GEN media pathways.""" + if num_vision_items <= 0: + raise ValueError(f"Cosmos3 num_vision_items must be positive, got {num_vision_items}.") B = text_mask.shape[0] S_text = text_mask.shape[1] text_lengths = text_mask.sum(dim=1).long() @@ -1249,17 +1253,33 @@ class Cosmos3VFMTransformer(nn.Module): real_len = int(text_lengths[b].item()) t_pos, t_offset = compute_mrope_position_ids_text(real_len, temporal_offset=0) media_temporal_offset = t_offset + self.temporal_modality_margin - v_pos, _ = compute_mrope_position_ids_vision( - t, - hp, - wp, - temporal_offset=media_temporal_offset, - fps=effective_fps, - base_fps=self.base_fps, - temporal_compression_factor=self.temporal_compression_factor, - enable_fps_modulation=self.enable_fps_modulation, - ) - gen_positions = [v_pos] + gen_positions = [] + if num_vision_items == 1 or share_vision_temporal_positions: + v_pos, _ = compute_mrope_position_ids_vision( + t, + hp, + wp, + temporal_offset=media_temporal_offset, + fps=effective_fps, + base_fps=self.base_fps, + temporal_compression_factor=self.temporal_compression_factor, + enable_fps_modulation=self.enable_fps_modulation, + ) + gen_positions.extend([v_pos] * num_vision_items) + else: + vision_offset: int | float = media_temporal_offset + for _ in range(num_vision_items): + v_pos, vision_offset = compute_mrope_position_ids_vision( + t, + hp, + wp, + temporal_offset=vision_offset, + fps=effective_fps, + base_fps=self.base_fps, + temporal_compression_factor=self.temporal_compression_factor, + enable_fps_modulation=self.enable_fps_modulation, + ) + gen_positions.append(v_pos) if action_frames > 0: a_pos, _ = compute_mrope_position_ids_action( action_frames, @@ -1316,16 +1336,21 @@ class Cosmos3VFMTransformer(nn.Module): *, s_gen: int, s_video: int, + s_control: int, s_action: int, s_sound: int, has_action: bool, has_sound: bool, + has_control: bool, ulysses_size: int, ) -> None: if ulysses_size <= 1 or s_gen % ulysses_size == 0: return - detail_parts = [f"video tokens {s_video}"] + detail_parts = [] + if has_control: + detail_parts.append(f"control tokens {s_control}") + detail_parts.append(f"video tokens {s_video}") if has_action: detail_parts.append(f"action tokens {s_action}") if has_sound: @@ -1335,7 +1360,7 @@ class Cosmos3VFMTransformer(nn.Module): "Adjust the spatial resolution, frame count, action chunk size, " "sound duration, or sound latent FPS so the combined media sequence is a " "multiple of ulysses_degree." - if has_action or has_sound + if has_control or has_action or has_sound else ( "Adjust the spatial resolution so that " "t * ceil(h/patch) * ceil(w/patch) is a multiple " @@ -1364,6 +1389,8 @@ class Cosmos3VFMTransformer(nn.Module): action_fps: float | None = None, sound_latents: torch.Tensor | None = None, noisy_frame_mask: torch.Tensor | None = None, + control_latents: list[torch.Tensor] | tuple[torch.Tensor, ...] | torch.Tensor | None = None, + transfer_share_vision_temporal_positions: bool = True, **kwargs, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """ @@ -1383,11 +1410,15 @@ class Cosmos3VFMTransformer(nn.Module): timestep embedding, predict velocity) and 0=conditioned (clean context, skip timestep embedding). None means all frames noisy (T2V mode). + control_latents: Optional transfer-control latents. Controls are + clean vision context and are packed before the noisy target. Returns: [B, C, t, h, w] velocity prediction, or tuple outputs in video, action, sound order when extra modalities are provided. """ + if kwargs: + raise TypeError(f"Unexpected Cosmos3 transformer kwargs: {sorted(kwargs)}") t, h, w = video_shape hp, wp, _, _ = self._pad_to_patch_size(h, w) text_lengths = text_mask.sum(dim=1) @@ -1400,6 +1431,15 @@ class Cosmos3VFMTransformer(nn.Module): ) has_action = action_latents is not None has_sound = sound_latents is not None + if control_latents is None: + control_latent_list: list[torch.Tensor] = [] + elif isinstance(control_latents, torch.Tensor): + control_latent_list = [control_latents] + else: + control_latent_list = list(control_latents) + has_control = len(control_latent_list) > 0 + if has_control and (has_action or has_sound): + raise ValueError("Cosmos3 transfer control latents cannot be combined with action or sound latents.") if has_action and not self.action_gen: raise ValueError( "Cosmos3 action generation was requested, but this transformer " @@ -1419,6 +1459,19 @@ class Cosmos3VFMTransformer(nn.Module): # Patchify latents and project to hidden space hidden_video = self.proj_in(self.patchify(hidden_states, t, h, w)) s_video = hidden_video.shape[1] + s_control = 0 + hidden_controls: list[torch.Tensor] = [] + for idx, control in enumerate(control_latent_list): + if control.shape != hidden_states.shape: + raise ValueError( + "Cosmos3 transfer control latent shape must match target latent shape: " + f"control[{idx}]={tuple(control.shape)}, target={tuple(hidden_states.shape)}." + ) + hidden_control = self.proj_in( + self.patchify(control.to(device=hidden_states.device, dtype=hidden_states.dtype), t, h, w) + ) + hidden_controls.append(hidden_control) + s_control += hidden_control.shape[1] s_action = 0 hidden_action = None s_sound = 0 @@ -1479,7 +1532,7 @@ class Cosmos3VFMTransformer(nn.Module): if hidden_sound is not None: hidden_sound = hidden_sound + time_embed.unsqueeze(1) - hidden_parts = [hidden_video] + hidden_parts = [*hidden_controls, hidden_video] if hidden_action is not None: hidden_parts.append(hidden_action) if hidden_sound is not None: @@ -1500,6 +1553,8 @@ class Cosmos3VFMTransformer(nn.Module): action_start_frame_offset=action_start_frame_offset, action_fps=action_fps, t_sound=s_sound, + num_vision_items=len(control_latent_list) + 1, + share_vision_temporal_positions=transfer_share_vision_temporal_positions, ) cached_kv_full = self.language_model(text_ids, freqs_und) self.cached_freqs_gen = freqs_gen @@ -1516,10 +1571,12 @@ class Cosmos3VFMTransformer(nn.Module): self._validate_gen_sequence_parallel( s_gen=hidden_gen.shape[1], s_video=s_video, + s_control=s_control, s_action=s_action, s_sound=s_sound, has_action=has_action, has_sound=has_sound, + has_control=has_control, ulysses_size=ulysses_size, ) freqs_cos, freqs_sin = self.cached_freqs_gen @@ -1553,19 +1610,27 @@ class Cosmos3VFMTransformer(nn.Module): # Final norm and project back to latent space hidden_gen = self.norm_moe_gen(hidden_gen) - if not has_action and not has_sound: + if not has_action and not has_sound and not has_control: return self.unpatchify(self.proj_out(hidden_gen), t, h, w) - split_sizes = [s_video] + split_sizes = [] + if has_control: + split_sizes.append(s_control) + split_sizes.append(s_video) if has_action: split_sizes.append(s_action) if has_sound: split_sizes.append(s_sound) split_hidden = hidden_gen.split(split_sizes, dim=1) - hidden_video = split_hidden[0] + split_idx = 0 + if has_control: + split_idx += 1 + hidden_video = split_hidden[split_idx] + split_idx += 1 video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w) + if has_control: + return video_pred outputs: list[torch.Tensor] = [video_pred] - split_idx = 1 if has_action: hidden_action = split_hidden[split_idx] split_idx += 1 diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 4bdbfce02..3e2927a27 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -128,6 +128,8 @@ class OmniOpenAIServingVideo: status_code=HTTPStatus.BAD_REQUEST.value, detail="Provide either an image reference or a video reference, not both.", ) + provided_fields = request.model_fields_set + fps_provided = self._request_fps_provided(request) vp = request.resolve_video_params() if input_image is not None and vp.width is not None and vp.height is not None: target_size = (vp.width, vp.height) @@ -147,10 +149,10 @@ class OmniOpenAIServingVideo: gen_params.height = vp.height if vp.num_frames is not None: gen_params.num_frames = vp.num_frames - if vp.fps is not None: + # Leave fps/frame_rate as None when the user did not provide fps. + if fps_provided and vp.fps is not None: gen_params.fps = vp.fps gen_params.frame_rate = float(vp.fps) - provided_fields = request.model_fields_set if "enable_frame_interpolation" in provided_fields: gen_params.enable_frame_interpolation = request.enable_frame_interpolation if "frame_interpolation_exp" in provided_fields: @@ -229,7 +231,9 @@ class OmniOpenAIServingVideo: audios = self._extract_audio_outputs(result, expected_count=len(videos)) actions = self._extract_action_outputs(result, expected_count=len(videos)) audio_sample_rate = self._resolve_audio_sample_rate(result) - output_fps = (vp.fps or self._resolve_fps(result) or 24) * self._resolve_video_fps_multiplier(result) + model_fps = self._resolve_fps(result) + output_fps_base = (vp.fps if fps_provided else None) or model_fps or vp.fps or 24 + output_fps = output_fps_base * self._resolve_video_fps_multiplier(result) return VideoGenerationArtifacts( videos=videos, audios=audios, @@ -344,6 +348,15 @@ class OmniOpenAIServingVideo: return int(multiplier) return 1 + @staticmethod + def _request_fps_provided(request: VideoGenerationRequest) -> bool: + if "fps" in request.model_fields_set and request.fps is not None: + return True + video_params = request.video_params + if video_params is None or "video_params" not in request.model_fields_set: + return False + return "fps" in video_params.model_fields_set and video_params.fps is not None + def _resolve_default_sampling_params(self) -> OmniDiffusionSamplingParams: default_sampling_params_list = getattr(self._engine_client, "default_sampling_params_list", None) if default_sampling_params_list: diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index aa4eeaba5..fbb37c36b 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -28,6 +28,29 @@ from vllm_omni.entrypoints.openai.protocol.videos import ( ) +class VideoFrames(list[Image.Image]): + """Decoded video frames plus source metadata.""" + + def __init__(self, frames: list[Image.Image] | None = None, *, fps: float | None = None) -> None: + super().__init__(frames or []) + self.fps = fps + self.frame_rate = fps + + +def positive_float(value: Any) -> float | None: + if value is None: + return None + if hasattr(value, "item") and not isinstance(value, (bytes, str)): + value = value.item() + try: + result = float(value) + except (TypeError, ValueError): + return None + if result <= 0: + return None + return result + + def _decode_image_bytes(image_bytes: bytes, *, source: str) -> Image.Image: try: return Image.open(BytesIO(image_bytes)).convert("RGB") @@ -41,7 +64,7 @@ def _decode_video_bytes( source: str, max_frames: int | None = None, keep: Literal["first", "last"] = "first", -) -> list[Image.Image]: +) -> VideoFrames: try: import av except ImportError as exc: # pragma: no cover - av is a serving dependency via media_utils @@ -56,8 +79,16 @@ def _decode_video_bytes( tail_frames: deque[Image.Image] | None = ( deque(maxlen=max_frames) if keep == "last" and max_frames is not None else None ) + fps: float | None = None try: with av.open(BytesIO(video_bytes)) as container: + video_stream = container.streams.video[0] if container.streams.video else None + if video_stream is not None: + fps = ( + positive_float(getattr(video_stream, "average_rate", None)) + or positive_float(getattr(video_stream, "base_rate", None)) + or positive_float(getattr(video_stream, "guessed_rate", None)) + ) for frame in container.decode(video=0): image = frame.to_image().convert("RGB") if tail_frames is not None: @@ -73,7 +104,7 @@ def _decode_video_bytes( frames = list(tail_frames) if not frames: raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.") - return frames + return VideoFrames(frames, fps=fps) def _decode_media_bytes( @@ -82,7 +113,7 @@ def _decode_media_bytes( source: str, max_video_frames: int | None = None, video_keep: Literal["first", "last"] = "first", -) -> Image.Image | list[Image.Image]: +) -> Image.Image | VideoFrames: try: return _decode_image_bytes(media_bytes, source=source) except InvalidInputReferenceError: @@ -138,7 +169,7 @@ def _decode_base64_video( source: str, max_frames: int | None = None, keep: Literal["first", "last"] = "first", -) -> list[Image.Image]: +) -> VideoFrames: if video_reference: if video_reference.startswith("data:video"): _, b64_data = video_reference.split(",", 1) @@ -158,7 +189,7 @@ async def decode_video_url( *, max_frames: int | None = None, keep: Literal["first", "last"] = "first", -) -> list[Image.Image]: +) -> VideoFrames: if video_url.startswith("data:video"): return _decode_base64_video( video_url, @@ -240,7 +271,7 @@ async def decode_input_reference( *, max_video_frames: int | None = None, video_keep: Literal["first", "last"] = "first", -) -> Image.Image | list[Image.Image] | None: +) -> Image.Image | VideoFrames | None: """Decode media input from multipart bytes, data URLs, or typed references.""" provided = sum(item is not None for item in (input_reference_bytes, image_reference, video_reference)) diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index a2b6bf722..73ac80642 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -326,6 +326,15 @@ class OmniDiffusionSamplingParams: @property def resolved_frame_rate(self) -> float | None: + """Frame rate the model should run at, or ``None`` if the user did not provide one. + + Precedence is ``frame_rate`` (the model-internal rate) over ``fps`` (the + request/output-encoding alias). Returns ``None`` when neither was set, which is + the model-agnostic "fps not provided" signal the video serving layer emits when + the request omits fps. Consumers MUST treat ``None`` as "use your own default" and + guard the read (e.g. ``req.sampling_params.resolved_frame_rate or 24.0``); never + feed it into arithmetic unguarded. + """ if self.frame_rate is not None: return float(self.frame_rate)