diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index e73dd7ff6..52762cbf7 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -1121,7 +1121,7 @@ def test_pipeline_init_passes_tokenizer_attrs_into_transformer( def test_preprocess_i2v_image_and_action_video_inputs() -> None: from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_pre_process_func - preprocess = get_cosmos3_pre_process_func(SimpleNamespace()) + preprocess = get_cosmos3_pre_process_func(SimpleNamespace(model_config={"guardrails": False}, tf_model_config=None)) i2v = SimpleNamespace( prompt={"prompt": "A slow camera push.", "multi_modal_data": {"image": Image.new("RGB", (320, 160))}}, sampling_params=make_sampling_params(height=None, width=None, extra_args={}), @@ -1155,6 +1155,130 @@ def test_preprocess_i2v_image_and_action_video_inputs() -> None: assert additional["condition_frame_indexes_vision"] == [0, 1] +def test_preprocess_v2v_decodes_uploaded_video_path(tmp_path) -> None: + """Serving may pass multipart uploads as /tmp/...mp4 path lists (#8073).""" + imageio = pytest.importorskip("imageio.v3") + + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_pre_process_func + + frames = [np.full((16, 32, 3), i * 40, dtype=np.uint8) for i in range(6)] + video_path = tmp_path / "vllm_omni_video_reference_test.mp4" + imageio.imwrite(video_path, frames, fps=4, codec="libx264") + + preprocess = get_cosmos3_pre_process_func(SimpleNamespace(model_config={"guardrails": False}, tf_model_config=None)) + request = SimpleNamespace( + prompt={"prompt": "Continue.", "multi_modal_data": {"video": [str(video_path)]}}, + sampling_params=make_sampling_params( + height=16, + width=32, + extra_args={"condition_frame_indexes_vision": [0, 1], "condition_video_keep": "first"}, + ), + ) + + additional = preprocess(request).prompt["additional_information"] + # condition_frame_indexes_vision=[0,1] => 5 pixel frames after VAE indexing. + assert tuple(additional["preprocessed_video"].shape) == (1, 3, 5, 16, 32) + assert additional["condition_frame_indexes_vision"] == [0, 1] + + +def test_decode_path_video_frames_honors_max_frames_and_keep(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: + pytest.importorskip("imageio.v3") + import imageio.v3 as iio + + from vllm_omni.diffusion.utils.video_decode import decode_path_video_frames + + source = tmp_path / "clip.mp4" + source.write_bytes(b"placeholder") + seen = {"n": 0} + + def fake_imiter(_path): + for idx in range(20): + seen["n"] += 1 + yield np.full((4, 4, 3), idx, dtype=np.uint8) + + monkeypatch.setattr(iio, "imiter", fake_imiter) + + first = decode_path_video_frames(source, max_frames=5, keep="first") + assert len(first) == 5 + assert seen["n"] == 5 + assert first[0][0, 0, 0] == 0 + assert first[-1][0, 0, 0] == 4 + + seen["n"] = 0 + last = decode_path_video_frames(source, max_frames=5, keep="last") + assert len(last) == 5 + assert seen["n"] == 20 + assert last[0][0, 0, 0] == 15 + assert last[-1][0, 0, 0] == 19 + + +def test_preprocess_v2v_video_path_honors_prompt_level_keep_and_indexes( + tmp_path, monkeypatch: pytest.MonkeyPatch +) -> None: + pytest.importorskip("imageio.v3") + import imageio.v3 as iio + + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 as cosmos_pipeline + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_pre_process_func + from vllm_omni.diffusion.utils.video_decode import decode_path_video_frames + + source = tmp_path / "vllm_omni_video_reference_prompt.mp4" + source.write_bytes(b"placeholder") + captured: dict[str, Any] = {} + + def fake_imiter(_path): + for idx in range(20): + yield np.full((16, 32, 3), idx, dtype=np.uint8) + + def spy_decode(*args, **kwargs): + frames = decode_path_video_frames(*args, **kwargs) + captured["keep"] = kwargs.get("keep") + captured["max_frames"] = kwargs.get("max_frames") + captured["n"] = len(frames) + captured["first"] = int(frames[0][0, 0, 0]) + captured["last"] = int(frames[-1][0, 0, 0]) + return frames + + monkeypatch.setattr(iio, "imiter", fake_imiter) + monkeypatch.setattr(cosmos_pipeline, "decode_path_video_frames", spy_decode) + + preprocess = get_cosmos3_pre_process_func(SimpleNamespace(model_config={"guardrails": False}, tf_model_config=None)) + last_keep = SimpleNamespace( + prompt={ + "prompt": "Continue.", + "condition_video_keep": "last", + "condition_frame_indexes_vision": [0, 1], + "multi_modal_data": {"video": [str(source)]}, + }, + sampling_params=make_sampling_params(height=16, width=32, extra_args={}), + ) + additional = preprocess(last_keep).prompt["additional_information"] + assert captured["keep"] == "last" + assert captured["max_frames"] == 5 + assert captured["n"] == 5 + assert captured["first"] == 15 + assert captured["last"] == 19 + assert tuple(additional["preprocessed_video"].shape) == (1, 3, 5, 16, 32) + assert additional["condition_frame_indexes_vision"] == [0, 1] + + wider_indexes = SimpleNamespace( + prompt={ + "prompt": "Continue.", + "condition_frame_indexes_vision": [0, 1, 2], + "multi_modal_data": {"video": [str(source)]}, + }, + sampling_params=make_sampling_params(height=16, width=32, extra_args={}), + ) + additional = preprocess(wider_indexes).prompt["additional_information"] + assert captured["keep"] == "first" + assert captured["max_frames"] == 9 + assert captured["n"] == 9 + assert captured["first"] == 0 + assert captured["last"] == 8 + assert tuple(additional["preprocessed_video"].shape) == (1, 3, 9, 16, 32) + assert additional["condition_frame_indexes_vision"] == [0, 1, 2] + + 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 ( @@ -1224,7 +1348,7 @@ def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest ) assert tuple(loaded.shape) == (3, 2, 8, 8) - preprocess = get_cosmos3_pre_process_func(SimpleNamespace()) + preprocess = get_cosmos3_pre_process_func(SimpleNamespace(model_config={"guardrails": False}, tf_model_config=None)) class FramesWithFps(list): fps = 12.5 diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 95ba4767f..7a7de7d70 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -72,6 +72,7 @@ from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep impor from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.video_decode import decode_path_video_frames, is_video_file_path from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.entrypoints.openai.video_api_utils import positive_float from vllm_omni.experimental.world_models.adapters.state_cosmos3_adapter import ( @@ -465,10 +466,39 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): return nested return video - def _video_payload_to_frames(video: Any) -> list[Any]: + def _decode_video_file_to_frames( + path: str | os.PathLike, + *, + max_frames: int | None, + keep: str, + ) -> list[PIL.Image.Image]: + rgb_frames = decode_path_video_frames(path, max_frames=max_frames, keep=keep) + return [PIL.Image.fromarray(frame).convert("RGB") for frame in rgb_frames] + + def _expand_video_payload_item( + item: Any, + *, + max_frames: int | None, + keep: str, + ) -> list[Any]: + if is_video_file_path(item): + return _decode_video_file_to_frames(item, max_frames=max_frames, keep=keep) + return [item] + + def _video_payload_to_frames( + video: Any, + *, + max_frames: int | None, + keep: str, + ) -> list[Any]: video = _unwrap_video_payload(video) + if is_video_file_path(video): + return _decode_video_file_to_frames(video, max_frames=max_frames, keep=keep) if isinstance(video, list): - return video + frames: list[Any] = [] + for item in video: + frames.extend(_expand_video_payload_item(item, max_frames=max_frames, keep=keep)) + return frames if isinstance(video, torch.Tensor): tensor = video.detach().cpu() if tensor.ndim == 5: @@ -531,11 +561,31 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): if "additional_information" not in prompt: prompt["additional_information"] = {} + extra = _extra_args(request) + transfer_requested = action_mode is None and has_transfer_hints(extra) + condition_frame_indexes_vision = normalize_condition_frame_indexes_vision( + extra.get("condition_frame_indexes_vision", prompt.get("condition_frame_indexes_vision")) + ) + condition_video_keep = normalize_condition_video_keep( + extra.get("condition_video_keep", prompt.get("condition_video_keep")) + ) + 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) + decode_extra = dict(extra) + decode_extra["condition_frame_indexes_vision"] = list(condition_frame_indexes_vision) + decode_extra["condition_video_keep"] = condition_video_keep + decode_spec = Cosmos3OmniDiffusersPipeline.reference_video_decode_spec( + num_frames=getattr(request.sampling_params, "num_frames", None), + extra_args=decode_extra, + ) + raw_video_frames = _video_payload_to_frames( + raw_video, + max_frames=decode_spec.max_frames, + keep=decode_spec.keep, + ) if not raw_video_frames: raise TypeError("Cosmos3 video input must be a non-empty list of PIL images or image paths.") @@ -544,8 +594,6 @@ 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) # Resolve missing H/W. if transfer_requested: @@ -610,22 +658,13 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): 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, + condition_video_keep, ) prompt["additional_information"]["condition_frame_indexes_vision"] = list( condition_frame_indexes_vision diff --git a/vllm_omni/diffusion/models/cosmos3/transfer.py b/vllm_omni/diffusion/models/cosmos3/transfer.py index 65bbbcef8..730a0d44e 100644 --- a/vllm_omni/diffusion/models/cosmos3/transfer.py +++ b/vllm_omni/diffusion/models/cosmos3/transfer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Cosmos3 transfer inference helpers. The reference Cosmos Framework transfer path accepts one or more control hints @@ -25,6 +25,8 @@ import PIL.Image import torch import torch.nn.functional as F +from vllm_omni.diffusion.utils.video_decode import decode_path_video_frames + TRANSFER_HINT_KEYS: tuple[str, ...] = ("edge", "blur", "depth", "seg", "wsm") _TRANSFER_HINT_COMMON_FIELDS = frozenset({"control_path", "control", "control_weight"}) _TRANSFER_HINT_FIELDS: dict[str, frozenset[str]] = { @@ -538,23 +540,9 @@ def _path_media_to_uint8_cthw(path: str | Path, max_frames: int | None) -> torch 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() + rgb_frames = decode_path_video_frames(media_path, max_frames=max_frames, keep="first") + stacked = [torch.from_numpy(frame).permute(2, 0, 1) for frame in rgb_frames] + return torch.stack(stacked, dim=1).contiguous() def media_to_uint8_cthw(value: Any, *, height: int, width: int, max_frames: int | None = None) -> torch.Tensor: diff --git a/vllm_omni/diffusion/utils/video_decode.py b/vllm_omni/diffusion/utils/video_decode.py new file mode 100644 index 000000000..888759230 --- /dev/null +++ b/vllm_omni/diffusion/utils/video_decode.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Decode video file paths to RGB frames. + +Path suffix checks and bounded first/last-N imageio decoding. Model-specific +conditioning and decode-budget stay in each pipeline's preprocess layer. + +Currently used by Cosmos3 V2V preprocess and transfer ``control_path``. +Unifying with ``OmniVideoBackend`` / ``_decode_video_bytes`` is a follow-up. +""" + +from __future__ import annotations + +import os +from collections import deque +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import PIL.Image + +VIDEO_EXTENSIONS = {".avi", ".m4v", ".mkv", ".mov", ".mp4", ".webm"} + + +def is_video_file_path(value: Any) -> bool: + """True when ``value`` is a filesystem path with a known video suffix. + + Serving may leave uploaded references as temp paths such as + ``/tmp/vllm_omni_video_reference_*.mp4`` instead of decoded frames. + """ + if not isinstance(value, str | os.PathLike): + return False + suffix = os.path.splitext(os.fspath(value))[1].lower() + return suffix in VIDEO_EXTENSIONS + + +def _frame_to_uint8_rgb(value: Any) -> np.ndarray: + if isinstance(value, PIL.Image.Image): + return np.array(value.convert("RGB"), dtype=np.uint8, copy=True) + 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"Video decode expected an RGB frame, got {type(value)!r}.") + + +def decode_path_video_frames( + path: str | Path, + *, + max_frames: int | None = None, + keep: Literal["first", "last"] = "first", +) -> list[np.ndarray]: + """Decode a video file to RGB uint8 frames (H, W, C). + + ``keep='first'`` stops after ``max_frames``. ``keep='last'`` still scans the + file and retains only the tail window. + """ + media_path = Path(path) + if not media_path.exists(): + raise FileNotFoundError(f"Video path does not exist: {media_path}") + if keep not in {"first", "last"}: + raise ValueError("Video keep must be either 'first' or 'last'.") + if max_frames is not None and int(max_frames) <= 0: + raise ValueError("Video max_frames must be positive.") + try: + import imageio.v3 as iio + except ImportError as exc: + raise ImportError( + "Video path decoding requires imageio. Install imageio[ffmpeg] or provide decoded frames." + ) from exc + + limit = None if max_frames is None else int(max_frames) + if keep == "last" and limit is not None: + window: deque[np.ndarray] = deque(maxlen=limit) + for frame in iio.imiter(media_path): + window.append(_frame_to_uint8_rgb(frame)) + frames = list(window) + else: + frames = [] + for frame in iio.imiter(media_path): + frames.append(_frame_to_uint8_rgb(frame)) + if limit is not None and len(frames) >= limit: + break + if not frames: + raise ValueError(f"Video path produced no frames: {media_path}") + return frames