diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index f58123317de..c4e181d5917 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -33,6 +33,7 @@ th { | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | +| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V | `nvidia/Cosmos3-Nano` | ✅︎ | | | | | `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 4e62cd6651d..3c58c1181f9 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -133,10 +133,12 @@ The following tables show which models support each feature: | **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ | | **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ | | **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | +| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | > Notes: > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT. > 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary. +> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload is not supported; use layerwise offload. ### VideoGen @@ -149,6 +151,8 @@ The following tables show which models support each feature: | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | +| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | + **Frame Interpolation Support** diff --git a/recipes/README.md b/recipes/README.md index 05b4788dd45..885710db1d4 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -36,6 +36,7 @@ recipes/ | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB | | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) | | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB | +| [`nvidia/Cosmos3-Nano.md`](./nvidia/Cosmos3-Nano.md) | Text-to-image, text-to-video, and image-to-video generation | 1x H200 141GB / B300 | | [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB | | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 | | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay | 1x A100 80GB | diff --git a/recipes/nvidia/Cosmos3-Nano.md b/recipes/nvidia/Cosmos3-Nano.md new file mode 100644 index 00000000000..5d5e524da58 --- /dev/null +++ b/recipes/nvidia/Cosmos3-Nano.md @@ -0,0 +1,215 @@ +# Cosmos3-Nano + +> Text-to-image, text-to-video, and image-to-video serving + +## Summary + +- Vendor: NVIDIA +- Model: `nvidia/Cosmos3-Nano` +- Task: Text-to-image (T2I), text-to-video (T2V), and image-to-video (I2V) generation +- Mode: Online serving with the OpenAI-compatible image/video APIs, plus offline generation via the `Omni` API +- Maintainer: Community + +## When to use this recipe + +Use this recipe to deploy `nvidia/Cosmos3-Nano` for image and video generation. +A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves all three modes; +the mode is selected per request: + +- **T2I** — `POST /v1/images/generations` (or a prompt carrying `modalities=["image"]`). +- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image. +- **I2V** — `POST /v1/videos/sync` with a reference image (`input_reference` file + upload, or `image_reference` JSON). + +## References + +- Model card (authoritative usage + example assets): +- Example inputs/outputs live in the repo's `assets/` (`example_t2v_prompt.json`, + `example_i2v_prompt.json`, `example_i2v_input.jpg`, `negative_prompt.json`). +- Prompt upsampling (recommended for quality): the model expects JSON-upsampled + structured prompts; see NVIDIA's `cosmos-framework` prompt-upsampling docs. +- Pipeline: [`vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py`](../../vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py) +- Smoke tests (canonical request formats): [`tests/e2e/accuracy/test_cosmos3_similarity.py`](../../tests/e2e/accuracy/test_cosmos3_similarity.py) + +## Hardware Support + +## GPU + +### 1x H200 141GB / B300 (Online serving) + +#### Environment + +- OS: Ubuntu 22.04+ +- Python: 3.12+ +- Driver / runtime: NVIDIA CUDA environment +- vLLM version: match the repository requirements from your current checkout +- vLLM-Omni version or commit: use the commit you are deploying from + +#### Command + +Safety guardrails are **on by default** (NVIDIA Open Model License). They load +the **gated** `nvidia/Cosmos-1.0-Guardrail` model, so to keep them on you must: + +1. `pip install cosmos-guardrail` +2. Accept the license at +3. Export a token with access: `export HF_TOKEN=hf_...` + +Then launch the recommended server: + +```bash +vllm serve nvidia/Cosmos3-Nano \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --init-timeout 1800 +``` + +To run **without** guardrails (you are responsible for license compliance), +add `--no-guardrails` (no token/`cosmos-guardrail` needed). For extra GPUs use +`--ulysses-degree N` (context parallel) or `--tensor-parallel-size N`; +`--enable-layerwise-offload` reduces VRAM on smaller GPUs. The pipeline +auto-resolves from `model_index.json`; pass +`--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly. + +#### Verification + +Best quality uses the JSON-upsampled prompts from `assets/` (download with +`hf download nvidia/Cosmos3-Nano assets/ --local-dir Cosmos3-Nano`). Minimal +self-contained examples: + +```bash +curl http://localhost:8000/v1/models + +# Text-to-image -> /v1/images/generations (1024x1024, 50 steps; base64 PNG) +curl -sS -X POST http://localhost:8000/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "model": "nvidia/Cosmos3-Nano", + "prompt": "A photorealistic red sports car on a city street at golden hour, cinematic lighting.", + "negative_prompt": "blurry, distorted, low quality", + "size": "1024x1024", "n": 1, "response_format": "b64_json", + "num_inference_steps": 50, "guidance_scale": 7.0, "seed": 42 + }' | python -c "import sys,json,base64; open('cosmos3_t2i.png','wb').write(base64.b64decode(json.load(sys.stdin)['data'][0]['b64_json']))" + +# Text-to-video -> /v1/videos/sync (720p, 189 frames @ 24fps; official params) +curl -sS -X POST http://localhost:8000/v1/videos/sync \ + -H "Accept: video/mp4" \ + -F "model=nvidia/Cosmos3-Nano" \ + -F "prompt=A robot arm is cleaning a plate in the kitchen" \ + -F "negative_prompt=blurry, distorted, low quality, jittery, deformed" \ + -F "size=1280x720" -F "num_frames=189" -F "fps=24" \ + -F "num_inference_steps=35" -F "guidance_scale=6.0" \ + -F "max_sequence_length=4096" -F "flow_shift=10.0" \ + -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \ + -F "seed=123" \ + -o cosmos3_t2v.mp4 + +# Image-to-video -> /v1/videos/sync with an uploaded reference image +curl -sS -X POST http://localhost:8000/v1/videos/sync \ + -H "Accept: video/mp4" \ + -F "model=nvidia/Cosmos3-Nano" \ + -F "prompt=The scene comes to life with smooth, natural motion." \ + -F "negative_prompt=blurry, distorted, low quality" \ + -F "size=1280x720" -F "num_frames=189" -F "fps=24" \ + -F "num_inference_steps=35" -F "guidance_scale=6.0" \ + -F "max_sequence_length=4096" -F "flow_shift=10.0" \ + -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \ + -F "seed=1111" \ + -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \ + -o cosmos3_i2v.mp4 +``` + +#### Notes + +- **Measured latency (1x B300, bf16, guardrails off):** + - T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s** + - T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s** + - I2V 1280×720, 189 frames @ 35 steps → ~**99 s** + - Guardrails-on overhead: ~8% on T2I, negligible on video. +- **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 GPU; + full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer) + ~33 GB on disk. +- **Determinism:** identical seed reproduces identical output on the same + hardware; outputs are not bit-identical across different GPU types. +- **Supported sizes (per model card):** 256p / 480p / 720p at 16:9, 4:3, 1:1, + 3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V 1280×720, + 189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`. +- **Key flags / params:** `--no-guardrails` (server) or + `extra_params={"guardrails":false}` (per request) toggles safety; + `use_resolution_template` / `use_duration_template` are off by default and only + needed when not using upsampled prompts that already encode resolution/duration. +- **Known limitations:** + - Guardrails-on requires `cosmos-guardrail` **and** access to the gated + `nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise + the server fails at pipeline build with a gated-repo / safety-checker error. + - A guardrail-blocked prompt currently returns HTTP 500 + (`"Guardrail blocked prompt"`). + - Video + audio, and action (policy / forward- / inverse-dynamics) modalities + are not part of this integration yet. + +### 1x GPU (Offline generation) + +#### Environment + +- OS: Ubuntu 22.04+ +- Python: 3.12+ +- Driver / runtime: NVIDIA CUDA environment +- vLLM-Omni version or commit: use the commit you are deploying from + +#### Command + +```python +# cosmos3_offline.py — run with: python cosmos3_offline.py +import torch +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + +def main(): + omni = Omni( + model="nvidia/Cosmos3-Nano", + model_class_name="Cosmos3OmniDiffusersPipeline", + trust_remote_code=True, + enforce_eager=True, + # Keep guardrails on by installing cosmos-guardrail + gated-repo access; + # this disables them for a quick local run. + model_config={"guardrails": False}, + ) + gen = torch.Generator(device="cpu").manual_seed(42) + + # Text-to-image (modalities=["image"]). For T2V use modalities=["video"] + # plus num_frames/fps; for I2V add multi_modal_data={"image": }. + outputs = omni.generate( + { + "prompt": "A photorealistic red sports car at golden hour, cinematic lighting.", + "negative_prompt": "blurry, distorted, low quality", + "modalities": ["image"], + }, + OmniDiffusionSamplingParams( + height=1024, width=1024, generator=gen, + guidance_scale=7.0, num_inference_steps=50, num_outputs_per_prompt=1, + ), + ) + outputs[0].request_output.images[0].save("cosmos3_t2i.png") + omni.close() + + +if __name__ == "__main__": + main() +``` + +#### Verification + +```bash +python cosmos3_offline.py +python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image', im.size, im.mode)" +``` + +#### Notes + +- Same `Cosmos3OmniDiffusersPipeline` as online; mode is chosen by + `prompt["modalities"]` (`["image"]` → T2I, `["video"]` → T2V) plus + `num_frames`/`fps`, and `multi_modal_data={"image": ...}` for I2V. For video, + frames are returned in `outputs[0].request_output.images` as an + `(B, F, H, W, 3)` array. +- The offline entry must be guarded by `if __name__ == "__main__":` — the engine + spawns workers with the `spawn` start method. diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py index 0fd07393d1b..146084a508c 100644 --- a/tests/diffusion/cache/test_cache_dit.py +++ b/tests/diffusion/cache/test_cache_dit.py @@ -25,6 +25,7 @@ cd_backend.enable_cache_for_helios, cd_backend.enable_cache_for_wan22, cd_backend.enable_cache_for_longcat_image, + cd_backend.enable_cache_for_cosmos3, ] SAMPLE_CACHE_CONFIG = DiffusionCacheConfig() @@ -47,6 +48,24 @@ def test_separate_cfg(mock_cache_dit, mock_block_adapter, enabler): assert adapter_kwargs["has_separate_cfg"] is True +@patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter") +@patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit") +def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter): + """Cosmos3 should cache only the repeated GEN pathway blocks.""" + mock_pipeline = Mock() + gen_layers = object() + mock_pipeline.transformer.gen_layers = gen_layers + + cd_backend.enable_cache_for_cosmos3(mock_pipeline, SAMPLE_CACHE_CONFIG) + + mock_cache_dit.enable_cache.assert_called_once() + adapter_kwargs = mock_block_adapter.call_args.kwargs + assert adapter_kwargs["transformer"] is mock_pipeline.transformer + assert adapter_kwargs["blocks"] == [gen_layers] + assert adapter_kwargs["has_separate_cfg"] is True + assert adapter_kwargs["check_forward_pattern"] is False + + # This test is skipped on ROCm since rocm_unquantized_gemm doesn't support CPU backend @pytest.mark.skipif( current_omni_platform.is_rocm(), diff --git a/tests/diffusion/models/cosmos3/__init__.py b/tests/diffusion/models/cosmos3/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/tests/diffusion/models/cosmos3/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py new file mode 100644 index 00000000000..31b40b6eee5 --- /dev/null +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -0,0 +1,482 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import sys +import types +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +from PIL import Image +from torch import nn + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class StubScheduler: + def __init__(self, timesteps: list[int] | None = None, *, flow_shift: float = 1.0) -> None: + self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64) + self.config = SimpleNamespace(num_train_timesteps=1000, flow_shift=flow_shift) + self.set_timesteps_calls: list[tuple[int, torch.device]] = [] + self.step_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + + def set_timesteps(self, num_steps: int, device: torch.device) -> None: + self.set_timesteps_calls.append((num_steps, device)) + self.timesteps = torch.arange(num_steps, 0, -1, dtype=torch.int64, device=device) + + def step(self, noise_pred: torch.Tensor, timestep: torch.Tensor, latents: torch.Tensor, **kwargs): + del kwargs + self.step_calls.append((noise_pred.clone(), timestep.clone(), latents.clone())) + return (latents + noise_pred,) + + +class _ModeLatentDist: + def __init__(self, latents: torch.Tensor) -> None: + self._latents = latents + + def mode(self) -> torch.Tensor: + return self._latents + + +class StubCosmos3VAE: + dtype = torch.float32 + + def __init__(self, z_dim: int = 2, *, temporal: int = 4, spatial: int = 8) -> None: + self.config = SimpleNamespace( + z_dim=z_dim, + scale_factor_temporal=temporal, + scale_factor_spatial=spatial, + latents_mean=[0.0] * z_dim, + latents_std=[1.0] * z_dim, + ) + + def encode(self, video: torch.Tensor): + latent_frames = (video.shape[2] - 1) // self.config.scale_factor_temporal + 1 + latent_height = video.shape[-2] // self.config.scale_factor_spatial + latent_width = video.shape[-1] // self.config.scale_factor_spatial + latents = torch.ones( + video.shape[0], + self.config.z_dim, + latent_frames, + latent_height, + latent_width, + dtype=video.dtype, + device=video.device, + ) + return SimpleNamespace(latent_dist=_ModeLatentDist(latents)) + + def decode(self, latents: torch.Tensor, return_dict: bool = False): + del return_dict + return (latents,) + + +class StubCosmos3Transformer(nn.Module): + def __init__( + self, + *, + latent_channel_size: int = 2, + ) -> None: + super().__init__() + self.latent_channel_size = latent_channel_size + self.cached_kv: Any | None = None + self.cached_freqs_gen: Any | None = None + self.calls: list[dict[str, Any]] = [] + self.reset_calls = 0 + + def reset_cache(self) -> None: + self.reset_calls += 1 + self.cached_kv = None + self.cached_freqs_gen = None + + def forward( + self, + *, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + text_ids: torch.Tensor, + text_mask: torch.Tensor, + **kwargs: Any, + ) -> torch.Tensor: + token = int(text_ids.reshape(-1)[0].item()) if text_ids.numel() else 0 + self.calls.append( + { + "token": token, + "timestep": timestep.clone(), + "text_mask": text_mask.clone(), + "cache_before": self.cached_kv, + "kwargs": dict(kwargs), + } + ) + if self.cached_kv is None: + marker = torch.tensor([token], dtype=torch.float32) + self.cached_kv = [(marker, marker + 100)] + self.cached_freqs_gen = (marker + 200, marker + 300) + return torch.full_like(hidden_states, float(token)) + + +def passthrough_progress_bar(iterable): + return iterable + + +@pytest.fixture(autouse=True) +def fake_cosmos3_guardrails(monkeypatch: pytest.MonkeyPatch): + module = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails") + module.is_guardrails_enabled = lambda od_config, sampling_params=None: False + module.ensure_initialized = lambda od_config: None + module.check_text_safety = lambda text: None + module.check_video_safety = lambda video: video + monkeypatch.setitem(sys.modules, module.__name__, module) + return module + + +@pytest.fixture +def make_cosmos3_pipeline(): + def _make(): + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + ) + + pipeline = object.__new__(Cosmos3OmniDiffusersPipeline) + nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace() + pipeline.device = torch.device("cpu") + pipeline.dtype = torch.float32 + pipeline.transformer = StubCosmos3Transformer(latent_channel_size=2) + pipeline.vae = StubCosmos3VAE(z_dim=2) + pipeline.vae_scale_factor_temporal = 4 + pipeline.vae_scale_factor_spatial = 8 + pipeline.scheduler = StubScheduler([9, 3], flow_shift=1.0) + pipeline._base_scheduler_config = pipeline.scheduler.config + pipeline._engine_init_flow_shift = 1.0 + pipeline._current_flow_shift = 1.0 + pipeline._guidance_scale = None + pipeline._num_timesteps = None + pipeline._cache_dit_requires_paired_cfg = False + pipeline.progress_bar = passthrough_progress_bar + return pipeline + + return _make + + +def make_sampling_params(**overrides: Any) -> SimpleNamespace: + values = { + "height": None, + "width": None, + "num_frames": None, + "num_inference_steps": None, + "guidance_scale": None, + "generator": None, + "seed": 123, + "num_outputs_per_prompt": 1, + "frame_rate": None, + "resolved_frame_rate": None, + "max_sequence_length": None, + "extra_args": {}, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _ids(value: int) -> torch.Tensor: + return torch.tensor([[value]], dtype=torch.long) + + +def _mask() -> torch.Tensor: + return torch.ones(1, 1, dtype=torch.long) + + +def test_pipeline_registered_and_exported() -> None: + from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS + from vllm_omni.diffusion.models import cosmos3 + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline + from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin + from vllm_omni.diffusion.registry import ( + _DIFFUSION_MODELS, + _DIFFUSION_POST_PROCESS_FUNCS, + _DIFFUSION_PRE_PROCESS_FUNCS, + ) + + assert issubclass(Cosmos3OmniDiffusersPipeline, nn.Module) + assert issubclass(Cosmos3OmniDiffusersPipeline, ProgressBarMixin) + assert Cosmos3OmniDiffusersPipeline.support_image_input is True + assert _DIFFUSION_MODELS["Cosmos3OmniDiffusersPipeline"] == ( + "cosmos3", + "pipeline_cosmos3", + "Cosmos3OmniDiffusersPipeline", + ) + assert _DIFFUSION_PRE_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_pre_process_func" + assert _DIFFUSION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_post_process_func" + assert "Cosmos3OmniDiffusersPipeline" in CUSTOM_DIT_ENABLERS + assert "Cosmos3OmniDiffusersPipeline" in cosmos3.__all__ + + +def test_preprocess_i2v_image_input() -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_pre_process_func + + preprocess = get_cosmos3_pre_process_func(SimpleNamespace()) + i2v = SimpleNamespace( + prompts=[{"prompt": "A slow camera push.", "multi_modal_data": {"image": Image.new("RGB", (320, 160))}}], + sampling_params=SimpleNamespace(height=None, width=None, extra_args={}), + ) + + result = preprocess(i2v) + assert (result.sampling_params.height, result.sampling_params.width) == (672, 1344) + assert tuple(result.prompts[0]["additional_information"]["preprocessed_image"].shape[-2:]) == (672, 1344) + + +def test_postprocess_handles_image_video_and_validation() -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_post_process_func + + func = get_cosmos3_post_process_func(SimpleNamespace()) + video = torch.zeros(1, 3, 1, 4, 4) + + assert func(video, output_type="latent") is video + assert func({"image": video})[0].size == (4, 4) + + with pytest.raises(ValueError, match="text-to-image postprocess expects"): + func({"image": torch.zeros(1, 3, 2, 4, 4)}) + with pytest.raises(ValueError, match="both image and video"): + func({"image": video, "video": video}) + + +def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline + + pipeline = make_cosmos3_pipeline() + captured: list[str] = [] + pipeline._tokenize_prompt = lambda text, *args, **kwargs: (captured.append(text) or _ids(len(captured)), _mask()) + + pipeline._format_and_tokenize_prompts( + "A robot", + "bad", + num_frames=48, + frame_rate=24, + height=720, + width=1280, + max_sequence_length=32, + sp=SimpleNamespace( + extra_args={ + "negative_metadata_mode": "inverse", + # Duration/resolution metadata templates are off by default + # (commit "Change resolution and duration templates to off by + # default"); enable them explicitly to exercise the formatting. + "use_duration_template": True, + "use_resolution_template": True, + } + ), + use_system_prompt=True, + is_t2i=False, + ) + assert "The video is 2.0 seconds long" in captured[0] + assert "The video is not 2.0 seconds long" in captured[1] + + remaps = { + "embed_tokens.weight": "transformer.language_model.embed_tokens.weight", + "model.embed_tokens.weight": "transformer.language_model.embed_tokens.weight", + "norm.weight": "transformer.language_model.norm.weight", + "norm_moe_gen.weight": "transformer.norm_moe_gen.weight", + "proj_in.weight": "transformer.proj_in.weight", + "proj_out.bias": "transformer.proj_out.bias", + "layers.3.self_attn.to_q.weight": "transformer.language_model.layers.3.self_attn.to_q.weight", + "layers.3.self_attn.to_out.weight": "transformer.language_model.layers.3.self_attn.to_out.weight", + "layers.3.self_attn.norm_q.weight": "transformer.language_model.layers.3.self_attn.norm_q.weight", + "layers.3.self_attn.add_q_proj.weight": "transformer.gen_layers.3.cross_attention.to_q.weight", + "layers.3.self_attn.to_add_out.weight": "transformer.gen_layers.3.cross_attention.to_out.weight", + "layers.3.self_attn.norm_added_q.weight": "transformer.gen_layers.3.cross_attention.norm_q.weight", + "transformer.model.layers.3.self_attn.add_k_proj.weight": ( + "transformer.gen_layers.3.cross_attention.to_k.weight" + ), + } + assert {key: Cosmos3OmniDiffusersPipeline._remap_ckpt_key(key) for key in remaps} == remaps + + +def test_prepare_latents_for_video_and_image(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + latents = pipeline._prepare_latents(16, 24, 5, torch.Generator(device="cpu").manual_seed(0)) + assert latents.shape == (1, 2, 2, 2, 3) + + pipeline._encode_conditioning_video = lambda *args, **kwargs: torch.full((1, 2, 2, 2, 3), 5.0) + i2v_latents, velocity_mask, image_latent = pipeline._prepare_latents_i2v( + torch.zeros(1, 3, 16, 24), 16, 24, 5, torch.Generator(device="cpu").manual_seed(0) + ) + torch.testing.assert_close(i2v_latents[:, :, 0], torch.full((1, 2, 2, 3), 5.0)) + assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]]]]] + assert image_latent.shape == (1, 2, 1, 2, 3) + + +def test_diffuse_covers_cfg_and_i2v_steps(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + latents = torch.zeros(1, 2, 1, 1, 1) + + result = pipeline.diffuse( + latents=latents, + timesteps=torch.tensor([900, 100]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=3.0, + shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0}, + guidance_interval=(500.0, 1000.0), + ) + assert [call["token"] for call in pipeline.transformer.calls] == [2, 1, 2] + torch.testing.assert_close(result, torch.full_like(latents, 6.0)) + + i2v = pipeline.diffuse( + latents=torch.zeros(1, 2, 2, 1, 1), + timesteps=torch.tensor([7]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=1.0, + shared_kwargs={"video_shape": (2, 1, 1), "fps": 24.0}, + velocity_mask=torch.tensor([[[[[0.0]], [[1.0]]]]]), + image_latent=torch.full((1, 2, 1, 1, 1), 7.0), + ) + torch.testing.assert_close(i2v[:, :, 0:1], torch.full((1, 2, 1, 1, 1), 7.0)) + + +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 + output is numerically identical to the skip path. + + Contrast with ``test_diffuse_covers_cfg_and_i2v_steps`` (no marker), where + the same inputs skip the out-of-interval uncond pass: calls == [2, 1, 2]. + """ + pipeline = make_cosmos3_pipeline() + # Marker normally set by ``enable_cache_for_cosmos3`` when cache-dit is on. + pipeline._cache_dit_requires_paired_cfg = True + latents = torch.zeros(1, 2, 1, 1, 1) + + result = pipeline.diffuse( + latents=latents, + timesteps=torch.tensor([900, 100]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=3.0, + shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0}, + guidance_interval=(500.0, 1000.0), + ) + + # t=900 is inside the interval (cond+uncond); t=100 is outside but the + # uncond pass is still issued -> paired cond/uncond at every step. + assert [call["token"] for call in pipeline.transformer.calls] == [2, 1, 2, 1] + # Identical result to the skip path: out-of-interval combine uses scale=1.0, + # so combine_cfg_noise(cond=2, uncond=1, 1.0) == 2 == the skipped cond value. + torch.testing.assert_close(result, torch.full_like(latents, 6.0)) + + +class TestForwardRouting: + def _install_forward_stubs(self, pipeline): + captured: dict[str, object] = {"diffuse_calls": [], "prepare_calls": []} + + def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, *args, **kwargs): + captured["format"] = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "num_frames": num_frames, + "frame_rate": frame_rate, + "height": height, + "width": width, + "is_t2i": kwargs["is_t2i"], + } + return _ids(2), _mask(), _ids(1), _mask() + + def fake_prepare(height, width, num_frames, generator): + captured["prepare_calls"].append((height, width, num_frames, generator.initial_seed())) + return torch.zeros(1, 2, 1, 1, 1) + + def fake_diffuse(**kwargs): + captured["diffuse_calls"].append(kwargs) + return kwargs["latents"] + len(captured["diffuse_calls"]) + + pipeline._format_and_tokenize_prompts = fake_format + pipeline._prepare_latents = fake_prepare + pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target) + pipeline.diffuse = fake_diffuse + pipeline._decode_latents = lambda latents: latents + return captured + + @pytest.mark.parametrize( + ("prompt", "sampling_params", "expected"), + [ + ( + {"prompt": "A painted robot", "modalities": ["image"]}, + make_sampling_params(num_outputs_per_prompt=2), + {"key": "image", "is_t2i": True, "flow": [3.0], "steps": [50, 50], "frames": 1}, + ), + ( + "A warehouse robot", + make_sampling_params(), + {"key": "video", "is_t2i": False, "flow": [1.0], "steps": [35], "frames": 189}, + ), + ], + ) + def test_forward_defaults_and_mode_selection( + self, + make_cosmos3_pipeline, + prompt, + sampling_params, + expected, + ) -> None: + pipeline = make_cosmos3_pipeline() + captured = self._install_forward_stubs(pipeline) + + output = pipeline.forward(SimpleNamespace(prompts=[prompt], sampling_params=sampling_params)) + + assert expected["key"] in output.output + assert captured["format"]["is_t2i"] is expected["is_t2i"] + assert captured["format"]["num_frames"] == expected["frames"] + assert captured["flow_shifts"] == expected["flow"] + assert [call[0] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"] + + def test_forward_i2v_route(self, make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + captured = self._install_forward_stubs(pipeline) + image_tensor = torch.zeros(1, 3, 16, 16) + velocity_mask = torch.ones(1, 1, 1, 1, 1) + + pipeline._prepare_latents_i2v = lambda *args, **kwargs: ( + torch.zeros(1, 2, 1, 1, 1), + velocity_mask, + torch.zeros(1, 2, 1, 1, 1), + ) + pipeline.forward( + SimpleNamespace( + prompts=[ + { + "prompt": "move", + "modalities": ["video"], + "additional_information": {"preprocessed_image": image_tensor}, + } + ], + sampling_params=make_sampling_params(height=16, width=16, num_frames=5), + ) + ) + assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask + + @pytest.mark.parametrize( + ("prompt", "sampling_params", "message"), + [ + (["one", "two"], make_sampling_params(), "single prompt"), + ([{"prompt": "one", "modalities": ["image", "video"]}], make_sampling_params(), "both image and video"), + ], + ) + def test_forward_rejects_invalid_public_requests( + self, + make_cosmos3_pipeline, + prompt, + sampling_params, + message, + ) -> None: + pipeline = make_cosmos3_pipeline() + + with pytest.raises(ValueError, match=message): + pipeline.forward(SimpleNamespace(prompts=prompt, sampling_params=sampling_params)) diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py new file mode 100644 index 00000000000..730079c116a --- /dev/null +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _tiny_cosmos3_config(**overrides): + config = { + "hidden_size": 8, + "num_hidden_layers": 0, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "head_dim": 4, + "intermediate_size": 16, + "vocab_size": 32, + "latent_patch_size": 1, + "latent_channel": 2, + "rope_scaling": {"mrope_section": [1, 1, 0]}, + } + config.update(overrides) + return config + + +def test_mrope_position_ids_cover_text_and_video() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import ( + compute_mrope_position_ids_text, + compute_mrope_position_ids_vision, + ) + + text_ids, text_offset = compute_mrope_position_ids_text(num_tokens=3, temporal_offset=5) + assert text_ids.tolist() == [[5, 6, 7], [5, 6, 7], [5, 6, 7]] + assert text_offset == 8 + + vision_ids, vision_offset = compute_mrope_position_ids_vision(2, 2, 3, temporal_offset=10, fps=None) + assert vision_ids.shape == (3, 12) + assert vision_ids[0].tolist() == [10] * 6 + [11] * 6 + assert vision_offset == 12 + + modulated_ids, modulated_offset = compute_mrope_position_ids_vision( + 2, + 1, + 1, + temporal_offset=10, + fps=12.0, + base_fps=24.0, + temporal_compression_factor=4, + ) + torch.testing.assert_close(modulated_ids[0], torch.tensor([10.0, 12.0])) + assert modulated_offset == 13 + + +@pytest.mark.parametrize( + ("key", "value"), + [ + ("qk_norm_for_diffusion", False), + ("qk_norm_for_text", False), + ("position_embedding_type", "rotary"), + ("unified_3d_mrope_reset_spatial_ids", False), + ("joint_attn_implementation", "one_way"), + ], +) +def test_validate_supported_config_rejects_unsupported_flags(key: str, value) -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer + + with pytest.raises(ValueError, match=f"{key}="): + Cosmos3VFMTransformer._validate_supported_config({key: value}) + Cosmos3VFMTransformer._validate_supported_config({}) + Cosmos3VFMTransformer._validate_supported_config(None) + + +def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer + + model = object.__new__(Cosmos3VFMTransformer) + nn.Module.__init__(model) + model.language_model = nn.Module() + model.language_model.layers = nn.ModuleList([nn.Linear(2, 2) for _ in range(2)]) + model.gen_layers = nn.ModuleList([nn.Linear(2, 2)]) + model.norm_moe_gen = nn.LayerNorm(2) + + matched = [ + name + for name, module in model.named_modules() + if any(condition(name, module) for condition in model._hsdp_shard_conditions) + ] + assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"] + assert Cosmos3VFMTransformer._layerwise_offload_blocks_attrs == ["gen_layers"] + assert Cosmos3VFMTransformer._repeated_blocks == ["Cosmos3GenDecoderLayer"] + + model.latent_patch_size = 2 + model.latent_channel_size = 3 + latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5) + torch.testing.assert_close(model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5), latents) + + +def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None)) + + output = transformer_cosmos3.Cosmos3VFMTransformer( + SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32) + )( + hidden_states=torch.zeros(1, 2, 1, 2, 2), + 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, + ) + + assert tuple(output.shape) == (1, 2, 1, 2, 2) + + +def test_compute_rope_freqs_places_text_and_video_positions() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer + + class FakeRotary: + def __init__(self) -> None: + self.position_ids: list[torch.Tensor] = [] + + def __call__(self, x, position_ids): + del x + self.position_ids.append(position_ids.detach().cpu()) + batch, seq = position_ids.shape[1], position_ids.shape[2] + return torch.zeros(batch, seq, 4), torch.ones(batch, seq, 4) + + rotary = FakeRotary() + model = object.__new__(Cosmos3VFMTransformer) + nn.Module.__init__(model) + model.language_model = SimpleNamespace(rotary_emb=rotary) + model.temporal_modality_margin = 100 + model.base_fps = 24.0 + model.temporal_compression_factor = 4 + model.enable_fps_modulation = False + + freqs_und, freqs_gen = model._compute_rope_freqs( + text_mask=torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.long), + t=2, + hp=1, + wp=1, + fps=24.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + text_pos, vision_pos = rotary.position_ids + assert text_pos[:, 0, :].tolist() == [[0, 1, 0], [0, 1, 0], [0, 1, 0]] + assert vision_pos[0, 0].tolist() == [102, 103] + assert freqs_und[0].shape == (2, 3, 1, 4) + assert freqs_gen[0].shape == (2, 2, 1, 4) diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index 2d44f6a6815..eb862a91ae0 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -10,6 +10,9 @@ from typing import Any import pytest +import torch + +from vllm_omni.diffusion.diffusion_engine import _move_tensor_tree_to_cpu @dataclass @@ -63,6 +66,80 @@ def update_from_output(self, sched_output, runner_output): return [req.request_id for req in sched_output.scheduled_new_reqs] +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +def test_move_tensor_tree_keeps_cpu_tensor_identity() -> None: + tensor = torch.arange(8, dtype=torch.float32) + + moved = _move_tensor_tree_to_cpu(tensor) + + assert moved is tensor + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +def test_move_tensor_tree_preserves_nested_structure_without_mutating_input() -> None: + tensor = torch.arange(4, dtype=torch.float32) + nested_tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3) + sentinel = object() + payload = { + "tensor": tensor, + "list": [nested_tensor, sentinel], + "tuple": ({"inner": tensor}, "metadata"), + "scalar": 3, + } + + moved = _move_tensor_tree_to_cpu(payload) + + assert moved is not payload + assert set(moved) == {"tensor", "list", "tuple", "scalar"} + assert moved["list"] is not payload["list"] + assert moved["tuple"] is not payload["tuple"] + assert moved["tuple"][0] is not payload["tuple"][0] + assert moved["tensor"] is tensor + assert moved["list"][0] is nested_tensor + assert moved["list"][1] is sentinel + assert moved["tuple"][0]["inner"] is tensor + assert moved["tuple"][1] == "metadata" + assert moved["scalar"] == 3 + assert payload["list"][0] is nested_tensor + assert payload["list"][1] is sentinel + assert payload["tuple"][0]["inner"] is tensor + assert payload["tuple"][1] == "metadata" + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +def test_move_tensor_tree_returns_non_tensor_values_unchanged() -> None: + value = object() + + moved = _move_tensor_tree_to_cpu(value) + + assert moved is value + + +@pytest.mark.diffusion +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_move_tensor_tree_moves_nested_cuda_tensors_to_cpu() -> None: + tensor = torch.arange(8, dtype=torch.float32, device="cuda") + other = torch.arange(4, dtype=torch.int64, device="cuda") + payload = {"tensor": tensor, "items": [other, ("keep", tensor)]} + + moved = _move_tensor_tree_to_cpu(payload) + + assert moved["tensor"].device.type == "cpu" + assert moved["items"][0].device.type == "cpu" + assert moved["items"][1][1].device.type == "cpu" + torch.testing.assert_close(moved["tensor"], tensor.cpu()) + torch.testing.assert_close(moved["items"][0], other.cpu()) + torch.testing.assert_close(moved["items"][1][1], tensor.cpu()) + assert moved["items"][1][0] == "keep" + + @pytest.mark.asyncio async def test_async_add_req_and_wait_for_response(): from vllm_omni.diffusion.diffusion_engine import DiffusionEngine diff --git a/tests/diffusion/test_diffusion_ipc.py b/tests/diffusion/test_diffusion_ipc.py new file mode 100644 index 00000000000..a99e42df782 --- /dev/null +++ b/tests/diffusion/test_diffusion_ipc.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import contextlib + +import pytest +import torch + +from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.ipc import ( + _SHM_TENSOR_THRESHOLD, + _pack_value_if_large, + _unpack_if_shm_handle, + pack_diffusion_output_shm, + unpack_diffusion_output_shm, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _large_numel(dtype: torch.dtype) -> int: + return (_SHM_TENSOR_THRESHOLD // torch.empty((), dtype=dtype).element_size()) + 1 + + +def _cleanup_shm_handle(value: object) -> None: + if isinstance(value, dict) and value.get("__tensor_shm__"): + with contextlib.suppress(FileNotFoundError): + _unpack_if_shm_handle(value) + + +def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None: + image = torch.arange(300_000, dtype=torch.float32) + video = torch.arange(300_000, dtype=torch.float32) * 2 + output = DiffusionOutput(output={"image": image, "video": video, "metadata": {"keep": "inline"}}) + + pack_diffusion_output_shm(output) + + assert output.output["image"]["__tensor_shm__"] is True + assert output.output["video"]["__tensor_shm__"] is True + assert output.output["metadata"] == {"keep": "inline"} + + unpack_diffusion_output_shm(output) + + torch.testing.assert_close(output.output["image"], image) + torch.testing.assert_close(output.output["video"], video) + assert output.output["metadata"] == {"keep": "inline"} + + +def test_diffusion_output_tuple_tensors_round_trip_through_shm() -> None: + # LTX2 / DreamID return (video, audio) tuples as DiffusionOutput.output. + video = torch.arange(300_000, dtype=torch.float32) + audio = torch.arange(300_000, dtype=torch.float32) * 3 + output = DiffusionOutput(output=(video, audio)) + + pack_diffusion_output_shm(output) + + assert isinstance(output.output, tuple) + assert output.output[0]["__tensor_shm__"] is True + assert output.output[1]["__tensor_shm__"] is True + + unpack_diffusion_output_shm(output) + + assert isinstance(output.output, tuple) + torch.testing.assert_close(output.output[0], video) + torch.testing.assert_close(output.output[1], audio) + + +def test_diffusion_output_list_tensors_round_trip_through_shm() -> None: + frames = [torch.arange(300_000, dtype=torch.float32), torch.arange(300_000, dtype=torch.float32) + 1] + output = DiffusionOutput(output=list(frames)) + + pack_diffusion_output_shm(output) + + assert isinstance(output.output, list) + assert all(isinstance(item, dict) and item["__tensor_shm__"] is True for item in output.output) + + unpack_diffusion_output_shm(output) + + assert isinstance(output.output, list) + torch.testing.assert_close(output.output[0], frames[0]) + torch.testing.assert_close(output.output[1], frames[1]) + + +def test_pack_value_keeps_tensor_at_threshold_inline() -> None: + tensor = torch.arange( + _SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(), + dtype=torch.float32, + ) + + packed = _pack_value_if_large(tensor) + + assert packed is tensor + + +def test_pack_value_packs_large_tensor_and_round_trips() -> None: + tensor = torch.arange(_large_numel(torch.float32), dtype=torch.float32) + packed = _pack_value_if_large(tensor) + + try: + assert isinstance(packed, dict) + assert packed["__tensor_shm__"] is True + assert packed["shape"] == [tensor.numel()] + assert packed["torch_dtype"] == "torch.float32" + + unpacked = _unpack_if_shm_handle(packed) + assert isinstance(unpacked, torch.Tensor) + torch.testing.assert_close(unpacked, tensor) + finally: + _cleanup_shm_handle(packed) + + +def test_pack_value_recurses_nested_dicts_and_lists_without_mutating_inline_values() -> None: + large = torch.arange(_large_numel(torch.float32), dtype=torch.float32) + small = torch.arange(8, dtype=torch.float32) + list_tensor = torch.arange(_large_numel(torch.float32), dtype=torch.float32) + payload = { + "media": { + "large": large, + "small": small, + }, + "list_value": [list_tensor], + "metadata": {"prompt": "keep inline"}, + } + + packed = _pack_value_if_large(payload) + + try: + assert packed is not payload + assert packed["media"] is not payload["media"] + assert packed["media"]["large"]["__tensor_shm__"] is True + assert packed["media"]["small"] is small + # Lists are recursed too: the large tensor inside is packed and a new + # list is returned, while the input payload is left untouched. + assert packed["list_value"] is not payload["list_value"] + assert packed["list_value"][0]["__tensor_shm__"] is True + assert payload["list_value"][0] is list_tensor + assert packed["metadata"] == {"prompt": "keep inline"} + + torch.testing.assert_close(_unpack_if_shm_handle(packed["media"]["large"]), large) + torch.testing.assert_close(_unpack_if_shm_handle(packed["list_value"][0]), list_tensor) + finally: + if isinstance(packed, dict): + _cleanup_shm_handle(packed.get("media", {}).get("large")) + list_value = packed.get("list_value") + if isinstance(list_value, list) and list_value: + _cleanup_shm_handle(list_value[0]) + + +def test_pack_value_preserves_dtype_shape_and_values_for_bfloat16() -> None: + tensor = torch.arange(_large_numel(torch.bfloat16), dtype=torch.float32).to(torch.bfloat16).reshape(1, -1) + packed = _pack_value_if_large(tensor) + + try: + assert isinstance(packed, dict) + assert packed["__tensor_shm__"] is True + assert packed["shape"] == list(tensor.shape) + assert packed["torch_dtype"] == "torch.bfloat16" + assert packed["numpy_dtype"] == "float32" + + unpacked = _unpack_if_shm_handle(packed) + assert isinstance(unpacked, torch.Tensor) + assert unpacked.dtype == torch.bfloat16 + torch.testing.assert_close(unpacked, tensor) + finally: + _cleanup_shm_handle(packed) + + +def test_pack_value_packs_non_contiguous_large_tensor_values() -> None: + tensor = torch.arange(_large_numel(torch.float32) * 2, dtype=torch.float32).reshape(-1, 2)[:, 0] + assert not tensor.is_contiguous() + + packed = _pack_value_if_large(tensor) + + try: + assert isinstance(packed, dict) + assert packed["__tensor_shm__"] is True + assert packed["shape"] == list(tensor.shape) + + unpacked = _unpack_if_shm_handle(packed) + assert isinstance(unpacked, torch.Tensor) + torch.testing.assert_close(unpacked, tensor) + finally: + _cleanup_shm_handle(packed) diff --git a/tests/e2e/accuracy/test_cosmos3_similarity.py b/tests/e2e/accuracy/test_cosmos3_similarity.py new file mode 100644 index 00000000000..ff2350096e0 --- /dev/null +++ b/tests/e2e/accuracy/test_cosmos3_similarity.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import base64 +import io +import json +import os +from pathlib import Path + +import pytest +import requests +import torch +from PIL import Image + +from tests.e2e.accuracy.helpers import model_output_dir +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServer + +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + +MODEL_ENV_VAR = "VLLM_TEST_COSMOS3_MODEL" +MODEL_ID = "cosmos3" +PROMPT = "A small warehouse robot moves a blue box across a clean floor." +NEGATIVE_PROMPT = "blurry, distorted, low quality" +SEED = 42 +WIDTH = HEIGHT = 256 +NUM_INFERENCE_STEPS = 2 + + +def _model_name() -> str: + model = os.environ.get(MODEL_ENV_VAR) + if not model: + pytest.skip(f"Set {MODEL_ENV_VAR} to run Cosmos3 full-model smoke tests.") + if not torch.cuda.is_available(): + pytest.skip("Cosmos3 full-model smoke tests require CUDA.") + return model + + +def _server_args() -> list[str]: + return [ + "--num-gpus", + "1", + "--model-class-name", + "Cosmos3OmniDiffusersPipeline", + "--stage-init-timeout", + "900", + "--init-timeout", + "1200", + ] + + +def _image_data_url(image: Image.Image) -> str: + buf = io.BytesIO() + image.save(buf, format="PNG") + return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode('ascii')}" + + +@pytest.mark.benchmark +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_cosmos3_t2i_serving_smoke(accuracy_artifact_root: Path) -> None: + output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) + with OmniServer(_model_name(), _server_args(), use_omni=True) as server: + response = requests.post( + f"http://{server.host}:{server.port}/v1/images/generations", + json={ + "model": server.model, + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "size": f"{WIDTH}x{HEIGHT}", + "n": 1, + "response_format": "b64_json", + "num_inference_steps": NUM_INFERENCE_STEPS, + "guidance_scale": 1.0, + "seed": SEED, + }, + timeout=1800, + ) + + response.raise_for_status() + payload = response.json() + assert len(payload["data"]) == 1 + image = Image.open(io.BytesIO(base64.b64decode(payload["data"][0]["b64_json"]))).convert("RGB") + image.save(output_dir / "cosmos3_t2i.png") + assert image.size == (WIDTH, HEIGHT) + + +@pytest.mark.parametrize( + ("name", "prompt", "num_frames", "image_reference"), + [ + ("t2v", PROMPT, "1", None), + ( + "i2v", + "The blue rectangle moves slowly forward.", + "5", + Image.new("RGB", (96, 64), color=(40, 80, 160)), + ), + ], +) +@pytest.mark.benchmark +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_cosmos3_video_serving_smoke( + accuracy_artifact_root: Path, + name: str, + prompt: str, + num_frames: str, + image_reference: Image.Image | None, +) -> None: + output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) + data = { + "model": "", + "prompt": prompt, + "negative_prompt": NEGATIVE_PROMPT, + "size": f"{WIDTH}x{HEIGHT}", + "num_frames": num_frames, + "fps": "1", + "num_inference_steps": str(NUM_INFERENCE_STEPS), + "guidance_scale": "1.0", + "seed": str(SEED), + } + if image_reference is not None: + data["image_reference"] = json.dumps({"image_url": _image_data_url(image_reference)}) + + with OmniServer(_model_name(), _server_args(), use_omni=True) as server: + data["model"] = server.model + response = requests.post(f"http://{server.host}:{server.port}/v1/videos/sync", data=data, timeout=1800) + + response.raise_for_status() + assert response.headers["content-type"].startswith("video/mp4") + assert response.content + (output_dir / f"cosmos3_{name}.mp4").write_bytes(response.content) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 2556ac71d34..d6a7dbaf74b 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -584,6 +584,7 @@ def test_generate_single_image(test_client): img_bytes = base64.b64decode(data["data"][0]["b64_json"]) img = Image.open(io.BytesIO(img_bytes)) assert img.size == (64, 64) # Our mock returns 64x64 images + assert test_client.app.state.engine_client.captured_prompt["modalities"] == ["image"] def test_generate_images_async_omni_sampling_params(async_omni_test_client): @@ -1011,6 +1012,42 @@ def test_with_custom_parameters(test_client): assert response.status_code == 200 +def test_flow_shift_forwarded_to_extra_args(test_client): + """flow_shift must reach the diffusion sampling params via extra_args. + + Regression: ``ImageGenerationRequest`` had no ``flow_shift`` field and the + single-stage handler never forwarded it, so a request like + ``{"flow_shift": 10.0}`` was silently dropped and Cosmos3 T2I always ran at + its hardcoded per-mode default shift. The pipeline reads + ``extra_args["flow_shift"]`` (via ``_get_sp_param``), so it must land there. + """ + response = test_client.post( + "/v1/images/generations", + json={ + "prompt": "a robot in a lab", + "size": "960x960", + "num_inference_steps": 50, + "guidance_scale": 4.0, + "flow_shift": 10.0, + }, + ) + assert response.status_code == 200 + captured = test_client.app.state.engine_client.captured_sampling_params_list[0] + assert captured.extra_args["flow_shift"] == 10.0 + + +def test_flow_shift_absent_when_not_requested(test_client): + """Omitting flow_shift must not inject an override, so the pipeline keeps + its per-mode default (e.g. Cosmos3 T2I shift=3.0).""" + response = test_client.post( + "/v1/images/generations", + json={"prompt": "a tree", "size": "1024x1024"}, + ) + assert response.status_code == 200 + captured = test_client.app.state.engine_client.captured_sampling_params_list[0] + assert "flow_shift" not in (captured.extra_args or {}) + + def test_invalid_size(test_client): """Test with invalid size parameter - rejected by Pydantic""" response = test_client.post( diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index a29f4493c28..36b19333980 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -243,6 +243,7 @@ def _fake_encode(video, fps, audio=None, audio_sample_rate=None, **kwargs): _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) engine = test_client.app.state.openai_serving_video._engine_client + assert engine.captured_prompt["modalities"] == ["video"] captured = engine.captured_sampling_params_list[0] assert captured.num_outputs_per_prompt == 1 assert captured.width == 640 @@ -1063,6 +1064,8 @@ def test_sync_t2v_returns_video_bytes(test_client, mocker: MockerFixture): assert float(response.headers["x-inference-time-s"]) >= 0 assert json.loads(response.headers["x-stage-durations"]) == {} assert float(response.headers["x-peak-memory-mb"]) == 0.0 + engine = test_client.app.state.openai_serving_video._engine_client + assert engine.captured_prompt["modalities"] == ["video"] def test_sync_t2v_returns_profiler_headers(test_client, mocker: MockerFixture): diff --git a/vllm_omni/diffusion/attention/backends/cudnn_attn.py b/vllm_omni/diffusion/attention/backends/cudnn_attn.py index f27fe18706f..357a2c0fd06 100644 --- a/vllm_omni/diffusion/attention/backends/cudnn_attn.py +++ b/vllm_omni/diffusion/attention/backends/cudnn_attn.py @@ -51,6 +51,7 @@ def __init__( ) -> None: self.causal = causal self.softmax_scale = softmax_scale + self.requires_gqa = num_heads != num_kv_heads def forward_cuda( self, @@ -84,6 +85,7 @@ def forward_cuda( dropout_p=0.0, is_causal=self.causal, scale=self.softmax_scale, + enable_gqa=self.requires_gqa, ) except RuntimeError as e: if "No available kernel" not in str(e): @@ -100,5 +102,6 @@ def forward_cuda( dropout_p=0.0, is_causal=self.causal, scale=self.softmax_scale, + enable_gqa=self.requires_gqa, ) return output.permute(0, 2, 1, 3) diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index 00585764292..6a2600e9fca 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/vllm_omni/diffusion/cache/cache_dit_backend.py @@ -1725,6 +1725,84 @@ def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool return refresh_cache_context +def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> Callable[[int], None]: + """Enable cache-dit for Cosmos3. + + Cosmos3 has a dual-pathway architecture (UND + GEN) but only the GEN + pathway (``gen_layers``) runs at every denoising step. The UND pathway + computes once and its K/V are cached by the pipeline itself; no cache-dit + needed there. We wrap only ``gen_layers`` via ``BlockAdapter``. + + Args: + pipeline: The Cosmos3 pipeline instance. + cache_config: DiffusionCacheConfig instance with cache configuration. + + Returns: + A refresh function that can be called to update cache context with new num_inference_steps. + """ + db_cache_config = _build_db_cache_config(cache_config) + + calibrator_config = None + if cache_config.enable_taylorseer: + taylorseer_order = cache_config.taylorseer_order + calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) + logger.info(f"TaylorSeer enabled with order={taylorseer_order}") + + logger.info( + f"Enabling cache-dit on Cosmos3 gen_layers: " + f"Fn={db_cache_config.Fn_compute_blocks}, " + f"Bn={db_cache_config.Bn_compute_blocks}, " + f"W={db_cache_config.max_warmup_steps}, " + ) + + cache_dit.enable_cache( + BlockAdapter( + transformer=pipeline.transformer, + blocks=[pipeline.transformer.gen_layers], + # Cosmos3 GEN blocks return only hidden_states. Per-layer UND K/V + # conditioning uses the transformer's cache-dit fallback path. + forward_pattern=[ForwardPattern.Pattern_3], + params_modifiers=[ + ParamsModifier( + cache_config=db_cache_config, + calibrator_config=calibrator_config, + ), + ], + check_forward_pattern=False, + has_separate_cfg=True, + ), + cache_config=db_cache_config, + calibrator_config=calibrator_config, + ) + + # The T2I denoising loop skips the unconditional forward outside the + # guidance interval as a speed optimization. cache-dit distinguishes the + # conditional vs unconditional passes purely by transformer-forward parity + # (has_separate_cfg=True above), so that skip would desync its per-generation + # step accounting. Still do both cond/uncond CFG steps when cache-dit is active. + # CFG is instead neutralized via scale=1.0 outside the interval. + pipeline._cache_dit_requires_paired_cfg = True + + def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + if cache_config.scm_steps_mask_policy is None: + cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) + else: + cache_dit.refresh_context( + pipeline.transformer, + cache_config=DBCacheConfig().reset( + num_inference_steps=num_inference_steps, + steps_computation_mask=cache_dit.steps_mask( + mask_policy=cache_config.scm_steps_mask_policy, + total_steps=num_inference_steps, + ), + steps_computation_policy=cache_config.scm_steps_policy, + ), + verbose=verbose, + ) + + return refresh_cache_context + + # Register custom cache-dit enablers after function definitions CUSTOM_DIT_ENABLERS.update( { @@ -1752,6 +1830,7 @@ def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool "HunyuanVideo15Pipeline": enable_cache_for_hunyuan_video_15, "HunyuanVideo15I2VPipeline": enable_cache_for_hunyuan_video_15, "HeliosPipeline": enable_cache_for_helios, + "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3, } ) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 1805b5dc349..4b4847e2540 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -83,6 +83,18 @@ def supports_audio_output(model_class_name: str) -> bool: return bool(getattr(model_cls, "support_audio_output", False)) +def _move_tensor_tree_to_cpu(value: object) -> object: + if isinstance(value, torch.Tensor): + return value.cpu() if value.device.type != "cpu" else value + if isinstance(value, dict): + return {key: _move_tensor_tree_to_cpu(item) for key, item in value.items()} + if isinstance(value, list): + return [_move_tensor_tree_to_cpu(item) for item in value] + if isinstance(value, tuple): + return tuple(_move_tensor_tree_to_cpu(item) for item in value) + return value + + def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool) -> int: model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) if model_cls is not None and hasattr(model_cls, "dummy_run_num_frames"): @@ -237,12 +249,8 @@ async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]: # post-processing to avoid device OOM — model weights may still # reside on the device and leave no headroom for intermediates. output_data = output.output - if ( - self.od_config.enable_cpu_offload - and isinstance(output_data, torch.Tensor) - and output_data.device.type != "cpu" - ): - output_data = output_data.cpu() + if self.od_config.enable_cpu_offload: + output_data = _move_tensor_tree_to_cpu(output_data) postprocess_start_time = time.perf_counter() if self.post_process_func is not None: diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index 6a96533fd40..ec5be5b833f 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/vllm_omni/diffusion/ipc.py @@ -85,16 +85,42 @@ def _pack_tensor_if_large(val: torch.Tensor) -> torch.Tensor | dict: return val +def _pack_value_if_large(val: object) -> object: + """Recursively replace large tensors with SHM handles. + + Walks the container shapes pipelines return as ``DiffusionOutput.output``: + bare tensors, dicts (e.g. Cosmos3 ``{"image"/"video": ...}``), and + tuples/lists (e.g. LTX2 and DreamID ``(video, audio)``). Other values pass + through unchanged. ``_unpack_if_shm_handle`` must mirror these shapes — keep + the two in sync. + """ + if isinstance(val, torch.Tensor): + return _pack_tensor_if_large(val) + if isinstance(val, dict): + return {key: _pack_value_if_large(value) for key, value in val.items()} + if isinstance(val, list): + return [_pack_value_if_large(item) for item in val] + if isinstance(val, tuple): + return tuple(_pack_value_if_large(item) for item in val) + return val + + def _unpack_if_shm_handle(val: object) -> object: - """Reconstruct a tensor from an SHM handle dict, or return as-is.""" + """Reconstruct tensors from SHM handles, mirroring ``_pack_value_if_large``.""" if isinstance(val, dict) and val.get("__tensor_shm__"): return _tensor_from_shm(val) + if isinstance(val, dict): + return {key: _unpack_if_shm_handle(value) for key, value in val.items()} + if isinstance(val, list): + return [_unpack_if_shm_handle(item) for item in val] + if isinstance(val, tuple): + return tuple(_unpack_if_shm_handle(item) for item in val) return val def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput: - if output.output is not None and isinstance(output.output, torch.Tensor): - output.output = _pack_tensor_if_large(output.output) + if output.output is not None: + output.output = _pack_value_if_large(output.output) if output.trajectory_latents is not None and isinstance(output.trajectory_latents, torch.Tensor): output.trajectory_latents = _pack_tensor_if_large(output.trajectory_latents) if output.trajectory_timesteps is not None and isinstance(output.trajectory_timesteps, torch.Tensor): diff --git a/vllm_omni/diffusion/models/cosmos3/__init__.py b/vllm_omni/diffusion/models/cosmos3/__init__.py new file mode 100644 index 00000000000..6df062b5c0d --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/__init__.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + get_cosmos3_post_process_func, + get_cosmos3_pre_process_func, +) +from .transformer_cosmos3 import Cosmos3VFMTransformer + +__all__ = [ + "Cosmos3OmniDiffusersPipeline", + "get_cosmos3_post_process_func", + "get_cosmos3_pre_process_func", + "Cosmos3VFMTransformer", +] diff --git a/vllm_omni/diffusion/models/cosmos3/guardrails.py b/vllm_omni/diffusion/models/cosmos3/guardrails.py new file mode 100644 index 00000000000..3458c2cc372 --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/guardrails.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cosmos3 guardrail hooks for vllm-omni. + +Thin adapter around the ``cosmos_guardrail`` package's ``CosmosSafetyChecker`` +(Blocklist + Qwen3Guard for text, RetinaFace face-blur for video). + +Enabled by default. Disable server-wide with ``--cosmos3-no-guardrails`` (which +sets ``od_config.model_config["guardrails"] = False``); per-request overrides +ride on ``sampling_params.extra_args["guardrails"]``. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.models.progress_bar import _is_rank_zero + +if TYPE_CHECKING: + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +logger = init_logger(__name__) + + +try: + from cosmos_guardrail import CosmosSafetyChecker + + _COSMOS_GUARDRAIL_AVAILABLE = True +except ImportError: + _COSMOS_GUARDRAIL_AVAILABLE = False + + class CosmosSafetyChecker: # type: ignore[no-redef] + # Raised at runtime (not import time) so guardrail-less inference + # continues to work when ``cosmos_guardrail`` is not installed and + # ``model_config["guardrails"]`` is False. + def __init__(self, *args, **kwargs): + raise ValueError( + f"You have disabled the safety checker for {self.__class__}. This is in violation of the " + "[NVIDIA Open Model License Agreement]" + "(https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license). " + f"Please ensure that you are compliant with the license agreement." + f"Please install cosmos-guardrail package to enable safety checks." + ) + + +TextGuardrailFn = Callable[[str], None] +VideoGuardrailFn = Callable[[np.ndarray], np.ndarray] + +_text_guardrail: TextGuardrailFn | None = None +_video_guardrail: VideoGuardrailFn | None = None + + +# --------------------------------------------------------------------------- +# Default guardrail builders +# --------------------------------------------------------------------------- +def _nn_models(runner: Any) -> list[torch.nn.Module]: + return [m for m in runner.models if isinstance(m, torch.nn.Module)] + + +def _build_text_guardrail(checker: Any) -> TextGuardrailFn: + def text_guardrail(prompt: str) -> None: + if not checker.check_text_safety(prompt): + # CosmosSafetyChecker logs the specific reason at CRITICAL. + raise ValueError("Guardrail blocked prompt") + + return text_guardrail + + +def _build_video_guardrail(checker: Any, offload_to_cpu: bool) -> VideoGuardrailFn: + video_models = _nn_models(checker.video_guardrail) + compute_device = "cuda" + + def video_guardrail(frames: np.ndarray) -> np.ndarray: + if offload_to_cpu: + for m in video_models: + m.to(compute_device) + try: + result = checker.check_video_safety(frames) + finally: + if offload_to_cpu: + for m in video_models: + m.to("cpu") + # ``check_video_safety`` returns ``None`` when the content safety + # filter blocks the frames. The face-blur postprocessor (the only + # video module enabled by default) does not block, so in practice + # ``result`` is always an ndarray here. + return result if result is not None else frames + + return video_guardrail + + +# --------------------------------------------------------------------------- +# Initialization +# --------------------------------------------------------------------------- +def _init_default_guardrails(offload_to_cpu: bool = False) -> None: + global _text_guardrail, _video_guardrail + if _text_guardrail is not None: + return + if _is_rank_zero(): + logger.info("Initializing Cosmos3 guardrails (offload_to_cpu=%s)...", offload_to_cpu) + + # Instantiation raises ValueError when ``cosmos_guardrail`` is not + # installed - this is the right moment to fail loudly because the + # caller has opted in to guardrails. + checker = CosmosSafetyChecker() + + # Place text models on their resting device permanently. Video models + # idle on CPU when offload is on and move to GPU per-call (handled in + # the video guardrail closure). + idle_device = "cpu" if offload_to_cpu else "cuda" + for m in _nn_models(checker.text_guardrail): + m.to(idle_device) + for m in _nn_models(checker.video_guardrail): + m.to(idle_device) + + _text_guardrail = _build_text_guardrail(checker) + _video_guardrail = _build_video_guardrail(checker, offload_to_cpu) + if _is_rank_zero(): + logger.info("Cosmos3 guardrails initialized.") + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- +def ensure_initialized(od_config: OmniDiffusionConfig) -> None: + if not is_guardrails_enabled(od_config): + return + model_config = od_config.model_config or {} + _init_default_guardrails(offload_to_cpu=bool(model_config.get("offload_guardrail_models", False))) + + +def check_text_safety(prompt: str) -> None: + if _text_guardrail is not None: + _text_guardrail(prompt) + + +def check_video_safety(video_tensor: torch.Tensor) -> torch.Tensor: + if _video_guardrail is None: + return video_tensor + + v = video_tensor.detach().cpu().float() + if v.dim() == 5: + v = v[0] + v = v.clamp(-1, 1) * 0.5 + 0.5 + frames_np = (v.permute(1, 2, 3, 0).numpy() * 255).round().astype(np.uint8) + + frames_np = _video_guardrail(frames_np) + + # Convert back to [-1, 1] to match the VAE output range. + result = torch.from_numpy(frames_np.copy()).float() / 127.5 - 1.0 + result = result.permute(3, 0, 1, 2) + if video_tensor.dim() == 5: + result = result.unsqueeze(0) + return result.to(video_tensor.device) + + +def is_guardrails_enabled( + od_config: OmniDiffusionConfig, + sampling_params: OmniDiffusionSamplingParams | None = None, +) -> bool: + """Resolve the active guardrail gate. + + Server-level ``od_config.model_config["guardrails"]`` decides whether the + guardrail models are loaded at all (eager load at pipeline build time). + When that is False, no per-request override can turn checks back on, + because the singletons in this module are never populated. + + When the server gate is on, ``sampling_params.extra_args["guardrails"]`` + may override on a per-request basis: ``False`` skips the check for that + request, anything else (or missing) keeps the default behavior. + """ + model_config = od_config.model_config or {} + if not bool(model_config.get("guardrails", True)): + return False + if sampling_params is not None: + per_request = (sampling_params.extra_args or {}).get("guardrails") + if per_request is not None: + return bool(per_request) + return True diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py new file mode 100644 index 00000000000..102b9216082 --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -0,0 +1,1218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cosmos3 text/image-to-video and text-to-image pipeline for vllm-omni. + +Single pipeline class supports T2V, I2V, and T2I; the mode is selected at +runtime by: + +* ``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`) + +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Iterable +from typing import Any, ClassVar + +import numpy as np +import PIL.Image +import torch +from diffusers import UniPCMultistepScheduler +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from torch import nn +from transformers import AutoTokenizer +from vllm.logger import init_logger +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_world_size, +) +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.interface import SupportImageInput +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.inputs.data import OmniDiffusionSamplingParams + +from .transformer_cosmos3 import Cosmos3VFMTransformer + +logger = init_logger(__name__) + +COSMOS3_DURATION_TEMPLATE = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS." +COSMOS3_RESOLUTION_TEMPLATE = "This video is of {height}x{width} resolution." +COSMOS3_IMAGE_RESOLUTION_TEMPLATE = "This image is of {height}x{width} resolution." +COSMOS3_INVERSE_DURATION_TEMPLATE = "The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS." +COSMOS3_INVERSE_RESOLUTION_TEMPLATE = "This video is not of {height}x{width} resolution." +COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE = "This image is not of {height}x{width} resolution." +COSMOS3_SYSTEM_PROMPT = "You are a helpful assistant who will generate videos from a given prompt." +COSMOS3_T2I_SYSTEM_PROMPT = "You are a helpful assistant who will generate images from a given prompt." + +COSMOS3_T2V_DEFAULT_HEIGHT = 720 +COSMOS3_T2V_DEFAULT_WIDTH = 1280 +COSMOS3_T2V_DEFAULT_NUM_FRAMES = 189 +COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS = 35 +COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE = 6.0 + +COSMOS3_T2I_DEFAULT_HEIGHT = 1024 +COSMOS3_T2I_DEFAULT_WIDTH = 1024 +COSMOS3_T2I_DEFAULT_NUM_INFERENCE_STEPS = 50 +COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE = 7.0 +COSMOS3_T2I_DEFAULT_FLOW_SHIFT = 3.0 +COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL: tuple[float, float] = (400.0, 1000.0) + +# Truncation cap on the prompt token count (shared by T2I and T2V). Prompts +# are tokenized to their natural length (no padding); this only bounds the +# UND pathway / GEN cross-attention cost for pathologically long prompts. +COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096 + + +# --------------------------------------------------------------------------- +# Post-process function (registered in registry.py) +# --------------------------------------------------------------------------- +def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): + """Pre-process function for both T2V and I2V. + + 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``. + """ + from .guardrails import check_text_safety, ensure_initialized, is_guardrails_enabled + + video_processor = VideoProcessor(vae_scale_factor=16) + # Eager-load guardrail models at pipeline build time when the server-level + # gate is on. Per-request overrides only decide whether the loaded models + # are *invoked* — they cannot turn checks on without a server-side preload. + if is_guardrails_enabled(od_config): + ensure_initialized(od_config) + + def _pil_to_rgb(value: Any) -> PIL.Image.Image: + if isinstance(value, str): + return PIL.Image.open(value).convert("RGB") + if isinstance(value, PIL.Image.Image): + return value.convert("RGB") + raise TypeError(f"Cosmos3 preprocessing expected PIL image or image path, got {type(value)!r}.") + + def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: + if is_guardrails_enabled(od_config, request.sampling_params): + for prompt in request.prompts: + text = prompt if isinstance(prompt, str) else prompt.get("prompt", "") + check_text_safety(text) + + for i, prompt in enumerate(request.prompts): + if isinstance(prompt, str): + continue + multi_modal_data = prompt.get("multi_modal_data", {}) or {} + raw_image = multi_modal_data.get("image") + if raw_image is None: + continue + + if "additional_information" not in prompt: + prompt["additional_information"] = {} + + image = _pil_to_rgb(raw_image) + + # Auto-calculate H/W from aspect ratio (720p max area) + if request.sampling_params.height is None or request.sampling_params.width is None: + max_area = 720 * 1280 + aspect_ratio = image.height / image.width + mod_value = 16 + height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value + width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value + if request.sampling_params.height is None: + request.sampling_params.height = height + if request.sampling_params.width is None: + request.sampling_params.width = width + + target_w = request.sampling_params.width + target_h = request.sampling_params.height + scale = max(target_w / image.width, target_h / image.height) + resize_w = int(np.ceil(scale * image.width)) + resize_h = int(np.ceil(scale * image.height)) + image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS) + left = (resize_w - target_w) // 2 + top = (resize_h - target_h) // 2 + image = image.crop((left, top, left + target_w, top + target_h)) + + prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess( + image, height=target_h, width=target_w + ) + request.prompts[i] = prompt + + return request + + return pre_process_func + + +def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig): + from .guardrails import check_video_safety, is_guardrails_enabled + + video_processor = VideoProcessor(vae_scale_factor=16) + + def post_process_func( + output: torch.Tensor | dict[str, torch.Tensor] | tuple, + output_type: str = "np", + sampling_params=None, + ): + if output_type == "latent": + return output + + if isinstance(output, dict): + if "image" in output and "video" in output: + raise ValueError("Cosmos3 output cannot contain both image and video payloads.") + if "image" in output: + video = output["image"] + elif "video" in output: + video = output["video"] + else: + raise ValueError("Cosmos3 postprocess expected an 'image' or 'video' output payload.") + else: + video = output + + if isinstance(output, dict) and "image" in output: + if video.ndim != 5 or video.shape[2] != 1: + raise ValueError( + "Cosmos3 text-to-image postprocess expects decoded output " + f"with shape [B, C, 1, H, W], got {tuple(video.shape)}." + ) + image = video.squeeze(2) # [B, 3, H, W] + if is_guardrails_enabled(od_config, sampling_params): + # check_video_safety expects a 5D tensor; re-add T axis. + checked = check_video_safety(image.unsqueeze(2)) + image = checked.squeeze(2) + return video_processor.postprocess(image, output_type="pil") + if is_guardrails_enabled(od_config, sampling_params): + video = check_video_safety(video) + return video_processor.postprocess_video(video, output_type=output_type) + + return post_process_func + + +# --------------------------------------------------------------------------- +# Pipeline +# --------------------------------------------------------------------------- +class Cosmos3OmniDiffusersPipeline( + nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin +): + """Cosmos3 text/image-to-video / text-to-image 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) + + Supports T2V, I2V, and T2I 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, + flow_shift=3.0, guidance_interval=[400, 1000]), the duration + template is suppressed, and post-process emits PIL images. + * **I2V** when the request supplies a preprocessed image via + ``multi_modal_data['image']`` (handled by + :func:`get_cosmos3_pre_process_func`) and the requested output modality + is not image. + Frame 0 of the initial latent is set to the VAE-encoded conditioning + image, frame-0 noise predictions are masked to zero, and the clean + image latent is re-injected at frame 0 after each scheduler step. + * **T2V** otherwise (default video generation). + """ + + support_image_input: ClassVar[bool] = True + color_format: ClassVar[str] = "RGB" + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ) -> None: + super().__init__() + if od_config.enable_cpu_offload: + raise ValueError( + "Cosmos3 has no separate text encoder, so CPU offloading " + "(transformer↔encoder swapping) is not supported. " + "Use --enable-layerwise-offload instead." + ) + self.od_config = od_config + self.device = get_local_device() + self.dtype = od_config.dtype + + model_path = od_config.model + local_files_only = os.path.exists(model_path) + + # --- Tokenizer --- + self.tokenizer = AutoTokenizer.from_pretrained( + model_path, + subfolder="text_tokenizer", + local_files_only=local_files_only, + ) + + # --- VAE --- + self.vae = DistributedAutoencoderKLWan.from_pretrained( + model_path, + subfolder="vae", + torch_dtype=self.dtype, + local_files_only=local_files_only, + ).to(self.device) + + if not hasattr(self.vae.config, "scale_factor_temporal"): + raise ValueError( + "Cosmos3 Diffusers VAE config must define scale_factor_temporal " + "so transformer mRoPE temporal positions can be computed correctly." + ) + self.vae_scale_factor_temporal = int(self.vae.config.scale_factor_temporal) + self.vae_scale_factor_spatial = getattr(self.vae.config, "scale_factor_spatial", 16) + + # --- Transformer (weights loaded later via weights_sources) --- + self.transformer = Cosmos3VFMTransformer( + od_config=od_config, + temporal_compression_factor=self.vae_scale_factor_temporal, + ) + + # --- Scheduler --- + # Load from checkpoint to preserve solver_order, timestep_spacing, + # beta_schedule, sigma bounds, flow_shift, etc. Only override + # flow_shift when explicitly requested by the user. + self.scheduler = UniPCMultistepScheduler.from_pretrained( + model_path, + subfolder="scheduler", + local_files_only=local_files_only, + ) + if od_config.flow_shift is not None: + self.scheduler = UniPCMultistepScheduler.from_config(self.scheduler.config, flow_shift=od_config.flow_shift) + self._cpu_scheduler_state() + + # --- Video processor for post-decode --- + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + + # --- Weight sources for DiffusersPipelineLoader --- + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=model_path, + subfolder=None, + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + allow_patterns_overrides=["transformer/*.safetensors"], + ), + ] + + # Snapshot the loaded scheduler config so we can rebuild the + # scheduler at request time when a per-request flow_shift override + # is supplied (T2I uses shift=3.0; T2V/I2V use the engine default). + self._base_scheduler_config = self.scheduler.config + self._engine_init_flow_shift = float(getattr(self.scheduler.config, "flow_shift", 1.0) or 1.0) + self._current_flow_shift = self._engine_init_flow_shift + + self._guidance_scale = None + self._num_timesteps = None + + # Set True by ``enable_cache_for_cosmos3`` when cache-dit is enabled on + # this pipeline. Tells the sequential-CFG loop to keep paired + # cond/uncond forwards so cache-dit's has_separate_cfg step accounting + # stays in sync. + self._cache_dit_requires_paired_cfg = False + + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) + + # -- Weight loading -------------------------------------------------------- + + @staticmethod + def _remap_ckpt_key(key: str) -> str | None: + """Remap a Diffusers transformer key to the model parameter namespace. + + Checkpoint keys arrive with a synthetic ``transformer.`` prefix from + ``weights_sources``. The source checkpoint itself uses the prefixless + Diffusers transformer namespace: top-level projections plus Qwen3-VL + backbone keys. UND and GEN components share each layer in the source + and are split into separate module lists here. Some sources wrap the + transformer namespace under ``model.``; that wrapper is structural and + is stripped before applying the Cosmos3 leaf-name remap. + + Returns the remapped name under ``transformer.``, or None to skip. + """ + k = key + # Strip the weights_sources prefix + if k.startswith("transformer."): + k = k[len("transformer.") :] + if k.startswith("model."): + k = k[len("model.") :] + + # Top-level generation components. + if k.startswith( + ( + "proj_in.", + "proj_out.", + "time_embedder.", + ) + ): + return f"transformer.{k}" + + # Skip lm_head + if k.startswith("lm_head."): + return None + + # embed_tokens / norm -> language_model.* + if k.startswith("embed_tokens."): + return f"transformer.language_model.{k}" + if k.startswith("norm."): + return f"transformer.language_model.{k}" + + # norm_moe_gen -> top level + if k.startswith("norm_moe_gen."): + return f"transformer.{k}" + + if not k.startswith("layers."): + return None + + parts = k.split(".", 2) # ['layers', '{i}', '{rest}'] + if len(parts) != 3: + return None + layer_idx = parts[1] + rest = parts[2] + + und_lp = f"transformer.language_model.layers.{layer_idx}" + gen_lp = f"transformer.gen_layers.{layer_idx}" + + _LAYER_MAP = { + # UND attention + "self_attn.to_q.": f"{und_lp}.self_attn.to_q.", + "self_attn.to_k.": f"{und_lp}.self_attn.to_k.", + "self_attn.to_v.": f"{und_lp}.self_attn.to_v.", + "self_attn.to_out.": f"{und_lp}.self_attn.to_out.", + "self_attn.norm_q.": f"{und_lp}.self_attn.norm_q.", + "self_attn.norm_k.": f"{und_lp}.self_attn.norm_k.", + # GEN attention + "self_attn.add_q_proj.": f"{gen_lp}.cross_attention.to_q.", + "self_attn.add_k_proj.": f"{gen_lp}.cross_attention.to_k.", + "self_attn.add_v_proj.": f"{gen_lp}.cross_attention.to_v.", + "self_attn.to_add_out.": f"{gen_lp}.cross_attention.to_out.", + "self_attn.norm_added_q.": f"{gen_lp}.cross_attention.norm_q.", + "self_attn.norm_added_k.": f"{gen_lp}.cross_attention.norm_k.", + # Norms + "input_layernorm.": f"{und_lp}.input_layernorm.", + "post_attention_layernorm.": f"{und_lp}.post_attention_layernorm.", + "input_layernorm_moe_gen.": f"{gen_lp}.input_layernorm.", + "post_attention_layernorm_moe_gen.": f"{gen_lp}.post_attention_layernorm.", + # UND MLP + "mlp.gate_proj.": f"{und_lp}.mlp.gate_proj.", + "mlp.up_proj.": f"{und_lp}.mlp.up_proj.", + "mlp.down_proj.": f"{und_lp}.mlp.down_proj.", + # GEN MLP + "mlp_moe_gen.gate_proj.": f"{gen_lp}.mlp.gate_proj.", + "mlp_moe_gen.up_proj.": f"{gen_lp}.mlp.up_proj.", + "mlp_moe_gen.down_proj.": f"{gen_lp}.mlp.down_proj.", + } + + for pattern, replacement in _LAYER_MAP.items(): + if rest.startswith(pattern): + suffix = rest[len(pattern) :] + return replacement + suffix + + return None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Stream-remap checkpoint weights and load via AutoWeightsLoader. + + Handles quantization, TP-aware weight_loader, and buffer loading. + Returns the set of loaded parameter names for strict validation. + """ + state = self.state_dict() + allowed = set(state.keys()) + tp_aware = {n for n, p in self.named_parameters() if hasattr(p, "weight_loader")} + + def _remapped_weights() -> Iterable[tuple[str, torch.Tensor]]: + total = kept = 0 + for name, tensor in weights: + total += 1 + remapped = self._remap_ckpt_key(name) + if remapped is not None and (remapped in allowed or remapped in tp_aware): + kept += 1 + yield remapped, tensor + if _is_rank_zero(): + logger.info( + "Cosmos3 weight remap: kept %d/%d tensors", + kept, + total, + ) + + loader = AutoWeightsLoader(self) + loaded = loader.load_weights(_remapped_weights()) + self.transformer.post_load_weights() + self.transformer.eval() + return loaded + + def predict_noise(self, **kwargs) -> torch.Tensor | tuple[torch.Tensor, ...]: + """Override CFGParallelMixin.predict_noise for Cosmos3. + + The transformer returns the raw video noise prediction. + """ + return self.transformer(**kwargs) + + @staticmethod + def _cfg_parallel_active() -> bool: + try: + return get_classifier_free_guidance_world_size() > 1 + except Exception: + return False + + def _cache_requires_paired_cfg(self) -> bool: + """Whether the sequential-CFG denoising loop must keep paired forwards. + + cache-dit wraps the GEN pathway with ``has_separate_cfg=True`` and + distinguishes the conditional vs unconditional passes purely by the + parity of its transformer-forward counter. The T2I ``guidance_interval`` + optimization that skips the uncond pass outside the interval would + desync that accounting (cond passes get mislabeled as uncond and the + per-generation step counter drifts). ``enable_cache_for_cosmos3`` sets + the marker below when it enables cache-dit on this pipeline; the loop + then keeps both passes and neutralizes CFG via scale=1.0 instead. + + Returns False when cache-dit is not active, preserving the skip speedup. + """ + return self._cache_dit_requires_paired_cfg + + @staticmethod + def _get_sp_param(sp: OmniDiffusionSamplingParams, key: str, default: Any = None) -> Any: + """Read a runtime control from sampling params. + + Order of precedence: + 1. ``sp.extra_args[key]`` - preferred path; the OpenAI image/video + endpoints surface custom controls here (see e.g. + ``serving_video.py`` writing ``extra_args['flow_shift']``). + 2. direct attribute on ``sp`` - backward compat for callers that + set attributes directly. + 3. ``default``. + + Skipping this helper would cause API-driven overrides like + ``request.flow_shift`` (forwarded as ``extra_args['flow_shift']``) to + be silently ignored. + """ + extra = sp.extra_args or {} + if extra.get(key) is not None: + return extra[key] + val = getattr(sp, key, None) + if val is not None: + return val + return default + + @staticmethod + def _is_t2i_request(req: OmniDiffusionRequest) -> bool: + """Detect text-to-image mode from request-level prompt modalities.""" + if not req.prompts: + return False + first_prompt = req.prompts[0] + modalities = first_prompt.get("modalities", []) if isinstance(first_prompt, dict) else [] + if modalities is None: + modalities = [] + if isinstance(modalities, str): + modalities = [modalities] + if "image" in modalities and "video" in modalities: + raise ValueError("Cosmos3 prompt modalities cannot request both image and video output.") + + accepted_modalities = ["image", "video", "text", "audio"] + if any(x not in accepted_modalities for x in modalities): + raise ValueError(f"Incorrect modality value in {modalities}, expected one of {accepted_modalities}.") + return "image" in modalities + + def _set_flow_shift(self, target_shift: float) -> None: + """Set the UniPC ``flow_shift`` to a concrete target value. + + The scheduler is rebuilt from the saved base config if + the target differs from the current shift. Tracking + ``self._current_flow_shift`` explicitly is required because the + previous mode may have rebuilt the scheduler - we cannot rely on + ``self.scheduler.config.flow_shift`` reflecting the last requested + target if a rebuild was skipped via the equality check. + """ + target = float(target_shift) + if target == float(self._current_flow_shift): + return + self.scheduler = UniPCMultistepScheduler.from_config(self._base_scheduler_config, flow_shift=target) + self._cpu_scheduler_state() + self._current_flow_shift = target + + def _cpu_scheduler_state(self) -> None: + # We need to move scheduler tensors to CPU, as unipc from diffusers assumes they are on CPU. + # However, after the creation they are on GPU due to "with target_device:" in diffusers_loader.py + for name, value in vars(self.scheduler).items(): + if isinstance(value, torch.Tensor) and value.device.type != "cpu": + setattr(self.scheduler, name, value.cpu()) + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale is not None and self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + # -- Prompt formatting ----------------------------------------------------- + + @staticmethod + def _apply_metadata_templates( + prompt: str, + num_frames: int, + frame_rate: float, + height: int, + width: int, + duration_template: str | None = COSMOS3_DURATION_TEMPLATE, + resolution_template: str | None = COSMOS3_RESOLUTION_TEMPLATE, + force_duration_template: bool = False, + ) -> str: + """ + Append duration and resolution metadata to a prompt. + """ + parts: list[str] = [] + head = prompt.rstrip(".").strip() + if head: + parts.append(head) + if duration_template is not None and (num_frames > 1 or force_duration_template): + duration = num_frames / frame_rate + parts.append(duration_template.format(duration=duration, fps=frame_rate).rstrip(".")) + if resolution_template is not None: + parts.append(resolution_template.format(height=height, width=width).rstrip(".")) + if not parts: + return "" + return ". ".join(parts) + "." + + # -- Tokenization -------------------------------------------------------- + + def _tokenize_prompt( + self, + text: str, + max_sequence_length: int, + use_system_prompt: bool = False, + system_prompt: str | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Tokenize a prompt using the Qwen2 chat template. + + Returns (input_ids, attention_mask) as [1, S] tensors on device. + """ + conversations = [] + if use_system_prompt: + conversations.append( + { + "role": "system", + "content": system_prompt or COSMOS3_SYSTEM_PROMPT, + } + ) + conversations.append({"role": "user", "content": text}) + + token_ids = self._normalize_token_ids( + self.tokenizer.apply_chat_template(conversations, tokenize=True, add_generation_prompt=True) + ) + original_token_count = len(token_ids) + if original_token_count > max_sequence_length and _is_rank_zero(): + logger.warning( + "Cosmos3 prompt token_ids shortened to max_sequence_length: " + "original_token_count=%d, max_sequence_length=%d, removed_token_count=%d", + original_token_count, + max_sequence_length, + original_token_count - max_sequence_length, + ) + token_ids = token_ids[:max_sequence_length] + token_ids.append(self.tokenizer.eos_token_id) # 151645 + token_ids.append(self.tokenizer.convert_tokens_to_ids("<|vision_start|>")) # 151652 + seq_len = len(token_ids) + + # No right-padding: the prompt is tokenized to its natural length. + # The UND pathway uses causal self-attention with no padding mask and + # the GEN cross-attention K/V is trimmed to the real text length, so + # padding to a fixed length only added dead compute and never changed + # the output. ``max_sequence_length`` is kept purely as a truncation + # cap (above). The mask is therefore all ones. + attention_mask = [1] * seq_len + + input_ids = torch.tensor([token_ids], dtype=torch.long, device=self.device) + attention_mask = torch.tensor([attention_mask], dtype=torch.long, device=self.device) + return input_ids, attention_mask + + @staticmethod + def _normalize_token_ids(tokenized_output: object) -> list[int]: + """Normalize tokenizer outputs into a flat ``list[int]``. + + Different Transformers/tokenizers versions can return ``list[int]``, + a mapping/BatchEncoding with ``input_ids``, tensors, or + ``tokenizers.Encoding`` objects from ``apply_chat_template``. + """ + token_ids = tokenized_output + while True: + if isinstance(token_ids, dict) and "input_ids" in token_ids: + token_ids = token_ids["input_ids"] + elif hasattr(token_ids, "input_ids"): + token_ids = token_ids.input_ids + elif hasattr(token_ids, "ids"): + token_ids = token_ids.ids + elif hasattr(token_ids, "tolist"): + token_ids = token_ids.tolist() + elif isinstance(token_ids, tuple): + token_ids = list(token_ids) + elif isinstance(token_ids, list) and len(token_ids) == 1: + first = token_ids[0] + if isinstance(first, list | tuple): + token_ids = list(first) + elif hasattr(first, "ids") or hasattr(first, "input_ids"): + token_ids = first + elif hasattr(first, "tolist"): + first_list = first.tolist() + if isinstance(first_list, list | tuple): + token_ids = list(first_list) + else: + break + else: + break + else: + break + + if not isinstance(token_ids, list): + raise TypeError( + "Cosmos3 tokenizer must return token IDs as a list-like value; " + f"got {type(token_ids).__name__}: {token_ids!r}" + ) + + normalized_ids = [] + for idx, token_id in enumerate(token_ids): + if hasattr(token_id, "item"): + token_id = token_id.item() + try: + normalized_ids.append(int(token_id)) + except (TypeError, ValueError) as exc: + raise TypeError( + "Cosmos3 tokenizer returned a non-integer token at " + f"index {idx}: {type(token_id).__name__}: {token_id!r}" + ) from exc + return normalized_ids + + # -- Latent preparation -------------------------------------------------- + + def _prepare_latents( + self, + height: int, + width: int, + num_frames: int, + generator: torch.Generator, + ) -> torch.Tensor: + num_channels_latents = self.transformer.latent_channel_size + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + shape = ( + 1, + num_channels_latents, + num_latent_frames, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + ) + return randn_tensor(shape, generator=generator, device=self.device, dtype=self.dtype) + + # -- VAE decode ---------------------------------------------------------- + + def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: + latents = latents.to(self.vae.dtype) + + if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"): + if not hasattr(self, "_latents_mean"): + self._latents_mean = ( + torch.tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1).to(self.device, self.vae.dtype) + ) + self._latents_std = ( + torch.tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1).to(self.device, self.vae.dtype) + ) + latents = (latents * self._latents_std) + self._latents_mean + else: + scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0) + latents = latents / scaling_factor + + video = self.vae.decode(latents, return_dict=False)[0] + return video + + # -- Prompt formatting + tokenization (shared by T2V and I2V) ------------ + + def _format_and_tokenize_prompts( + self, + prompt: str, + negative_prompt: str, + num_frames: int, + frame_rate: float, + height: int, + width: int, + max_sequence_length: int, + sp: OmniDiffusionSamplingParams, + use_system_prompt: bool = False, + is_t2i: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Format prompts with metadata templates and tokenize. + + Returns (cond_ids, cond_mask, uncond_ids, uncond_mask). + + For T2I (``is_t2i=True``) the duration template is suppressed (no FPS + or duration concept for a single image) and the image-flavored + resolution template is used. + """ + # Route cosmos3-specific controls through ``_get_sp_param`` so they + # are picked up from ``extra_args`` (OpenAI endpoint path) as well + # as from direct attributes. + use_duration_template = bool(self._get_sp_param(sp, "use_duration_template", False)) and not is_t2i + dur_tmpl = COSMOS3_DURATION_TEMPLATE if use_duration_template else None + if bool(self._get_sp_param(sp, "use_resolution_template", False)): + res_tmpl = COSMOS3_IMAGE_RESOLUTION_TEMPLATE if is_t2i else COSMOS3_RESOLUTION_TEMPLATE + else: + res_tmpl = None + prompt = self._apply_metadata_templates( + prompt, + num_frames, + frame_rate, + height, + width, + duration_template=dur_tmpl, + resolution_template=res_tmpl, + ) + if _is_rank_zero(): + logger.info("Final prompt: '%s'", prompt) + + # Negative prompt: inverse templates ("not {duration}...", "not {height}x{width}..."). + # Applied whenever the matching positive template is enabled; an empty + # negative_prompt yields output that starts with the template, not a dot. + inv_dur = COSMOS3_INVERSE_DURATION_TEMPLATE if dur_tmpl else None + if res_tmpl is None: + inv_res = None + elif is_t2i: + inv_res = COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE + else: + inv_res = COSMOS3_INVERSE_RESOLUTION_TEMPLATE + negative_prompt = self._apply_metadata_templates( + negative_prompt, + num_frames, + frame_rate, + height, + width, + duration_template=inv_dur, + resolution_template=inv_res, + force_duration_template=True, + ) + + default_sys_prompt = COSMOS3_T2I_SYSTEM_PROMPT if is_t2i else COSMOS3_SYSTEM_PROMPT + sys_prompt = self._get_sp_param(sp, "system_prompt", default_sys_prompt) or default_sys_prompt + cond_ids, cond_mask = self._tokenize_prompt( + prompt, max_sequence_length, use_system_prompt, system_prompt=sys_prompt + ) + uncond_ids, uncond_mask = self._tokenize_prompt( + negative_prompt, max_sequence_length, use_system_prompt, system_prompt=sys_prompt + ) + return cond_ids, cond_mask, uncond_ids, uncond_mask + + # -- I2V latent preparation --------------------------------------------- + + def _encode_conditioning_video( + self, + image_tensor: torch.Tensor, + num_frames: int, + height: int, + width: int, + ) -> torch.Tensor: + """VAE-encode a conditioning image as a full-length video. + + The WAN VAE has temporal compression (factor 4), so encoding a single + frame produces degenerate temporal features. We fill the entire + pixel-space video with the conditioning image (repeating it across all + frames) so the temporal encoder sees plausible content everywhere. + The caller keeps only the conditioned latent frame(s) and replaces + the rest with noise. + """ + # image_tensor: [1, 3, H, W] -> [1, 3, num_frames, H, W] + video = image_tensor.unsqueeze(2).expand(-1, -1, num_frames, -1, -1).contiguous() + video = video.to(device=self.device, dtype=self.vae.dtype) + + latent = self.vae.encode(video).latent_dist.mode() + + # Normalize (inverse of _decode_latents denormalization) + if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"): + latents_mean = ( + torch.tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype) + ) + latents_std = torch.tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype) + latent = (latent - latents_mean) / latents_std + else: + scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0) + latent = latent * scaling_factor + + return latent.to(self.dtype) + + def _prepare_latents_i2v( + self, + image_tensor: torch.Tensor, + height: int, + width: int, + num_frames: int, + generator: torch.Generator, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Prepare initial latents with frame 0 conditioned on the input image. + + Returns: + latents: [1, C, T_lat, H_lat, W_lat] with frame 0 = image, rest = noise + velocity_mask: [1, 1, T_lat, 1, 1] with frame 0 = 0, rest = 1 + image_latent: [1, C, 1, H_lat, W_lat] clean frame 0 for re-injection + """ + C = self.transformer.latent_channel_size + T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + H_lat = height // self.vae_scale_factor_spatial + W_lat = width // self.vae_scale_factor_spatial + + noise = randn_tensor( + (1, C, T_lat, H_lat, W_lat), + generator=generator, + device=self.device, + dtype=self.dtype, + ) + + cond_latent = self._encode_conditioning_video(image_tensor, num_frames, height, width) + image_latent = cond_latent[:, :, 0:1, :, :] + + condition_mask = torch.zeros(1, 1, T_lat, 1, 1, device=self.device, dtype=self.dtype) + condition_mask[:, :, 0, :, :] = 1.0 + latents = condition_mask * cond_latent + (1.0 - condition_mask) * noise + velocity_mask = 1.0 - condition_mask + return latents, velocity_mask, image_latent + + # -- Denoising loop (shared by T2V and I2V) ----------------------------- + + def diffuse( + 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, + shared_kwargs: dict, + *, + velocity_mask: torch.Tensor | None = None, + image_latent: torch.Tensor | None = None, + condition_latents: torch.Tensor | None = None, + guidance_interval: tuple[float, float] | None = None, + ) -> torch.Tensor: + """Denoising loop with 3-mode CFG support (parallel, sequential, none). + + Cosmos3's UND pathway is text-dependent, so CFG needs separate K/V + caches for conditional and unconditional text. + + Two modes: + 1. CFG parallel (multi-GPU): each rank handles one condition via + predict_noise_maybe_with_cfg; caching is rank-local. + 2. Sequential CFG (single-GPU or cfg_size=1): two separate + forward passes with explicit cache swapping. We cannot + batch B=2 because different text lengths would cause the + shorter branch to attend to padding in cross-attention. + + I2V conditioning (when both arguments are supplied): + * ``velocity_mask`` zeros frame-0 noise predictions before stepping. + * ``image_latent`` is re-injected into frame 0 after each scheduler + step, since UniPC's predictor-corrector update rescales the + sample (sigma-dependent), so even zero velocity does not preserve + frame 0. + + ``guidance_interval`` (T2I) restricts CFG to + timesteps inside the closed interval ``[lo, hi]``. The interval is + compared against the raw scheduler timestep value; works for both + the [0, 1000] discrete scale and normalized flow-matching scales. + Outside the interval the cond/uncond delta is zeroed so all ranks + continue to execute identical control flow (CFG-Parallel safe). + """ + do_cfg = guidance_scale > 1.0 + cfg_parallel = self._cfg_parallel_active() and do_cfg + self.transformer.reset_cache() + + def _cfg_active_at(t: torch.Tensor) -> bool: + if guidance_interval is None: + return True + t_scalar = float(t.item()) if torch.is_tensor(t) else float(t) + lo, hi = guidance_interval + return lo <= t_scalar <= hi + + def _step( + noise_pred: torch.Tensor, + t: torch.Tensor, + latents: torch.Tensor, + ) -> torch.Tensor: + if isinstance(noise_pred, tuple): + raise ValueError("Cosmos3 noise prediction must be a single tensor; got a tuple.") + if velocity_mask is not None: + noise_pred = noise_pred * velocity_mask + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + if condition_latents is not None and velocity_mask is not None: + latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents + elif image_latent is not None: + latents[:, :, 0:1, :, :] = image_latent + return latents + + if cfg_parallel: + for t in self.progress_bar(timesteps): + timestep = t.unsqueeze(0) + # Out-of-interval steps run with effective scale 1.0 so the + # combined output equals the cond branch (uncond is dropped). + # All ranks still execute both branches; no CFG-Parallel + # divergence. + step_scale = guidance_scale if _cfg_active_at(t) else 1.0 + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=step_scale, + positive_kwargs=dict( + hidden_states=latents, + timestep=timestep, + text_ids=cond_ids, + text_mask=cond_mask, + **shared_kwargs, + ), + negative_kwargs=dict( + hidden_states=latents, + timestep=timestep, + text_ids=uncond_ids, + text_mask=uncond_mask, + **shared_kwargs, + ), + cfg_normalize=False, + ) + latents = _step(noise_pred, t, latents) + + elif do_cfg: + cond_cache: tuple = (None, None) + uncond_cache: tuple = (None, None) + + keep_uncond_for_cache = self._cache_requires_paired_cfg() + + for t in self.progress_bar(timesteps): + timestep = t.unsqueeze(0) + cfg_active = _cfg_active_at(t) + + self.transformer.cached_kv, self.transformer.cached_freqs_gen = cond_cache + noise_cond = self.transformer( + hidden_states=latents, + timestep=timestep, + text_ids=cond_ids, + text_mask=cond_mask, + **shared_kwargs, + ) + if cond_cache[0] is None: + cond_cache = (self.transformer.cached_kv, self.transformer.cached_freqs_gen) + + if cfg_active or keep_uncond_for_cache: + self.transformer.cached_kv, self.transformer.cached_freqs_gen = uncond_cache + noise_uncond = self.transformer( + hidden_states=latents, + timestep=timestep, + text_ids=uncond_ids, + text_mask=uncond_mask, + **shared_kwargs, + ) + if uncond_cache[0] is None: + uncond_cache = (self.transformer.cached_kv, self.transformer.cached_freqs_gen) + # Outside the interval, scale=1.0 makes the combined result + # equal to noise_cond; the uncond pass is computed only to + # preserve cache-dit's cond/uncond parity. + step_scale = guidance_scale if cfg_active else 1.0 + noise_pred = self.combine_cfg_noise(noise_cond, noise_uncond, step_scale, cfg_normalize=False) + else: + noise_pred = noise_cond + + latents = _step(noise_pred, t, latents) + + else: + for t in self.progress_bar(timesteps): + timestep = t.unsqueeze(0) + noise_pred = self.transformer( + hidden_states=latents, + timestep=timestep, + text_ids=cond_ids, + text_mask=cond_mask, + **shared_kwargs, + ) + latents = _step(noise_pred, t, latents) + + return latents + + # -- Forward (main generation entry point) ------------------------------- + + def forward( + self, + req: OmniDiffusionRequest, + ) -> DiffusionOutput: + pipeline_start = time.time() + + # --- Parse request --- + if len(req.prompts) > 1: + raise ValueError("Cosmos3OmniDiffusersPipeline currently supports a single prompt per request.") + + prompt_data = req.prompts[0] + if isinstance(prompt_data, str): + prompt = prompt_data + negative_prompt = None + image_tensor = 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") + + sp = req.sampling_params + is_t2i = self._is_t2i_request(req) + if negative_prompt is None: + negative_prompt = "" + + # T2I and T2V share the same model + forward path; only defaults + # differ: + # T2I: 1024x1024, 50 steps, shift=3.0, guidance_interval=[400, 1000] + # T2V: 720x1280, 35 steps, shift=engine-init, no interval + if is_t2i: + height = sp.height or COSMOS3_T2I_DEFAULT_HEIGHT + width = sp.width or COSMOS3_T2I_DEFAULT_WIDTH + num_frames = 1 + num_inference_steps = sp.num_inference_steps or COSMOS3_T2I_DEFAULT_NUM_INFERENCE_STEPS + guidance_scale = sp.guidance_scale if sp.guidance_scale else COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE + default_flow_shift = COSMOS3_T2I_DEFAULT_FLOW_SHIFT + default_guidance_interval: tuple[float, float] | None = COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL + batch_size = max(1, int(sp.num_outputs_per_prompt or 1)) + else: + height = sp.height or COSMOS3_T2V_DEFAULT_HEIGHT + width = sp.width or COSMOS3_T2V_DEFAULT_WIDTH + num_frames = sp.num_frames or COSMOS3_T2V_DEFAULT_NUM_FRAMES + num_inference_steps = sp.num_inference_steps or COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS + guidance_scale = sp.guidance_scale if sp.guidance_scale else COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE + # Fall back to the engine-init shift, NOT None: passing None + # to ``_set_flow_shift`` would leak a prior T2I rebuild + # (shift=3.0) into a subsequent video request. + default_flow_shift = self._engine_init_flow_shift + default_guidance_interval = None + batch_size = 1 # Existing video pipeline assumes B=1. + + # Runtime controls: prefer ``extra_args`` (OpenAI endpoints write + # there) over direct attrs. + flow_shift_target = float(self._get_sp_param(sp, "flow_shift", default_flow_shift)) + guidance_interval = self._get_sp_param(sp, "guidance_interval", default_guidance_interval) + + frame_rate = self._get_sp_param(sp, "resolved_frame_rate") or self._get_sp_param(sp, "frame_rate") or 24.0 + 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 + + # Always resolve to a concrete target shift for this request, then + # update the scheduler. This is what guarantees mode-to-mode + # transitions restore the right schedule (no T2I to T2V leak). + 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) + + # --- Format prompts & tokenize (B=1; reused across loop iterations + # for T2I num_outputs_per_prompt > 1) --- + cond_ids, cond_mask, uncond_ids, uncond_mask = self._format_and_tokenize_prompts( + prompt, + negative_prompt, + num_frames, + frame_rate, + height, + width, + max_sequence_length, + sp, + use_system_prompt, + is_t2i=is_t2i, + ) + + # --- Prepare latents (T2I, T2V, or I2V) --- + # T2I shares _prepare_latents with T2V; the math collapses cleanly + # at num_frames=1 ((1-1)//4 + 1 = 1 latent frame). For T2I with + # ``num_outputs_per_prompt > 1`` we loop the diffusion below; + # batching B=N together would require expanding text K/V (UND + # pathway is text-only and cached) and is left as a future + # optimization. + if image_tensor is not None and not is_t2i: + latents, velocity_mask, image_latent = self._prepare_latents_i2v( + image_tensor, + height, + width, + num_frames, + generator, + ) + condition_latents = None + else: + latents = self._prepare_latents(height, width, num_frames, generator) + velocity_mask = None + image_latent = None + condition_latents = None + + T_latent = latents.shape[2] + H_latent = latents.shape[3] + W_latent = latents.shape[4] + video_shape = (T_latent, H_latent, W_latent) + + # --- Denoising loop --- + shared_kwargs = dict(video_shape=video_shape, fps=frame_rate) + if velocity_mask is not None: + shared_kwargs["noisy_frame_mask"] = velocity_mask + + def _run_diffusion(start_latents): + self.scheduler.set_timesteps(num_inference_steps, device=self.device) + return self.diffuse( + latents=start_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, + shared_kwargs=shared_kwargs, + velocity_mask=velocity_mask, + image_latent=image_latent, + condition_latents=condition_latents, + guidance_interval=guidance_interval, + ) + + if is_t2i and batch_size > 1: + # Generate N independent images by re-running the full diffusion + # loop with different noise seeds. The first sample reuses + # ``latents`` already drawn from ``generator``; subsequent + # samples draw fresh noise from the same generator (state + # advances per call), giving distinct outputs from a single + # user-provided seed. Batched B=N would be more efficient but + # requires expanding cached UND text K/V to match. + samples = [_run_diffusion(latents)] + for _ in range(batch_size - 1): + next_latents = self._prepare_latents(height, width, num_frames, generator) + samples.append(_run_diffusion(next_latents)) + latents = torch.cat(samples, dim=0) + else: + latents = _run_diffusion(latents) + + # --- Decode --- + if _is_rank_zero(): + logger.info("Decoding video...") + decode_start = time.time() + 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) + + return DiffusionOutput(output={"image": video} if is_t2i else {"video": video}) diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py new file mode 100644 index 00000000000..7b3848a089a --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -0,0 +1,1173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cosmos3 VFM Transformer for vllm-omni. + +Implements the Mixture-of-Transformers architecture with two pathways: +- Understanding (UND): causal self-attention on text tokens (Qwen3-VL backbone) +- Generation (GEN): cross-attention where visual Q attends to [K_und, K_gen] + +Ported from the TRT-LLM integration (tekit branch user/shreyasm/cosmos3). +""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, +) + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput +from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available +from vllm_omni.diffusion.layers.norm import RMSNorm + +logger = init_logger(__name__) + + +def _get_ulysses_state() -> tuple[int, int, dist.ProcessGroup | None]: + """Return (ulysses_size, ulysses_rank, ulysses_pg) from vllm-omni parallel state. + + Returns (1, 0, None) when sequence parallelism is not active. + """ + from vllm_omni.diffusion.distributed.parallel_state import ( + get_sp_group, + get_ulysses_parallel_rank, + get_ulysses_parallel_world_size, + ) + + size = get_ulysses_parallel_world_size() + if size <= 1: + return 1, 0, None + return size, get_ulysses_parallel_rank(), get_sp_group().ulysses_group + + +def _is_sp_active() -> bool: + """Check whether sequence parallelism is active in the current forward context. + + Follows the Bagel pattern: read ``forward_context.sp_active`` which returns + True when ``sequence_parallel_size > 1`` even without ``_sp_plan`` hooks. + """ + + if not is_forward_context_available(): + return False + return get_forward_context().sp_active + + +def _tf_config_get(config: Any, key: str, default: Any) -> Any: + """Read a value from TransformerConfig, dict, or simple namespace.""" + if config is None: + return default + if hasattr(config, "get"): + return config.get(key, default) + return getattr(config, key, default) + + +# --------------------------------------------------------------------------- +# Rotary Position Embeddings (mRoPE) +# --------------------------------------------------------------------------- +def compute_mrope_position_ids_text( + num_tokens: int, + temporal_offset: int, +) -> tuple[torch.Tensor, int]: + """Generate 3D mRoPE position IDs for text tokens. + + Text tokens: all three axes (t, h, w) share the same monotonically + increasing position IDs. + """ + ids = torch.arange(num_tokens, dtype=torch.long) + temporal_offset + mrope_ids = ids.unsqueeze(0).expand(3, -1).contiguous() + return mrope_ids, temporal_offset + num_tokens + + +def compute_mrope_position_ids_vision( + grid_t: int, + grid_h: int, + grid_w: int, + temporal_offset: int | float, + fps: float | None = None, + base_fps: float = 24.0, + temporal_compression_factor: int = 4, + base_temporal_compression_factor: int | None = None, + enable_fps_modulation: bool = True, +) -> tuple[torch.Tensor, int | float]: + """Generate 3D mRoPE position IDs for vision tokens. + + Creates a (t, h, w) position grid with spatial indices reset per segment + (Qwen3VL-style). Flattened in t-major order. + """ + fps_modulation = enable_fps_modulation and fps is not None + + if fps_modulation: + tps = fps / temporal_compression_factor + effective_base_tcf = ( + base_temporal_compression_factor + if base_temporal_compression_factor is not None + else temporal_compression_factor + ) + base_tps = base_fps / effective_base_tcf + frame_indices = torch.arange(grid_t, dtype=torch.float32) + t_index = (frame_indices / tps * base_tps + temporal_offset).view(-1, 1).expand(-1, grid_h * grid_w).flatten() + else: + t_index = torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten() + int( + temporal_offset + ) + + h_index = torch.arange(grid_h, dtype=torch.long).view(1, -1, 1).expand(grid_t, -1, grid_w).flatten() + w_index = torch.arange(grid_w, dtype=torch.long).view(1, 1, -1).expand(grid_t, grid_h, -1).flatten() + + if fps_modulation: + mrope_ids = torch.stack([t_index, h_index.to(torch.float32), w_index.to(torch.float32)], dim=0) + else: + mrope_ids = torch.stack([t_index, h_index, w_index], dim=0) + + next_offset = math.floor(mrope_ids.max().item()) + 1 + return mrope_ids, next_offset + + +class Qwen3VLTextRotaryEmbedding(nn.Module): + """Multi-dimensional rotary position embedding for Qwen3-VL.""" + + def __init__( + self, + *, + head_dim: int, + rope_theta: float, + mrope_section: list[int], + ) -> None: + super().__init__() + self.head_dim = head_dim + self.rope_theta = rope_theta + self.mrope_section = mrope_section + inv_freq = 1.0 / ( + rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.int64).to(dtype=torch.float) / head_dim) + ) + self.attention_scaling = 1.0 + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def apply_interleaved_mrope(self, freqs: torch.Tensor, mrope_section: list[int]) -> torch.Tensor: + """Reorganize from chunked [TTT...HHH...WWW] to interleaved [THTHW...].""" + freqs_t = freqs[0] + for dim, offset in enumerate((1, 2), start=1): + length = mrope_section[dim] * 3 + idx = slice(offset, length, 3) + freqs_t[..., idx] = freqs[dim, ..., idx] + return freqs_t + + @torch.no_grad() + def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + if position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + inv_freq_expanded = ( + self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1).to(x.device) + ) + position_ids_expanded = position_ids[:, :, None, :].float() + + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) + freqs = self.apply_interleaved_mrope(freqs, self.mrope_section) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +# --------------------------------------------------------------------------- +# RoPE application (Qwen3/Llama style) +# --------------------------------------------------------------------------- +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _apply_rotary_pos_emb( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Qwen3-style RoPE: (x * cos) + (rotate_half(x) * sin). + + Args: + q: [B, S, h, D] + k: [B, S, H_kv, D] + cos: [1, S, 1, D] or broadcastable + sin: [1, S, 1, D] or broadcastable + """ + q_embed = (q * cos) + (_rotate_half(q) * sin) + k_embed = (k * cos) + (_rotate_half(k) * sin) + return q_embed, k_embed + + +# --------------------------------------------------------------------------- +# Timestep Embedder +# --------------------------------------------------------------------------- +class TimestepEmbedder(nn.Module): + """Embeds scalar timesteps into vector representations.""" + + def __init__( + self, + hidden_size: int, + frequency_embedding_size: int = 256, + max_period: int = 10000, + target_dtype: torch.dtype = torch.bfloat16, + ) -> None: + super().__init__() + # Following diffusers naming pattern here for checkpoint compatibility. + self.linear_1 = nn.Linear(frequency_embedding_size, hidden_size, bias=True) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True) + self.frequency_embedding_size = frequency_embedding_size + self.hidden_size = hidden_size + + half = frequency_embedding_size // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=target_dtype) / half) + self.register_buffer("freqs", freqs, persistent=False) + + def forward(self, t: torch.Tensor) -> torch.Tensor: + args = t[:, None] * self.freqs[None] + t_freq = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + return self.linear_2(self.act(self.linear_1(t_freq))) + + +# --------------------------------------------------------------------------- +# GatedMLP (replaces TRT-LLM GatedMLP) +# --------------------------------------------------------------------------- +class Cosmos3GatedMLP(nn.Module): + def __init__( + self, + hidden_size: int = 4096, + intermediate_size: int = 12288, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_proj", + ) + self.up_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +# --------------------------------------------------------------------------- +# Attention Modules +# --------------------------------------------------------------------------- +class Cosmos3CausalAttention(nn.Module): + """Understanding pathway: causal self-attention on text tokens. + + Returns (output, K, V) where K/V are post-norm, post-RoPE for the + generation pathway's cross-attention. + """ + + def __init__( + self, + *, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + rms_norm_eps: float, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_attention_heads + self.num_kv_heads = num_key_value_heads + self.head_dim = head_dim + + tp_size = get_tensor_model_parallel_world_size() + self.num_heads_local = self.num_heads // tp_size + self.num_kv_heads_local = self.num_kv_heads // tp_size + + self.to_q = ColumnParallelLinear( + hidden_size, + self.num_heads * self.head_dim, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_q", + ) + self.to_k = ColumnParallelLinear( + hidden_size, + self.num_kv_heads * self.head_dim, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_k", + ) + self.to_v = ColumnParallelLinear( + hidden_size, + self.num_kv_heads * self.head_dim, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_v", + ) + self.to_out = RowParallelLinear( + self.num_heads * self.head_dim, + hidden_size, + bias=False, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_out", + ) + + self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps) + self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps) + + # skip_sequence_parallel=True because the UND pathway is + # computed once and replicated across SP ranks. + # Only the GEN pathway is sequence-sharded. + self.attn = FrameworkAttention( + num_heads=self.num_heads, + head_size=self.head_dim, + causal=True, + softmax_scale=1.0 / (self.head_dim**0.5), + num_kv_heads=self.num_kv_heads, + skip_sequence_parallel=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, S, _ = hidden_states.shape + + q = self.to_q(hidden_states).view(B, S, self.num_heads_local, self.head_dim) + k = self.to_k(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim) + v = self.to_v(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim) + + # Per-head QK norm + q = F.rms_norm(q, (q.shape[-1],), self.norm_q.weight, self.norm_q.variance_epsilon) + k = F.rms_norm(k, (k.shape[-1],), self.norm_k.weight, self.norm_k.variance_epsilon) + + # Qwen3-style RoPE + q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin) + + out = self.attn(q, k, v).reshape(B, S, -1) + return self.to_out(out), k, v + + +class Cosmos3CrossAttention(nn.Module): + """Generation pathway: full attention where visual Q attends to all K/V. + + * **Non-SP path**: explicit ``cat([k_und, k_gen])``. Text conditioning is + always present because K/V are physically concatenated. + + * **SP path** (Ulysses active): ``k_und``/``v_und`` are passed as + ``joint_key``/``joint_value`` in ``AttentionMetadata``. The Ulysses + wrapper head-slices the replicated UND K/V and performs all-to-all on the + sharded GEN Q/K/V so every query sees the full context. + """ + + def __init__( + self, + *, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + rms_norm_eps: float, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_attention_heads + self.num_kv_heads = num_key_value_heads + self.head_dim = head_dim + + tp_size = get_tensor_model_parallel_world_size() + self.num_heads_local = self.num_heads // tp_size + self.num_kv_heads_local = self.num_kv_heads // tp_size + + self.to_q = ColumnParallelLinear( + hidden_size, + self.num_heads * self.head_dim, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_q", + ) + self.to_k = ColumnParallelLinear( + hidden_size, + self.num_kv_heads * self.head_dim, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_k", + ) + self.to_v = ColumnParallelLinear( + hidden_size, + self.num_kv_heads * self.head_dim, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_v", + ) + self.to_out = RowParallelLinear( + self.num_heads * self.head_dim, + hidden_size, + bias=False, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_out", + ) + + self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps) + self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps) + + self.attn = FrameworkAttention( + num_heads=self.num_heads, + head_size=self.head_dim, + causal=False, + softmax_scale=1.0 / (self.head_dim**0.5), + num_kv_heads=self.num_kv_heads, + ) + + # TODO(follow-up): collapse _forward_local and _forward_sp into a single + # joint-based path when NoParallelAttention can process joint_key/value. + # Currently the non-SP path must concatenate the replicated UND K/V explicitly, + # while the SP path passes it as joint_*. + + # -- Non-SP path: explicit K/V concatenation + framework Attention -------- + + def _forward_local( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_und: torch.Tensor, + v_und: torch.Tensor, + ) -> torch.Tensor: + B, S_gen = q.shape[:2] + k_all = torch.cat([k_und, k], dim=1) + v_all = torch.cat([v_und, v], dim=1) + + out = self.attn(q, k_all, v_all) + return out.reshape(B, S_gen, -1) + + # -- SP path: framework Attention with joint_key/value ------------------- + + def _forward_sp( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_und: torch.Tensor, + v_und: torch.Tensor, + ) -> torch.Tensor: + B, S_gen = q.shape[:2] + + # Zero-length joint_query satisfies the Ulysses contract + # (joint_query, joint_key, joint_value must all be present) without + # adding text tokens to Q. joint_len=0 keeps post_attention on the + # standard reverse-all-to-all path (no joint-output splitting). + joint_q = q.new_empty(B, 0, self.num_heads_local, self.head_dim) + + attn_metadata = AttentionMetadata( + joint_query=joint_q, + joint_key=k_und, + joint_value=v_und, + joint_strategy="front", + ) + out = self.attn(q, k, v, attn_metadata) + return out.reshape(B, S_gen, -1) + + # -- Public forward: routes to the appropriate path ---------------------- + + def forward( + self, + hidden_states: torch.Tensor, + k_und: torch.Tensor, + v_und: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ) -> torch.Tensor: + """ + Args: + hidden_states: [B, S_gen_local, hidden_size] (may be sequence-sharded) + k_und: [B, S_und, H_kv_local, D] pre-computed UND keys (TP-sharded, post-norm/RoPE) + v_und: [B, S_und, H_kv_local, D] pre-computed UND values (TP-sharded) + freqs_cos: [B, S_gen_local, 1, D] + freqs_sin: [B, S_gen_local, 1, D] + """ + B, S_gen, _ = hidden_states.shape + + q = self.to_q(hidden_states).view(B, S_gen, self.num_heads_local, self.head_dim) + k = self.to_k(hidden_states).view(B, S_gen, self.num_kv_heads_local, self.head_dim) + v = self.to_v(hidden_states).view(B, S_gen, self.num_kv_heads_local, self.head_dim) + + # Per-head QK norm + q = F.rms_norm(q, (q.shape[-1],), self.norm_q.weight, self.norm_q.variance_epsilon) + k = F.rms_norm(k, (k.shape[-1],), self.norm_k.weight, self.norm_k.variance_epsilon) + + # Qwen3-style RoPE + q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin) + + if _is_sp_active(): + out = self._forward_sp(q, k, v, k_und, v_und) + else: + out = self._forward_local(q, k, v, k_und, v_und) + + return self.to_out(out) + + +# --------------------------------------------------------------------------- +# Decoder Layers +# --------------------------------------------------------------------------- +class Cosmos3UndDecoderLayer(nn.Module): + """Understanding pathway decoder layer: causal self-attention + MLP.""" + + def __init__( + self, + *, + hidden_size: int, + intermediate_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + rms_norm_eps: float, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.self_attn = Cosmos3CausalAttention( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + rms_norm_eps=rms_norm_eps, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) + self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) + self.mlp = Cosmos3GatedMLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + + def forward( + self, + hidden_states: torch.Tensor, + freqs: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Returns (hidden_states, K, V) where K/V are for GEN cross-attention.""" + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + cos, sin = freqs + attn_out, k, v = self.self_attn(hidden_states, cos, sin) + hidden_states = residual + attn_out + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + self.mlp(hidden_states) + + return hidden_states, k, v + + +class Cosmos3GenDecoderLayer(nn.Module): + """Generation pathway decoder layer: cross-attention (to UND K/V) + MLP.""" + + def __init__( + self, + *, + layer_idx: int | None = None, + hidden_size: int, + intermediate_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + rms_norm_eps: float, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.layer_idx = layer_idx + self.cross_attention = Cosmos3CrossAttention( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + rms_norm_eps=rms_norm_eps, + quant_config=quant_config, + prefix=f"{prefix}.cross_attention", + ) + self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) + self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) + self.mlp = Cosmos3GatedMLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + + def forward( + self, + hidden_states: torch.Tensor, + *, + k_und: torch.Tensor | None = None, + v_und: torch.Tensor | None = None, + freqs_cos: torch.Tensor | None = None, + freqs_sin: torch.Tensor | None = None, + cached_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None, + freqs_gen: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + if cached_kv is not None: + if self.layer_idx is None: + raise ValueError("Cosmos3 GEN layer requires layer_idx when cached_kv is provided.") + k_und, v_und = cached_kv[self.layer_idx] + if freqs_gen is not None: + freqs_cos, freqs_sin = freqs_gen + if k_und is None or v_und is None or freqs_cos is None or freqs_sin is None: + raise ValueError("Cosmos3 GEN layer requires k_und/v_und/freqs_cos/freqs_sin.") + + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + hidden_states = self.cross_attention( + hidden_states, k_und=k_und, v_und=v_und, freqs_cos=freqs_cos, freqs_sin=freqs_sin + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + self.mlp(hidden_states) + + return hidden_states + + +# --------------------------------------------------------------------------- +# Language Model (Understanding pathway) +# --------------------------------------------------------------------------- +class Cosmos3LanguageModel(nn.Module): + """Understanding pathway: a standard causal LM that processes text tokens. + + Returns per-layer K/V tensors for the generation pathway's cross-attention. + The UND pathway is independent of the denoising step, so its K/V can be + computed once and reused across all sampling steps. + """ + + def __init__( + self, + *, + hidden_size: int, + intermediate_size: int, + num_hidden_layers: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + vocab_size: int, + rms_norm_eps: float, + rope_theta: float, + mrope_section: list[int], + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.embed_tokens = nn.Embedding(vocab_size, hidden_size) + self.rotary_emb = Qwen3VLTextRotaryEmbedding( + head_dim=head_dim, + rope_theta=rope_theta, + mrope_section=mrope_section, + ) + self.layers = nn.ModuleList( + [ + Cosmos3UndDecoderLayer( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + rms_norm_eps=rms_norm_eps, + quant_config=quant_config, + prefix=f"{prefix}.layers.{i}", + ) + for i in range(num_hidden_layers) + ] + ) + # TODO: Not used right now, will be used in the future for prompt upsampler. + self.norm = RMSNorm(hidden_size, eps=rms_norm_eps) + + def forward( + self, + text_ids: torch.Tensor, + freqs: tuple[torch.Tensor, torch.Tensor], + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + """ + Args: + text_ids: [B, S] token IDs + freqs: (cos, sin) each [B, S, 1, D] + + Returns: + List of (K, V) per layer, each [B, S, H_kv, D]. + + No padding mask is applied: with right-padding + causal self-attention, + real query positions only attend to real keys, and the caller trims pad + K/V via ``max_real_len`` before the GEN cross-attention sees them. + """ + hidden = self.embed_tokens(text_ids) + + cached_kv: list[tuple[torch.Tensor, torch.Tensor]] = [] + for layer in self.layers: + hidden, k, v = layer(hidden, freqs) + cached_kv.append((k, v)) + + return cached_kv + + +# --------------------------------------------------------------------------- +# Main Transformer +# --------------------------------------------------------------------------- +class Cosmos3GenSPPrepare(nn.Module): + """Module boundary used by _sp_plan to shard GEN states and RoPE together.""" + + def forward( + self, + hidden_gen: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return hidden_gen, freqs_cos, freqs_sin + + +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. + + Layerwise offloading uses ``gen_layers`` as the block container. + + Sequence parallelism uses ``_sp_plan`` to shard/gather the GEN pathway at + module boundaries. ``Cosmos3CrossAttention`` checks + ``forward_context.sp_active`` at runtime and routes to the framework + ``Attention`` layer (with Ulysses all-to-all) or plain SDPA accordingly. + """ + + _repeated_blocks = ["Cosmos3GenDecoderLayer"] + + _layerwise_offload_blocks_attrs = ["gen_layers"] + + packed_modules_mapping = {} + + @staticmethod + def _is_transformer_block(name: str, module) -> bool: + return ("gen_layers" in name or "language_model.layers" in name) and name.split(".")[-1].isdigit() + + _hsdp_shard_conditions = [_is_transformer_block] + + _sp_plan = { + "gen_sp_prepare": { + 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True), + 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + 2: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + }, + "gen_sp_gather": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + + @staticmethod + def _validate_supported_config(model_config: Any) -> None: + """Fail loudly when a checkpoint requests an unsupported architecture.""" + expected_values = { + "qk_norm_for_diffusion": True, + "qk_norm_for_text": True, + "position_embedding_type": "unified_3d_mrope", + "unified_3d_mrope_reset_spatial_ids": True, + "joint_attn_implementation": "two_way", + } + for key, expected in expected_values.items(): + actual = _tf_config_get(model_config, key, expected) + if actual != expected: + raise ValueError(f"Unsupported Cosmos3 transformer config: {key}={actual!r}; expected {expected!r}.") + + def __init__( + self, + od_config: OmniDiffusionConfig, + *, + temporal_compression_factor: int | None = None, + ) -> None: + super().__init__() + model_config = od_config.tf_model_config + self._validate_supported_config(model_config) + rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {} + + self.hidden_size = int(_tf_config_get(model_config, "hidden_size", 4096)) + self.num_hidden_layers = int(_tf_config_get(model_config, "num_hidden_layers", 36)) + self.num_attention_heads = int(_tf_config_get(model_config, "num_attention_heads", 32)) + self.num_key_value_heads = int(_tf_config_get(model_config, "num_key_value_heads", 8)) + self.head_dim = int(_tf_config_get(model_config, "head_dim", 128)) + self.intermediate_size = int(_tf_config_get(model_config, "intermediate_size", 12288)) + self.vocab_size = int(_tf_config_get(model_config, "vocab_size", 151936)) + self.rms_norm_eps = float(_tf_config_get(model_config, "rms_norm_eps", 1e-6)) + self.rope_theta = float(_tf_config_get(model_config, "rope_theta", 5_000_000)) + self.mrope_section = list(rope_scaling.get("mrope_section", [24, 20, 20])) + self.latent_patch_size = int(_tf_config_get(model_config, "latent_patch_size", 2)) + self.latent_channel_size = int(_tf_config_get(model_config, "latent_channel", 48)) + self.timestep_scale = float(_tf_config_get(model_config, "timestep_scale", 0.001)) + self.base_fps = float(_tf_config_get(model_config, "base_fps", 24.0)) + if temporal_compression_factor is None: + temporal_compression_factor = _tf_config_get(model_config, "temporal_compression_factor", 4) + self.temporal_compression_factor = int(temporal_compression_factor) + self.enable_fps_modulation = bool(_tf_config_get(model_config, "enable_fps_modulation", True)) + self.temporal_modality_margin = int( + _tf_config_get( + model_config, + "unified_3d_mrope_temporal_modality_margin", + 15000, + ) + ) + self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel_size + + dtype = od_config.dtype + quant_config = getattr(od_config, "quantization_config", None) if od_config else None + + self.language_model = Cosmos3LanguageModel( + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads, + num_key_value_heads=self.num_key_value_heads, + head_dim=self.head_dim, + vocab_size=self.vocab_size, + rms_norm_eps=self.rms_norm_eps, + rope_theta=self.rope_theta, + mrope_section=self.mrope_section, + quant_config=quant_config, + prefix="language_model", + ) + + # Video projection layers are small; not worth quantizing. + self.proj_in = nn.Linear(self.patch_latent_dim, self.hidden_size) + self.proj_out = nn.Linear(self.hidden_size, self.patch_latent_dim) + self.time_embedder = TimestepEmbedder(self.hidden_size, target_dtype=dtype) + + self.gen_layers = nn.ModuleList( + [ + Cosmos3GenDecoderLayer( + layer_idx=i, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + num_attention_heads=self.num_attention_heads, + num_key_value_heads=self.num_key_value_heads, + head_dim=self.head_dim, + rms_norm_eps=self.rms_norm_eps, + quant_config=quant_config, + prefix=f"gen_layers.{i}", + ) + for i in range(self.num_hidden_layers) + ] + ) + + self.norm_moe_gen = RMSNorm(self.hidden_size, eps=self.rms_norm_eps) + self.gen_sp_prepare = Cosmos3GenSPPrepare() + self.gen_sp_gather = nn.Identity() + + # Cached state (populated on first forward, reused across denoising steps) + self.cached_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None + self.cached_freqs_gen: tuple[torch.Tensor, torch.Tensor] | None = None + + @property + def device(self) -> torch.device: + return next(self.parameters()).device + + # -- Patchify / Unpatchify ----------------------------------------------- + + def _pad_to_patch_size(self, h: int, w: int) -> tuple[int, int, int, int]: + """Returns (hp, wp, H_padded, W_padded).""" + p = self.latent_patch_size + H_padded = ((h + p - 1) // p) * p + W_padded = ((w + p - 1) // p) * p + return H_padded // p, W_padded // p, H_padded, W_padded + + def patchify(self, latents: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor: + """[B, C, t, h, w] -> [B, t*hp*wp, p*p*C], padding h/w if needed.""" + B = latents.shape[0] + p = self.latent_patch_size + C = self.latent_channel_size + hp, wp, H_padded, W_padded = self._pad_to_patch_size(h, w) + + if H_padded != h or W_padded != w: + latents = F.pad(latents, (0, W_padded - w, 0, H_padded - h)) + + x = latents.reshape(B, C, t, hp, p, wp, p) + x = x.permute(0, 2, 3, 5, 4, 6, 1) # [B, t, hp, wp, p, p, C] + return x.reshape(B, t * hp * wp, p * p * C) + + def unpatchify(self, tokens: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor: + """[B, t*hp*wp, p*p*C] -> [B, C, t, h, w], cropping padding if needed.""" + B = tokens.shape[0] + p = self.latent_patch_size + C = self.latent_channel_size + hp, wp, H_padded, W_padded = self._pad_to_patch_size(h, w) + + x = tokens.reshape(B, t, hp, wp, p, p, C) + x = x.permute(0, 6, 1, 2, 4, 3, 5) # [B, C, t, hp, p, wp, p] + x = x.reshape(B, C, t, H_padded, W_padded) + + if H_padded != h or W_padded != w: + x = x[:, :, :, :h, :w] + return x + + # -- RoPE computation ---------------------------------------------------- + + def _compute_rope_freqs( + self, + text_mask: torch.Tensor, + t: int, + hp: int, + wp: int, + fps: float | None, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]: + """Compute mRoPE cos/sin for UND text and GEN media pathways.""" + B = text_mask.shape[0] + S_text = text_mask.shape[1] + text_lengths = text_mask.sum(dim=1).long() + effective_fps = fps if fps is not None and t > 1 else None + + text_pos_list = [] + gen_pos_list = [] + for b in range(B): + 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, + ) + if real_len < S_text: + t_pos = torch.cat( + [t_pos, torch.zeros(3, S_text - real_len, dtype=t_pos.dtype)], + dim=1, + ) + text_pos_list.append(t_pos) + gen_pos_list.append(v_pos) + + text_pos_ids = torch.stack(text_pos_list, dim=1).to(device) # [3, B, S_text] + gen_pos_ids = torch.stack(gen_pos_list, dim=1).to(device) # [3, B, S_gen] + + rotary_emb = self.language_model.rotary_emb + _dummy = torch.tensor([], dtype=dtype, device=device) + cos_und, sin_und = rotary_emb(_dummy, position_ids=text_pos_ids) + cos_gen, sin_gen = rotary_emb(_dummy, position_ids=gen_pos_ids) + + freqs_und = (cos_und.unsqueeze(2), sin_und.unsqueeze(2)) # (B, S, 1, D) + freqs_gen = (cos_gen.unsqueeze(2), sin_gen.unsqueeze(2)) + return freqs_und, freqs_gen + + # -- Cache management ---------------------------------------------------- + + def reset_cache(self) -> None: + self.cached_kv = None + self.cached_freqs_gen = None + + @staticmethod + def _validate_gen_sequence_parallel( + *, + s_gen: int, + ulysses_size: int, + ) -> None: + if ulysses_size <= 1 or s_gen % ulysses_size == 0: + return + + adjust_detail = ( + "Adjust the spatial resolution so that t * ceil(h/patch) * ceil(w/patch) is a multiple of ulysses_degree." + ) + raise ValueError( + f"GEN sequence length ({s_gen} video tokens) must be divisible by " + f"ulysses_degree ({ulysses_size}). {adjust_detail}" + ) + + # -- Forward ------------------------------------------------------------- + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + text_ids: torch.Tensor, + text_mask: torch.Tensor, + video_shape: tuple[int, int, int], + fps: float | None = None, + noisy_frame_mask: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + """ + Args: + hidden_states: [B, C, t, h, w] noisy latents + timestep: [B] diffusion timestep + text_ids: [B, S_text] tokenized text + text_mask: [B, S_text] attention mask (1=real, 0=pad) + video_shape: (t, h, w) in latent space + fps: video frame rate for temporal mRoPE modulation + 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). + + Returns: + [B, C, t, h, w] velocity prediction. + """ + t, h, w = video_shape + hp, wp, _, _ = self._pad_to_patch_size(h, w) + text_lengths = text_mask.sum(dim=1) + min_real_len = int(text_lengths.min().item()) + max_real_len = int(text_lengths.max().item()) + if min_real_len != max_real_len: + raise ValueError( + f"Cosmos3 requires identical real text lengths within a batch " + f"(got min={min_real_len}, max={max_real_len})." + ) + + # Query Ulysses state at runtime + ulysses_size, _, _ = _get_ulysses_state() + + # Patchify latents and project to hidden space + hidden_video = self.proj_in(self.patchify(hidden_states, t, h, w)) + + # Timestep embedding (fp32 for precision). + # For I2V: only add to noisy tokens, not conditioned ones. + # Conditioned frames are clean context and should not receive + # the diffusion timestep signal. + with torch.autocast("cuda", enabled=True, dtype=torch.float32): + time_embed = self.time_embedder(timestep * self.timestep_scale) + time_embed = time_embed.to(hidden_states.dtype) + + if noisy_frame_mask is not None: + # Build per-token mask from per-frame mask. + # noisy_frame_mask: [B, 1, t, 1, 1] → token mask: [B, t*hp*wp, 1] + token_noisy_mask = ( + noisy_frame_mask[:, 0, :, 0, 0] # [B, t] + .unsqueeze(-1) # [B, t, 1] + .expand(-1, -1, hp * wp) # [B, t, hp*wp] + .reshape(hidden_video.shape[0], -1, 1) # [B, t*hp*wp, 1] + ) + hidden_video = hidden_video + time_embed.unsqueeze(1) * token_noisy_mask + else: + hidden_video = hidden_video + time_embed.unsqueeze(1) + + hidden_gen = hidden_video + + # Run UND pathway once and cache K/V (replicated across all ranks) + if self.cached_kv is None: + freqs_und, freqs_gen = self._compute_rope_freqs( + text_mask, + t, + hp, + wp, + fps, + hidden_states.device, + hidden_states.dtype, + ) + cached_kv_full = self.language_model(text_ids, freqs_und) + self.cached_freqs_gen = freqs_gen + + # Trim to real text length (remove padding). K/V stay replicated; + # the framework Attention layer head-slices them via joint_key/value. + self.cached_kv = [(k[:, :max_real_len], v[:, :max_real_len]) for k, v in cached_kv_full] + + # Run GEN layers. UND K/V (replicated) is passed to each layer; + # the Cosmos3CrossAttention forwards them as joint_key/value so the + # framework Attention handles the Ulysses head-slicing internally. + if self.cached_kv is None or self.cached_freqs_gen is None: + raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.") + self._validate_gen_sequence_parallel( + s_gen=hidden_gen.shape[1], + ulysses_size=ulysses_size, + ) + freqs_cos, freqs_sin = self.cached_freqs_gen + hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) + freqs_gen = (freqs_cos, freqs_sin) + + if len(self.gen_layers) == len(self.cached_kv): + for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True): + hidden_gen = layer( + hidden_gen, + k_und=k_und, + v_und=v_und, + freqs_cos=freqs_cos, + freqs_sin=freqs_sin, + ) + # Cache-dit's block wrapper may return a tuple; unwrap it. + if isinstance(hidden_gen, tuple): + hidden_gen = hidden_gen[0] + else: + # Cache-dit patches gen_layers to a grouped wrapper. + for layer in self.gen_layers: + hidden_gen = layer( + hidden_gen, + cached_kv=self.cached_kv, + freqs_gen=freqs_gen, + ) + if isinstance(hidden_gen, tuple): + hidden_gen = hidden_gen[0] + + hidden_gen = self.gen_sp_gather(hidden_gen) + + # Final norm and project back to latent space + hidden_gen = self.norm_moe_gen(hidden_gen) + return self.unpatchify(self.proj_out(hidden_gen), t, h, w) + + def post_load_weights(self) -> None: + """Post-load processing: ensure correct dtypes.""" + self.time_embedder.to(torch.float32) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index fe9ed8fcf9e..55dfaa5b451 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -266,6 +266,11 @@ "pipeline_omnivoice", "OmniVoicePipeline", ), + "Cosmos3OmniDiffusersPipeline": ( + "cosmos3", + "pipeline_cosmos3", + "Cosmos3OmniDiffusersPipeline", + ), "DiffusersAdapterPipeline": ( "diffusers_adapter", "pipeline_diffusers_adapter", @@ -498,6 +503,7 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - "OmniVoicePipeline": "get_omnivoice_post_process_func", "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func", "SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func", + "Cosmos3OmniDiffusersPipeline": "get_cosmos3_post_process_func", "HiDreamImagePipeline": "get_hidream_image_post_process_func", } @@ -522,6 +528,7 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_pre_process_func", "HunyuanImage3ForCausalMM": "get_hunyuan_image_3_pre_process_func", "MagiHumanPipeline": "get_magi_human_pre_process_func", + "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func", } diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 5163de016de..dbfbe15e92a 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1958,6 +1958,7 @@ def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list: "max_num_seqs": kwargs.get("max_num_seqs") or 1, "parallel_config": parallel_config, "model_class_name": kwargs.get("model_class_name", None), + "model_config": kwargs.get("model_config", None), "additional_config": kwargs.get("additional_config", None), "step_execution": kwargs.get("step_execution", False), "vae_use_slicing": kwargs.get("vae_use_slicing", False), diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 4be487415b1..d59f3dd480e 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -96,6 +96,12 @@ def cmd(args: argparse.Namespace) -> None: if hasattr(args, "model_tag") and args.model_tag is not None: args.model = args.model_tag + if getattr(args, "no_guardrails", False): + existing = getattr(args, "model_config", None) + model_config = dict(existing) if isinstance(existing, dict) else {} + model_config["guardrails"] = False + args.model_config = model_config + if args.headless: run_headless(args) else: @@ -642,6 +648,15 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgu help="Maximum length for TTS voice style instructions (overrides stage config, default: 500).", ) + # Disable safety guardrails for this server (currently only applicable for Cosmos3) + # TODO: drop once --model-config-override lands (3/N config refactor) + omni_config_group.add_argument( + "--no-guardrails", + dest="no_guardrails", + action="store_true", + help="Disable Cosmos3 text/video safety guardrails for this server.", + ) + # Enable diffusion pipeline profiling omni_config_group.add_argument( "--enable-diffusion-pipeline-profiler", diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 7afba4d4553..ca7182edb43 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1631,6 +1631,10 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request) extra_body["guidance_scale"] = request.guidance_scale if request.true_cfg_scale is not None: extra_body["true_cfg_scale"] = request.true_cfg_scale + if request.flow_shift is not None: + extra_body["flow_shift"] = request.flow_shift + if request.extra_params is not None: + extra_body["extra_params"] = request.extra_params if request.generator_device is not None: extra_body["generator_device"] = request.generator_device if request.lora is not None: @@ -1661,17 +1665,19 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request) return ImageGenerationResponse(created=int(time.time()), data=image_data) # Build params - pass through user values directly - prompt: OmniTextPrompt = {"prompt": request.prompt} + prompt: OmniTextPrompt = {"prompt": request.prompt, "modalities": ["image"]} if request.negative_prompt is not None: prompt["negative_prompt"] = request.negative_prompt gen_params = OmniDiffusionSamplingParams(num_outputs_per_prompt=request.n) - extra_args = {} + extra_args = dict(request.extra_params or {}) if request.use_system_prompt is not None: extra_args["use_system_prompt"] = request.use_system_prompt if request.system_prompt is not None: extra_args["system_prompt"] = request.system_prompt if request.bot_task is not None: extra_args["bot_task"] = request.bot_task + if request.flow_shift is not None: + extra_args["flow_shift"] = request.flow_shift if extra_args: gen_params.extra_args = extra_args # Parse per-request LoRA (compatible with chat's extra_body.lora shape). @@ -1846,7 +1852,7 @@ async def edit_images( try: # 2. Build prompt & images params cot_output = None - prompt: OmniTextPrompt = {"prompt": prompt} + prompt: OmniTextPrompt = {"prompt": prompt, "modalities": ["image"]} if negative_prompt is not None: prompt["negative_prompt"] = negative_prompt input_images_list = [] diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index 34fee760b0e..7f82151bfae 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -122,6 +122,13 @@ def validate_use_system_prompt(cls, v): le=20.0, description="True CFG scale (model-specific parameter, may be ignored if not supported)", ) + flow_shift: float | None = Field( + default=None, description="Scheduler flow_shift (sigma shift) for flow-matching diffusion models." + ) + extra_params: dict[str, Any] | None = Field( + default=None, + description="Optional model-specific parameters passed directly to the model's extra_args.", + ) seed: int | None = Field(default=None, description="Random seed for reproducibility") generator_device: str | None = Field( default=None, diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 2247f739452..58c61a921cb 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -2666,6 +2666,7 @@ def _prepare_diffusion_image_request( gen_prompt: OmniTextPrompt = { "prompt": prompt, "negative_prompt": negative_prompt, + "modalities": ["image"], } if pil_images: if len(pil_images) == 1: @@ -2998,6 +2999,7 @@ async def _create_diffusion_chat_completion( gen_prompt: OmniTextPrompt = { "prompt": prompt, "negative_prompt": negative_prompt, + "modalities": ["image"], } gen_params = OmniDiffusionSamplingParams( height=height, diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 043ccd98322..b6ed49996fe 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -96,7 +96,7 @@ async def _run_and_extract( reference_image: ReferenceImage | None = None, ) -> VideoGenerationArtifacts: """Run the generation pipeline and extract video/audio/profiler outputs.""" - prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt) + prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt, modalities=["video"]) if request.negative_prompt is not None: prompt["negative_prompt"] = request.negative_prompt @@ -259,18 +259,11 @@ async def generate_video_bytes( @staticmethod def _resolve_video_fps_multiplier(result: Any) -> int: - custom_output = getattr(result, "custom_output", None) + custom_output = OmniOpenAIServingVideo._extract_custom_output(result) if isinstance(custom_output, dict): multiplier = custom_output.get("video_fps_multiplier") if multiplier is not None: return int(multiplier) - request_output = getattr(result, "request_output", None) - if request_output is not None: - custom_output = getattr(request_output, "custom_output", None) - if isinstance(custom_output, dict): - multiplier = custom_output.get("video_fps_multiplier") - if multiplier is not None: - return int(multiplier) return 1 def _resolve_default_sampling_params(self) -> OmniDiffusionSamplingParams: @@ -443,6 +436,24 @@ def _resolve_audio_sample_rate(self, result: Any) -> int: return 24000 + @staticmethod + def _extract_custom_output(result: Any) -> dict[str, Any]: + custom_output = getattr(result, "custom_output", None) + if isinstance(custom_output, dict): + return custom_output + + request_output = getattr(result, "request_output", None) + if isinstance(request_output, dict): + custom_output = request_output.get("custom_output") + if custom_output is None: + custom_output = request_output.get("_custom_output") + elif request_output is not None: + custom_output = getattr(request_output, "custom_output", None) + if custom_output is None: + custom_output = getattr(request_output, "_custom_output", None) + + return custom_output if isinstance(custom_output, dict) else {} + @staticmethod def _resolve_fps(result: Any) -> int | None: """Extract fps from multimodal_output if the model reported it.""" diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 1b80f4b1b77..a2b6bf722d5 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -33,6 +33,9 @@ class OmniTextPrompt(TextPrompt): """ negative_prompt: NotRequired[str] + # Using modalities field to differentiate between different tasks for the same pipeline + # for example Cosmos3OmniDiffusersPipeline handles t2i and t2v in the same pipeline. + modalities: NotRequired[list[str]] prompt_embeds: NotRequired[torch.Tensor] negative_prompt_embeds: NotRequired[torch.Tensor] additional_information: NotRequired[dict[str, Any]]