Add Cosmos3 model support (#3454)
Signed-off-by: Maciej Bala <mbala@nvidia.com> Signed-off-by: MaciejBalaNV <mbala@nvidia.com> Signed-off-by: lishunyang12 <lishunyang12@163.com> Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com> Co-authored-by: lishunyang12 <lishunyang12@163.com>
This commit is contained in:
29 files changed
+4095
-21
No files matched your search
@@ -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` | ✅︎ | ✅︎ | | |
|
||||
|
||||
@@ -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**
|
||||
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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): <https://huggingface.co/nvidia/Cosmos3-Nano>
|
||||
- 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 <https://huggingface.co/nvidia/Cosmos-1.0-Guardrail>
|
||||
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": <PIL.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.
|
||||
+19
@@ -25,6 +25,7 @@ SEPARATE_CFG_ENABLERS = [
|
||||
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(),
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -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))
|
||||
@@ -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)
|
||||
@@ -10,6 +10,9 @@ from types import SimpleNamespace
|
||||
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 @@ class MockScheduler:
|
||||
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
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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(
|
||||
|
||||
@@ -243,6 +243,7 @@ def test_t2v_video_generation_form(test_client, mocker: MockerFixture):
|
||||
_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):
|
||||
|
||||
@@ -51,6 +51,7 @@ class CuDNNAttentionImpl(AttentionImpl):
|
||||
) -> 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 @@ class CuDNNAttentionImpl(AttentionImpl):
|
||||
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 @@ class CuDNNAttentionImpl(AttentionImpl):
|
||||
dropout_p=0.0,
|
||||
is_causal=self.causal,
|
||||
scale=self.softmax_scale,
|
||||
enable_gqa=self.requires_gqa,
|
||||
)
|
||||
return output.permute(0, 2, 1, 3)
|
||||
+79
@@ -1725,6 +1725,84 @@ def enable_cache_for_hunyuan_video_15(pipeline: Any, cache_config: Any) -> Calla
|
||||
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 @@ CUSTOM_DIT_ENABLERS.update(
|
||||
"HunyuanVideo15Pipeline": enable_cache_for_hunyuan_video_15,
|
||||
"HunyuanVideo15I2VPipeline": enable_cache_for_hunyuan_video_15,
|
||||
"HeliosPipeline": enable_cache_for_helios,
|
||||
"Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -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 @@ class DiffusionEngine:
|
||||
# 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:
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -266,6 +266,11 @@ _DIFFUSION_MODELS = {
|
||||
"pipeline_omnivoice",
|
||||
"OmniVoicePipeline",
|
||||
),
|
||||
"Cosmos3OmniDiffusersPipeline": (
|
||||
"cosmos3",
|
||||
"pipeline_cosmos3",
|
||||
"Cosmos3OmniDiffusersPipeline",
|
||||
),
|
||||
"DiffusersAdapterPipeline": (
|
||||
"diffusers_adapter",
|
||||
"pipeline_diffusers_adapter",
|
||||
@@ -498,6 +503,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
|
||||
"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 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
|
||||
"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",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -1958,6 +1958,7 @@ class AsyncOmniEngine:
|
||||
"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),
|
||||
|
||||
@@ -96,6 +96,12 @@ class OmniServeCommand(CLISubcommand):
|
||||
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 @@ class OmniServeCommand(CLISubcommand):
|
||||
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",
|
||||
|
||||
@@ -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 = []
|
||||
|
||||
@@ -122,6 +122,13 @@ class ImageGenerationRequest(BaseModel):
|
||||
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,
|
||||
|
||||
@@ -2666,6 +2666,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
gen_prompt: OmniTextPrompt = {
|
||||
"prompt": prompt,
|
||||
"negative_prompt": negative_prompt,
|
||||
"modalities": ["image"],
|
||||
}
|
||||
if pil_images:
|
||||
if len(pil_images) == 1:
|
||||
@@ -2998,6 +2999,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
gen_prompt: OmniTextPrompt = {
|
||||
"prompt": prompt,
|
||||
"negative_prompt": negative_prompt,
|
||||
"modalities": ["image"],
|
||||
}
|
||||
gen_params = OmniDiffusionSamplingParams(
|
||||
height=height,
|
||||
|
||||
@@ -96,7 +96,7 @@ class OmniOpenAIServingVideo:
|
||||
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 @@ class OmniOpenAIServingVideo:
|
||||
|
||||
@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 @@ class OmniOpenAIServingVideo:
|
||||
|
||||
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."""
|
||||
|
||||
@@ -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]]
|
||||
|
||||
Reference in new issue
Block a user