Cosmos3 video to video generation (#4266)
Signed-off-by: Maciej Bala <mbala@nvidia.com>
This commit is contained in:
MaciejBalaNV authored and GitHub committed
1 parent
9c9d4ed4a8
commit
2f25195750
14 files changed
+1037
-94
No files matched your search
@@ -33,7 +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, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
|
||||
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `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` | ✅︎ | ✅︎ | | |
|
||||
|
||||
@@ -67,8 +67,9 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `input_reference` | file | null | Uploaded reference image for image-to-video requests |
|
||||
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` |
|
||||
| `input_reference` | file | null | Uploaded reference image or video for image-to-video/video-to-video requests |
|
||||
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` or `video_reference` |
|
||||
| `video_reference` | string | null | JSON-encoded reference video payload; do not combine with `input_reference` or `image_reference` |
|
||||
| `width` | integer | model default | Output video width |
|
||||
| `height` | integer | model default | Output video height |
|
||||
| `num_frames` | integer | model default | Number of generated frames |
|
||||
@@ -80,6 +81,8 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
|
||||
| `flow_shift` | number | null | Scheduler flow-shift value |
|
||||
| `true_cfg_scale` | number | null | True CFG scale when supported by the model |
|
||||
| `seed` | integer | null | Random seed for reproducibility |
|
||||
| `generate_sound` | boolean | false | Request model-generated audio for video models that support sound generation |
|
||||
| `sound_duration` | number | null | Duration in seconds for generated audio; defaults to generated video duration |
|
||||
| `negative_prompt` | string | null | Text describing what to avoid in the generated video |
|
||||
| `enable_frame_interpolation` | boolean | null | Enable post-generation frame interpolation |
|
||||
| `frame_interpolation_exp` | integer | null | Interpolation exponent; `1=2x`, `2=4x`, and so on |
|
||||
@@ -123,6 +126,41 @@ curl -s http://localhost:8091/v1/videos \
|
||||
-F "fps=16"
|
||||
```
|
||||
|
||||
### Video-to-Video
|
||||
|
||||
For models that support video conditioning, upload the reference video with
|
||||
`input_reference`:
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/videos \
|
||||
-F "prompt=continue this motion with consistent subjects and lighting" \
|
||||
-F "input_reference=@input.mp4;type=video/mp4" \
|
||||
-F "width=1280" \
|
||||
-F "height=720" \
|
||||
-F "num_frames=80" \
|
||||
-F "fps=16"
|
||||
```
|
||||
|
||||
You can also pass a JSON-safe video URL or `data:video/...;base64,...` payload
|
||||
through `video_reference`. Do not send `video_reference` together with
|
||||
`input_reference` or `image_reference`.
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/videos \
|
||||
-F "prompt=continue this motion with consistent subjects and lighting" \
|
||||
-F 'video_reference={"video_url":"https://example.com/input.mp4"}' \
|
||||
-F "width=1280" \
|
||||
-F "height=720" \
|
||||
-F "num_frames=80" \
|
||||
-F "fps=16"
|
||||
```
|
||||
|
||||
JSON references currently support `image_url`/`video_url`; `file_id` references
|
||||
are not implemented yet. Models may expose additional V2V controls through
|
||||
`extra_params`. For example, Cosmos3 supports
|
||||
`condition_frame_indexes_vision` and `condition_video_keep` to select which
|
||||
decoded reference frames are used as clean conditioning.
|
||||
|
||||
### Synchronous Generation
|
||||
|
||||
```bash
|
||||
@@ -146,7 +184,10 @@ export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
|
||||
|
||||
## Model-Specific Examples
|
||||
|
||||
For complete text-to-video and image-to-video walkthroughs, see:
|
||||
For complete text-to-video, image-to-video, and model-specific video-to-video
|
||||
walkthroughs, see:
|
||||
|
||||
- [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md)
|
||||
- [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md)
|
||||
- [Cosmos3 recipes](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md)
|
||||
for model-specific video-to-video examples and conditioning controls
|
||||
@@ -0,0 +1,14 @@
|
||||
{
|
||||
"prompt": "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. The robotic arm is in the process of pouring a transparent liquid into the glass. To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. The lighting is bright, casting soft shadows on the table. The robotic arm's movements are smooth and controlled, demonstrating precision in its task. As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. The other robotic arm on the right side also stays stationary throughout the video. The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper.",
|
||||
"negative_prompt": "",
|
||||
"vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4",
|
||||
"height": 720,
|
||||
"width": 1280,
|
||||
"num_frames": 189,
|
||||
"num_inference_steps": 35,
|
||||
"guidance_scale": 6.0,
|
||||
"flow_shift": 10.0,
|
||||
"fps": 24,
|
||||
"condition_frame_indexes_vision": [0, 1],
|
||||
"condition_video_keep": "first"
|
||||
}
|
||||
+2
-2
@@ -36,8 +36,8 @@ 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 |
|
||||
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
|
||||
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
|
||||
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
|
||||
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
|
||||
| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
|
||||
| [`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 |
|
||||
|
||||
@@ -1,34 +1,43 @@
|
||||
# Cosmos3-Nano
|
||||
|
||||
> Text-to-image, text-to-video, and image-to-video serving
|
||||
> Text-to-image, text-to-video, image-to-video, and video-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, with optional synchronized audio (video + sound), action policy
|
||||
- Task: Text-to-image (T2I), text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) generation, with optional synchronized audio (video + sound), action policy
|
||||
- Mode: Online serving with the OpenAI-compatible image/video APIs, plus offline generation via the `Omni` API
|
||||
- Maintainer: Community
|
||||
|
||||
## 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:
|
||||
A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves these 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.
|
||||
- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image/video.
|
||||
- **I2V** — `POST /v1/videos/sync` with a reference image (`input_reference` file
|
||||
upload, or `image_reference` JSON).
|
||||
- **V2V** — `POST /v1/videos/sync` with a reference video (`input_reference`
|
||||
video upload, or `video_reference` JSON). Cosmos3 conditions on selected
|
||||
reference-video latent frames; use `extra_params.condition_frame_indexes_vision`
|
||||
and `extra_params.condition_video_keep` to choose which prefix/tail frames guide
|
||||
generation.
|
||||
- **T2VS / I2VS** — add `generate_sound=true` (and optional `sound_duration`) to a
|
||||
T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into
|
||||
the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples.
|
||||
- **Action** — pass `extra_params={"action_mode": ...}` to drive Physical-AI tasks:
|
||||
- `forward_dynamics` — given a first frame **and** an action trajectory, roll out
|
||||
the resulting video. Synchronous: `POST /v1/videos/sync`.
|
||||
- `policy` — given a first frame and a language instruction, **predict** the action
|
||||
trajectory (and a rollout video). Use the async `POST /v1/videos` endpoint and
|
||||
read the predicted action from the top-level `action` field
|
||||
- `forward_dynamics` — given a first frame or video **and** an action trajectory,
|
||||
roll out the resulting video. Synchronous: `POST /v1/videos/sync`.
|
||||
- `policy` — given a first frame or video and a language instruction,
|
||||
**predict** the action trajectory (and a rollout video). Use the async
|
||||
`POST /v1/videos` endpoint and read the predicted action from the top-level
|
||||
`action` field.
|
||||
- `inverse_dynamics` — given a video, **recover** the action trajectory. Use
|
||||
the async `POST /v1/videos` endpoint and read the recovered action from
|
||||
the top-level `action` field
|
||||
(`{data, shape, dtype, raw_action_dim, domain_id}`).
|
||||
|
||||
Action requests also take `domain_name` (e.g. `av`, `bridge_orig_lerobot`,
|
||||
@@ -38,8 +47,9 @@ the mode is selected per request:
|
||||
**`nvidia/Cosmos3-Nano-Policy-DROID`** is served the same way
|
||||
(`domain_name=droid_lerobot`).
|
||||
|
||||
`inverse_dynamics` (recover the action from a given video) is supported by the
|
||||
pipeline; **online inference of inverse dynamics will be added in a follow-up MR.**
|
||||
Action requests can use `input_reference` or `video_reference` for video input.
|
||||
`policy` and `forward_dynamics` can also use an image reference; `inverse_dynamics`
|
||||
requires a video reference.
|
||||
|
||||
## References
|
||||
|
||||
@@ -142,6 +152,36 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
-F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
|
||||
-o cosmos3_i2v.mp4
|
||||
|
||||
# Video-to-video -> /v1/videos/sync with an uploaded reference video.
|
||||
# By default Cosmos3 conditions on latent indexes [0, 1]. For the default
|
||||
# temporal VAE stride this decodes only the first 5 input frames.
|
||||
# The model works best when the prompt describes the actual situation happening in the video.
|
||||
# Generic prompts may create sub-standard generations.
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
-H "Accept: video/mp4" \
|
||||
-F "model=nvidia/Cosmos3-Nano" \
|
||||
-F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
|
||||
-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,"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}' \
|
||||
-F "seed=2222" \
|
||||
-F "input_reference=@/path/to/reference.mp4;type=video/mp4" \
|
||||
-o cosmos3_v2v.mp4
|
||||
|
||||
# V2V can also use a JSON-safe URL/data-URL video reference. Do not combine
|
||||
# video_reference with input_reference or image_reference.
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
-H "Accept: video/mp4" \
|
||||
-F "model=nvidia/Cosmos3-Nano" \
|
||||
-F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
|
||||
-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={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"last"}' \
|
||||
-F 'video_reference={"video_url":"https://example.com/reference.mp4"}' \
|
||||
-o cosmos3_v2v_from_url.mp4
|
||||
|
||||
# Text-to-video-with-sound
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
@@ -190,6 +230,21 @@ VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
|
||||
# poll until status == completed, then:
|
||||
curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
|
||||
curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_policy.mp4
|
||||
|
||||
# Action — inverse dynamics (video -> recovered action trajectory).
|
||||
# Asynchronous: use the job metadata to read the recovered action.
|
||||
VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
|
||||
-H "Accept: application/json" \
|
||||
--form-string "model=nvidia/Cosmos3-Nano" \
|
||||
--form-string "prompt=Recover the robot action trajectory from this clip." \
|
||||
-F "input_reference=@motion_clip.mp4;type=video/mp4" \
|
||||
-F "size=640x480" -F "num_frames=17" -F "fps=5" \
|
||||
-F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \
|
||||
--form-string 'extra_params={"action_mode":"inverse_dynamics","domain_name":"bridge_orig_lerobot","raw_action_dim":10,"action_chunk_size":16}' \
|
||||
-F "seed=0" | jq -r '.id')
|
||||
# poll until status == completed, then:
|
||||
curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
|
||||
curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_inverse_dynamics.mp4
|
||||
```
|
||||
|
||||
#### Notes
|
||||
@@ -206,24 +261,27 @@ curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_polic
|
||||
- **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`.
|
||||
3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V/V2V
|
||||
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. The
|
||||
per-request flag only takes effect when the server was launched **with**
|
||||
guardrails enabled (it cannot re-enable them on a `--no-guardrails` server).
|
||||
`use_resolution_template` / `use_duration_template` are off by default and only
|
||||
needed when not using upsampled prompts that already encode resolution/duration.
|
||||
For V2V, `condition_frame_indexes_vision` selects the clean conditioned latent
|
||||
frame indexes (default `[0, 1]`), and `condition_video_keep` selects whether the
|
||||
API decodes the first or last needed reference frames (`"first"` by default).
|
||||
- **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"`).
|
||||
- Action `forward_dynamics` (sync `/v1/videos/sync`) and `policy` (async
|
||||
`/v1/videos`, returns the predicted action under the top-level `action`
|
||||
field) are supported online. **Online inference of inverse dynamics will be
|
||||
added in a follow-up MR.**
|
||||
- Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported
|
||||
online. Use async `POST /v1/videos` when you need the predicted/recovered
|
||||
action payload under the top-level `action` field; sync `/v1/videos/sync`
|
||||
returns raw MP4 bytes and does not expose action metadata in the response body.
|
||||
|
||||
### 1x GPU (Offline generation)
|
||||
|
||||
@@ -287,8 +345,9 @@ python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image
|
||||
|
||||
- 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
|
||||
`num_frames`/`fps`, `multi_modal_data={"image": ...}` for I2V, and
|
||||
`multi_modal_data={"video": [<PIL frames>]}` or a video tensor/array for V2V.
|
||||
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.
|
||||
@@ -1,12 +1,12 @@
|
||||
# Cosmos3-Super
|
||||
|
||||
> Frontier 64B world model: text-to-image, text-to-video, image-to-video (+ optional audio)
|
||||
> Frontier 64B world model: text-to-image, text-to-video, image-to-video, video-to-video (+ optional audio)
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: NVIDIA
|
||||
- Model: `nvidia/Cosmos3-Super` (64B; also `Cosmos3-Super-Text2Image`, `Cosmos3-Super-Image2Video`)
|
||||
- Task: T2I, T2V, I2V generation, with optional synchronized audio (video + sound)
|
||||
- Task: T2I, T2V, I2V, V2V generation, with optional synchronized audio (video + sound)
|
||||
- Mode: Online serving with the OpenAI-compatible image/video APIs
|
||||
- Maintainer: Community
|
||||
|
||||
@@ -16,8 +16,8 @@ Use this recipe to deploy the 64B `nvidia/Cosmos3-Super` for the highest-quality
|
||||
Cosmos3 generation. It shares the same `Cosmos3OmniDiffusersPipeline` and request
|
||||
formats as [Cosmos3-Nano](./Cosmos3-Nano.md) — only the checkpoint size and the
|
||||
recommended parallelism differ. Mode is selected per request (T2I →
|
||||
`/v1/images/generations`; T2V/I2V → `/v1/videos/sync`; add `generate_sound=true`
|
||||
for audio).
|
||||
`/v1/images/generations`; T2V/I2V/V2V → `/v1/videos/sync`; add
|
||||
`generate_sound=true` for audio).
|
||||
|
||||
## References
|
||||
|
||||
@@ -62,8 +62,9 @@ disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs.
|
||||
#### Verification
|
||||
|
||||
Requests are identical to Nano (see [`Cosmos3-Nano.md`](./Cosmos3-Nano.md) for full
|
||||
T2I/T2V/I2V/T2VS curls); official params: `size=1280x720, num_frames=189, fps=24,
|
||||
num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0, max_sequence_length=4096`.
|
||||
T2I/T2V/I2V/V2V/T2VS curls); official params: `size=1280x720, num_frames=189,
|
||||
fps=24, num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0,
|
||||
max_sequence_length=4096`.
|
||||
|
||||
```bash
|
||||
curl http://localhost:8000/v1/models
|
||||
@@ -84,6 +85,15 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
-F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
|
||||
-o cosmos3_super_i2v.mp4
|
||||
|
||||
# V2V — add an uploaded reference video. condition_video_keep can be "first" or "last".
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
-F "model=nvidia/Cosmos3-Super" -F "prompt=Continue the same scene with smooth natural motion." \
|
||||
-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={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}' \
|
||||
-F "seed=2222" -F "input_reference=@/path/to/reference.mp4;type=video/mp4" \
|
||||
-o cosmos3_super_v2v.mp4
|
||||
|
||||
# T2V + sound — add generate_sound/sound_duration (output muxes AAC 48 kHz stereo)
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
-F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \
|
||||
@@ -103,10 +113,12 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
- T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo**
|
||||
- (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
|
||||
- **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk.
|
||||
- Same generation defaults, supported sizes, and `generate_sound`/`sound_duration`
|
||||
semantics as Nano, including the **action** modality: `forward_dynamics`
|
||||
(sync `/v1/videos/sync`) and `policy` (async `/v1/videos`, predicted action under
|
||||
the top-level `action` field) — see the Cosmos3-Nano recipe for the request shape.
|
||||
Online inference of inverse dynamics will be added in a follow-up MR. Verified on
|
||||
the 64B Super under `--cfg-parallel-size 2`: async `policy` returns the predicted
|
||||
action (`[16, 10]`) and the rollout video reliably.
|
||||
- Same generation defaults, supported sizes, V2V reference-video controls
|
||||
(`condition_frame_indexes_vision`, `condition_video_keep`), and
|
||||
`generate_sound`/`sound_duration`
|
||||
semantics as Nano, including the **action** modality: `forward_dynamics`,
|
||||
`policy`, and `inverse_dynamics` — see the Cosmos3-Nano recipe for the request
|
||||
shapes. Use async `/v1/videos` when you need predicted/recovered action metadata
|
||||
under the top-level `action` field. Verified on the 64B Super under
|
||||
`--cfg-parallel-size 2`: async `policy` returns the predicted action (`[16, 10]`)
|
||||
and the rollout video reliably.
|
||||
@@ -374,6 +374,19 @@ def test_preprocess_i2v_image_and_action_video_inputs() -> None:
|
||||
assert tuple(additional["preprocessed_image"].shape) == (1, 3, 16, 32)
|
||||
assert tuple(additional["preprocessed_video"].shape) == (1, 3, 3, 16, 32)
|
||||
|
||||
frames = [Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue", "yellow", "purple", "black")]
|
||||
v2v = SimpleNamespace(
|
||||
prompts=[{"prompt": "Continue.", "multi_modal_data": {"video": frames}}],
|
||||
sampling_params=SimpleNamespace(
|
||||
height=16,
|
||||
width=32,
|
||||
extra_args={"condition_frame_indexes_vision": [0, 1], "condition_video_keep": "last"},
|
||||
),
|
||||
)
|
||||
additional = preprocess(v2v).prompts[0]["additional_information"]
|
||||
assert tuple(additional["preprocessed_video"].shape) == (1, 3, 5, 16, 32)
|
||||
assert additional["condition_frame_indexes_vision"] == [0, 1]
|
||||
|
||||
|
||||
def test_postprocess_handles_image_video_audio_and_validation() -> None:
|
||||
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_post_process_func
|
||||
@@ -464,6 +477,19 @@ def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline)
|
||||
assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]]]]]
|
||||
assert image_latent.shape == (1, 2, 1, 2, 3)
|
||||
|
||||
pipeline._encode_video_tensor = lambda *args, **kwargs: torch.full((1, 2, 3, 2, 3), 6.0)
|
||||
v2v_latents, v2v_velocity_mask, v2v_condition = pipeline._prepare_latents_v2v(
|
||||
torch.zeros(1, 3, 5, 16, 24),
|
||||
16,
|
||||
24,
|
||||
9,
|
||||
torch.Generator(device="cpu").manual_seed(0),
|
||||
[0, 1],
|
||||
)
|
||||
torch.testing.assert_close(v2v_latents[:, :, 0:2], torch.full((1, 2, 2, 2, 3), 6.0))
|
||||
assert v2v_velocity_mask.tolist() == [[[[[0.0]], [[0.0]], [[1.0]]]]]
|
||||
assert v2v_condition.shape == (1, 2, 3, 2, 3)
|
||||
|
||||
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
|
||||
pipeline._sound_tokenizer = SimpleNamespace(
|
||||
sample_rate=10,
|
||||
@@ -670,6 +696,34 @@ class TestForwardRouting:
|
||||
)
|
||||
assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask
|
||||
|
||||
video_tensor = torch.zeros(1, 3, 5, 16, 16)
|
||||
v2v_condition = torch.full((1, 2, 2, 1, 1), 4.0)
|
||||
v2v_mask = torch.tensor([[[[[0.0]], [[1.0]]]]])
|
||||
pipeline._prepare_latents_v2v = lambda *args, **kwargs: (
|
||||
torch.zeros(1, 2, 2, 1, 1),
|
||||
v2v_mask,
|
||||
v2v_condition,
|
||||
)
|
||||
pipeline.forward(
|
||||
SimpleNamespace(
|
||||
prompts=[
|
||||
{
|
||||
"prompt": "continue",
|
||||
"modalities": ["video"],
|
||||
"additional_information": {
|
||||
"preprocessed_video": video_tensor,
|
||||
"condition_frame_indexes_vision": [0],
|
||||
},
|
||||
}
|
||||
],
|
||||
sampling_params=make_sampling_params(height=16, width=16, num_frames=5),
|
||||
)
|
||||
)
|
||||
assert captured["flow_shifts"][-1] == 10.0
|
||||
assert captured["format"]["negative_prompt"] == ""
|
||||
assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is v2v_mask
|
||||
assert captured["diffuse_calls"][-1]["condition_latents"] is v2v_condition
|
||||
|
||||
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
|
||||
sound_latents = torch.zeros(1, 3, 4)
|
||||
pipeline._resolve_sound_target_samples = lambda *args: (20, 2.0, 10)
|
||||
|
||||
@@ -13,12 +13,14 @@ import threading
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
from PIL import Image
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
from vllm_omni.entrypoints.openai import api_server
|
||||
from vllm_omni.entrypoints.openai.api_server import router
|
||||
from vllm_omni.entrypoints.openai.protocol.videos import (
|
||||
@@ -84,7 +86,8 @@ class BlockingVideoHandler:
|
||||
if self.stage_configs is None:
|
||||
self.stage_configs = stage_configs
|
||||
|
||||
async def generate_video_bytes(self, request, reference_id, *, reference_image=None):
|
||||
async def generate_video_bytes(self, request, reference_id, *, reference_image=None, reference_video=None):
|
||||
del request, reference_id, reference_image, reference_video
|
||||
self.started.set()
|
||||
try:
|
||||
await asyncio.Future()
|
||||
@@ -130,6 +133,30 @@ def _make_test_image_data_url(size=(64, 64)) -> str:
|
||||
return f"data:image/png;base64,{encoded}"
|
||||
|
||||
|
||||
def _make_test_video_bytes(size=(32, 24), num_frames=3) -> bytes:
|
||||
width, height = size
|
||||
frames = np.zeros((num_frames, height, width, 3), dtype=np.uint8)
|
||||
for idx in range(num_frames):
|
||||
frames[idx, :, :, 0] = idx * 40
|
||||
frames[idx, :, :, 1] = 128
|
||||
frames[idx, :, :, 2] = 255 - idx * 40
|
||||
return mux_video_audio_bytes(frames, fps=8, video_codec_options={"preset": "ultrafast", "threads": "0"})
|
||||
|
||||
|
||||
def _make_test_video_data_url(size=(32, 24), num_frames=3) -> str:
|
||||
encoded = base64.b64encode(_make_test_video_bytes(size, num_frames)).decode("utf-8")
|
||||
return f"data:video/mp4;base64,{encoded}"
|
||||
|
||||
|
||||
def _cosmos3_stage_configs():
|
||||
return [
|
||||
SimpleNamespace(
|
||||
stage_type="diffusion",
|
||||
engine_args=SimpleNamespace(model_class_name="Cosmos3OmniDiffusersPipeline"),
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def _wait_for_status(client: TestClient, video_id: str, status: str, timeout_s: float = 2.0):
|
||||
deadline = time.time() + timeout_s
|
||||
last_payload = None
|
||||
@@ -332,6 +359,104 @@ def test_i2v_video_generation_with_image_reference_form(test_client, mocker: Moc
|
||||
assert input_image.size == (40, 24)
|
||||
|
||||
|
||||
def test_v2v_video_generation_form(test_client, mocker: MockerFixture):
|
||||
video_bytes = _make_test_video_bytes((32, 24), num_frames=3)
|
||||
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
|
||||
return_value=b"fake-video",
|
||||
)
|
||||
response = test_client.post(
|
||||
"/v1/videos",
|
||||
data={"prompt": "Continue this motion."},
|
||||
files={"input_reference": ("input.mp4", video_bytes, "video/mp4")},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
video_id = response.json()["id"]
|
||||
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
|
||||
|
||||
engine = test_client.app.state.openai_serving_video._engine_client
|
||||
prompt = engine.captured_prompt
|
||||
assert "multi_modal_data" in prompt
|
||||
assert "video" in prompt["multi_modal_data"]
|
||||
input_video = prompt["multi_modal_data"]["video"]
|
||||
assert len(input_video) == 3
|
||||
assert all(isinstance(frame, Image.Image) for frame in input_video)
|
||||
assert input_video[0].size == (32, 24)
|
||||
|
||||
|
||||
def test_v2v_video_generation_with_video_reference_form(test_client, mocker: MockerFixture):
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
|
||||
return_value=b"fake-video",
|
||||
)
|
||||
response = test_client.post(
|
||||
"/v1/videos",
|
||||
data={
|
||||
"prompt": "Continue this motion.",
|
||||
"video_reference": json.dumps({"video_url": _make_test_video_data_url((32, 24), 2)}),
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
video_id = response.json()["id"]
|
||||
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
|
||||
|
||||
engine = test_client.app.state.openai_serving_video._engine_client
|
||||
input_video = engine.captured_prompt["multi_modal_data"]["video"]
|
||||
assert len(input_video) == 2
|
||||
assert input_video[0].size == (32, 24)
|
||||
|
||||
|
||||
def test_decode_video_bytes_can_keep_last_frames():
|
||||
from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes
|
||||
|
||||
frames = _decode_video_bytes(
|
||||
_make_test_video_bytes((32, 24), num_frames=6),
|
||||
source="input_reference",
|
||||
max_frames=2,
|
||||
keep="last",
|
||||
)
|
||||
|
||||
assert len(frames) == 2
|
||||
red_means = [np.asarray(frame)[:, :, 0].mean() for frame in frames]
|
||||
assert red_means[0] > 100
|
||||
assert red_means[1] > red_means[0]
|
||||
|
||||
|
||||
def test_cosmos3_reference_video_limit_uses_v2v_condition_frames():
|
||||
request = VideoGenerationRequest(
|
||||
prompt="Continue this motion.",
|
||||
num_frames=189,
|
||||
extra_params={"condition_frame_indexes_vision": [0, 2]},
|
||||
)
|
||||
|
||||
spec = api_server._reference_video_decode_spec(request, _cosmos3_stage_configs())
|
||||
assert spec.max_frames == 9
|
||||
assert spec.keep == "first"
|
||||
|
||||
|
||||
def test_cosmos3_reference_video_limit_preserves_action_frames():
|
||||
request = VideoGenerationRequest(
|
||||
prompt="Predict the action.",
|
||||
num_frames=17,
|
||||
extra_params={"action_mode": "inverse_dynamics", "action_chunk_size": 16},
|
||||
)
|
||||
|
||||
assert api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 17
|
||||
|
||||
|
||||
def test_cosmos3_reference_video_limit_caps_condition_frames_to_output_frames():
|
||||
request = VideoGenerationRequest(
|
||||
prompt="Continue this motion.",
|
||||
num_frames=5,
|
||||
extra_params={"condition_frame_indexes_vision": [0, 20]},
|
||||
)
|
||||
|
||||
assert api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 5
|
||||
|
||||
|
||||
def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
|
||||
fps_values = []
|
||||
|
||||
@@ -794,7 +919,20 @@ def test_rejects_input_reference_and_image_reference_together(test_client):
|
||||
files={"input_reference": ("input.png", _make_test_image_bytes(), "image/png")},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert "either input_reference or image_reference" in response.json()["detail"].lower()
|
||||
assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
|
||||
|
||||
|
||||
def test_rejects_image_reference_and_video_reference_together(test_client):
|
||||
response = test_client.post(
|
||||
"/v1/videos",
|
||||
data={
|
||||
"prompt": "bad refs",
|
||||
"image_reference": '{"image_url": "https://example.com/cat.png"}',
|
||||
"video_reference": '{"video_url": "https://example.com/cat.mp4"}',
|
||||
},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
|
||||
|
||||
|
||||
def test_invalid_seconds_returns_422(test_client):
|
||||
@@ -856,6 +994,18 @@ def test_unsupported_image_reference_file_id_returns_400(test_client):
|
||||
assert response.json()["detail"] == "Invalid image_reference: file_id is not supported yet."
|
||||
|
||||
|
||||
def test_unsupported_video_reference_file_id_returns_400(test_client):
|
||||
response = test_client.post(
|
||||
"/v1/videos",
|
||||
data={
|
||||
"prompt": "unsupported ref",
|
||||
"video_reference": '{"file_id": "file-123"}',
|
||||
},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert response.json()["detail"] == "Invalid video_reference: file_id is not supported yet."
|
||||
|
||||
|
||||
def test_invalid_uploaded_input_reference_returns_400(test_client):
|
||||
response = test_client.post(
|
||||
"/v1/videos",
|
||||
@@ -863,7 +1013,7 @@ def test_invalid_uploaded_input_reference_returns_400(test_client):
|
||||
files={"input_reference": ("input.png", b"not-an-image", "image/png")},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert response.json()["detail"] == "Invalid input_reference: provided content is not a valid image."
|
||||
assert response.json()["detail"] == "Invalid input_reference: provided content is not a valid image or video."
|
||||
|
||||
|
||||
def test_video_request_validation():
|
||||
@@ -880,6 +1030,8 @@ def test_video_request_validation():
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
VideoGenerationRequest(prompt="test", image_reference={"file_id": "file-1", "image_url": "https://example.com"})
|
||||
with pytest.raises(ValueError):
|
||||
VideoGenerationRequest(prompt="test", video_reference={"file_id": "file-1", "video_url": "https://example.com"})
|
||||
with pytest.raises(ValueError):
|
||||
VideoGenerationRequest(prompt="test", frame_interpolation_exp=0)
|
||||
with pytest.raises(ValueError):
|
||||
@@ -1265,6 +1417,23 @@ def test_sync_i2v_with_image_reference(test_client, mocker: MockerFixture):
|
||||
assert response.content == b"ref-video"
|
||||
|
||||
|
||||
def test_sync_v2v_returns_video_bytes(test_client, mocker: MockerFixture):
|
||||
video_bytes = _make_test_video_bytes((32, 24), num_frames=3)
|
||||
_mock_encode_video_bytes(mocker, b"v2v-video-data")
|
||||
response = test_client.post(
|
||||
"/v1/videos/sync",
|
||||
data={"prompt": "Continue this motion."},
|
||||
files={"input_reference": ("input.mp4", video_bytes, "video/mp4")},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.content == b"v2v-video-data"
|
||||
engine = test_client.app.state.openai_serving_video._engine_client
|
||||
input_video = engine.captured_prompt["multi_modal_data"]["video"]
|
||||
assert len(input_video) == 3
|
||||
assert input_video[0].size == (32, 24)
|
||||
|
||||
|
||||
def test_sync_missing_handler_returns_503():
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
@@ -1297,7 +1466,7 @@ def test_sync_rejects_both_references(test_client):
|
||||
files={"input_reference": ("input.png", _make_test_image_bytes(), "image/png")},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert "either input_reference or image_reference" in response.json()["detail"].lower()
|
||||
assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
|
||||
|
||||
|
||||
def test_sync_generation_error_returns_500(test_client, mocker: MockerFixture):
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
# 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.
|
||||
"""Cosmos3 text/image/video-to-video and text-to-image pipeline for vllm-omni.
|
||||
|
||||
Single pipeline class supports T2V, I2V, and T2I; the mode is selected at
|
||||
Single pipeline class supports T2V, I2V, V2V, and T2I; the mode is selected at
|
||||
runtime by:
|
||||
|
||||
* ``prompt["modalities"]`` contains ``"image"``: **T2I** (text-to-image).
|
||||
@@ -10,6 +10,8 @@ runtime by:
|
||||
(text-to-video).
|
||||
* ``multi_modal_data['image']`` present on the prompt: **I2V**
|
||||
(handled by :func:`get_cosmos3_pre_process_func`)
|
||||
* ``multi_modal_data['video']`` present on the prompt with no action mode:
|
||||
**V2V** (handled by :func:`get_cosmos3_pre_process_func`)
|
||||
|
||||
"""
|
||||
|
||||
@@ -40,7 +42,10 @@ from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
)
|
||||
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.interface import (
|
||||
ReferenceVideoDecodeSpec,
|
||||
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
|
||||
@@ -63,6 +68,17 @@ from .transformer_cosmos3 import Cosmos3VFMTransformer, resolve_sound_gen
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION = (0, 1)
|
||||
COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP = "first"
|
||||
# Mirrors the WAN VAE's temporal compression. Authoritative value is
|
||||
# ``self.vae.config.scale_factor_temporal`` at runtime; this constant exists so
|
||||
# off-line / API code that runs before the pipeline is constructed can compute
|
||||
# pixel-frame budgets without instantiating the VAE.
|
||||
COSMOS3_VAE_TEMPORAL_COMPRESSION = 4
|
||||
COSMOS3_DEFAULT_CONDITION_PIXEL_FRAMES = (
|
||||
max(COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION) * COSMOS3_VAE_TEMPORAL_COMPRESSION + 1
|
||||
)
|
||||
COSMOS3_V2V_DEFAULT_FLOW_SHIFT = 10.0
|
||||
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."
|
||||
@@ -91,16 +107,55 @@ COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL: tuple[float, float] = (400.0, 1000.0)
|
||||
COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096
|
||||
|
||||
|
||||
def _normalize_condition_frame_indexes_vision(value: Any) -> tuple[int, ...]:
|
||||
"""Normalize Cosmos3 vision-conditioning latent frame indexes."""
|
||||
if value is None:
|
||||
return COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION
|
||||
if isinstance(value, str):
|
||||
value = [item.strip() for item in value.split(",") if item.strip()]
|
||||
elif isinstance(value, int):
|
||||
value = [value]
|
||||
|
||||
if not isinstance(value, Iterable):
|
||||
raise TypeError(
|
||||
"Cosmos3 condition_frame_indexes_vision must be an int, comma-separated string, "
|
||||
f"or iterable of ints; got {type(value)!r}."
|
||||
)
|
||||
|
||||
indexes = tuple(sorted({int(index) for index in value}))
|
||||
if not indexes:
|
||||
raise ValueError("Cosmos3 condition_frame_indexes_vision must contain at least one index.")
|
||||
if any(index < 0 for index in indexes):
|
||||
raise ValueError(f"Cosmos3 condition_frame_indexes_vision must be non-negative, got {indexes}.")
|
||||
return indexes
|
||||
|
||||
|
||||
def _condition_pixel_frame_count(
|
||||
condition_frame_indexes_vision: Iterable[int],
|
||||
temporal_compression: int = COSMOS3_VAE_TEMPORAL_COMPRESSION,
|
||||
) -> int:
|
||||
return max(condition_frame_indexes_vision) * int(temporal_compression) + 1
|
||||
|
||||
|
||||
def _normalize_condition_video_keep(value: Any) -> str:
|
||||
keep = str(value or COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP).strip().lower()
|
||||
if keep not in {"first", "last"}:
|
||||
raise ValueError("Cosmos3 condition_video_keep must be either 'first' or 'last'.")
|
||||
return keep
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Post-process function (registered in registry.py)
|
||||
# ---------------------------------------------------------------------------
|
||||
def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
|
||||
"""Pre-process function for both T2V and I2V.
|
||||
"""Pre-process function for T2V, I2V, V2V, and action inputs.
|
||||
|
||||
For T2V (no image in ``multi_modal_data``), the request is returned
|
||||
unchanged after the optional guardrails check. For I2V (image present),
|
||||
the conditioning image is loaded, aspect-resized + center-cropped, and
|
||||
stored back on the prompt as ``additional_information.preprocessed_image``.
|
||||
For V2V (video present without action mode), source frames are cropped to
|
||||
the target size and stored as ``additional_information.preprocessed_video``.
|
||||
"""
|
||||
from .guardrails import check_text_safety, ensure_initialized, is_guardrails_enabled
|
||||
|
||||
@@ -137,7 +192,27 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
|
||||
return PIL.Image.open(value).convert("RGB")
|
||||
if isinstance(value, PIL.Image.Image):
|
||||
return value.convert("RGB")
|
||||
raise TypeError(f"Cosmos3 action preprocessing expected PIL image or image path, got {type(value)!r}.")
|
||||
if isinstance(value, np.ndarray):
|
||||
array = value
|
||||
if array.ndim == 3 and array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4):
|
||||
array = np.transpose(array, (1, 2, 0))
|
||||
if np.issubdtype(array.dtype, np.floating):
|
||||
if array.min() < 0.0 or array.max() > 1.0:
|
||||
array = np.clip(array, -1.0, 1.0) * 0.5 + 0.5
|
||||
array = (np.clip(array, 0.0, 1.0) * 255.0).round().astype(np.uint8)
|
||||
return PIL.Image.fromarray(array).convert("RGB")
|
||||
if isinstance(value, torch.Tensor):
|
||||
tensor = value.detach().cpu()
|
||||
if tensor.ndim == 3 and tensor.shape[0] in (3, 4):
|
||||
tensor = tensor.permute(1, 2, 0)
|
||||
if tensor.is_floating_point():
|
||||
if tensor.min().item() < 0.0 or tensor.max().item() > 1.0:
|
||||
tensor = tensor.clamp(-1.0, 1.0) * 0.5 + 0.5
|
||||
tensor = (tensor.clamp(0.0, 1.0) * 255.0).round().to(torch.uint8)
|
||||
return PIL.Image.fromarray(tensor.numpy()).convert("RGB")
|
||||
raise TypeError(
|
||||
f"Cosmos3 preprocessing expected PIL image, numpy array, torch tensor, or path, got {type(value)!r}."
|
||||
)
|
||||
|
||||
def _resize_and_pad_action_image(image: PIL.Image.Image, target_h: int, target_w: int) -> PIL.Image.Image:
|
||||
scale = min(target_w / image.width, target_h / image.height, 1.0)
|
||||
@@ -170,6 +245,61 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
|
||||
processed = [_preprocess_action_image(_pil_to_rgb(frame), target_h, target_w).squeeze(0) for frame in frames]
|
||||
return torch.stack(processed, dim=1).unsqueeze(0).contiguous()
|
||||
|
||||
def _preprocess_condition_image(image: PIL.Image.Image, target_h: int, target_w: int) -> torch.Tensor:
|
||||
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))
|
||||
return video_processor.preprocess(image, height=target_h, width=target_w)
|
||||
|
||||
def _video_payload_to_frames(video: Any) -> list[Any]:
|
||||
if isinstance(video, list):
|
||||
return video
|
||||
if isinstance(video, torch.Tensor):
|
||||
tensor = video.detach().cpu()
|
||||
if tensor.ndim == 5:
|
||||
if tensor.shape[0] != 1:
|
||||
raise TypeError("Cosmos3 video preprocessing supports only batch size 1.")
|
||||
tensor = tensor[0]
|
||||
if tensor.ndim == 4 and tensor.shape[0] in (3, 4) and tensor.shape[-1] not in (3, 4):
|
||||
return [tensor[:, i] for i in range(tensor.shape[1])]
|
||||
if tensor.ndim == 4 and tensor.shape[-1] in (3, 4):
|
||||
return [tensor[i] for i in range(tensor.shape[0])]
|
||||
if isinstance(video, np.ndarray):
|
||||
array = video
|
||||
if array.ndim == 5:
|
||||
if array.shape[0] != 1:
|
||||
raise TypeError("Cosmos3 video preprocessing supports only batch size 1.")
|
||||
array = array[0]
|
||||
if array.ndim == 4 and array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4):
|
||||
return [array[:, i] for i in range(array.shape[1])]
|
||||
if array.ndim == 4 and array.shape[-1] in (3, 4):
|
||||
return [array[i] for i in range(array.shape[0])]
|
||||
raise TypeError("Cosmos3 video input must be a non-empty list of frames or a single video tensor/array.")
|
||||
|
||||
def _select_video_frames(frames: list[Any], max_frames: int, keep: str) -> list[Any]:
|
||||
if not frames:
|
||||
raise ValueError("Cosmos3 video input must contain at least one frame.")
|
||||
if keep == "last":
|
||||
return frames[-max_frames:]
|
||||
return frames[:max_frames]
|
||||
|
||||
def _preprocess_condition_video(
|
||||
frames: list[Any],
|
||||
target_h: int,
|
||||
target_w: int,
|
||||
max_frames: int,
|
||||
keep: str,
|
||||
) -> torch.Tensor:
|
||||
selected = _select_video_frames(frames, max_frames, keep)
|
||||
processed = [
|
||||
_preprocess_condition_image(_pil_to_rgb(frame), target_h, target_w).squeeze(0) for frame in selected
|
||||
]
|
||||
return torch.stack(processed, dim=1).unsqueeze(0).contiguous()
|
||||
|
||||
def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
|
||||
action_mode = _request_action_mode(request)
|
||||
if is_guardrails_enabled(od_config, request.sampling_params):
|
||||
@@ -183,16 +313,23 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
|
||||
multi_modal_data = prompt.get("multi_modal_data", {}) or {}
|
||||
raw_image = multi_modal_data.get("image")
|
||||
raw_video = multi_modal_data.get("video")
|
||||
if raw_image is None and not (action_mode is not None and raw_video is not None):
|
||||
if raw_image is None and raw_video is None:
|
||||
continue
|
||||
if raw_image is not None and raw_video is not None and action_mode is None:
|
||||
raise ValueError("Cosmos3 non-action generation accepts either image or video input, not both.")
|
||||
|
||||
if "additional_information" not in prompt:
|
||||
prompt["additional_information"] = {}
|
||||
|
||||
raw_video_frames: list[Any] | None = None
|
||||
if raw_video is not None:
|
||||
raw_video_frames = _video_payload_to_frames(raw_video)
|
||||
if not raw_video_frames:
|
||||
raise TypeError("Cosmos3 video input must be a non-empty list of PIL images or image paths.")
|
||||
|
||||
if raw_image is None:
|
||||
if not isinstance(raw_video, list) or not raw_video:
|
||||
raise TypeError("Cosmos3 action video input must be a non-empty list of PIL images or image paths.")
|
||||
image = _pil_to_rgb(raw_video[0])
|
||||
assert raw_video_frames is not None # raw_image and raw_video can't both be None here
|
||||
image = _pil_to_rgb(raw_video_frames[0])
|
||||
else:
|
||||
image = _pil_to_rgb(raw_image)
|
||||
|
||||
@@ -219,23 +356,38 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
|
||||
int(target_h),
|
||||
int(target_w),
|
||||
)
|
||||
else:
|
||||
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
|
||||
elif raw_video is None:
|
||||
prompt["additional_information"]["preprocessed_image"] = _preprocess_condition_image(
|
||||
image,
|
||||
int(target_h),
|
||||
int(target_w),
|
||||
)
|
||||
if action_mode is not None and raw_video is not None:
|
||||
if not isinstance(raw_video, list):
|
||||
raise TypeError("Cosmos3 action video input must be a list of PIL images or image paths.")
|
||||
else:
|
||||
assert raw_video_frames is not None
|
||||
extra = _extra_args(request)
|
||||
condition_frame_indexes_vision = _normalize_condition_frame_indexes_vision(
|
||||
extra.get(
|
||||
"condition_frame_indexes_vision",
|
||||
prompt.get("condition_frame_indexes_vision"),
|
||||
)
|
||||
)
|
||||
keep = _normalize_condition_video_keep(
|
||||
extra.get("condition_video_keep", prompt.get("condition_video_keep"))
|
||||
)
|
||||
max_frames = _condition_pixel_frame_count(condition_frame_indexes_vision)
|
||||
prompt["additional_information"]["preprocessed_video"] = _preprocess_condition_video(
|
||||
raw_video_frames,
|
||||
int(target_h),
|
||||
int(target_w),
|
||||
max_frames,
|
||||
keep,
|
||||
)
|
||||
prompt["additional_information"]["condition_frame_indexes_vision"] = list(
|
||||
condition_frame_indexes_vision
|
||||
)
|
||||
if action_mode is not None and raw_video_frames is not None:
|
||||
prompt["additional_information"]["preprocessed_video"] = _preprocess_action_video(
|
||||
raw_video,
|
||||
raw_video_frames,
|
||||
int(target_h),
|
||||
int(target_w),
|
||||
)
|
||||
@@ -345,13 +497,13 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig):
|
||||
class Cosmos3OmniDiffusersPipeline(
|
||||
nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin
|
||||
):
|
||||
"""Cosmos3 text/image-to-video / text-to-image pipeline.
|
||||
"""Cosmos3 text/image/video-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
|
||||
Supports T2V, I2V, V2V, and T2I from the same class. Mode is selected at
|
||||
runtime:
|
||||
|
||||
* **T2I** when ``prompt["modalities"]`` contains ``"image"``. Latent
|
||||
@@ -365,12 +517,45 @@ class Cosmos3OmniDiffusersPipeline(
|
||||
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.
|
||||
* **V2V** when the request supplies a preprocessed video via
|
||||
``multi_modal_data['video']`` without an action mode. Explicit latent
|
||||
frame indexes are kept clean with ``noisy_frame_mask`` and re-injected
|
||||
after each scheduler step.
|
||||
* **T2V** otherwise (default video generation).
|
||||
"""
|
||||
|
||||
support_image_input: ClassVar[bool] = True
|
||||
color_format: ClassVar[str] = "RGB"
|
||||
|
||||
@classmethod
|
||||
def reference_video_decode_spec(
|
||||
cls,
|
||||
*,
|
||||
num_frames: int | None = None,
|
||||
extra_args: dict[str, Any] | None = None,
|
||||
) -> ReferenceVideoDecodeSpec:
|
||||
extra_args = extra_args if isinstance(extra_args, dict) else {}
|
||||
action_mode = normalize_action_mode(extra_args.get("action_mode"))
|
||||
if action_mode is not None:
|
||||
if num_frames is not None:
|
||||
return ReferenceVideoDecodeSpec(max_frames=int(num_frames), keep="first")
|
||||
action_chunk_size = extra_args.get("action_chunk_size")
|
||||
if action_chunk_size is not None:
|
||||
try:
|
||||
max_frames = int(action_chunk_size) + 1
|
||||
except (TypeError, ValueError):
|
||||
max_frames = None
|
||||
if max_frames is not None and max_frames > 0:
|
||||
return ReferenceVideoDecodeSpec(max_frames=max_frames, keep="first")
|
||||
return ReferenceVideoDecodeSpec(max_frames=None, keep="first")
|
||||
|
||||
condition_indexes = _normalize_condition_frame_indexes_vision(extra_args.get("condition_frame_indexes_vision"))
|
||||
max_frames = _condition_pixel_frame_count(condition_indexes)
|
||||
if num_frames is not None:
|
||||
max_frames = min(max_frames, int(num_frames))
|
||||
keep = _normalize_condition_video_keep(extra_args.get("condition_video_keep"))
|
||||
return ReferenceVideoDecodeSpec(max_frames=max_frames, keep=keep)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
@@ -1188,6 +1373,70 @@ class Cosmos3OmniDiffusersPipeline(
|
||||
velocity_mask = 1.0 - condition_mask
|
||||
return latents, velocity_mask, image_latent
|
||||
|
||||
def _prepare_latents_v2v(
|
||||
self,
|
||||
video_tensor: torch.Tensor,
|
||||
height: int,
|
||||
width: int,
|
||||
num_frames: int,
|
||||
generator: torch.Generator,
|
||||
condition_frame_indexes_vision: Iterable[int] | int | str | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Prepare V2V latents with explicit clean conditioned latent frames."""
|
||||
del height, width
|
||||
if video_tensor.ndim == 4:
|
||||
video_tensor = video_tensor.unsqueeze(0)
|
||||
if video_tensor.ndim != 5 or video_tensor.shape[0] != 1 or video_tensor.shape[1] != 3:
|
||||
raise ValueError(f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}.")
|
||||
if video_tensor.shape[2] < 1:
|
||||
raise ValueError("Cosmos3 V2V video tensor must contain at least one frame.")
|
||||
|
||||
C = self.transformer.latent_channel_size
|
||||
T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1
|
||||
H_lat = video_tensor.shape[-2] // self.vae_scale_factor_spatial
|
||||
W_lat = video_tensor.shape[-1] // self.vae_scale_factor_spatial
|
||||
indexes = _normalize_condition_frame_indexes_vision(condition_frame_indexes_vision)
|
||||
out_of_range = [index for index in indexes if index >= T_lat]
|
||||
if out_of_range:
|
||||
raise ValueError(
|
||||
"Cosmos3 condition_frame_indexes_vision contains indexes outside the latent video: "
|
||||
f"indexes={indexes}, latent_frames={T_lat}."
|
||||
)
|
||||
|
||||
noise = randn_tensor(
|
||||
(1, C, T_lat, H_lat, W_lat),
|
||||
generator=generator,
|
||||
device=self.device,
|
||||
dtype=self.dtype,
|
||||
)
|
||||
condition_pixel_frames = _condition_pixel_frame_count(indexes, self.vae_scale_factor_temporal)
|
||||
condition_video = video_tensor[:, :, :condition_pixel_frames]
|
||||
if condition_video.shape[2] < condition_pixel_frames:
|
||||
pad = condition_video[:, :, -1:].repeat(1, 1, condition_pixel_frames - condition_video.shape[2], 1, 1)
|
||||
condition_video = torch.cat([condition_video, pad], dim=2)
|
||||
|
||||
cond_prefix_latent = self._encode_video_tensor(condition_video)
|
||||
expected_prefix = (1, C, max(indexes) + 1, H_lat, W_lat)
|
||||
if (
|
||||
cond_prefix_latent.shape[0] != expected_prefix[0]
|
||||
or cond_prefix_latent.shape[1] != expected_prefix[1]
|
||||
or cond_prefix_latent.shape[2] < expected_prefix[2]
|
||||
or cond_prefix_latent.shape[3:] != expected_prefix[3:]
|
||||
):
|
||||
raise ValueError(
|
||||
"Cosmos3 V2V condition latent shape mismatch: "
|
||||
f"encoded={tuple(cond_prefix_latent.shape)}, expected at least {expected_prefix}."
|
||||
)
|
||||
|
||||
condition_mask = torch.zeros(1, 1, T_lat, 1, 1, device=self.device, dtype=self.dtype)
|
||||
condition_latents = torch.zeros_like(noise)
|
||||
for index in indexes:
|
||||
condition_mask[:, :, index, :, :] = 1.0
|
||||
condition_latents[:, :, index : index + 1] = cond_prefix_latent[:, :, index : index + 1]
|
||||
latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise
|
||||
velocity_mask = 1.0 - condition_mask
|
||||
return latents, velocity_mask, condition_latents
|
||||
|
||||
def _prepare_latents_action_video(
|
||||
self,
|
||||
video_tensor: torch.Tensor,
|
||||
@@ -1567,19 +1816,20 @@ class Cosmos3OmniDiffusersPipeline(
|
||||
prompt = prompt_data
|
||||
negative_prompt = None
|
||||
image_tensor = None
|
||||
action_video_tensor = None
|
||||
video_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")
|
||||
action_video_tensor = additional_info.get("preprocessed_video")
|
||||
video_tensor = additional_info.get("preprocessed_video")
|
||||
|
||||
sp = req.sampling_params
|
||||
is_t2i = self._is_t2i_request(req)
|
||||
sound_enabled = self._is_sound_request(prompt_data, sp)
|
||||
action_mode = self._get_action_mode(prompt_data, sp)
|
||||
action_enabled = action_mode is not None
|
||||
action_video_tensor = video_tensor if action_enabled else None
|
||||
if action_enabled and is_t2i:
|
||||
raise ValueError("Cosmos3 action generation is supported only for video outputs.")
|
||||
if action_enabled and sound_enabled:
|
||||
@@ -1603,6 +1853,11 @@ class Cosmos3OmniDiffusersPipeline(
|
||||
)
|
||||
if negative_prompt is None:
|
||||
negative_prompt = ""
|
||||
if image_tensor is not None and video_tensor is not None and not action_enabled:
|
||||
raise ValueError("Cosmos3 non-action generation accepts either image or video input, not both.")
|
||||
if video_tensor is not None and is_t2i:
|
||||
raise ValueError("Cosmos3 video-to-video generation is supported only for video outputs.")
|
||||
is_v2v = video_tensor is not None and not is_t2i and not action_enabled
|
||||
|
||||
# T2I and T2V share the same model + forward path; only defaults
|
||||
# differ:
|
||||
@@ -1626,7 +1881,7 @@ class Cosmos3OmniDiffusersPipeline(
|
||||
# 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_flow_shift = COSMOS3_V2V_DEFAULT_FLOW_SHIFT if is_v2v else self._engine_init_flow_shift
|
||||
default_guidance_interval = None
|
||||
batch_size = 1 # Existing video pipeline assumes B=1.
|
||||
|
||||
@@ -1775,6 +2030,23 @@ class Cosmos3OmniDiffusersPipeline(
|
||||
generator,
|
||||
)
|
||||
image_latent = condition_latents[:, :, 0:1]
|
||||
elif is_v2v:
|
||||
condition_frame_indexes_vision = _normalize_condition_frame_indexes_vision(
|
||||
self._get_sp_param(
|
||||
sp,
|
||||
"condition_frame_indexes_vision",
|
||||
self._get_prompt_param(prompt_data, "condition_frame_indexes_vision", None),
|
||||
)
|
||||
)
|
||||
latents, velocity_mask, condition_latents = self._prepare_latents_v2v(
|
||||
video_tensor,
|
||||
height,
|
||||
width,
|
||||
num_frames,
|
||||
generator,
|
||||
condition_frame_indexes_vision,
|
||||
)
|
||||
image_latent = None
|
||||
elif image_tensor is not None and not is_t2i:
|
||||
latents, velocity_mask, image_latent = self._prepare_latents_i2v(
|
||||
image_tensor,
|
||||
|
||||
@@ -2,10 +2,12 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
ClassVar,
|
||||
Literal,
|
||||
Protocol,
|
||||
runtime_checkable,
|
||||
)
|
||||
@@ -23,6 +25,12 @@ class SupportImageInput(Protocol):
|
||||
color_format: ClassVar[str] = "RGB" # Default color format
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ReferenceVideoDecodeSpec:
|
||||
max_frames: int | None = None
|
||||
keep: Literal["first", "last"] = "first"
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SupportAudioInput(Protocol):
|
||||
support_audio_input: ClassVar[bool] = True
|
||||
|
||||
@@ -91,6 +91,7 @@ from vllm.utils import random_uuid
|
||||
from vllm.utils.system_utils import decorate_logs
|
||||
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
|
||||
|
||||
from vllm_omni.diffusion.models.interface import ReferenceVideoDecodeSpec
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
|
||||
from vllm_omni.entrypoints.openai.image_api_utils import (
|
||||
@@ -125,7 +126,7 @@ from vllm_omni.entrypoints.openai.serving_audio_generate import OmniOpenAIServin
|
||||
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
|
||||
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
|
||||
from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler
|
||||
from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage
|
||||
from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage, ReferenceVideo
|
||||
from vllm_omni.entrypoints.openai.serving_video_stream import OmniStreamingVideoHandler
|
||||
from vllm_omni.entrypoints.openai.stage_params import (
|
||||
build_stage_sampling_params_list,
|
||||
@@ -2531,6 +2532,84 @@ def _parse_form_json(value: str | None, expected_type: type | None = None) -> An
|
||||
return parsed
|
||||
|
||||
|
||||
def _config_get(config: Any, key: str, default: Any = None) -> Any:
|
||||
if isinstance(config, dict):
|
||||
return config.get(key, default)
|
||||
if hasattr(config, "get"):
|
||||
try:
|
||||
return config.get(key, default)
|
||||
except Exception:
|
||||
pass
|
||||
return getattr(config, key, default)
|
||||
|
||||
|
||||
def _stage_engine_args(stage_cfg: Any) -> Any:
|
||||
return _config_get(stage_cfg, "engine_args", {}) or {}
|
||||
|
||||
|
||||
def _diffusion_model_classes(stage_configs: list[Any] | None) -> list[type]:
|
||||
if not stage_configs:
|
||||
return []
|
||||
|
||||
from vllm_omni.diffusion.registry import DiffusionModelRegistry
|
||||
|
||||
model_classes: list[type] = []
|
||||
for stage_cfg in stage_configs:
|
||||
if get_stage_type(stage_cfg) != "diffusion":
|
||||
continue
|
||||
model_class_name = _config_get(_stage_engine_args(stage_cfg), "model_class_name")
|
||||
if not model_class_name:
|
||||
continue
|
||||
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
|
||||
if model_cls is not None:
|
||||
model_classes.append(model_cls)
|
||||
return model_classes
|
||||
|
||||
|
||||
def _normalize_reference_video_decode_spec(spec: ReferenceVideoDecodeSpec) -> ReferenceVideoDecodeSpec:
|
||||
max_frames = spec.max_frames
|
||||
if max_frames is not None:
|
||||
try:
|
||||
max_frames = int(max_frames)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Invalid reference video decode spec: max_frames must be an integer.",
|
||||
) from exc
|
||||
if max_frames <= 0:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Invalid reference video decode spec: max_frames must be positive.",
|
||||
)
|
||||
|
||||
keep = str(spec.keep or "first").strip().lower()
|
||||
if keep not in {"first", "last"}:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Invalid reference video decode spec: keep must be either 'first' or 'last'.",
|
||||
)
|
||||
return ReferenceVideoDecodeSpec(max_frames=max_frames, keep=cast(Literal["first", "last"], keep))
|
||||
|
||||
|
||||
def _reference_video_decode_spec(
|
||||
req: VideoGenerationRequest,
|
||||
stage_configs: list[Any] | None,
|
||||
) -> ReferenceVideoDecodeSpec:
|
||||
video_params = req.resolve_video_params()
|
||||
extra_params = req.extra_params if isinstance(req.extra_params, dict) else {}
|
||||
for model_cls in _diffusion_model_classes(stage_configs):
|
||||
resolver = getattr(model_cls, "reference_video_decode_spec", None)
|
||||
if resolver is None:
|
||||
continue
|
||||
try:
|
||||
spec = resolver(num_frames=video_params.num_frames, extra_args=extra_params)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc
|
||||
if spec is not None:
|
||||
return _normalize_reference_video_decode_spec(spec)
|
||||
return ReferenceVideoDecodeSpec(max_frames=video_params.num_frames, keep="first")
|
||||
|
||||
|
||||
def video_response_from_request(model_name: str, req: VideoGenerationRequest) -> VideoResponse:
|
||||
resp = VideoResponse(
|
||||
model=model_name,
|
||||
@@ -2597,6 +2676,7 @@ async def _run_video_generation_job(
|
||||
request: VideoGenerationRequest,
|
||||
video_id: str,
|
||||
reference_image: ReferenceImage | None = None,
|
||||
reference_video: ReferenceVideo | None = None,
|
||||
app_state: Any | None = None,
|
||||
) -> None:
|
||||
job = await VIDEO_STORE.get(video_id)
|
||||
@@ -2609,7 +2689,10 @@ async def _run_video_generation_job(
|
||||
output_path = None
|
||||
try:
|
||||
video_bytes, stage_durations, peak_memory_mb, action = await handler.generate_video_bytes(
|
||||
request, video_id, reference_image=reference_image
|
||||
request,
|
||||
video_id,
|
||||
reference_image=reference_image,
|
||||
reference_video=reference_video,
|
||||
)
|
||||
|
||||
file_name = f"{video_id}.{job.file_extension}"
|
||||
@@ -2676,6 +2759,7 @@ async def _parse_video_form(
|
||||
prompt: str = Form(...),
|
||||
input_reference: UploadFile | None = File(default=None),
|
||||
image_reference: str | None = Form(default=None),
|
||||
video_reference: str | None = Form(default=None),
|
||||
model: str | None = Form(default=None),
|
||||
seconds: SecondStr | None = Form(default=None),
|
||||
size: SizeStr | None = Form(default=None),
|
||||
@@ -2700,7 +2784,7 @@ async def _parse_video_form(
|
||||
frame_interpolation_model_path: str | None = Form(default=None),
|
||||
lora: str | None = Form(default=None),
|
||||
extra_params: str | None = Form(default=None),
|
||||
) -> tuple[VideoGenerationRequest, "OmniOpenAIServingVideo", str, ReferenceImage | None]:
|
||||
) -> tuple[VideoGenerationRequest, "OmniOpenAIServingVideo", str, ReferenceImage | None, ReferenceVideo | None]:
|
||||
"""FastAPI dependency that parses video form data, validates inputs,
|
||||
resolves the handler, and decodes any reference image.
|
||||
|
||||
@@ -2708,11 +2792,15 @@ async def _parse_video_form(
|
||||
"""
|
||||
input_reference_bytes = await input_reference.read() if input_reference is not None else None
|
||||
parsed_image_reference = _parse_form_json(image_reference)
|
||||
parsed_video_reference = _parse_form_json(video_reference)
|
||||
|
||||
if parsed_image_reference is not None and input_reference_bytes is not None:
|
||||
provided_references = sum(
|
||||
item is not None for item in (parsed_image_reference, parsed_video_reference, input_reference_bytes)
|
||||
)
|
||||
if provided_references > 1:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Provide either input_reference or image_reference, not both.",
|
||||
detail="Provide only one of input_reference, image_reference, or video_reference.",
|
||||
)
|
||||
|
||||
request_data: dict[str, Any] = {
|
||||
@@ -2721,6 +2809,7 @@ async def _parse_video_form(
|
||||
"seconds": seconds,
|
||||
"size": size,
|
||||
"image_reference": parsed_image_reference,
|
||||
"video_reference": parsed_video_reference,
|
||||
"user": user,
|
||||
"width": width,
|
||||
"height": height,
|
||||
@@ -2774,13 +2863,28 @@ async def _parse_video_form(
|
||||
detail=f"Video generation setup failed: {str(e)}",
|
||||
)
|
||||
|
||||
decode_spec = ReferenceVideoDecodeSpec()
|
||||
if parsed_video_reference is not None or input_reference_bytes is not None:
|
||||
stage_configs = (
|
||||
handler.stage_configs
|
||||
or app_stage_configs
|
||||
or getattr(getattr(handler, "_engine_client", None), "stage_configs", None)
|
||||
)
|
||||
decode_spec = _reference_video_decode_spec(request, stage_configs)
|
||||
try:
|
||||
image_data = await decode_input_reference(request.image_reference, input_reference_bytes)
|
||||
media_data = await decode_input_reference(
|
||||
request.image_reference,
|
||||
request.video_reference,
|
||||
input_reference_bytes,
|
||||
max_video_frames=decode_spec.max_frames,
|
||||
video_keep=decode_spec.keep,
|
||||
)
|
||||
except InvalidInputReferenceError as exc:
|
||||
raise HTTPException(400, detail=str(exc) or "Invalid input reference.") from exc
|
||||
|
||||
reference_image = ReferenceImage(data=image_data) if image_data is not None else None
|
||||
return request, handler, effective_model_name, reference_image
|
||||
reference_image = ReferenceImage(data=media_data) if isinstance(media_data, Image.Image) else None
|
||||
reference_video = ReferenceVideo(data=media_data) if isinstance(media_data, list) else None
|
||||
return request, handler, effective_model_name, reference_image, reference_video
|
||||
|
||||
|
||||
@router.post(
|
||||
@@ -2794,18 +2898,31 @@ async def _parse_video_form(
|
||||
)
|
||||
async def create_video(
|
||||
raw_request: Request,
|
||||
ctx: tuple[VideoGenerationRequest, OmniOpenAIServingVideo, str, ReferenceImage | None] = Depends(_parse_video_form),
|
||||
ctx: tuple[
|
||||
VideoGenerationRequest,
|
||||
OmniOpenAIServingVideo,
|
||||
str,
|
||||
ReferenceImage | None,
|
||||
ReferenceVideo | None,
|
||||
] = Depends(_parse_video_form),
|
||||
) -> VideoResponse:
|
||||
"""Create an asynchronous video generation job.
|
||||
|
||||
Accepts multipart form-data (see ``_parse_video_form`` for parameters),
|
||||
persists a queued job record, and starts generation in the background.
|
||||
"""
|
||||
request, handler, effective_model_name, reference_image = ctx
|
||||
request, handler, effective_model_name, reference_image, reference_video = ctx
|
||||
ref = video_response_from_request(effective_model_name, request)
|
||||
await VIDEO_STORE.upsert(ref.id, ref)
|
||||
task = asyncio.create_task(
|
||||
_run_video_generation_job(handler, request, ref.id, reference_image, app_state=raw_request.app.state)
|
||||
_run_video_generation_job(
|
||||
handler,
|
||||
request,
|
||||
ref.id,
|
||||
reference_image,
|
||||
reference_video,
|
||||
app_state=raw_request.app.state,
|
||||
)
|
||||
)
|
||||
await VIDEO_TASKS.upsert(ref.id, task)
|
||||
return ref
|
||||
@@ -2822,7 +2939,13 @@ async def create_video(
|
||||
)
|
||||
async def create_video_sync(
|
||||
raw_request: Request,
|
||||
ctx: tuple[VideoGenerationRequest, OmniOpenAIServingVideo, str, ReferenceImage | None] = Depends(_parse_video_form),
|
||||
ctx: tuple[
|
||||
VideoGenerationRequest,
|
||||
OmniOpenAIServingVideo,
|
||||
str,
|
||||
ReferenceImage | None,
|
||||
ReferenceVideo | None,
|
||||
] = Depends(_parse_video_form),
|
||||
) -> Response:
|
||||
"""Synchronous video generation endpoint.
|
||||
|
||||
@@ -2833,13 +2956,18 @@ async def create_video_sync(
|
||||
Metadata is returned via response headers ``X-Request-Id``,
|
||||
``X-Model``, and ``X-Inference-Time-S``.
|
||||
"""
|
||||
request, handler, effective_model_name, reference_image = ctx
|
||||
request, handler, effective_model_name, reference_image, reference_video = ctx
|
||||
request_id = f"video_sync-{random_uuid()}"
|
||||
raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id)
|
||||
started_at = time.perf_counter()
|
||||
try:
|
||||
video_bytes, stage_durations, peak_memory_mb, _action = await asyncio.wait_for(
|
||||
handler.generate_video_bytes(request, request_id, reference_image=reference_image),
|
||||
handler.generate_video_bytes(
|
||||
request,
|
||||
request_id,
|
||||
reference_image=reference_image,
|
||||
reference_video=reference_video,
|
||||
),
|
||||
timeout=VIDEO_SYNC_TIMEOUT_S,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
|
||||
@@ -73,6 +73,19 @@ class UrlImageReference(BaseModel):
|
||||
ImageReference = UrlImageReference | FileImageReference
|
||||
|
||||
|
||||
class FileVideoReference(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
file_id: str
|
||||
|
||||
|
||||
class UrlVideoReference(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
video_url: str
|
||||
|
||||
|
||||
VideoReference = UrlVideoReference | FileVideoReference
|
||||
|
||||
|
||||
class VideoGenerationRequest(BaseModel):
|
||||
"""
|
||||
OpenAI-style video generation request.
|
||||
@@ -99,6 +112,10 @@ class VideoGenerationRequest(BaseModel):
|
||||
default=None,
|
||||
description="Optional JSON-safe image reference that guides generation. Provide either image_url or file_id.",
|
||||
)
|
||||
video_reference: VideoReference | None = Field(
|
||||
default=None,
|
||||
description="Optional JSON-safe video reference that guides generation. Provide either video_url or file_id.",
|
||||
)
|
||||
|
||||
# Video params block for extensibility
|
||||
video_params: VideoParams | None = Field(default=None, description="Optional video-specific parameters")
|
||||
|
||||
@@ -39,6 +39,13 @@ class ReferenceImage:
|
||||
data: Image.Image
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReferenceVideo:
|
||||
"""Reference video frames for video-conditioned generation."""
|
||||
|
||||
data: list[Image.Image]
|
||||
|
||||
|
||||
@dataclass
|
||||
class VideoGenerationArtifacts:
|
||||
"""Normalized outputs and profiler metadata extracted from one request."""
|
||||
@@ -96,6 +103,7 @@ class OmniOpenAIServingVideo:
|
||||
reference_id: str,
|
||||
*,
|
||||
reference_image: ReferenceImage | None = None,
|
||||
reference_video: ReferenceVideo | None = None,
|
||||
) -> VideoGenerationArtifacts:
|
||||
"""Run the generation pipeline and extract video/audio/profiler outputs."""
|
||||
prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt, modalities=["video"])
|
||||
@@ -105,6 +113,12 @@ class OmniOpenAIServingVideo:
|
||||
gen_params = self._resolve_default_sampling_params()
|
||||
|
||||
input_image = None if reference_image is None else reference_image.data
|
||||
input_video = None if reference_video is None else reference_video.data
|
||||
if input_image is not None and input_video is not None:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Provide either an image reference or a video reference, not both.",
|
||||
)
|
||||
vp = request.resolve_video_params()
|
||||
if input_image is not None and vp.width is not None and vp.height is not None:
|
||||
target_size = (vp.width, vp.height)
|
||||
@@ -112,6 +126,8 @@ class OmniOpenAIServingVideo:
|
||||
input_image = input_image.resize(target_size, Image.Resampling.LANCZOS)
|
||||
if input_image is not None:
|
||||
prompt["multi_modal_data"] = {"image": input_image}
|
||||
elif input_video is not None:
|
||||
prompt["multi_modal_data"] = {"video": input_video}
|
||||
if vp.width is not None and vp.height is not None:
|
||||
gen_params.width = vp.width
|
||||
gen_params.height = vp.height
|
||||
@@ -211,8 +227,14 @@ class OmniOpenAIServingVideo:
|
||||
reference_id: str,
|
||||
*,
|
||||
reference_image: ReferenceImage | None = None,
|
||||
reference_video: ReferenceVideo | None = None,
|
||||
) -> VideoGenerationResponse:
|
||||
artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image)
|
||||
artifacts = await self._run_and_extract(
|
||||
request,
|
||||
reference_id,
|
||||
reference_image=reference_image,
|
||||
reference_video=reference_video,
|
||||
)
|
||||
|
||||
video_codec_options = {"preset": "ultrafast", "threads": "0"}
|
||||
if request.extra_params is not None and isinstance(request.extra_params, dict):
|
||||
@@ -252,9 +274,15 @@ class OmniOpenAIServingVideo:
|
||||
reference_id: str,
|
||||
*,
|
||||
reference_image: ReferenceImage | None = None,
|
||||
reference_video: ReferenceVideo | None = None,
|
||||
) -> tuple[bytes, dict[str, float], float, VideoAction | None]:
|
||||
"""Generate a video and return raw MP4 bytes, bypassing base64 encoding."""
|
||||
artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image)
|
||||
artifacts = await self._run_and_extract(
|
||||
request,
|
||||
reference_id,
|
||||
reference_image=reference_image,
|
||||
reference_video=reference_video,
|
||||
)
|
||||
if len(artifacts.videos) > 1:
|
||||
logger.warning(
|
||||
"Video request %s generated %d outputs; returning only the first.",
|
||||
|
||||
@@ -8,8 +8,9 @@ from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
from collections import deque
|
||||
from io import BytesIO
|
||||
from typing import Any
|
||||
from typing import Any, Literal
|
||||
|
||||
import httpx
|
||||
import numpy as np
|
||||
@@ -19,8 +20,11 @@ from PIL import Image, UnidentifiedImageError
|
||||
from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
|
||||
from vllm_omni.entrypoints.openai.protocol.videos import (
|
||||
FileImageReference,
|
||||
FileVideoReference,
|
||||
ImageReference,
|
||||
UrlImageReference,
|
||||
UrlVideoReference,
|
||||
VideoReference,
|
||||
)
|
||||
|
||||
|
||||
@@ -31,6 +35,70 @@ def _decode_image_bytes(image_bytes: bytes, *, source: str) -> Image.Image:
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid image.") from exc
|
||||
|
||||
|
||||
def _decode_video_bytes(
|
||||
video_bytes: bytes,
|
||||
*,
|
||||
source: str,
|
||||
max_frames: int | None = None,
|
||||
keep: Literal["first", "last"] = "first",
|
||||
) -> list[Image.Image]:
|
||||
try:
|
||||
import av
|
||||
except ImportError as exc: # pragma: no cover - av is a serving dependency via media_utils
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: video decoding requires PyAV.") from exc
|
||||
|
||||
if keep not in {"first", "last"}:
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: video frame selection must be 'first' or 'last'.")
|
||||
if max_frames is not None and max_frames <= 0:
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: max video frames must be positive.")
|
||||
|
||||
frames: list[Image.Image] = []
|
||||
tail_frames: deque[Image.Image] | None = (
|
||||
deque(maxlen=max_frames) if keep == "last" and max_frames is not None else None
|
||||
)
|
||||
try:
|
||||
with av.open(BytesIO(video_bytes)) as container:
|
||||
for frame in container.decode(video=0):
|
||||
image = frame.to_image().convert("RGB")
|
||||
if tail_frames is not None:
|
||||
tail_frames.append(image)
|
||||
else:
|
||||
frames.append(image)
|
||||
if keep == "first" and max_frames is not None and len(frames) >= max_frames:
|
||||
break
|
||||
except Exception as exc:
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.") from exc
|
||||
|
||||
if tail_frames is not None:
|
||||
frames = list(tail_frames)
|
||||
if not frames:
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.")
|
||||
return frames
|
||||
|
||||
|
||||
def _decode_media_bytes(
|
||||
media_bytes: bytes,
|
||||
*,
|
||||
source: str,
|
||||
max_video_frames: int | None = None,
|
||||
video_keep: Literal["first", "last"] = "first",
|
||||
) -> Image.Image | list[Image.Image]:
|
||||
try:
|
||||
return _decode_image_bytes(media_bytes, source=source)
|
||||
except InvalidInputReferenceError:
|
||||
try:
|
||||
return _decode_video_bytes(
|
||||
media_bytes,
|
||||
source=source,
|
||||
max_frames=max_video_frames,
|
||||
keep=video_keep,
|
||||
)
|
||||
except InvalidInputReferenceError as video_exc:
|
||||
raise InvalidInputReferenceError(
|
||||
f"Invalid {source}: provided content is not a valid image or video."
|
||||
) from video_exc
|
||||
|
||||
|
||||
def _decode_base64_image(input_reference: str, *, source: str) -> Image.Image:
|
||||
if input_reference:
|
||||
if input_reference.startswith("data:image"):
|
||||
@@ -64,23 +132,96 @@ async def decode_image_url(image_url: str) -> Image.Image:
|
||||
raise InvalidInputReferenceError("Invalid image_reference.image_url: must be an http(s) URL or data URL.")
|
||||
|
||||
|
||||
def _decode_base64_video(
|
||||
video_reference: str,
|
||||
*,
|
||||
source: str,
|
||||
max_frames: int | None = None,
|
||||
keep: Literal["first", "last"] = "first",
|
||||
) -> list[Image.Image]:
|
||||
if video_reference:
|
||||
if video_reference.startswith("data:video"):
|
||||
_, b64_data = video_reference.split(",", 1)
|
||||
else:
|
||||
b64_data = video_reference
|
||||
|
||||
try:
|
||||
video_bytes = base64.b64decode(b64_data)
|
||||
except (binascii.Error, ValueError) as exc: # pragma: no cover - malformed base64
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: video data is not valid base64.") from exc
|
||||
return _decode_video_bytes(video_bytes, source=source, max_frames=max_frames, keep=keep)
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: video data is empty.")
|
||||
|
||||
|
||||
async def decode_video_url(
|
||||
video_url: str,
|
||||
*,
|
||||
max_frames: int | None = None,
|
||||
keep: Literal["first", "last"] = "first",
|
||||
) -> list[Image.Image]:
|
||||
if video_url.startswith("data:video"):
|
||||
return _decode_base64_video(
|
||||
video_url,
|
||||
source="video_reference.video_url",
|
||||
max_frames=max_frames,
|
||||
keep=keep,
|
||||
)
|
||||
|
||||
if video_url.startswith(("http://", "https://")):
|
||||
async with httpx.AsyncClient(timeout=60) as client:
|
||||
try:
|
||||
response = await client.get(video_url)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPError as exc:
|
||||
raise InvalidInputReferenceError(
|
||||
"Invalid video_reference.video_url: failed to download video."
|
||||
) from exc
|
||||
return _decode_video_bytes(
|
||||
response.content,
|
||||
source="video_reference.video_url",
|
||||
max_frames=max_frames,
|
||||
keep=keep,
|
||||
)
|
||||
|
||||
raise InvalidInputReferenceError("Invalid video_reference.video_url: must be an http(s) URL or data URL.")
|
||||
|
||||
|
||||
async def decode_input_reference(
|
||||
image_reference: ImageReference | None,
|
||||
video_reference: VideoReference | None,
|
||||
input_reference_bytes: bytes | None,
|
||||
) -> Image.Image | None:
|
||||
"""Decode image input from multipart bytes, base64/data URL, or image_reference."""
|
||||
*,
|
||||
max_video_frames: int | None = None,
|
||||
video_keep: Literal["first", "last"] = "first",
|
||||
) -> Image.Image | list[Image.Image] | None:
|
||||
"""Decode media input from multipart bytes, data URLs, or typed references."""
|
||||
|
||||
if input_reference_bytes is not None and image_reference is not None:
|
||||
raise InvalidInputReferenceError("Provide either input_reference or image_reference, not both.")
|
||||
provided = sum(item is not None for item in (input_reference_bytes, image_reference, video_reference))
|
||||
if provided > 1:
|
||||
raise InvalidInputReferenceError("Provide only one of input_reference, image_reference, or video_reference.")
|
||||
|
||||
if isinstance(input_reference_bytes, bytes):
|
||||
return _decode_image_bytes(input_reference_bytes, source="input_reference")
|
||||
return _decode_media_bytes(
|
||||
input_reference_bytes,
|
||||
source="input_reference",
|
||||
max_video_frames=max_video_frames,
|
||||
video_keep=video_keep,
|
||||
)
|
||||
|
||||
if isinstance(image_reference, UrlImageReference):
|
||||
return await decode_image_url(image_reference.image_url)
|
||||
elif isinstance(image_reference, FileImageReference):
|
||||
raise InvalidInputReferenceError("Invalid image_reference: file_id is not supported yet.")
|
||||
|
||||
if isinstance(video_reference, UrlVideoReference):
|
||||
return await decode_video_url(
|
||||
video_reference.video_url,
|
||||
max_frames=max_video_frames,
|
||||
keep=video_keep,
|
||||
)
|
||||
elif isinstance(video_reference, FileVideoReference):
|
||||
raise InvalidInputReferenceError("Invalid video_reference: file_id is not supported yet.")
|
||||
|
||||
return None
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user