diff --git a/vllm_omni/diffusion/models/cosmos3/action.py b/vllm_omni/diffusion/models/cosmos3/action.py index 80a1bc7a9..27f6d36c2 100644 --- a/vllm_omni/diffusion/models/cosmos3/action.py +++ b/vllm_omni/diffusion/models/cosmos3/action.py @@ -1,6 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Action-token helpers for Cosmos3 UVA/action generation.""" +"""Action-token helpers for Cosmos3 action generation. + +These helpers cover action modes that use action tokens as an auxiliary output +stream: ``policy``, ``forward_dynamics``, and ``inverse_dynamics``. The pipeline +returns predicted actions through ``custom_output`` rather than as a normal +``modalities`` entry. +""" from __future__ import annotations diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 3b8a9bd0d..694ae4e8e 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -1,18 +1,29 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Cosmos3 text/image/video-to-video and text-to-image pipeline for vllm-omni. +"""Cosmos3 text/image/video/sound/action pipeline for vllm-omni. -Single pipeline class supports T2V, I2V, V2V, and T2I; the mode is selected at -runtime by: +One pipeline class serves the Cosmos3 family modes. Output modality is selected +mainly by ``prompt["modalities"]``: -* ``prompt["modalities"]`` contains ``"image"``: **T2I** (text-to-image). -* ``prompt["modalities"]`` contains ``"video"`` or is omitted: **T2V** - (text-to-video). -* ``multi_modal_data['image']`` present on the prompt: **I2V** - (handled by :func:`get_cosmos3_pre_process_func`) -* ``multi_modal_data['video']`` present on the prompt with no action mode: - **V2V** (handled by :func:`get_cosmos3_pre_process_func`) +* ``"image"`` selects T2I (text-to-image) and forces a single visual frame. +* ``"video"`` or omitted modalities select video generation. +* ``"audio"`` is accepted for compatibility but does not request sound by + itself; sound is enabled with ``generate_sound`` or ``sound_gen``. +Video generation is further specialized by inputs and extra args: + +* no image/video input: T2V (text-to-video). +* ``multi_modal_data["image"]``: I2V (image-to-video). +* ``multi_modal_data["video"]`` with no action/transfer mode: V2V + (video-to-video). +* transfer hints (``edge``, ``blur``, ``depth``, ``seg``, or ``wsm``): control + transfer video generation. +* ``action_mode``: action-capable video generation. RoboLab/OpenPI observation + payloads in ``extra_args["robot_obs"]`` or ``extra_args["observation"]`` + bypass normal video output and return action-only custom output. + +Generated sound is video-only, cannot be combined with action or transfer, and +is produced from sound latents rather than from ``multi_modal_data["audio"]``. """ from __future__ import annotations @@ -156,14 +167,20 @@ COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096 # Post-process function (registered in registry.py) # --------------------------------------------------------------------------- def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): - """Pre-process function for T2V, I2V, V2V, and action inputs. + """Build the request preprocessor for Cosmos3 image/video inputs. - For T2V (no image in ``multi_modal_data``), the request is returned - unchanged after the optional guardrails check. For I2V (image present), - the conditioning image is loaded, aspect-resized + center-cropped, and - stored back on the prompt as ``additional_information.preprocessed_image``. - For V2V (video present without action mode), source frames are cropped to - the target size and stored as ``additional_information.preprocessed_video``. + For plain T2V (no image or video in ``multi_modal_data``), the request is + returned unchanged after the optional guardrail check. For I2V, the + conditioning image is loaded, aspect-resized, center-cropped, and stored as + ``additional_information.preprocessed_image``. For V2V, source frames are + cropped to the target size and stored as + ``additional_information.preprocessed_video``. + + Action modes reuse image/video preprocessing but use action-specific resize + and padding rules. Transfer requests store + ``additional_information.preprocessed_transfer_video`` for optional input + video conditioning. Cosmos3 sound generation is not driven by + ``multi_modal_data["audio"]``; it is enabled later from sampling params. """ from .guardrails import check_text_safety, ensure_initialized, is_guardrails_enabled @@ -470,6 +487,13 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig): + """Build the postprocessor for Cosmos3 image, video, and video+audio output. + + The pipeline returns image payloads as ``{"image": tensor}`` and video + payloads as ``{"video": tensor}``. Sound-enabled video returns the same + video payload plus ``audio`` and ``audio_sample_rate``. Image output with + audio is rejected because Cosmos3 sound generation is video-only. + """ from .guardrails import check_video_safety, is_guardrails_enabled video_processor = VideoProcessor(vae_scale_factor=16) @@ -563,6 +587,13 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig): def get_cosmos3_action_post_process_func(od_config: OmniDiffusionConfig): + """Build the custom-output postprocessor for Cosmos3 action predictions. + + Action modes return predicted action tensors in ``custom_output`` alongside + normal video output. RoboLab/OpenPI policy serving marks action-only output + and carries observation metadata used here to map model-space actions back + to the requested robot action representation. + """ del od_config def action_post_process_func(action: Any, custom_output: dict[str, Any] | None = None, sampling_params=None): @@ -598,14 +629,15 @@ def get_cosmos3_ir_op_priority_func(od_config: OmniDiffusionConfig): class Cosmos3OmniDiffusersPipeline( nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin ): - """Cosmos3 text/image/video-to-video / text-to-image pipeline. + """Cosmos3 text/image/video/sound/action pipeline. Architecture: Mixture-of-Transformers with Qwen3-VL backbone. - Understanding pathway: causal self-attention on text (runs once, K/V cached) - - Generation pathway: cross-attention on noisy visual latents (runs each step) + - Generation pathway: cross-attention on visual latents and optional + transfer-control, action, and sound latents (runs each step) - Supports T2V, I2V, V2V, and T2I from the same class. Mode is selected at - runtime: + Supports T2V, I2V, V2V, T2I, transfer, sound-enabled video, and action + generation from the same class. Mode is selected at runtime: * **T2I** when ``prompt["modalities"]`` contains ``"image"``. Latent T-dim is forced to 1, T2I-specific scheduler defaults are applied (50 steps, @@ -622,6 +654,17 @@ class Cosmos3OmniDiffusersPipeline( ``multi_modal_data['video']`` without an action mode. Explicit latent frame indexes are kept clean with ``noisy_frame_mask`` and re-injected after each scheduler step. + * **Transfer** when ``edge``, ``blur``, ``depth``, ``seg``, or ``wsm`` hints + are supplied. Transfer is video-output only and cannot be combined with + sound or action generation. + * **Sound-enabled video** when ``generate_sound`` or ``sound_gen`` is true. + Sound is generated from sound latents, not from ``multi_modal_data['audio']``; + T2I, transfer, and action+sound are rejected. + * **Action generation** when ``action_mode`` is provided. ``policy`` and + ``forward_dynamics`` require an image or video input; ``inverse_dynamics`` + requires video input. Action predictions are returned in ``custom_output``. + RoboLab/OpenPI observations in ``extra_args['robot_obs']`` or + ``extra_args['observation']`` return action-only custom output. * **T2V** otherwise (default video generation). """ @@ -1386,7 +1429,14 @@ class Cosmos3OmniDiffusersPipeline( @staticmethod def _is_t2i_request(req: OmniDiffusionRequest) -> bool: - """Detect text-to-image mode from request-level prompt modalities.""" + """Return whether request-level modalities select image output. + + Only ``"image"`` switches Cosmos3 into T2I. ``"video"`` and omitted + modalities select video output. ``"text"`` and ``"audio"`` are accepted + compatibility values for callers that share prompt schemas, but they do + not select text or audio output in this pipeline. ``"image"`` and + ``"video"`` cannot be requested together. + """ if not req.prompts: return False first_prompt = req.prompts[0] diff --git a/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py b/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py index 66937b33c..5d6ff8dd9 100644 --- a/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py +++ b/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py @@ -1,6 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Cosmos3 sound tokenizer integration.""" +"""Cosmos3 sound tokenizer integration. + +The tokenizer decodes model-generated sound latents for video+audio output. It +does not encode or condition on ``multi_modal_data["audio"]``; request-side +sound generation is enabled by ``generate_sound`` or ``sound_gen``. +""" from __future__ import annotations diff --git a/vllm_omni/diffusion/models/cosmos3/transfer.py b/vllm_omni/diffusion/models/cosmos3/transfer.py index a7930a905..42cbea559 100644 --- a/vllm_omni/diffusion/models/cosmos3/transfer.py +++ b/vllm_omni/diffusion/models/cosmos3/transfer.py @@ -6,6 +6,9 @@ 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. + +Transfer is a video-output mode. It cannot be combined with T2I, sound +generation, or action generation. """ from __future__ import annotations diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 3eef62621..99e0eacf0 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -956,7 +956,8 @@ class Cosmos3VFMTransformer(nn.Module): """Cosmos3 VFM Transformer: UND language model + GEN denoising layers. The UND pathway runs once per generation (K/V cached). The GEN pathway - runs at each denoising step. + runs at each denoising step over the target video/image latent stream and + optional transfer-control, action, and sound latent streams. Layerwise offloading uses ``gen_layers`` as the block container. @@ -1408,14 +1409,16 @@ class Cosmos3VFMTransformer(nn.Module): sound_latents: Optional [B, C_sound, T_sound] noisy sound latents. noisy_frame_mask: Optional [B, 1, t, 1, 1] mask where 1=noisy (add timestep embedding, predict velocity) and 0=conditioned (clean - context, skip timestep embedding). None means all frames noisy - (T2V mode). + context, skip timestep embedding). None means all target vision + frames are noisy, as in T2I/T2V. 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. + tuple outputs in video, action, sound order when action/sound streams + are provided. Transfer-control streams condition the video prediction + and are not returned. """ if kwargs: raise TypeError(f"Unexpected Cosmos3 transformer kwargs: {sorted(kwargs)}") diff --git a/vllm_omni/model_extras/cosmos3.py b/vllm_omni/model_extras/cosmos3.py index 6303f0f11..a17e8193f 100644 --- a/vllm_omni/model_extras/cosmos3.py +++ b/vllm_omni/model_extras/cosmos3.py @@ -12,6 +12,14 @@ def build_text_to_image_prompt( height: int | None = None, width: int | None = None, ) -> dict[str, Any]: + """Build a Cosmos3 T2I prompt by selecting image output modality. + + Cosmos3 uses the same pipeline class for T2I and video modes. The OpenAI + image endpoint selects T2I by adding ``modalities=["image"]``. ``height`` + and ``width`` are accepted for registry compatibility; sizing travels + through sampling params instead of the prompt payload. + """ + del height, width text_prompt: dict[str, Any] = { "prompt": prompt, "modalities": ["image"],