Merge remote-tracking branch 'remotes/origin/main' into dev/rebase-v0170
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
This commit is contained in:
Taichang Zhou committed
commit
0940c5a05e
91 files changed
+5235
-729
No files matched your search
@@ -0,0 +1,38 @@
|
||||
steps:
|
||||
- label: ":docker: Build image"
|
||||
key: image-build
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
|
||||
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L2 Test
|
||||
- label: "Upload Ready Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L3 Test
|
||||
- label: "Upload Merge Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L4 Test
|
||||
- label: "Upload Nightly Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
@@ -68,6 +68,7 @@ steps:
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- |
|
||||
huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk
|
||||
agents:
|
||||
|
||||
@@ -214,9 +214,10 @@ steps:
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py -k "not NoAsyncChunk"
|
||||
'
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
|
||||
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
||||
</h3>
|
||||
|
||||
<p align="center">
|
||||
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
|
||||
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://drive.google.com/file/d/12k_DOo2xoEROvd4xQxYNuEi6C2GECoRx/view?usp=drive_link"><b>Slides</b></a> |
|
||||
</p>
|
||||
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@ COPY . .
|
||||
|
||||
# Install system dependencies
|
||||
RUN apt-get update && \
|
||||
apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
|
||||
@@ -19,7 +19,7 @@ WORKDIR ${COMMON_WORKDIR}
|
||||
|
||||
# Step 1: Setup - Install system dependencies
|
||||
RUN apt-get update && \
|
||||
apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
|
||||
@@ -14,6 +14,7 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO
|
||||
RUN apt clean && apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends --fix-missing \
|
||||
curl \
|
||||
espeak-ng \
|
||||
ffmpeg \
|
||||
git \
|
||||
libsndfile1 \
|
||||
|
||||
@@ -69,20 +69,8 @@ With `--log-stats` enabled, the server will output detailed metrics logs after e
|
||||
| rx_decode_time_ms | 111.865 | 31.706 |
|
||||
| in_flight_time_ms | 2.015 | 2.819 |
|
||||
|
||||
### DiffusionStats
|
||||
|
||||
| Field | Value |
|
||||
|---------------------------------|------------|
|
||||
| batch_id | 1 |
|
||||
| batch_size | 1 |
|
||||
| diffusion_engine_exec_time_ms | 16,363.419 |
|
||||
| diffusion_engine_total_time_ms | 16,234.849 |
|
||||
| image_num | 1.000 |
|
||||
| num_inference_steps | 50.000 |
|
||||
| postprocess_time_ms | 67.685 |
|
||||
| preprocess_time_ms | 60.106 |
|
||||
| resolution | 640.000 |
|
||||
| stage_gen_time_ms | 16,363.714 |
|
||||
These logs include:
|
||||
|
||||
- **Overall summary**: total requests, wall time, average tokens/sec, etc.
|
||||
|
||||
@@ -92,9 +80,6 @@ With `--log-stats` enabled, the server will output detailed metrics logs after e
|
||||
|
||||
- **Transfer table**: data transfer and timing for each edge.
|
||||
|
||||
- **DiffusionStats**:preprocessing time, model execution time, postprocessing time, generated images, resolution, and inference steps.
|
||||
|
||||
|
||||
You can use these logs to monitor system health, debug performance, and analyze request-level metrics as described above.
|
||||
|
||||
|
||||
@@ -145,11 +130,6 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
|
||||
| `image_num` | Number of images generated (for diffusion/image stages). |
|
||||
| `resolution` | Image resolution (for diffusion/image stages). |
|
||||
| `postprocess_time_ms` | Diffusion/image: post-processing time in ms. |
|
||||
| `preprocess_time_ms` | Diffusion/image: pre-processing time in ms (before model execution). |
|
||||
| `diffusion_engine_exec_time_ms` | Diffusion/image: model execution time in ms, excluding pre-processing and post-processing. |
|
||||
| `diffusion_engine_total_time_ms` | Diffusion/image: total diffusion engine time in ms, including pre-processing, model execution, and post-processing. |
|
||||
| `num_inference_steps` | Number of inference steps for diffusion generation. |
|
||||
|
||||
|
||||
---
|
||||
|
||||
@@ -163,8 +143,6 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
|
||||
| `in_flight_time_ms` | In-flight time in ms. |
|
||||
|
||||
|
||||
|
||||
|
||||
### Expectation of the Numbers (Verification)
|
||||
|
||||
**Formulas:**
|
||||
@@ -193,14 +171,3 @@ For each edge:
|
||||
- 1->2: tx_time_ms (**18.790**) + rx_decode_time_ms (**31.706**) + in_flight_time_ms (**2.819**) = **53.315**
|
||||
|
||||
192.581 + 53.315 = **245.896** = transfers_total_time_ms, which matches the calculation (difference is due to rounding)
|
||||
|
||||
**diffusion_engine_total_time_ms**
|
||||
|
||||
- preprocess_time_ms: Time spent on input preprocessing
|
||||
- diffusion_engine_exec_time_ms: Pure model execution time
|
||||
- postprocess_time_ms: Time spent on output postprocessing
|
||||
- diffusion_engine_total_time_ms: Total time spent in the diffusion engine
|
||||
|
||||
so `diffusion_engine_total_time_ms` = `diffusion_engine_exec_time_ms` + `preprocess_time_ms` + `postprocess_time_ms`
|
||||
= **16,363.419** + **60.106** + **67.685**
|
||||
= **16,491.210**
|
||||
@@ -115,7 +115,7 @@ python image_edit.py \
|
||||
--image cherry_blossom.jpg \
|
||||
--prompt "Let this mascot dance under the moon, surrounded by floating stars" \
|
||||
--output output_image_edit.png \
|
||||
--num_inference_steps 10
|
||||
--num-inference-steps 10
|
||||
```
|
||||
|
||||
## Advanced Usage
|
||||
|
||||
@@ -104,7 +104,7 @@ Content-Type: application/json
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `ref_audio` | string | null | Reference audio (URL or base64 data URL) |
|
||||
| `ref_audio` | string | null | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`) |
|
||||
| `ref_text` | string | null | Transcript of reference audio |
|
||||
| `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) |
|
||||
|
||||
@@ -125,6 +125,81 @@ Lists available voices for the loaded model.
|
||||
"voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
|
||||
}
|
||||
```
|
||||
```
|
||||
POST /v1/audio/voices
|
||||
Content-Type: multipart/form-data
|
||||
```
|
||||
|
||||
Upload a new voice sample for voice cloning in Base task TTS requests.
|
||||
|
||||
**Form Parameters:**
|
||||
|
||||
| Parameter | Type | Required | Description |
|
||||
|-----------|------|----------|-------------|
|
||||
| `audio_sample` | file | Yes | Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4) |
|
||||
| `consent` | string | Yes | Consent recording ID |
|
||||
| `name` | string | Yes | Name for the new voice |
|
||||
|
||||
**Response Example:**
|
||||
|
||||
```json
|
||||
{
|
||||
"success": true,
|
||||
"voice": {
|
||||
"name": "custom_voice_1",
|
||||
"consent": "user_consent_id",
|
||||
"created_at": 1738660000,
|
||||
"mime_type": "audio/wav",
|
||||
"file_size": 1024000
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Usage Example:**
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
-F "audio_sample=@/path/to/voice_sample.wav" \
|
||||
-F "consent=user_consent_id" \
|
||||
-F "name=custom_voice_1"
|
||||
```
|
||||
|
||||
|
||||
```bash
|
||||
DELETE /v1/audio/voices/{name}
|
||||
```
|
||||
|
||||
Delete an uploaded voice sample.
|
||||
|
||||
**Path Parameters:**
|
||||
|
||||
| Parameter | Type | Required | Description |
|
||||
|-----------|------|----------|-------------|
|
||||
| `name` | string | Yes | Name of the voice to delete |
|
||||
|
||||
**Response Example:**
|
||||
|
||||
```json
|
||||
{
|
||||
"success": true,
|
||||
"message": "Voice 'custom_voice_1' deleted successfully"
|
||||
}
|
||||
```
|
||||
|
||||
**Error Response (404 Not Found):**
|
||||
|
||||
```json
|
||||
{
|
||||
"success": false,
|
||||
"error": "Voice 'unknown_voice' not found"
|
||||
}
|
||||
```
|
||||
|
||||
**Usage Example:**
|
||||
|
||||
```bash
|
||||
curl -X DELETE http://localhost:8091/v1/audio/voices/custom_voice_1
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
@@ -185,6 +260,25 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
upload voice
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
-F "audio_sample=@/path/to/voice_sample.wav" \
|
||||
-F "consent=user_consent_id" \
|
||||
-F "name=custom_voice_1"
|
||||
```
|
||||
|
||||
use upload voice
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, this is a cloned voice",
|
||||
"task_type": "Base",
|
||||
"voice": "custom_voice_1"
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | Task Type | Description |
|
||||
|
||||
@@ -136,9 +136,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--height 768 \
|
||||
--width 1360 \
|
||||
--seed 42 \
|
||||
--cfg_scale 4.0 \
|
||||
--num_images_per_prompt 1 \
|
||||
--num_inference_steps 4 \
|
||||
--cfg-scale 4.0 \
|
||||
--num-images-per-prompt 1 \
|
||||
--num-inference-steps 4 \
|
||||
--output outputs/flux2_klein_4b.png
|
||||
```
|
||||
|
||||
@@ -152,9 +152,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--height 768 \
|
||||
--width 1360 \
|
||||
--seed 42 \
|
||||
--cfg_scale 4.0 \
|
||||
--num_images_per_prompt 1 \
|
||||
--num_inference_steps 4 \
|
||||
--cfg-scale 4.0 \
|
||||
--num-images-per-prompt 1 \
|
||||
--num-inference-steps 4 \
|
||||
--output outputs/flux2_klein_4b_gguf.png
|
||||
```
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ Qwen3 TTS provides multiple task variants for speech generation:
|
||||
|
||||
- **CustomVoice**: Generate speech with a known speaker identity (speaker ID) and optional instruction.
|
||||
- **VoiceDesign**: Generate speech from text plus a descriptive instruction that designs a new voice.
|
||||
- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection.
|
||||
- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection. The `ref_audio` field accepts a local file path, HTTP URL, or base64 data URL.
|
||||
|
||||
## Setup
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
|
||||
@@ -112,7 +112,7 @@ The Python client supports the following command-line arguments:
|
||||
- `--voice`: Speaker/voice name (default: `vivian`). Options: `vivian`, `ryan`, `aiden`, etc.
|
||||
- `--language`: Language. Options: `Auto`, `Chinese`, `English`, `Japanese`, `Korean`, `German`, `French`, `Russian`, `Portuguese`, `Spanish`, `Italian`
|
||||
- `--instructions`: Voice style/emotion instructions
|
||||
- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task)
|
||||
- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task). Local paths are automatically base64-encoded by the client before sending to the server.
|
||||
- `--ref-text`: Reference audio transcript for voice cloning (Base task)
|
||||
- `--response-format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`, `aac`, `opus`
|
||||
- `--output` (or `-o`): Output audio file path (default: `tts_output.wav`)
|
||||
@@ -220,7 +220,7 @@ Lists available voices for the loaded model:
|
||||
"task_type": "CustomVoice",
|
||||
"language": "Auto",
|
||||
"instructions": "Optional style instructions",
|
||||
"ref_audio": "URL or base64 for voice cloning",
|
||||
"ref_audio": "HTTP URL, base64 data URL, or file:// URI for voice cloning",
|
||||
"ref_text": "Reference audio transcript",
|
||||
"x_vector_only_mode": false,
|
||||
"max_new_tokens": 2048
|
||||
@@ -261,7 +261,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
|
||||
|
||||
| Parameter | Type | Required | Description |
|
||||
| -------------------- | ------ | -------- | -------------------------------------------- |
|
||||
| `ref_audio` | string | **Yes** | Reference audio (URL or base64 data URL) |
|
||||
| `ref_audio` | string | **Yes** | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`) |
|
||||
| `ref_text` | string | No | Transcript of reference audio (for ICL mode) |
|
||||
| `x_vector_only_mode` | bool | No | Use speaker embedding only (no ICL) |
|
||||
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
import argparse
|
||||
import os
|
||||
from typing import cast
|
||||
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
|
||||
from vllm_omni.inputs.data import OmniPromptType
|
||||
|
||||
|
||||
def parse_args():
|
||||
@@ -58,6 +57,7 @@ def parse_args():
|
||||
choices=[1, 2, 3],
|
||||
help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.")
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
@@ -88,108 +88,80 @@ def main():
|
||||
|
||||
from PIL import Image
|
||||
|
||||
if args.modality == "img2img":
|
||||
from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
print(f"[Info] Running in {args.modality} mode (Stage 1 only, cfg_parallel_size={args.cfg_parallel_size})")
|
||||
omni_kwargs = {}
|
||||
if args.stage_configs_path:
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path
|
||||
|
||||
client = OmniDiffusion(
|
||||
model=model_name,
|
||||
parallel_config={"cfg_parallel_size": args.cfg_parallel_size},
|
||||
)
|
||||
omni_kwargs.update(
|
||||
{
|
||||
"log_stats": args.log_stats,
|
||||
"init_sleep_seconds": args.init_sleep_seconds,
|
||||
"batch_timeout": args.batch_timeout,
|
||||
"init_timeout": args.init_timeout,
|
||||
"shm_threshold_bytes": args.shm_threshold_bytes,
|
||||
"worker_backend": args.worker_backend,
|
||||
"ray_address": args.ray_address,
|
||||
}
|
||||
)
|
||||
|
||||
if args.image_path:
|
||||
if os.path.exists(args.image_path):
|
||||
loaded_image = Image.open(args.image_path).convert("RGB")
|
||||
prompts = [
|
||||
{
|
||||
"prompt": cast(str, p),
|
||||
"multi_modal_data": {"image": loaded_image},
|
||||
}
|
||||
for p in prompts
|
||||
]
|
||||
else:
|
||||
print(f"[Warning] Image path {args.image_path} does not exist.")
|
||||
omni = Omni(model=model_name, **omni_kwargs)
|
||||
|
||||
result = client.generate(
|
||||
prompts,
|
||||
OmniDiffusionSamplingParams(
|
||||
seed=52,
|
||||
need_kv_receive=False,
|
||||
num_inference_steps=args.steps,
|
||||
extra_args={
|
||||
"cfg_text_scale": args.cfg_text_scale,
|
||||
"cfg_img_scale": args.cfg_img_scale,
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
# Ensure result is a list for iteration
|
||||
if not isinstance(result, list):
|
||||
omni_outputs = [result]
|
||||
else:
|
||||
omni_outputs = result
|
||||
|
||||
else:
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni_kwargs = {}
|
||||
if args.stage_configs_path:
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path
|
||||
|
||||
omni_kwargs.update(
|
||||
{
|
||||
"log_stats": args.log_stats,
|
||||
"init_sleep_seconds": args.init_sleep_seconds,
|
||||
"batch_timeout": args.batch_timeout,
|
||||
"init_timeout": args.init_timeout,
|
||||
"shm_threshold_bytes": args.shm_threshold_bytes,
|
||||
"worker_backend": args.worker_backend,
|
||||
"ray_address": args.ray_address,
|
||||
formatted_prompts = []
|
||||
for p in prompts:
|
||||
if args.modality == "img2img":
|
||||
if not args.image_path or not os.path.exists(args.image_path):
|
||||
raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}")
|
||||
loaded_image = Image.open(args.image_path).convert("RGB")
|
||||
final_prompt_text = f"<|fim_middle|><|im_start|>{p}<|im_end|>"
|
||||
prompt_dict = {
|
||||
"prompt": final_prompt_text,
|
||||
"multi_modal_data": {"img2img": loaded_image},
|
||||
"modalities": ["img2img"],
|
||||
}
|
||||
)
|
||||
|
||||
omni = Omni(model=model_name, **omni_kwargs)
|
||||
|
||||
formatted_prompts = []
|
||||
for p in args.prompts:
|
||||
if args.modality == "img2text":
|
||||
if args.image_path:
|
||||
loaded_image = Image.open(args.image_path).convert("RGB")
|
||||
final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
|
||||
prompt_dict = {
|
||||
"prompt": final_prompt_text,
|
||||
"multi_modal_data": {"image": loaded_image},
|
||||
"modalities": ["text"],
|
||||
}
|
||||
formatted_prompts.append(prompt_dict)
|
||||
elif args.modality == "text2text":
|
||||
final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
|
||||
prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
|
||||
formatted_prompts.append(prompt_dict)
|
||||
else:
|
||||
# text2img
|
||||
final_prompt_text = f"<|im_start|>{p}<|im_end|>"
|
||||
prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
|
||||
if args.negative_prompt is not None:
|
||||
prompt_dict["negative_prompt"] = args.negative_prompt
|
||||
formatted_prompts.append(prompt_dict)
|
||||
|
||||
params_list = omni.default_sampling_params_list
|
||||
if args.modality == "text2img":
|
||||
params_list[0].max_tokens = 1 # type: ignore # The first stage is a SamplingParam (vllm)
|
||||
if len(params_list) > 1:
|
||||
diffusion_params = params_list[1]
|
||||
diffusion_params.num_inference_steps = args.steps # type: ignore
|
||||
extra = {
|
||||
"cfg_text_scale": args.cfg_text_scale,
|
||||
"cfg_img_scale": args.cfg_img_scale,
|
||||
if args.negative_prompt is not None:
|
||||
prompt_dict["negative_prompt"] = args.negative_prompt
|
||||
formatted_prompts.append(prompt_dict)
|
||||
elif args.modality == "img2text":
|
||||
if args.image_path:
|
||||
loaded_image = Image.open(args.image_path).convert("RGB")
|
||||
final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
|
||||
prompt_dict = {
|
||||
"prompt": final_prompt_text,
|
||||
"multi_modal_data": {"image": loaded_image},
|
||||
"modalities": ["text"],
|
||||
}
|
||||
if args.negative_prompt is not None:
|
||||
extra["negative_prompt"] = args.negative_prompt
|
||||
diffusion_params.extra_args = extra # type: ignore
|
||||
formatted_prompts.append(prompt_dict)
|
||||
elif args.modality == "text2text":
|
||||
final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
|
||||
prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
|
||||
formatted_prompts.append(prompt_dict)
|
||||
else:
|
||||
final_prompt_text = f"<|im_start|>{p}<|im_end|>"
|
||||
prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
|
||||
if args.negative_prompt is not None:
|
||||
prompt_dict["negative_prompt"] = args.negative_prompt
|
||||
formatted_prompts.append(prompt_dict)
|
||||
|
||||
omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
|
||||
params_list = omni.default_sampling_params_list
|
||||
if args.modality in ("text2img", "img2img"):
|
||||
params_list[0].max_tokens = 1 # type: ignore
|
||||
if len(params_list) > 1:
|
||||
diffusion_params = params_list[1]
|
||||
diffusion_params.num_inference_steps = args.steps # type: ignore
|
||||
diffusion_params.cfg_parallel_size = args.cfg_parallel_size # type: ignore
|
||||
if args.seed is not None:
|
||||
diffusion_params.seed = args.seed # type: ignore
|
||||
extra = {
|
||||
"cfg_text_scale": args.cfg_text_scale,
|
||||
"cfg_img_scale": args.cfg_img_scale,
|
||||
}
|
||||
if args.negative_prompt is not None:
|
||||
extra["negative_prompt"] = args.negative_prompt
|
||||
diffusion_params.extra_args = extra # type: ignore
|
||||
|
||||
omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
|
||||
|
||||
for i, req_output in enumerate(omni_outputs):
|
||||
images = getattr(req_output, "images", None)
|
||||
|
||||
@@ -29,8 +29,8 @@ Accelerated with cache_dit:
|
||||
python image_edit.py \
|
||||
--image input.png \
|
||||
--prompt "Edit description" \
|
||||
--cache_backend cache_dit \
|
||||
--cache_dit_enable_taylorseer
|
||||
--cache-backend cache_dit \
|
||||
--cache-dit-enable-taylorseer
|
||||
|
||||
Layered RGBA output:
|
||||
python image_edit.py \
|
||||
@@ -66,40 +66,39 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--model", default="Qwen/Qwen-Image-Edit", help="Model name or local path.")
|
||||
parser.add_argument("--image", type=str, nargs="+", required=True, help="Input image file(s).")
|
||||
parser.add_argument("--prompt", type=str, required=True, help="Edit description prompt.")
|
||||
parser.add_argument("--negative_prompt", type=str, default=None)
|
||||
parser.add_argument("--negative-prompt", type=str, default=None)
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
parser.add_argument("--cfg_scale", type=float, default=4.0)
|
||||
parser.add_argument("--guidance_scale", type=float, default=1.0)
|
||||
parser.add_argument("--cfg-scale", type=float, default=4.0)
|
||||
parser.add_argument("--guidance-scale", type=float, default=1.0)
|
||||
parser.add_argument("--output", type=str, default="output.png")
|
||||
parser.add_argument("--num_outputs_per_prompt", type=int, default=1)
|
||||
parser.add_argument("--num_inference_steps", type=int, default=50)
|
||||
parser.add_argument("--cache_backend", type=str, default=None, choices=["cache_dit", "tea_cache"])
|
||||
parser.add_argument("--ulysses_degree", type=int, default=1)
|
||||
parser.add_argument("--ring_degree", type=int, default=1)
|
||||
parser.add_argument("--tensor_parallel_size", type=int, default=1)
|
||||
parser.add_argument("--num-outputs-per-prompt", type=int, default=1)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=50)
|
||||
parser.add_argument("--cache-backend", type=str, default=None, choices=["cache_dit", "tea_cache"])
|
||||
parser.add_argument("--ulysses-degree", type=int, default=1)
|
||||
parser.add_argument("--ring-degree", type=int, default=1)
|
||||
parser.add_argument("--tensor-parallel-size", type=int, default=1)
|
||||
parser.add_argument("--layers", type=int, default=4)
|
||||
parser.add_argument("--resolution", type=int, default=640)
|
||||
parser.add_argument("--color-format", type=str, default="RGB")
|
||||
|
||||
# Acceleration + Optimization Options
|
||||
parser.add_argument("--cache_dit_fn_compute_blocks", type=int, default=1)
|
||||
parser.add_argument("--cache_dit_bn_compute_blocks", type=int, default=0)
|
||||
parser.add_argument("--cache_dit_max_warmup_steps", type=int, default=4)
|
||||
parser.add_argument("--cache_dit_residual_diff_threshold", type=float, default=0.24)
|
||||
parser.add_argument("--cache_dit_max_continuous_cached_steps", type=int, default=3)
|
||||
parser.add_argument("--cache_dit_enable_taylorseer", action="store_true", default=False)
|
||||
parser.add_argument("--cache_dit_taylorseer_order", type=int, default=1)
|
||||
parser.add_argument("--cache-dit-fn-compute-blocks", type=int, default=1)
|
||||
parser.add_argument("--cache-dit-bn-compute-blocks", type=int, default=0)
|
||||
parser.add_argument("--cache-dit-max-warmup-steps", type=int, default=4)
|
||||
parser.add_argument("--cache-dit-residual-diff-threshold", type=float, default=0.24)
|
||||
parser.add_argument("--cache-dit-max-continuous-cached-steps", type=int, default=3)
|
||||
parser.add_argument("--cache-dit-enable-taylorseer", action="store_true", default=False)
|
||||
parser.add_argument("--cache-dit-taylorseer-order", type=int, default=1)
|
||||
parser.add_argument(
|
||||
"--cache_dit_scm_steps_mask_policy", type=str, default=None, choices=[None, "slow", "medium", "fast", "ultra"]
|
||||
"--cache-dit-scm-steps-mask-policy", type=str, default=None, choices=[None, "slow", "medium", "fast", "ultra"]
|
||||
)
|
||||
parser.add_argument("--cache_dit_scm_steps_policy", type=str, default="dynamic", choices=["dynamic", "static"])
|
||||
parser.add_argument("--tea_cache_rel_l1_thresh", type=float, default=0.2)
|
||||
parser.add_argument("--cfg_parallel_size", type=int, default=1, choices=[1, 2])
|
||||
parser.add_argument("--enforce_eager", action="store_true")
|
||||
parser.add_argument("--vae_use_slicing", action="store_true")
|
||||
parser.add_argument("--vae_use_tiling", action="store_true")
|
||||
parser.add_argument("--cache-dit-scm-steps-policy", type=str, default="dynamic", choices=["dynamic", "static"])
|
||||
parser.add_argument("--tea-cache-rel-l1-thresh", type=float, default=0.2)
|
||||
parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2])
|
||||
parser.add_argument("--enforce-eager", action="store_true")
|
||||
parser.add_argument("--vae-use-slicing", action="store_true")
|
||||
parser.add_argument("--vae-use-tiling", action="store_true")
|
||||
parser.add_argument("--enable-cpu-offload", action="store_true")
|
||||
parser.add_argument("--log-stats", "--log_stats", dest="log_stats", action="store_true", default=False)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
@@ -156,7 +155,6 @@ async def main():
|
||||
cache_config=cache_config,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
log_stats=args.log_stats,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
diffusion_load_format="dummy",
|
||||
custom_pipeline_args={"pipeline_class": "custom_pipeline.CustomPipeline"},
|
||||
|
||||
@@ -3,4 +3,4 @@ python image_edit.py \
|
||||
--image cherry_blossom.jpg \
|
||||
--prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
|
||||
--output output_image_edit.png \
|
||||
--num_inference_steps 10
|
||||
--num-inference-steps 10
|
||||
@@ -320,11 +320,6 @@ def parse_args() -> argparse.Namespace:
|
||||
action="store_true",
|
||||
help="Enable layerwise (blockwise) offloading on DiT modules.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
help="Enable vLLM-Omni statistics logging.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -388,7 +383,6 @@ def main():
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
log_stats=args.log_stats,
|
||||
)
|
||||
print("Pipeline loaded")
|
||||
|
||||
|
||||
@@ -22,7 +22,7 @@ Usage:
|
||||
python image_to_video.py --model /path/to/LTX-2 \
|
||||
--model-class-name LTX2ImageToVideoPipeline \
|
||||
--image input.jpg --prompt "A cinematic dolly shot of a boat" \
|
||||
--num-frames 121 --num_inference_steps 40 --guidance_scale 4.0 \
|
||||
--num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
|
||||
--frame-rate 24 --fps 24 --output ltx2_i2v.mp4
|
||||
"""
|
||||
|
||||
@@ -176,11 +176,6 @@ def parse_args() -> argparse.Namespace:
|
||||
"Default 1 means pure sharding (no replication). "
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
help="Enable vLLM-Omni statistics logging.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -283,7 +278,6 @@ def main():
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
log_stats=args.log_stats,
|
||||
model_class_name=model_class_name,
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
|
||||
@@ -64,6 +64,51 @@ If media file paths are not provided, the script will use default assets. Suppor
|
||||
- `multi_audios`: Multiple audio inputs
|
||||
- `mixed_modalities`: Combination of video, image, and audio inputs
|
||||
|
||||
### Async-chunk (offline)
|
||||
|
||||
For true stage-level concurrency -- where downstream stages (Talker, Code2Wav)
|
||||
start **before** the upstream stage (Thinker) finishes -- use the async_chunk
|
||||
example. This requires:
|
||||
|
||||
1. A stage config YAML with ``async_chunk: true`` (e.g.
|
||||
``qwen3_omni_moe_async_chunk.yaml``).
|
||||
2. Hardware that matches the config (e.g. 2x H100 for the default 3-stage
|
||||
config).
|
||||
|
||||
The async_chunk example uses ``AsyncOmni`` instead of the synchronous ``Omni``
|
||||
class, which enables the async orchestrator to receive stage-0 intermediate
|
||||
outputs and trigger downstream stages early. Chunk data flows directly between
|
||||
stage workers via the in-worker ``OmniChunkTransferAdapter`` / connector,
|
||||
**not** through the orchestrator.
|
||||
|
||||
#### Single prompt
|
||||
```bash
|
||||
cd examples/offline_inference/qwen3_omni
|
||||
bash run_single_prompt_async_chunk.sh
|
||||
```
|
||||
|
||||
#### Multiple prompts with concurrency control
|
||||
```bash
|
||||
bash run_multiple_prompts_async_chunk.sh --max-in-flight 4
|
||||
```
|
||||
|
||||
#### Text-only output (skip audio generation)
|
||||
```bash
|
||||
python end2end_async_chunk.py --query-type text --modalities text
|
||||
```
|
||||
|
||||
#### Custom stage config
|
||||
```bash
|
||||
python end2end_async_chunk.py \
|
||||
--query-type use_audio \
|
||||
--stage-configs-path /path/to/your_async_chunk.yaml
|
||||
```
|
||||
|
||||
> **Note**: The synchronous ``end2end.py`` (using ``Omni``) is still the
|
||||
> recommended entry point for non-async-chunk workflows. Only use the
|
||||
> async_chunk example when you need the stage-level concurrency semantics
|
||||
> described in PR #962 / #1151.
|
||||
|
||||
### FAQ
|
||||
|
||||
If you encounter error about backend of librosa, try to install ffmpeg with command below.
|
||||
|
||||
@@ -0,0 +1,614 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline inference with async_chunk enabled via AsyncOmni.
|
||||
|
||||
This script uses AsyncOmni (the async orchestrator) to run offline inference
|
||||
with async_chunk semantics: downstream stages (Talker, Code2Wav) start
|
||||
*before* upstream stages finish, consuming chunks as they arrive via
|
||||
the in-worker OmniChunkTransferAdapter / connector.
|
||||
|
||||
Compared to the synchronous ``end2end.py`` (which uses ``Omni``), this
|
||||
entry point achieves true stage-level concurrency -- stage-1/2 are
|
||||
actively processing while stage-0 is still generating.
|
||||
|
||||
Usage
|
||||
-----
|
||||
python end2end_async_chunk.py --query-type use_audio \
|
||||
--stage-configs-path <path-to-async-chunk-yaml>
|
||||
|
||||
See ``--help`` for all options.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from typing import NamedTuple
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
import librosa
|
||||
from PIL import Image
|
||||
from vllm import SamplingParams
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Query builders (reuse the patterns from end2end.py)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
default_system = (
|
||||
"You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
|
||||
"Group, capable of perceiving auditory and visual inputs, as well as "
|
||||
"generating text and speech."
|
||||
)
|
||||
|
||||
|
||||
class QueryResult(NamedTuple):
|
||||
inputs: dict
|
||||
limit_mm_per_prompt: dict[str, int]
|
||||
|
||||
|
||||
def get_text_query(question: str = None) -> QueryResult:
|
||||
if question is None:
|
||||
question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
|
||||
prompt = (
|
||||
f"<|im_start|>system\n{default_system}<|im_end|>\n"
|
||||
"<|im_start|>user\n"
|
||||
f"{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
return QueryResult(inputs={"prompt": prompt}, limit_mm_per_prompt={})
|
||||
|
||||
|
||||
def get_audio_query(
|
||||
question: str = None,
|
||||
audio_path: str | None = None,
|
||||
sampling_rate: int = 16000,
|
||||
) -> QueryResult:
|
||||
if question is None:
|
||||
question = "What is the content of this audio?"
|
||||
prompt = (
|
||||
f"<|im_start|>system\n{default_system}<|im_end|>\n"
|
||||
"<|im_start|>user\n<|audio_start|><|audio_pad|><|audio_end|>"
|
||||
f"{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
if audio_path:
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
|
||||
audio_data = (audio_signal.astype(np.float32), sr)
|
||||
else:
|
||||
audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
|
||||
return QueryResult(
|
||||
inputs={
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"audio": audio_data},
|
||||
},
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
)
|
||||
|
||||
|
||||
def get_image_query(question: str = None, image_path: str | None = None) -> QueryResult:
|
||||
if question is None:
|
||||
question = "What is the content of this image?"
|
||||
prompt = (
|
||||
f"<|im_start|>system\n{default_system}<|im_end|>\n"
|
||||
"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
|
||||
f"{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
if image_path:
|
||||
if not os.path.exists(image_path):
|
||||
raise FileNotFoundError(f"Image file not found: {image_path}")
|
||||
pil_image = Image.open(image_path)
|
||||
image_data = convert_image_mode(pil_image, "RGB")
|
||||
else:
|
||||
image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
|
||||
return QueryResult(
|
||||
inputs={
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": image_data},
|
||||
},
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
)
|
||||
|
||||
|
||||
def get_video_query(
|
||||
question: str = None,
|
||||
video_path: str | None = None,
|
||||
num_frames: int = 16,
|
||||
) -> QueryResult:
|
||||
if question is None:
|
||||
question = "Why is this video funny?"
|
||||
prompt = (
|
||||
f"<|im_start|>system\n{default_system}<|im_end|>\n"
|
||||
"<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
|
||||
f"{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
if video_path:
|
||||
if not os.path.exists(video_path):
|
||||
raise FileNotFoundError(f"Video file not found: {video_path}")
|
||||
video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
|
||||
else:
|
||||
video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
|
||||
return QueryResult(
|
||||
inputs={
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"video": video_frames},
|
||||
},
|
||||
limit_mm_per_prompt={"video": 1},
|
||||
)
|
||||
|
||||
|
||||
query_map = {
|
||||
"text": get_text_query,
|
||||
"use_audio": get_audio_query,
|
||||
"use_image": get_image_query,
|
||||
"use_video": get_video_query,
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Core async routine
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def clone_prompt_for_request(template: dict) -> dict:
|
||||
"""Shallow-clone prompt dict so concurrent requests own independent containers."""
|
||||
cloned = dict(template)
|
||||
for key in ("multi_modal_data", "mm_processor_kwargs", "additional_information"):
|
||||
value = template.get(key)
|
||||
if isinstance(value, dict):
|
||||
cloned[key] = dict(value)
|
||||
elif isinstance(value, list):
|
||||
cloned[key] = list(value)
|
||||
return cloned
|
||||
|
||||
|
||||
def _default_async_chunk_stage_configs_path() -> str | None:
|
||||
"""Best-effort default stage config for running Qwen3-Omni with async_chunk.
|
||||
|
||||
When this example is executed from within the repository, we resolve the
|
||||
default YAML path relative to this file. When installed elsewhere, the
|
||||
file may not exist and callers should pass --stage-configs-path explicitly.
|
||||
"""
|
||||
repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
|
||||
candidate = os.path.join(
|
||||
repo_root,
|
||||
"vllm_omni",
|
||||
"model_executor",
|
||||
"stage_configs",
|
||||
"qwen3_omni_moe_async_chunk.yaml",
|
||||
)
|
||||
return candidate if os.path.exists(candidate) else None
|
||||
|
||||
|
||||
async def run_single_request(
|
||||
async_omni: AsyncOmni,
|
||||
prompt: dict,
|
||||
request_id: str,
|
||||
sampling_params_list: list[SamplingParams] | None,
|
||||
output_dir: str,
|
||||
output_modalities: list[str] | None = None,
|
||||
stream_audio_to_disk: bool = False,
|
||||
) -> dict:
|
||||
"""Run one request through AsyncOmni and collect outputs.
|
||||
|
||||
Returns a dict with timing information and saved file paths.
|
||||
"""
|
||||
t_start = time.perf_counter()
|
||||
text_parts: list[str] = []
|
||||
audio_chunks: list[torch.Tensor] = []
|
||||
audio_sr: int | None = None
|
||||
first_audio_ts: float | None = None
|
||||
audio_list_consumed: int = 0
|
||||
audio_last_tensor: torch.Tensor | None = None
|
||||
stage_0_first_output_ts: float | None = None
|
||||
|
||||
samplerate = 24000
|
||||
wav_file = os.path.join(output_dir, f"output_{request_id}.wav")
|
||||
sf_writer: sf.SoundFile | None = None
|
||||
audio_samples_written: int = 0
|
||||
|
||||
try:
|
||||
async for omni_output in async_omni.generate(
|
||||
prompt=prompt,
|
||||
request_id=request_id,
|
||||
sampling_params_list=sampling_params_list,
|
||||
output_modalities=output_modalities,
|
||||
):
|
||||
if not isinstance(omni_output.request_output, list):
|
||||
outputs_list = [omni_output.request_output]
|
||||
else:
|
||||
outputs_list = omni_output.request_output
|
||||
|
||||
for output in outputs_list:
|
||||
if omni_output.final_output_type == "text":
|
||||
if stage_0_first_output_ts is None:
|
||||
stage_0_first_output_ts = time.perf_counter()
|
||||
text_output = output.outputs[0].text
|
||||
if output.finished:
|
||||
text_parts.append(text_output)
|
||||
elif omni_output.final_output_type == "audio":
|
||||
mm_out = output.outputs[0].multimodal_output
|
||||
if mm_out and "audio" in mm_out:
|
||||
if first_audio_ts is None:
|
||||
first_audio_ts = time.perf_counter()
|
||||
if audio_sr is None and "sr" in mm_out:
|
||||
sr_val = mm_out["sr"]
|
||||
audio_sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
|
||||
samplerate = audio_sr
|
||||
audio_data = mm_out["audio"]
|
||||
if isinstance(audio_data, list):
|
||||
new_chunks = audio_data[audio_list_consumed:]
|
||||
audio_list_consumed = len(audio_data)
|
||||
elif isinstance(audio_data, torch.Tensor):
|
||||
new_chunks = [audio_data]
|
||||
audio_last_tensor = audio_data
|
||||
else:
|
||||
new_chunks = []
|
||||
|
||||
if stream_audio_to_disk and new_chunks:
|
||||
if sf_writer is None:
|
||||
sf_writer = sf.SoundFile(
|
||||
wav_file,
|
||||
mode="w",
|
||||
samplerate=samplerate,
|
||||
channels=1,
|
||||
subtype="FLOAT",
|
||||
)
|
||||
for chunk in new_chunks:
|
||||
chunk_np = chunk.float().detach().cpu().numpy().flatten()
|
||||
sf_writer.write(chunk_np)
|
||||
audio_samples_written += len(chunk_np)
|
||||
else:
|
||||
audio_chunks.extend(new_chunks)
|
||||
finally:
|
||||
if sf_writer is not None:
|
||||
sf_writer.close()
|
||||
|
||||
t_end = time.perf_counter()
|
||||
result = {
|
||||
"request_id": request_id,
|
||||
"e2e_latency_s": t_end - t_start,
|
||||
"saved_files": [],
|
||||
}
|
||||
|
||||
# Save text output
|
||||
if text_parts:
|
||||
text_file = os.path.join(output_dir, f"{request_id}.txt")
|
||||
with open(text_file, "w", encoding="utf-8") as f:
|
||||
f.write("\n".join(text_parts))
|
||||
result["saved_files"].append(text_file)
|
||||
print(
|
||||
f"[Request {request_id}] Text saved to {text_file} "
|
||||
f"(stage-0 first output at {stage_0_first_output_ts - t_start:.3f}s)"
|
||||
)
|
||||
|
||||
# Save audio output
|
||||
if stream_audio_to_disk and audio_samples_written > 0:
|
||||
result["saved_files"].append(wav_file)
|
||||
result["audio_duration_s"] = audio_samples_written / samplerate
|
||||
result["num_audio_chunks"] = audio_list_consumed
|
||||
ttfa = (first_audio_ts - t_start) if first_audio_ts else None
|
||||
result["time_to_first_audio_s"] = ttfa
|
||||
ttfa_str = f"{ttfa:.3f}s" if ttfa is not None else "N/A"
|
||||
print(
|
||||
f"[Request {request_id}] Audio streamed to {wav_file} "
|
||||
f"(duration={result['audio_duration_s']:.2f}s, "
|
||||
f"TTFA={ttfa_str}, "
|
||||
f"e2e={result['e2e_latency_s']:.3f}s)"
|
||||
)
|
||||
elif audio_chunks or audio_last_tensor is not None:
|
||||
if audio_chunks:
|
||||
if len(audio_chunks) > 1:
|
||||
audio_tensor = torch.cat(audio_chunks, dim=-1)
|
||||
else:
|
||||
audio_tensor = audio_chunks[0]
|
||||
else:
|
||||
audio_tensor = audio_last_tensor
|
||||
audio_numpy = audio_tensor.float().detach().cpu().numpy()
|
||||
if audio_numpy.ndim > 1:
|
||||
audio_numpy = audio_numpy.flatten()
|
||||
sf.write(wav_file, audio_numpy, samplerate=samplerate, format="WAV")
|
||||
result["saved_files"].append(wav_file)
|
||||
result["audio_duration_s"] = len(audio_numpy) / samplerate
|
||||
result["num_audio_chunks"] = len(audio_chunks)
|
||||
ttfa = (first_audio_ts - t_start) if first_audio_ts else None
|
||||
result["time_to_first_audio_s"] = ttfa
|
||||
ttfa_str = f"{ttfa:.3f}s" if ttfa is not None else "N/A"
|
||||
print(
|
||||
f"[Request {request_id}] Audio saved to {wav_file} "
|
||||
f"({len(audio_chunks)} chunks, "
|
||||
f"duration={result['audio_duration_s']:.2f}s, "
|
||||
f"TTFA={ttfa_str}, "
|
||||
f"e2e={result['e2e_latency_s']:.3f}s)"
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
async def run_all(args):
|
||||
"""Main async entry: build prompts, create AsyncOmni, run requests."""
|
||||
# Build query
|
||||
query_func = query_map[args.query_type]
|
||||
if args.query_type == "use_video":
|
||||
query_result = query_func(
|
||||
video_path=getattr(args, "video_path", None),
|
||||
num_frames=getattr(args, "num_frames", 16),
|
||||
)
|
||||
elif args.query_type == "use_image":
|
||||
query_result = query_func(image_path=getattr(args, "image_path", None))
|
||||
elif args.query_type == "use_audio":
|
||||
query_result = query_func(
|
||||
audio_path=getattr(args, "audio_path", None),
|
||||
sampling_rate=getattr(args, "sampling_rate", 16000),
|
||||
)
|
||||
else:
|
||||
query_result = query_func()
|
||||
|
||||
# Build prompt list
|
||||
if args.txt_prompts is not None:
|
||||
assert args.query_type == "text", "txt-prompts is only supported for text query type"
|
||||
with open(args.txt_prompts, encoding="utf-8") as f:
|
||||
lines = [ln.strip() for ln in f if ln.strip()]
|
||||
prompts = [get_text_query(ln).inputs for ln in lines]
|
||||
print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
|
||||
else:
|
||||
prompts = [clone_prompt_for_request(query_result.inputs) for _ in range(args.num_prompts)]
|
||||
|
||||
# Inject output modalities if specified
|
||||
output_modalities = None
|
||||
if args.modalities is not None:
|
||||
output_modalities = args.modalities.split(",")
|
||||
for prompt in prompts:
|
||||
prompt["modalities"] = output_modalities
|
||||
|
||||
# Create AsyncOmni
|
||||
print(f"[Info] Creating AsyncOmni with stage_configs_path={args.stage_configs_path}")
|
||||
async_omni = None
|
||||
try:
|
||||
async_omni = AsyncOmni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
# Use default sampling params from stage config (they are pre-configured
|
||||
# in the YAML for each stage).
|
||||
sampling_params_list = None
|
||||
|
||||
output_dir = args.output_dir
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Run requests with concurrency control
|
||||
semaphore = asyncio.Semaphore(args.max_in_flight)
|
||||
request_timeout = getattr(args, "request_timeout_s", None)
|
||||
stream_audio = getattr(args, "stream_audio_to_disk", False)
|
||||
|
||||
async def _run_one(idx: int, prompt: dict):
|
||||
async with semaphore:
|
||||
request_id = f"req_{idx}_{uuid.uuid4().hex[:8]}"
|
||||
coro = run_single_request(
|
||||
async_omni=async_omni,
|
||||
prompt=prompt,
|
||||
request_id=request_id,
|
||||
sampling_params_list=sampling_params_list,
|
||||
output_dir=output_dir,
|
||||
output_modalities=output_modalities,
|
||||
stream_audio_to_disk=stream_audio,
|
||||
)
|
||||
if request_timeout and request_timeout > 0:
|
||||
return await asyncio.wait_for(coro, timeout=request_timeout)
|
||||
return await coro
|
||||
|
||||
wall_start = time.perf_counter()
|
||||
tasks = [_run_one(i, p) for i, p in enumerate(prompts)]
|
||||
all_results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
wall_end = time.perf_counter()
|
||||
|
||||
# Print summary
|
||||
print("\n" + "=" * 60)
|
||||
print("Summary")
|
||||
print("=" * 60)
|
||||
success_count = 0
|
||||
total_audio_dur = 0.0
|
||||
for r in all_results:
|
||||
if isinstance(r, Exception):
|
||||
print(f" [ERROR] {type(r).__name__}: {r}")
|
||||
else:
|
||||
success_count += 1
|
||||
total_audio_dur += r.get("audio_duration_s", 0.0)
|
||||
print(f" [{r['request_id']}] e2e={r['e2e_latency_s']:.3f}s files={r['saved_files']}")
|
||||
wall_time = wall_end - wall_start
|
||||
print(f"\nTotal: {success_count}/{len(prompts)} succeeded")
|
||||
print(f"Wall time: {wall_time:.3f}s")
|
||||
if total_audio_dur > 0:
|
||||
print(f"Total audio duration: {total_audio_dur:.2f}s")
|
||||
print(f"Real-time factor: {total_audio_dur / wall_time:.2f}x")
|
||||
print("=" * 60)
|
||||
finally:
|
||||
if async_omni is not None:
|
||||
async_omni.shutdown()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(
|
||||
description=(
|
||||
"Offline inference with async_chunk enabled via AsyncOmni. "
|
||||
"Downstream stages start before upstream stages finish, "
|
||||
"achieving true stage-level concurrency."
|
||||
)
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
help="Model name or path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
type=str,
|
||||
default="use_audio",
|
||||
choices=query_map.keys(),
|
||||
help="Query type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=_default_async_chunk_stage_configs_path(),
|
||||
help=(
|
||||
"Path to an async_chunk stage config YAML. "
|
||||
"If not set, uses the model's default config "
|
||||
"(make sure it has async_chunk: true)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable writing detailed statistics.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-init-timeout",
|
||||
type=int,
|
||||
default=300,
|
||||
help="Timeout for initializing a single stage (seconds).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default="output_audio_async_chunk",
|
||||
help="Directory to save output files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-prompts",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of prompts to generate (duplicated from query).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--txt-prompts",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a .txt file with one prompt per line.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-in-flight",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Maximum concurrent requests (default: 1).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--request-timeout-s",
|
||||
type=float,
|
||||
default=None,
|
||||
help=(
|
||||
"Per-request timeout in seconds. When set, a request that "
|
||||
"exceeds this duration is cancelled and reported as an error. "
|
||||
"Default: None (no timeout)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-timeout-s",
|
||||
type=float,
|
||||
default=None,
|
||||
help=(
|
||||
"Global timeout for the entire batch in seconds. When set, "
|
||||
"the whole run_all() is cancelled if it exceeds this duration. "
|
||||
"Default: None (no global timeout)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stream-audio-to-disk",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help=(
|
||||
"Write audio chunks to WAV incrementally instead of "
|
||||
"accumulating in memory. Useful for very long audio or "
|
||||
"high --max-in-flight to reduce memory footprint."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--modalities",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Comma-separated output modalities filter (e.g. 'text', 'audio', 'text,audio').",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-path",
|
||||
"-a",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to local audio file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image-path",
|
||||
"-i",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to local image file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video-path",
|
||||
"-v",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to local video file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-frames",
|
||||
type=int,
|
||||
default=16,
|
||||
help="Number of frames to extract from video.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sampling-rate",
|
||||
type=int,
|
||||
default=16000,
|
||||
help="Sampling rate for audio loading.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
|
||||
async def _main():
|
||||
batch_timeout = getattr(args, "batch_timeout_s", None)
|
||||
if batch_timeout and batch_timeout > 0:
|
||||
await asyncio.wait_for(run_all(args), timeout=batch_timeout)
|
||||
else:
|
||||
await run_all(args)
|
||||
|
||||
try:
|
||||
asyncio.run(_main())
|
||||
except asyncio.TimeoutError:
|
||||
print(
|
||||
f"\n[TIMEOUT] Batch exceeded --batch-timeout-s="
|
||||
f"{args.batch_timeout_s}s. AsyncOmni shutdown was handled "
|
||||
f"by finally block."
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
print("\nInterrupted by user. AsyncOmni shutdown was handled by finally block.")
|
||||
@@ -0,0 +1,23 @@
|
||||
#!/bin/bash
|
||||
# Run multiple Qwen3-Omni requests with async_chunk enabled.
|
||||
#
|
||||
# Uses AsyncOmni with --max-in-flight to control request-level
|
||||
# concurrency (each request still gets true stage-level concurrency
|
||||
# via async_chunk).
|
||||
#
|
||||
# Usage:
|
||||
# bash run_multiple_prompts_async_chunk.sh
|
||||
# bash run_multiple_prompts_async_chunk.sh --max-in-flight 4
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
|
||||
|
||||
python "${SCRIPT_DIR}/end2end_async_chunk.py" \
|
||||
--query-type text \
|
||||
--txt-prompts "${SCRIPT_DIR}/text_prompts_10.txt" \
|
||||
--stage-configs-path "${REPO_ROOT}/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml" \
|
||||
--output-dir output_audio_async_chunk \
|
||||
--max-in-flight 2 \
|
||||
"$@"
|
||||
@@ -0,0 +1,26 @@
|
||||
#!/bin/bash
|
||||
# Run a single Qwen3-Omni request with async_chunk enabled.
|
||||
#
|
||||
# This uses AsyncOmni (async orchestrator) so that downstream stages
|
||||
# (Talker, Code2Wav) start *before* stage-0 (Thinker) finishes,
|
||||
# achieving true stage-level concurrency via chunk-level streaming.
|
||||
#
|
||||
# Prerequisites:
|
||||
# - An async_chunk stage config YAML (e.g. qwen3_omni_moe_async_chunk.yaml)
|
||||
# - Hardware matching the config (e.g. 2x H100 for the default 3-stage config)
|
||||
#
|
||||
# Usage:
|
||||
# bash run_single_prompt_async_chunk.sh
|
||||
# bash run_single_prompt_async_chunk.sh --query-type text --modalities text
|
||||
# bash run_single_prompt_async_chunk.sh --stage-configs-path /path/to/custom.yaml
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
|
||||
|
||||
python "${SCRIPT_DIR}/end2end_async_chunk.py" \
|
||||
--query-type use_audio \
|
||||
--stage-configs-path "${REPO_ROOT}/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml" \
|
||||
--output-dir output_audio_async_chunk \
|
||||
"$@"
|
||||
@@ -79,10 +79,16 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
|
||||
task_type = "CustomVoice"
|
||||
model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
|
||||
if use_batch_sample:
|
||||
texts = ["其实我真的有发现，我是一个特别善于观察别人情绪的人。", "She said she would be here by noon."]
|
||||
instructs = ["", "Very happy."]
|
||||
languages = ["Chinese", "English"]
|
||||
speakers = ["Vivian", "Ryan"]
|
||||
texts = [
|
||||
"其实我真的有发现，我是一个特别善于观察别人情绪的人。",
|
||||
"She said she would be here by noon.",
|
||||
"I like you very much.",
|
||||
"Really, you do?",
|
||||
"Yes, absolutely.",
|
||||
]
|
||||
instructs = ["", "Very happy.", "Very happy.", "Very happy.", "Very happy."]
|
||||
languages = ["Chinese", "English", "English", "English", "English"]
|
||||
speakers = ["Vivian", "Ryan", "Ryan", "Ryan", "Ryan"]
|
||||
inputs = []
|
||||
for text, instruct, language, speaker in zip(texts, instructs, languages, speakers):
|
||||
additional_information = {
|
||||
|
||||
@@ -94,11 +94,6 @@ def parse_args() -> argparse.Namespace:
|
||||
action="store_true",
|
||||
help="Enable cache-dit summary logging after diffusion forward passes.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
help="Enable vLLM-Omni statistics logging.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ulysses-degree",
|
||||
type=int,
|
||||
@@ -304,7 +299,6 @@ def main():
|
||||
"parallel_config": parallel_config,
|
||||
"enforce_eager": args.enforce_eager,
|
||||
"enable_cpu_offload": args.enable_cpu_offload,
|
||||
"log_stats": args.log_stats,
|
||||
**lora_args,
|
||||
**quant_kwargs,
|
||||
}
|
||||
|
||||
@@ -127,11 +127,6 @@ def parse_args() -> argparse.Namespace:
|
||||
default=1,
|
||||
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
help="Enable vLLM-Omni statistics logging.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -181,7 +176,6 @@ def main():
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
log_stats=args.log_stats,
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
)
|
||||
|
||||
@@ -184,6 +184,59 @@ sudo apt install ffmpeg
|
||||
|
||||
## API Reference
|
||||
|
||||
### Voices Endpoint
|
||||
|
||||
#### GET /v1/audio/voices
|
||||
|
||||
List all available voices/speakers from the loaded model, including both built-in model voices and uploaded custom voices.
|
||||
|
||||
**Response Example:**
|
||||
```json
|
||||
{
|
||||
"voices": ["vivian", "ryan", "custom_voice_1"],
|
||||
"uploaded_voices": [
|
||||
{
|
||||
"name": "custom_voice_1",
|
||||
"consent": "user_consent_id",
|
||||
"created_at": 1738660000,
|
||||
"file_size": 1024000,
|
||||
"mime_type": "audio/wav"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
#### POST /v1/audio/voices
|
||||
|
||||
Upload a new voice sample for voice cloning in Base task TTS requests.
|
||||
|
||||
**Form Parameters:**
|
||||
- `audio_sample` (required): Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4)
|
||||
- `consent` (required): Consent recording ID
|
||||
- `name` (required): Name for the new voice
|
||||
|
||||
**Response Example:**
|
||||
```json
|
||||
{
|
||||
"success": true,
|
||||
"voice": {
|
||||
"name": "custom_voice_1",
|
||||
"consent": "user_consent_id",
|
||||
"created_at": 1738660000,
|
||||
"mime_type": "audio/wav",
|
||||
"file_size": 1024000
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Usage Example:**
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/audio/voices \
|
||||
-F "audio_sample=@/path/to/voice_sample.wav" \
|
||||
-F "consent=user_consent_id" \
|
||||
-F "name=custom_voice_1"
|
||||
```
|
||||
|
||||
### Endpoint
|
||||
|
||||
```
|
||||
@@ -193,20 +246,6 @@ Content-Type: application/json
|
||||
|
||||
This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters.
|
||||
|
||||
### Voices Endpoint
|
||||
|
||||
```
|
||||
GET /v1/audio/voices
|
||||
```
|
||||
|
||||
Lists available voices for the loaded model:
|
||||
|
||||
```json
|
||||
{
|
||||
"voices": ["aiden", "dylan", "eric", "one_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
|
||||
}
|
||||
```
|
||||
|
||||
### Request Body
|
||||
|
||||
```json
|
||||
|
||||
@@ -51,6 +51,7 @@ dev = [
|
||||
"mooncake-transfer-engine==0.3.8.post1",
|
||||
"av", # for ComfyUI tests
|
||||
"openpyxl>=3.0.0", # for nightly CI
|
||||
"pyttsx3>=2.99"
|
||||
]
|
||||
|
||||
docs = [
|
||||
|
||||
+200
-94
@@ -12,6 +12,7 @@ if "VLLM_TARGET_DEVICE" not in os.environ:
|
||||
|
||||
import concurrent.futures
|
||||
import gc
|
||||
import multiprocessing
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
@@ -217,102 +218,193 @@ def generate_synthetic_audio(
|
||||
sample_rate: int = 48000, # Default use 48000Hz.
|
||||
save_to_file: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
""" "Generate synthetic audio with rain."""
|
||||
"""
|
||||
Generate TTS speech with pyttsx3 and return base64 string.
|
||||
"""
|
||||
import tempfile
|
||||
|
||||
import pyttsx3
|
||||
import soundfile as sf
|
||||
|
||||
# Initialize audio data array
|
||||
num_samples = int(sample_rate * duration)
|
||||
def _pick_voice(engine: pyttsx3.Engine) -> str | None:
|
||||
voices = engine.getProperty("voices")
|
||||
if not voices:
|
||||
return None
|
||||
|
||||
preferred_tokens = (
|
||||
"natural",
|
||||
"jenny",
|
||||
"sonia",
|
||||
"susan",
|
||||
"zira",
|
||||
"aria",
|
||||
"hazel",
|
||||
"samantha",
|
||||
"ava",
|
||||
"allison",
|
||||
"female",
|
||||
"woman",
|
||||
"english-us",
|
||||
"en-us",
|
||||
"english",
|
||||
)
|
||||
discouraged_tokens = (
|
||||
"espeak",
|
||||
"robot",
|
||||
"mbrola",
|
||||
"microsoft david",
|
||||
"male",
|
||||
"man",
|
||||
)
|
||||
|
||||
best_voice = voices[0]
|
||||
best_score = float("-inf")
|
||||
for voice in voices:
|
||||
voice_text = f"{getattr(voice, 'id', '')} {getattr(voice, 'name', '')}".lower()
|
||||
voice_languages = " ".join(
|
||||
lang.decode(errors="ignore") if isinstance(lang, bytes) else str(lang)
|
||||
for lang in getattr(voice, "languages", [])
|
||||
).lower()
|
||||
combined_text = f"{voice_text} {voice_languages}"
|
||||
score = 0
|
||||
for idx, token in enumerate(preferred_tokens):
|
||||
if token in combined_text:
|
||||
score += 20 - idx
|
||||
for token in discouraged_tokens:
|
||||
if token in combined_text:
|
||||
score -= 10
|
||||
if "english" in combined_text or "en_" in combined_text or "en-" in combined_text:
|
||||
score += 4
|
||||
if "en-us" in combined_text or "english-us" in combined_text:
|
||||
score += 4
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_voice = voice
|
||||
|
||||
return best_voice.id
|
||||
|
||||
def _resample_audio(audio: np.ndarray, src_sr: int, dst_sr: int) -> np.ndarray:
|
||||
if src_sr == dst_sr or len(audio) == 0:
|
||||
return audio.astype(np.float32)
|
||||
|
||||
src_len = audio.shape[0]
|
||||
dst_len = max(1, int(round(src_len * float(dst_sr) / float(src_sr))))
|
||||
src_idx = np.arange(src_len, dtype=np.float32)
|
||||
dst_idx = np.linspace(0, src_len - 1, dst_len, dtype=np.float32)
|
||||
|
||||
resampled_channels: list[np.ndarray] = []
|
||||
for ch in range(audio.shape[1]):
|
||||
resampled_channels.append(np.interp(dst_idx, src_idx, audio[:, ch]).astype(np.float32))
|
||||
return np.stack(resampled_channels, axis=1)
|
||||
|
||||
def _match_channels(audio: np.ndarray, target_channels: int) -> np.ndarray:
|
||||
current_channels = audio.shape[1]
|
||||
if current_channels == target_channels:
|
||||
return audio.astype(np.float32)
|
||||
if target_channels == 1:
|
||||
return np.mean(audio, axis=1, keepdims=True, dtype=np.float32)
|
||||
if current_channels == 1:
|
||||
return np.repeat(audio, target_channels, axis=1).astype(np.float32)
|
||||
|
||||
collapsed = np.mean(audio, axis=1, keepdims=True, dtype=np.float32)
|
||||
return np.repeat(collapsed, target_channels, axis=1).astype(np.float32)
|
||||
|
||||
def _trim_silence(audio: np.ndarray, threshold: float = 0.01) -> np.ndarray:
|
||||
if len(audio) == 0:
|
||||
return audio
|
||||
energy = np.max(np.abs(audio), axis=1)
|
||||
voiced = np.where(energy > threshold)[0]
|
||||
if len(voiced) == 0:
|
||||
return audio
|
||||
start = max(0, int(voiced[0]) - int(sample_rate * 0.02))
|
||||
end = min(len(audio), int(voiced[-1]) + int(sample_rate * 0.04) + 1)
|
||||
return audio[start:end]
|
||||
|
||||
def _enhance_speech(audio: np.ndarray) -> np.ndarray:
|
||||
if len(audio) == 0:
|
||||
return audio.astype(np.float32)
|
||||
enhanced = audio.astype(np.float32).copy()
|
||||
enhanced -= np.mean(enhanced, axis=0, keepdims=True, dtype=np.float32)
|
||||
if len(enhanced) > 1:
|
||||
preemphasis = enhanced.copy()
|
||||
preemphasis[1:] = enhanced[1:] - 0.94 * enhanced[:-1]
|
||||
enhanced = 0.7 * enhanced + 0.3 * preemphasis
|
||||
# Mild dynamic-range compression for ASR/TTS robustness.
|
||||
enhanced = np.sign(enhanced) * np.sqrt(np.abs(enhanced))
|
||||
# Light fade to avoid clicks after trimming/repeating.
|
||||
fade = min(len(enhanced) // 4, max(1, int(sample_rate * 0.01)))
|
||||
if fade > 1:
|
||||
ramp_in = np.linspace(0.0, 1.0, fade, dtype=np.float32)
|
||||
ramp_out = np.linspace(1.0, 0.0, fade, dtype=np.float32)
|
||||
enhanced[:fade] *= ramp_in[:, None]
|
||||
enhanced[-fade:] *= ramp_out[:, None]
|
||||
peak = float(np.max(np.abs(enhanced)))
|
||||
if peak > 1e-8:
|
||||
enhanced = enhanced / peak * 0.95
|
||||
return enhanced.astype(np.float32)
|
||||
|
||||
phrase_text = "test"
|
||||
num_samples = int(sample_rate * max(1, duration))
|
||||
audio_data = np.zeros((num_samples, num_channels), dtype=np.float32)
|
||||
|
||||
# Configure parameters based on rain intensity
|
||||
drop_density = 10 # Number of raindrops per second
|
||||
drop_volume = 0.15 # Volume of individual raindrops
|
||||
background_volume = 0.02 # Volume of background rain noise
|
||||
engine = pyttsx3.init()
|
||||
engine.setProperty("rate", 112)
|
||||
engine.setProperty("volume", 1.0)
|
||||
selected_voice = _pick_voice(engine)
|
||||
if selected_voice is not None:
|
||||
engine.setProperty("voice", selected_voice)
|
||||
|
||||
# Pink noise sounds more natural than white noise for rain
|
||||
white_noise = np.random.randn(num_samples)
|
||||
pink_noise = np.convolve(white_noise, np.ones(8) / 8, mode="same")
|
||||
pink_noise = pink_noise / np.max(np.abs(pink_noise)) if np.max(np.abs(pink_noise)) > 0 else pink_noise
|
||||
bg_noise = pink_noise * background_volume
|
||||
temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
|
||||
temp_wav.close()
|
||||
|
||||
# Add background noise to all channels
|
||||
for ch in range(num_channels):
|
||||
audio_data[:, ch] += bg_noise
|
||||
try:
|
||||
engine.save_to_file(phrase_text, temp_wav.name)
|
||||
engine.runAndWait()
|
||||
engine.stop()
|
||||
|
||||
# Total number of raindrops = density × duration × channels for stereo effect
|
||||
total_drops = int(drop_density * duration * num_channels)
|
||||
ready = False
|
||||
for _ in range(50):
|
||||
if os.path.exists(temp_wav.name) and os.path.getsize(temp_wav.name) > 44:
|
||||
ready = True
|
||||
break
|
||||
time.sleep(0.1)
|
||||
|
||||
for _ in range(total_drops):
|
||||
# Random timing for raindrop
|
||||
drop_time = random.uniform(0, duration)
|
||||
if not ready:
|
||||
raise RuntimeError("pyttsx3 did not produce a WAV file in time.")
|
||||
|
||||
# Random duration of raindrop sound (0.01-0.05 seconds)
|
||||
drop_duration = random.uniform(0.01, 0.05)
|
||||
tts_audio, tts_sr = sf.read(temp_wav.name, dtype="float32", always_2d=True)
|
||||
finally:
|
||||
if os.path.exists(temp_wav.name):
|
||||
os.unlink(temp_wav.name)
|
||||
|
||||
# Random frequency gives variation in raindrop pitch
|
||||
drop_freq = random.uniform(500, 5000) # Hz
|
||||
if len(tts_audio) == 0:
|
||||
raise RuntimeError("pyttsx3 produced an empty WAV file.")
|
||||
|
||||
# Random channel selection for stereo positioning
|
||||
channel = random.randint(0, num_channels - 1)
|
||||
tts_audio = _resample_audio(tts_audio, tts_sr, sample_rate)
|
||||
tts_audio = _match_channels(tts_audio, num_channels)
|
||||
tts_audio = _trim_silence(tts_audio, threshold=0.012)
|
||||
tts_audio = _enhance_speech(tts_audio)
|
||||
|
||||
# Calculate sample positions for this raindrop
|
||||
start_sample = int(drop_time * sample_rate)
|
||||
drop_samples = int(drop_duration * sample_rate)
|
||||
end_sample = min(start_sample + drop_samples, num_samples)
|
||||
lead_silence = min(int(sample_rate * 0.02), num_samples // 8)
|
||||
pause_samples = int(sample_rate * 0.18)
|
||||
start = lead_silence
|
||||
phrase_len = tts_audio.shape[0]
|
||||
|
||||
if start_sample < end_sample:
|
||||
# Generate the raindrop sound
|
||||
num_drop_samples = end_sample - start_sample
|
||||
t = np.arange(num_drop_samples) / sample_rate
|
||||
while start < num_samples:
|
||||
take = min(phrase_len, num_samples - start)
|
||||
audio_data[start : start + take] = tts_audio[:take]
|
||||
start += phrase_len + pause_samples
|
||||
|
||||
# Basic sine wave for raindrop sound
|
||||
drop_sound = drop_volume * np.sin(2 * math.pi * drop_freq * t)
|
||||
|
||||
# Apply envelope for natural attack and decay
|
||||
envelope = np.ones(num_drop_samples)
|
||||
attack_samples = int(num_drop_samples * 0.1) # 10% of samples for attack
|
||||
decay_samples = num_drop_samples - attack_samples
|
||||
|
||||
if attack_samples > 0:
|
||||
# Linear attack: volume increases from 0 to 1
|
||||
envelope[:attack_samples] = np.linspace(0, 1, attack_samples)
|
||||
|
||||
if decay_samples > 0:
|
||||
# Exponential decay for natural sound fade
|
||||
decay = np.exp(-8 * t[attack_samples:] / drop_duration)
|
||||
envelope[attack_samples:] = decay
|
||||
|
||||
# Apply envelope to raindrop sound
|
||||
drop_sound *= envelope
|
||||
|
||||
# Add raindrop sound to selected channel
|
||||
audio_data[start_sample:end_sample, channel] += drop_sound
|
||||
|
||||
# Step 3: Add simple reverb effect for realism
|
||||
# Reverb simulates sound reflections in environment
|
||||
if duration > 2:
|
||||
# Single delay reverb (100ms delay)
|
||||
delay_samples = int(0.1 * sample_rate)
|
||||
if delay_samples < num_samples:
|
||||
for ch in range(num_channels):
|
||||
delayed = np.zeros(num_samples)
|
||||
delayed[delay_samples:] = audio_data[:-delay_samples, ch] * 0.3
|
||||
audio_data[:, ch] += delayed
|
||||
|
||||
# Step 4: Normalize audio to prevent clipping
|
||||
# Find maximum amplitude and scale to 80% of maximum volume
|
||||
max_amp = np.max(np.abs(audio_data))
|
||||
max_amp = float(np.max(np.abs(audio_data)))
|
||||
if max_amp > 0:
|
||||
audio_data = audio_data / max_amp * 0.8
|
||||
audio_data = audio_data / max_amp * 0.95
|
||||
|
||||
audio_array = audio_data.copy()
|
||||
result = {
|
||||
"np_array": audio_array,
|
||||
audio_bytes: bytes | None = None
|
||||
output_path: str | None = None
|
||||
result: dict[str, Any] = {
|
||||
"np_array": audio_data.copy(),
|
||||
}
|
||||
|
||||
# Handle file saving
|
||||
audio_bytes = None
|
||||
|
||||
if save_to_file:
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
output_path = f"audio_{num_channels}ch_{timestamp}.wav"
|
||||
@@ -617,24 +709,38 @@ def convert_audio_to_text(audio_data):
|
||||
return text
|
||||
|
||||
|
||||
def convert_audio_file_to_text(output_path):
|
||||
def _whisper_transcribe_in_current_process(output_path: str) -> str:
|
||||
import whisper
|
||||
|
||||
model = whisper.load_model("small")
|
||||
text = model.transcribe(
|
||||
output_path,
|
||||
temperature=0.0,
|
||||
word_timestamps=True,
|
||||
condition_on_previous_text=False,
|
||||
)["text"]
|
||||
del model
|
||||
if torch.cuda.is_available():
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
model = whisper.load_model("small", device=device)
|
||||
try:
|
||||
text = model.transcribe(
|
||||
output_path,
|
||||
temperature=0.0,
|
||||
word_timestamps=True,
|
||||
condition_on_previous_text=False,
|
||||
)["text"]
|
||||
finally:
|
||||
# Sync GPU so in-flight ops finish before we free the model; otherwise
|
||||
# freed memory may not show up until those ops complete.
|
||||
if torch is not None and torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
del model
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
if text:
|
||||
return text
|
||||
else:
|
||||
return ""
|
||||
if torch is not None and torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return text or ""
|
||||
|
||||
|
||||
def convert_audio_file_to_text(output_path: str) -> str:
|
||||
"""Convert an audio file to text in an isolated subprocess."""
|
||||
# Import locally to avoid impacting test module import time.
|
||||
ctx = multiprocessing.get_context("spawn")
|
||||
with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor:
|
||||
future = executor.submit(_whisper_transcribe_in_current_process, output_path)
|
||||
return future.result()
|
||||
|
||||
|
||||
def merge_base64_and_convert_to_text(base64_list):
|
||||
@@ -1105,7 +1211,7 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any],
|
||||
assert response.success, "The request failed."
|
||||
e2e_latency = response.e2e_latency
|
||||
if e2e_latency is not None:
|
||||
print(f"the avg e2e is: {e2e_latency}")
|
||||
print(f"the avg e2e latency is: {e2e_latency}")
|
||||
|
||||
modalities = request_config.get("modalities", ["text", "audio"])
|
||||
|
||||
|
||||
@@ -156,13 +156,22 @@ class TestDiffusionParallelConfigHSDP:
|
||||
|
||||
def test_hsdp_cannot_use_with_tp(self):
|
||||
"""Test that HSDP and Tensor Parallelism cannot be used together."""
|
||||
with pytest.raises(ValueError, match="cannot be used with Tensor Parallelism"):
|
||||
with pytest.raises(ValueError, match="cannot be used with TP or DP"):
|
||||
DiffusionParallelConfig(
|
||||
tensor_parallel_size=2,
|
||||
use_hsdp=True,
|
||||
hsdp_shard_size=4,
|
||||
)
|
||||
|
||||
def test_hsdp_cannot_use_with_dp(self):
|
||||
"""Test that HSDP and Data Parallelism cannot be used together."""
|
||||
with pytest.raises(ValueError, match="cannot be used with TP or DP"):
|
||||
DiffusionParallelConfig(
|
||||
data_parallel_size=2,
|
||||
use_hsdp=True,
|
||||
hsdp_shard_size=4,
|
||||
)
|
||||
|
||||
def test_from_dict_with_hsdp(self):
|
||||
"""Test creating config from dict with HSDP settings."""
|
||||
config = DiffusionParallelConfig.from_dict(
|
||||
@@ -177,6 +186,158 @@ class TestDiffusionParallelConfigHSDP:
|
||||
assert config.hsdp_shard_size == 2 # auto: 4 // 2
|
||||
|
||||
|
||||
class TestStandaloneHSDPDetection:
|
||||
"""Tests for standalone HSDP detection and dit_parallel_size calculation.
|
||||
|
||||
These tests verify the logic used in initialize_model_parallel() to detect
|
||||
standalone HSDP mode and calculate effective parallel sizes.
|
||||
|
||||
Standalone HSDP is when all non-HSDP parallelism dimensions are 1.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def compute_standalone_hsdp_params(
|
||||
data_parallel_size: int = 1,
|
||||
cfg_parallel_size: int = 1,
|
||||
sequence_parallel_size: int = 1,
|
||||
pipeline_parallel_size: int = 1,
|
||||
tensor_parallel_size: int = 1,
|
||||
fully_shard_degree: int = 1,
|
||||
hsdp_replicate_size: int = 1,
|
||||
) -> dict:
|
||||
"""Compute standalone HSDP detection parameters.
|
||||
|
||||
This mirrors the logic in initialize_model_parallel().
|
||||
"""
|
||||
dit_parallel_size = (
|
||||
data_parallel_size
|
||||
* cfg_parallel_size
|
||||
* sequence_parallel_size
|
||||
* pipeline_parallel_size
|
||||
* tensor_parallel_size
|
||||
)
|
||||
|
||||
# Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
|
||||
is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1
|
||||
|
||||
# For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size)
|
||||
if is_standalone_hsdp:
|
||||
effective_dit_parallel_size = fully_shard_degree * hsdp_replicate_size
|
||||
else:
|
||||
effective_dit_parallel_size = dit_parallel_size
|
||||
|
||||
effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size
|
||||
|
||||
return {
|
||||
"original_dit_parallel_size": dit_parallel_size,
|
||||
"is_standalone_hsdp": is_standalone_hsdp,
|
||||
"effective_dit_parallel_size": effective_dit_parallel_size,
|
||||
"effective_dp_size": effective_dp_size,
|
||||
}
|
||||
|
||||
def test_standalone_hsdp_basic(self):
|
||||
"""Test basic standalone HSDP detection (shard_size=4, replicate=1)."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
fully_shard_degree=4,
|
||||
hsdp_replicate_size=1,
|
||||
)
|
||||
assert result["original_dit_parallel_size"] == 1
|
||||
assert result["is_standalone_hsdp"] is True
|
||||
assert result["effective_dit_parallel_size"] == 4
|
||||
assert result["effective_dp_size"] == 4
|
||||
|
||||
def test_standalone_hsdp_with_replicate(self):
|
||||
"""Test standalone HSDP with replication (shard_size=4, replicate=2)."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
fully_shard_degree=4,
|
||||
hsdp_replicate_size=2,
|
||||
)
|
||||
assert result["original_dit_parallel_size"] == 1
|
||||
assert result["is_standalone_hsdp"] is True
|
||||
assert result["effective_dit_parallel_size"] == 8 # 4 * 2
|
||||
assert result["effective_dp_size"] == 8
|
||||
|
||||
def test_combined_hsdp_sp_not_standalone(self):
|
||||
"""Test HSDP combined with SP is NOT detected as standalone.
|
||||
|
||||
This is a regression test for the bug where the condition
|
||||
`dit_parallel_size == fully_shard_degree` incorrectly matched
|
||||
combined modes like SP=4 + HSDP=4.
|
||||
"""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
sequence_parallel_size=4,
|
||||
fully_shard_degree=4,
|
||||
hsdp_replicate_size=1,
|
||||
)
|
||||
assert result["original_dit_parallel_size"] == 4
|
||||
assert result["is_standalone_hsdp"] is False
|
||||
# Should NOT override dp_size for combined mode
|
||||
assert result["effective_dp_size"] == 1 # original data_parallel_size
|
||||
|
||||
def test_combined_hsdp_cfg_not_standalone(self):
|
||||
"""Test HSDP combined with CFG is NOT detected as standalone."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
cfg_parallel_size=2,
|
||||
fully_shard_degree=4,
|
||||
hsdp_replicate_size=1,
|
||||
)
|
||||
assert result["original_dit_parallel_size"] == 2
|
||||
assert result["is_standalone_hsdp"] is False
|
||||
assert result["effective_dp_size"] == 1
|
||||
|
||||
def test_combined_hsdp_pp_not_standalone(self):
|
||||
"""Test HSDP combined with PP is NOT detected as standalone."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
pipeline_parallel_size=2,
|
||||
fully_shard_degree=4,
|
||||
hsdp_replicate_size=1,
|
||||
)
|
||||
assert result["original_dit_parallel_size"] == 2
|
||||
assert result["is_standalone_hsdp"] is False
|
||||
assert result["effective_dp_size"] == 1
|
||||
|
||||
def test_no_hsdp_not_standalone(self):
|
||||
"""Test that no HSDP (fully_shard_degree=1) is NOT standalone."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
fully_shard_degree=1,
|
||||
)
|
||||
assert result["original_dit_parallel_size"] == 1
|
||||
assert result["is_standalone_hsdp"] is False
|
||||
assert result["effective_dp_size"] == 1
|
||||
|
||||
def test_combined_multiple_parallelism_not_standalone(self):
|
||||
"""Test HSDP combined with multiple parallelism is NOT standalone."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
sequence_parallel_size=2,
|
||||
cfg_parallel_size=2,
|
||||
fully_shard_degree=4,
|
||||
hsdp_replicate_size=1,
|
||||
)
|
||||
assert result["original_dit_parallel_size"] == 4 # 2 * 2
|
||||
assert result["is_standalone_hsdp"] is False
|
||||
assert result["effective_dp_size"] == 1
|
||||
|
||||
def test_standalone_hsdp_large_shard(self):
|
||||
"""Test standalone HSDP with large shard size."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
fully_shard_degree=8,
|
||||
hsdp_replicate_size=1,
|
||||
)
|
||||
assert result["is_standalone_hsdp"] is True
|
||||
assert result["effective_dit_parallel_size"] == 8
|
||||
assert result["effective_dp_size"] == 8
|
||||
|
||||
def test_standalone_hsdp_large_replicate(self):
|
||||
"""Test standalone HSDP with large replicate size."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
fully_shard_degree=4,
|
||||
hsdp_replicate_size=4,
|
||||
)
|
||||
assert result["is_standalone_hsdp"] is True
|
||||
assert result["effective_dit_parallel_size"] == 16 # 4 * 4
|
||||
assert result["effective_dp_size"] == 16
|
||||
|
||||
|
||||
class TestHSDPShardConditions:
|
||||
"""Tests for _hsdp_shard_conditions matching logic."""
|
||||
|
||||
|
||||
@@ -96,6 +96,7 @@ class TestWorkerWrapperBaseInitialization:
|
||||
local_rank=0,
|
||||
rank=0,
|
||||
od_config=mock_od_config,
|
||||
skip_load_model=False,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
|
||||
OmniKVCacheConfig,
|
||||
OmniKVTransferManager,
|
||||
)
|
||||
from vllm_omni.distributed.omni_connectors.utils.kv_utils import normalize_layer_kv
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.cache]
|
||||
@@ -192,8 +193,7 @@ def test_normalize_layer_kv_rejects_invalid_inputs(kv_config, common_constants,
|
||||
else:
|
||||
layer_kv = (torch.randn(2, block_size, num_heads, head_dim), "not-a-tensor")
|
||||
|
||||
manager = OmniKVTransferManager(kv_config)
|
||||
normalized = manager._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0)
|
||||
normalized = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0)
|
||||
assert normalized is None
|
||||
|
||||
|
||||
|
||||
@@ -13,7 +13,7 @@ stage_args:
|
||||
model_arch: BagelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.35
|
||||
gpu_memory_utilization: 0.45
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: text
|
||||
@@ -46,7 +46,7 @@ stage_args:
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: dit
|
||||
gpu_memory_utilization: 0.55
|
||||
gpu_memory_utilization: 0.45
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: image
|
||||
|
||||
@@ -13,7 +13,7 @@ stage_args:
|
||||
model_arch: BagelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.35
|
||||
gpu_memory_utilization: 0.45
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: text
|
||||
@@ -45,7 +45,7 @@ stage_args:
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: dit
|
||||
gpu_memory_utilization: 0.55
|
||||
gpu_memory_utilization: 0.45
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: image
|
||||
|
||||
@@ -21,6 +21,7 @@ stage_args:
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
enable_prefix_caching: false
|
||||
mm_processor_cache_gb: 0
|
||||
is_comprehension: true
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import gc
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
@@ -20,11 +21,18 @@ models = ["riverclouds/qwen_image_random"]
|
||||
|
||||
|
||||
def inference(model_name: str, offload: bool = True):
|
||||
gc.collect()
|
||||
current_omni_platform.empty_cache()
|
||||
device_index = torch.cuda.current_device()
|
||||
torch.cuda.reset_peak_memory_stats(device=device_index)
|
||||
monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
|
||||
monitor.start()
|
||||
m = Omni(model=model_name, enable_cpu_offload=offload)
|
||||
m = Omni(
|
||||
model=model_name,
|
||||
# TODO: we might want to add overlapped feature e2e tests
|
||||
# cache_backend="cache_dit",
|
||||
enable_cpu_offload=offload,
|
||||
)
|
||||
torch.cuda.reset_peak_memory_stats(device=device_index)
|
||||
height = 256
|
||||
width = 256
|
||||
@@ -42,6 +50,10 @@ def inference(model_name: str, offload: bool = True):
|
||||
peak = monitor.peak_used_mb
|
||||
monitor.stop()
|
||||
|
||||
del m
|
||||
gc.collect()
|
||||
current_omni_platform.empty_cache()
|
||||
|
||||
return peak
|
||||
|
||||
|
||||
@@ -51,9 +63,9 @@ def inference(model_name: str, offload: bool = True):
|
||||
@pytest.mark.parametrize("model_name", models)
|
||||
def test_cpu_offload_diffusion_model(model_name: str):
|
||||
try:
|
||||
no_offload_peak_memory = inference(model_name, offload=False)
|
||||
cleanup_dist_env_and_memory()
|
||||
offload_peak_memory = inference(model_name, offload=True)
|
||||
cleanup_dist_env_and_memory()
|
||||
no_offload_peak_memory = inference(model_name, offload=False)
|
||||
except Exception:
|
||||
pytest.fail("Inference failed")
|
||||
print(f"Offload peak memory: {offload_peak_memory} MB")
|
||||
|
||||
@@ -38,6 +38,8 @@ def run_inference(
|
||||
m = Omni(
|
||||
model=model_name,
|
||||
enable_layerwise_offload=layerwise_offload,
|
||||
# TODO: we might want to add overlapped feature e2e tests
|
||||
# cache_backend="cache_dit",
|
||||
boundary_ratio=0.875,
|
||||
flow_shift=5.0,
|
||||
)
|
||||
|
||||
@@ -10,12 +10,29 @@ from tests.conftest import (
|
||||
generate_synthetic_audio,
|
||||
generate_synthetic_image,
|
||||
generate_synthetic_video,
|
||||
modify_stage_config,
|
||||
)
|
||||
from tests.utils import hardware_test
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
models = ["Qwen/Qwen2.5-Omni-7B"]
|
||||
|
||||
|
||||
def get_cuda_graph_config():
|
||||
path = modify_stage_config(
|
||||
str(Path(__file__).parent.parent / "stage_configs" / "qwen2_5_omni_ci.yaml"),
|
||||
updates={
|
||||
"stage_args": {
|
||||
0: {
|
||||
"engine_args.enforce_eager": "true",
|
||||
},
|
||||
1: {"engine_args.enforce_eager": "true"},
|
||||
},
|
||||
},
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
# CI stage config optimized for 24GB GPU (L4/RTX3090) or NPU
|
||||
if current_omni_platform.is_npu():
|
||||
stage_config = str(Path(__file__).parent / "stage_configs" / "npu" / "qwen2_5_omni_ci.yaml")
|
||||
@@ -23,7 +40,7 @@ elif current_omni_platform.is_rocm():
|
||||
# ROCm stage config optimized for MI325 GPU
|
||||
stage_config = str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml")
|
||||
else:
|
||||
stage_config = str(Path(__file__).parent.parent / "stage_configs" / "qwen2_5_omni_ci.yaml")
|
||||
stage_config = get_cuda_graph_config()
|
||||
|
||||
# Create parameter combinations for model and stage config
|
||||
test_params = [(model, stage_config) for model in models]
|
||||
|
||||
@@ -13,18 +13,35 @@ import pytest
|
||||
|
||||
from tests.conftest import (
|
||||
generate_synthetic_video,
|
||||
modify_stage_config,
|
||||
)
|
||||
from tests.utils import hardware_test
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
|
||||
|
||||
|
||||
def get_cuda_graph_config():
|
||||
path = modify_stage_config(
|
||||
str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml"),
|
||||
updates={
|
||||
"stage_args": {
|
||||
0: {
|
||||
"engine_args.enforce_eager": "true",
|
||||
},
|
||||
1: {"engine_args.enforce_eager": "true"},
|
||||
},
|
||||
},
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
# CI stage config for 2xH100-80G GPUs or AMD GPU MI325
|
||||
if current_omni_platform.is_rocm():
|
||||
# ROCm stage config optimized for MI325 GPU
|
||||
stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
|
||||
else:
|
||||
stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")]
|
||||
stage_configs = [get_cuda_graph_config()]
|
||||
|
||||
# Create parameter combinations for model and stage config
|
||||
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
|
||||
|
||||
@@ -114,7 +114,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
"messages": messages,
|
||||
"stream": True,
|
||||
"key_words": {
|
||||
"audio": ["water", "chirping", "crackling", "rain"],
|
||||
"audio": ["test"],
|
||||
"image": ["square", "quadrate"],
|
||||
},
|
||||
}
|
||||
@@ -126,13 +126,14 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_text_to_text_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Test text input processing and text/audio output generation via OpenAI API.
|
||||
Deploy Setting: default yaml
|
||||
Input Modal: text
|
||||
Output Modal: text + audio
|
||||
Output Modal: text
|
||||
Datasets: few requests
|
||||
"""
|
||||
messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
|
||||
|
||||
@@ -22,7 +22,7 @@ from tests.utils import hardware_test
|
||||
|
||||
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
|
||||
|
||||
AUDIO_KEY = ["water", "chirping", "crackling", "rain"]
|
||||
AUDIO_KEY = ["test"]
|
||||
IMAGE_KEY = ["square", "quadrate"]
|
||||
VIDEO_KEY = ["sphere", "globe", "circle", "round", "ball"]
|
||||
|
||||
@@ -90,12 +90,19 @@ def get_max_batch_size(size_type="few"):
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_text_to_audio_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Input Modal: text
|
||||
Output Modal: audio
|
||||
Input Setting: stream=True
|
||||
Datasets: single request
|
||||
"""
|
||||
messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"modalities": ["audio"],
|
||||
"stream": True,
|
||||
"key_words": {"text": ["beijing"]},
|
||||
}
|
||||
|
||||
@@ -107,6 +114,12 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_text_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Input Modal: text
|
||||
Output Modal: text, audio
|
||||
Input Setting: stream=False
|
||||
Datasets: few requests
|
||||
"""
|
||||
messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
|
||||
|
||||
request_config = {
|
||||
@@ -115,7 +128,7 @@ def test_text_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
"key_words": {"text": ["beijing"]},
|
||||
}
|
||||
|
||||
openai_client.send_request(request_config)
|
||||
openai_client.send_request(request_config, request_num=get_max_batch_size())
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@@ -123,6 +136,12 @@ def test_text_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_image_to_text_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Input Modal: image
|
||||
Output Modal: text
|
||||
Input Setting: stream=True
|
||||
Datasets: single request
|
||||
"""
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
|
||||
|
||||
@@ -130,6 +149,7 @@ def test_image_to_text_001(omni_server, openai_client) -> None:
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"modalities": ["text"],
|
||||
"stream": True,
|
||||
"key_words": {"image": IMAGE_KEY},
|
||||
}
|
||||
|
||||
@@ -141,6 +161,12 @@ def test_image_to_text_001(omni_server, openai_client) -> None:
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_image_to_audio_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Input Modal: image
|
||||
Output Modal: audio
|
||||
Input Setting: stream=False
|
||||
Datasets: single request
|
||||
"""
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
|
||||
|
||||
@@ -159,6 +185,12 @@ def test_image_to_audio_001(omni_server, openai_client) -> None:
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_image_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Input Modal: image
|
||||
Output Modal: text, audio
|
||||
Input Setting: stream=False
|
||||
Datasets: few requests
|
||||
"""
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
|
||||
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
|
||||
@@ -177,6 +209,12 @@ def test_image_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_video_to_text_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Input Modal: video
|
||||
Output Modal: text
|
||||
Input Setting: stream=False
|
||||
Datasets: single request
|
||||
"""
|
||||
video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(video_data_url=video_data_url)
|
||||
|
||||
@@ -195,6 +233,12 @@ def test_video_to_text_001(omni_server, openai_client) -> None:
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_video_to_audio_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Input Modal: video
|
||||
Output Modal: audio
|
||||
Input Setting: stream=False
|
||||
Datasets: single request
|
||||
"""
|
||||
video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(video_data_url=video_data_url)
|
||||
|
||||
@@ -213,6 +257,12 @@ def test_video_to_audio_001(omni_server, openai_client) -> None:
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_video_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Input Modal: video
|
||||
Output Modal: text, audio
|
||||
Input Setting: stream=False
|
||||
Datasets: few requests
|
||||
"""
|
||||
video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
|
||||
|
||||
messages = dummy_messages_from_mix_data(video_data_url=video_data_url)
|
||||
@@ -231,6 +281,12 @@ def test_video_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Input Modal: text, audio
|
||||
Output Modal: text, audio
|
||||
Input Setting: stream=False
|
||||
Datasets: single request
|
||||
"""
|
||||
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(
|
||||
audio_data_url=audio_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_audio")
|
||||
@@ -239,7 +295,6 @@ def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"modalities": ["text"],
|
||||
"key_words": {"audio": AUDIO_KEY},
|
||||
}
|
||||
|
||||
@@ -251,6 +306,12 @@ def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Input Modal: text, image
|
||||
Output Modal: text, audio
|
||||
Input Setting: stream=False
|
||||
Datasets: single request
|
||||
"""
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
|
||||
|
||||
messages = dummy_messages_from_mix_data(
|
||||
@@ -260,19 +321,23 @@ def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"modalities": ["audio"],
|
||||
"key_words": {"image": IMAGE_KEY},
|
||||
}
|
||||
|
||||
openai_client.send_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="There is a known issue with oom error.")
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Input Modal: text, video
|
||||
Output Modal: text, audio
|
||||
Input Setting: stream=True
|
||||
Datasets: single requests
|
||||
"""
|
||||
video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(1280, 720, 30)['base64']}"
|
||||
|
||||
messages = dummy_messages_from_mix_data(
|
||||
@@ -282,10 +347,11 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"stream": True,
|
||||
"key_words": {"video": VIDEO_KEY},
|
||||
}
|
||||
|
||||
openai_client.send_request(request_config, request_num=get_max_batch_size())
|
||||
openai_client.send_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="There is a known issue with shape mismatch error.")
|
||||
@@ -294,6 +360,12 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Input Modal: text, audio, image, video
|
||||
Output Modal: text, audio
|
||||
Input Setting: stream=True
|
||||
Datasets: few requests
|
||||
"""
|
||||
video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
|
||||
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
|
||||
|
||||
@@ -18,6 +18,7 @@ stage_args:
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
enable_prefix_caching: false
|
||||
mm_processor_cache_gb: 0
|
||||
is_comprehension: true
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
|
||||
@@ -22,6 +22,7 @@ stage_args:
|
||||
max_num_batched_tokens: 32768
|
||||
max_model_len: 32768
|
||||
enable_prefix_caching: false
|
||||
mm_processor_cache_gb: 0
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 1
|
||||
load_format: dummy
|
||||
|
||||
@@ -22,6 +22,7 @@ stage_args:
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
enable_prefix_caching: false
|
||||
mm_processor_cache_gb: 0
|
||||
is_comprehension: true
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
|
||||
@@ -20,6 +20,7 @@ stage_args:
|
||||
engine_output_type: latent # Output hidden states for talker
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
mm_processor_cache_gb: 0
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 1
|
||||
load_format: dummy
|
||||
|
||||
@@ -1,11 +1,15 @@
|
||||
# tests/entrypoints/openai/test_serving_speech.py
|
||||
import logging
|
||||
import os
|
||||
from inspect import Signature, signature
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from fastapi import FastAPI
|
||||
from fastapi import FastAPI, HTTPException, UploadFile
|
||||
from fastapi.params import File, Form
|
||||
from fastapi.testclient import TestClient
|
||||
from pydantic import ValidationError
|
||||
from pytest_mock import MockerFixture
|
||||
@@ -206,10 +210,52 @@ def test_app(mocker: MockerFixture):
|
||||
# Add list_voices endpoint
|
||||
async def list_voices():
|
||||
speakers = sorted(speech_server.supported_speakers) if speech_server.supported_speakers else []
|
||||
return {"voices": speakers}
|
||||
uploaded_voices = []
|
||||
if hasattr(speech_server, "uploaded_speakers"):
|
||||
for voice_name, info in speech_server.uploaded_speakers.items():
|
||||
uploaded_voices.append(
|
||||
{
|
||||
"name": info.get("name", voice_name),
|
||||
"consent": info.get("consent", ""),
|
||||
"created_at": info.get("created_at", 0),
|
||||
"file_size": info.get("file_size", 0),
|
||||
"mime_type": info.get("mime_type", ""),
|
||||
}
|
||||
)
|
||||
return {"voices": speakers, "uploaded_voices": uploaded_voices}
|
||||
|
||||
app.add_api_route("/v1/audio/voices", list_voices, methods=["GET"])
|
||||
|
||||
# Add upload_voice endpoint
|
||||
async def upload_voice(audio_sample: UploadFile = File(...), consent: str = Form(...), name: str = Form(...)):
|
||||
try:
|
||||
result = await speech_server.upload_voice(audio_sample, consent, name)
|
||||
return {"success": True, "voice": result}
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except Exception as e:
|
||||
logger.exception(f"Failed to upload voice: {e}")
|
||||
raise HTTPException(status_code=500, detail=f"Failed to upload voice: {str(e)}")
|
||||
|
||||
app.add_api_route("/v1/audio/voices", upload_voice, methods=["POST"])
|
||||
|
||||
# Add delete_voice endpoint
|
||||
async def delete_voice(name: str):
|
||||
try:
|
||||
success = await speech_server.delete_voice(name)
|
||||
if not success:
|
||||
raise HTTPException(status_code=404, detail=f"Voice '{name}' not found")
|
||||
return {"success": True, "message": f"Voice '{name}' deleted successfully"}
|
||||
except HTTPException:
|
||||
raise
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except Exception as e:
|
||||
logger.exception(f"Failed to delete voice '{name}': {e}")
|
||||
raise HTTPException(status_code=500, detail=f"Failed to delete voice: {str(e)}")
|
||||
|
||||
app.add_api_route("/v1/audio/voices/{name}", delete_voice, methods=["DELETE"])
|
||||
|
||||
return app
|
||||
|
||||
|
||||
@@ -284,6 +330,149 @@ class TestSpeechAPI:
|
||||
assert response.status_code == 200
|
||||
assert "voices" in response.json()
|
||||
|
||||
def test_upload_voice_success(self, client, tmp_path):
|
||||
"""Test successful voice upload."""
|
||||
# Create a mock audio file
|
||||
audio_content = b"fake audio content" * 1000 # ~17KB
|
||||
files = {
|
||||
"audio_sample": ("test.wav", audio_content, "audio/wav"),
|
||||
}
|
||||
data = {
|
||||
"consent": "user_consent_123",
|
||||
"name": "test_voice",
|
||||
}
|
||||
|
||||
response = client.post("/v1/audio/voices", files=files, data=data)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["success"] is True
|
||||
assert "voice" in result
|
||||
voice_info = result["voice"]
|
||||
assert voice_info["name"] == "test_voice"
|
||||
assert voice_info["consent"] == "user_consent_123"
|
||||
assert "created_at" in voice_info
|
||||
assert voice_info["mime_type"] == "audio/wav"
|
||||
assert voice_info["file_size"] == len(audio_content)
|
||||
response = client.delete("/v1/audio/voices/test_voice")
|
||||
|
||||
def test_upload_voice_file_too_large(self, client):
|
||||
"""Test voice upload with file exceeding size limit."""
|
||||
# Create a file larger than 10MB
|
||||
audio_content = b"x" * (11 * 1024 * 1024) # 11MB
|
||||
files = {
|
||||
"audio_sample": ("test.wav", audio_content, "audio/wav"),
|
||||
}
|
||||
data = {
|
||||
"consent": "user_consent_123",
|
||||
"name": "test_voice",
|
||||
}
|
||||
|
||||
response = client.post("/v1/audio/voices", files=files, data=data)
|
||||
assert response.status_code == 400
|
||||
result = response.json()
|
||||
assert "detail" in result
|
||||
assert "10MB" in result["detail"]
|
||||
|
||||
def test_upload_voice_invalid_mime_type(self, client):
|
||||
"""Test voice upload with invalid MIME type."""
|
||||
audio_content = b"fake audio content"
|
||||
files = {
|
||||
"audio_sample": ("test.txt", audio_content, "text/plain"),
|
||||
}
|
||||
data = {
|
||||
"consent": "user_consent_123",
|
||||
"name": "test_voice",
|
||||
}
|
||||
|
||||
response = client.post("/v1/audio/voices", files=files, data=data)
|
||||
assert response.status_code == 400
|
||||
result = response.json()
|
||||
assert "detail" in result
|
||||
assert "MIME type" in result["detail"]
|
||||
|
||||
def test_upload_voice_name_collision(self, client):
|
||||
"""Test voice upload with duplicate name."""
|
||||
# First upload
|
||||
audio_content = b"fake audio content"
|
||||
files = {
|
||||
"audio_sample": ("test.wav", audio_content, "audio/wav"),
|
||||
}
|
||||
data = {
|
||||
"consent": "user_consent_123",
|
||||
"name": "test_voice",
|
||||
}
|
||||
|
||||
response = client.post("/v1/audio/voices", files=files, data=data)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Second upload with same name
|
||||
response = client.post("/v1/audio/voices", files=files, data=data)
|
||||
assert response.status_code == 400
|
||||
result = response.json()
|
||||
assert "detail" in result
|
||||
assert "already exists" in result["detail"]
|
||||
response = client.delete("/v1/audio/voices/test_voice")
|
||||
|
||||
def test_upload_voice_missing_parameters(self, client):
|
||||
"""Test voice upload with missing required parameters."""
|
||||
audio_content = b"fake audio content"
|
||||
files = {
|
||||
"audio_sample": ("test.wav", audio_content, "audio/wav"),
|
||||
}
|
||||
|
||||
# Missing consent
|
||||
data = {"name": "test_voice5"}
|
||||
response = client.post("/v1/audio/voices", files=files, data=data)
|
||||
assert response.status_code == 422 # Validation error
|
||||
|
||||
# Missing name
|
||||
data = {"consent": "user_consent_123"}
|
||||
response = client.post("/v1/audio/voices", files=files, data=data)
|
||||
assert response.status_code == 422 # Validation error
|
||||
|
||||
# Missing file
|
||||
data = {
|
||||
"consent": "user_consent_123",
|
||||
"name": "test_voice6",
|
||||
}
|
||||
response = client.post("/v1/audio/voices", data=data)
|
||||
assert response.status_code == 422 # Validation error
|
||||
|
||||
def test_delete_voice_success(self, client):
|
||||
"""Test successful voice deletion."""
|
||||
# First upload a voice
|
||||
audio_content = b"fake audio content"
|
||||
files = {
|
||||
"audio_sample": ("test.wav", audio_content, "audio/wav"),
|
||||
}
|
||||
data = {
|
||||
"consent": "user_consent_123",
|
||||
"name": "test_voice7",
|
||||
}
|
||||
|
||||
response = client.post("/v1/audio/voices", files=files, data=data)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Then delete it
|
||||
response = client.delete("/v1/audio/voices/test_voice7")
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["success"] is True
|
||||
assert "deleted successfully" in result["message"]
|
||||
|
||||
# Verify it's gone by trying to delete again
|
||||
response = client.delete("/v1/audio/voices/test_voice7")
|
||||
assert response.status_code == 404
|
||||
result = response.json()
|
||||
assert "not found" in result["detail"]
|
||||
|
||||
def test_delete_voice_not_found(self, client):
|
||||
"""Test deleting a non-existent voice."""
|
||||
response = client.delete("/v1/audio/voices/nonexistent")
|
||||
assert response.status_code == 404
|
||||
result = response.json()
|
||||
assert "not found" in result["detail"]
|
||||
|
||||
|
||||
class TestTTSMethods:
|
||||
"""Unit tests for TTS validation and parameter building."""
|
||||
@@ -438,6 +627,121 @@ class TestTTSMethods:
|
||||
# Verify speakers are normalized to lowercase
|
||||
assert server.supported_speakers == {"ryan", "vivian", "aiden"}
|
||||
|
||||
def test_build_tts_params_with_uploaded_voice(self, speech_server):
|
||||
"""Test _build_tts_params auto-sets ref_audio for uploaded voices."""
|
||||
# Mock an uploaded speaker
|
||||
speech_server.uploaded_speakers = {
|
||||
"custom_voice": {
|
||||
"name": "custom_voice",
|
||||
"file_path": "/tmp/voice_samples/custom_voice_consent_123.wav",
|
||||
"mime_type": "audio/wav",
|
||||
}
|
||||
}
|
||||
speech_server.supported_speakers = {"ryan", "vivian", "custom_voice"}
|
||||
|
||||
# Mock _get_uploaded_audio_data to return base64 data
|
||||
with patch.object(speech_server, "_get_uploaded_audio_data") as mock_get_audio:
|
||||
mock_get_audio.return_value = "data:audio/wav;base64,ZmFrZWF1ZGlv"
|
||||
|
||||
req = OpenAICreateSpeechRequest(input="Hello", voice="custom_voice", task_type="Base")
|
||||
|
||||
params = speech_server._build_tts_params(req)
|
||||
|
||||
# Verify ref_audio was auto-set
|
||||
assert "ref_audio" in params
|
||||
assert params["ref_audio"] == ["data:audio/wav;base64,ZmFrZWF1ZGlv"]
|
||||
assert "x_vector_only_mode" in params
|
||||
assert params["x_vector_only_mode"] == [True]
|
||||
mock_get_audio.assert_called_once_with("custom_voice")
|
||||
|
||||
def test_build_tts_params_without_uploaded_voice(self, speech_server):
|
||||
"""Test _build_tts_params does not auto-set ref_audio for non-uploaded voices."""
|
||||
# No uploaded speakers
|
||||
speech_server.uploaded_speakers = {}
|
||||
speech_server.supported_speakers = {"ryan", "vivian"}
|
||||
|
||||
req = OpenAICreateSpeechRequest(input="Hello", voice="ryan", task_type="Base")
|
||||
|
||||
params = speech_server._build_tts_params(req)
|
||||
|
||||
# Verify ref_audio was NOT auto-set
|
||||
assert "ref_audio" not in params
|
||||
assert "x_vector_only_mode" not in params
|
||||
|
||||
def test_build_tts_params_with_explicit_ref_audio(self, speech_server):
|
||||
"""Test _build_tts_params uses explicit ref_audio even for uploaded voices."""
|
||||
# Mock an uploaded speaker
|
||||
speech_server.uploaded_speakers = {
|
||||
"custom_voice": {
|
||||
"name": "custom_voice",
|
||||
"file_path": "/tmp/voice_samples/custom_voice_consent_123.wav",
|
||||
"mime_type": "audio/wav",
|
||||
}
|
||||
}
|
||||
speech_server.supported_speakers = {"ryan", "vivian", "custom_voice"}
|
||||
|
||||
req = OpenAICreateSpeechRequest(
|
||||
input="Hello", voice="custom_voice", task_type="Base", ref_audio="data:audio/wav;base64,ZXhwbGljaXQ="
|
||||
)
|
||||
|
||||
params = speech_server._build_tts_params(req)
|
||||
|
||||
# _build_tts_params should NOT auto-set ref_audio when explicit ref_audio
|
||||
# is provided (request.ref_audio is not None skips the auto-set branch).
|
||||
# The explicit ref_audio is resolved later in create_speech() via
|
||||
# _resolve_ref_audio(), not in _build_tts_params().
|
||||
assert "ref_audio" not in params
|
||||
# x_vector_only_mode should not be set when explicit ref_audio is provided
|
||||
assert "x_vector_only_mode" not in params
|
||||
|
||||
def test_get_uploaded_audio_data(self, speech_server):
|
||||
"""Test _get_uploaded_audio_data function."""
|
||||
# Mock file operations
|
||||
with (
|
||||
patch("builtins.open", create=True) as mock_open,
|
||||
patch("base64.b64encode") as mock_b64encode,
|
||||
patch("pathlib.Path.exists") as mock_exists,
|
||||
):
|
||||
mock_exists.return_value = True
|
||||
mock_b64encode.return_value = b"ZmFrZWF1ZGlv"
|
||||
|
||||
# Setup mock file
|
||||
mock_file = MagicMock()
|
||||
mock_file.read.return_value = b"fakeaudio"
|
||||
mock_open.return_value.__enter__.return_value = mock_file
|
||||
|
||||
# Setup uploaded speaker
|
||||
speech_server.uploaded_speakers = {
|
||||
"test_voice": {"name": "test_voice", "file_path": "/tmp/test.wav", "mime_type": "audio/wav"}
|
||||
}
|
||||
result = speech_server._get_uploaded_audio_data("test_voice")
|
||||
|
||||
assert result == "data:audio/wav;base64,ZmFrZWF1ZGlv"
|
||||
mock_open.assert_called_once_with(Path("/tmp/test.wav"), "rb")
|
||||
mock_b64encode.assert_called_once_with(b"fakeaudio")
|
||||
|
||||
def test_get_uploaded_audio_data_missing_file(self, speech_server):
|
||||
"""Test _get_uploaded_audio_data when file is missing."""
|
||||
with patch("pathlib.Path.exists") as mock_exists:
|
||||
mock_exists.return_value = False
|
||||
|
||||
# Setup uploaded speaker
|
||||
speech_server.uploaded_speakers = {
|
||||
"test_voice": {"name": "test_voice", "file_path": "/tmp/test.wav", "mime_type": "audio/wav"}
|
||||
}
|
||||
|
||||
result = speech_server._get_uploaded_audio_data("test_voice")
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_get_uploaded_audio_data_voice_not_found(self, speech_server):
|
||||
"""Test _get_uploaded_audio_data when voice is not in uploaded_speakers."""
|
||||
speech_server.uploaded_speakers = {}
|
||||
|
||||
result = speech_server._get_uploaded_audio_data("nonexistent")
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_max_instructions_length_default(self, speech_server):
|
||||
"""Test default max instructions length (500) when no config provided."""
|
||||
# Fixture creates server with no CLI override and no TTS stage
|
||||
@@ -544,6 +848,70 @@ class TestTTSMethods:
|
||||
assert "max 10 characters" in error
|
||||
|
||||
|
||||
class TestFileValidationFunctions:
|
||||
"""Unit tests for file validation helper functions."""
|
||||
|
||||
def test_sanitize_filename(self):
|
||||
"""Test _sanitize_filename function."""
|
||||
from vllm_omni.entrypoints.openai.serving_speech import _sanitize_filename
|
||||
|
||||
# Test normal filenames
|
||||
assert _sanitize_filename("test.wav") == "test.wav"
|
||||
assert _sanitize_filename("test-file.mp3") == "test-file.mp3"
|
||||
assert _sanitize_filename("test_file.flac") == "test_file.flac"
|
||||
|
||||
# Test path traversal attempts
|
||||
assert _sanitize_filename("../../../etc/passwd") == "passwd"
|
||||
assert _sanitize_filename("/absolute/path/file.wav") == "file.wav"
|
||||
|
||||
# Test special characters
|
||||
assert _sanitize_filename("file with spaces.wav") == "file_with_spaces.wav"
|
||||
assert _sanitize_filename("file&with&special&chars.wav") == "file_with_special_chars.wav"
|
||||
assert _sanitize_filename("file@with#special$chars%.wav") == "file_with_special_chars_.wav"
|
||||
|
||||
# Test empty filename
|
||||
assert _sanitize_filename("") == "file"
|
||||
|
||||
# Test very long filename
|
||||
long_name = "a" * 300
|
||||
sanitized = _sanitize_filename(long_name)
|
||||
assert len(sanitized) == 255
|
||||
assert sanitized.startswith("a")
|
||||
|
||||
def test_validate_path_within_directory(self, tmp_path):
|
||||
"""Test _validate_path_within_directory function."""
|
||||
from vllm_omni.entrypoints.openai.serving_speech import _validate_path_within_directory
|
||||
|
||||
# Create test directory structure
|
||||
base_dir = tmp_path / "uploads"
|
||||
base_dir.mkdir()
|
||||
|
||||
# Valid paths within directory
|
||||
valid_file = base_dir / "test.wav"
|
||||
valid_subdir_file = base_dir / "subdir" / "test.wav"
|
||||
valid_subdir_file.parent.mkdir()
|
||||
|
||||
assert _validate_path_within_directory(valid_file, base_dir) is True
|
||||
assert _validate_path_within_directory(valid_subdir_file, base_dir) is True
|
||||
|
||||
# Invalid paths outside directory
|
||||
outside_file = tmp_path / "outside.wav"
|
||||
assert _validate_path_within_directory(outside_file, base_dir) is False
|
||||
|
||||
# Test with symlink (should fail)
|
||||
if hasattr(os, "symlink"):
|
||||
link_target = tmp_path / "target.wav"
|
||||
link_target.touch()
|
||||
symlink = base_dir / "link.wav"
|
||||
os.symlink(link_target, symlink)
|
||||
# Symlinks to outside should be rejected
|
||||
assert _validate_path_within_directory(symlink, base_dir) is False
|
||||
|
||||
# Test with non-existent file (should still validate path)
|
||||
non_existent = base_dir / "nonexistent.wav"
|
||||
assert _validate_path_within_directory(non_existent, base_dir) is True
|
||||
|
||||
|
||||
class TestStreamingProtocolValidation:
|
||||
"""Unit tests for the stream field validators in OpenAICreateSpeechRequest."""
|
||||
|
||||
|
||||
@@ -11,12 +11,18 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
MODEL = "riverclouds/qwen_image_random"
|
||||
|
||||
|
||||
def _noop_inline_engine(self, model, stage_config, kwargs):
|
||||
self._inline_diffusion = False
|
||||
self._inline_engine = None
|
||||
|
||||
|
||||
def test_default_stage_config_includes_cache_backend(monkeypatch):
|
||||
"""Ensure cache_backend/cache_config are preserved in default diffusion stage."""
|
||||
monkeypatch.setattr(utils_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
|
||||
monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
|
||||
monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
|
||||
monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
|
||||
monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
|
||||
|
||||
omni = AsyncOmni(
|
||||
model=MODEL,
|
||||
@@ -47,6 +53,7 @@ def test_default_cache_config_used_when_missing(monkeypatch):
|
||||
monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
|
||||
monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
|
||||
monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
|
||||
monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
|
||||
|
||||
omni = AsyncOmni(
|
||||
model=MODEL,
|
||||
@@ -65,6 +72,7 @@ def test_default_stage_devices_from_sequence_parallel(monkeypatch):
|
||||
monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
|
||||
monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
|
||||
monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
|
||||
monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
|
||||
|
||||
omni = AsyncOmni(
|
||||
model=MODEL,
|
||||
|
||||
@@ -85,7 +85,7 @@ stage_args:
|
||||
engine_output_type: audio # Final output: audio waveform
|
||||
gpu_memory_utilization: 0.1
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 1000000
|
||||
max_num_batched_tokens: 100000
|
||||
hf_config_name: thinker_config
|
||||
engine_input_source: [1]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
|
||||
|
||||
@@ -111,11 +111,12 @@ class DiffusionParallelConfig:
|
||||
# 1. Standalone: when other parallelism is all 1, HSDP determines world_size
|
||||
# 2. Combined: HSDP overlays on top of other parallelism
|
||||
if self.use_hsdp:
|
||||
if self.tensor_parallel_size > 1:
|
||||
if self.tensor_parallel_size > 1 or self.data_parallel_size > 1:
|
||||
raise ValueError(
|
||||
"HSDP (use_hsdp=True) cannot be used with Tensor Parallelism "
|
||||
f"(tensor_parallel_size={self.tensor_parallel_size}). "
|
||||
"Set tensor_parallel_size=1 when using HSDP."
|
||||
"HSDP (use_hsdp=True) cannot be used with TP or DP "
|
||||
f"(tensor_parallel_size={self.tensor_parallel_size}, "
|
||||
f"data_parallel_size={self.data_parallel_size}). "
|
||||
"Set tensor_parallel_size=1 and data_parallel_size=1 when using HSDP."
|
||||
)
|
||||
if self.hsdp_shard_size == -1:
|
||||
# Auto-calculate: use other_parallel_world_size as shard_size
|
||||
@@ -618,6 +619,10 @@ class DiffusionOutput:
|
||||
|
||||
post_process_func: Callable[..., Any] | None = None
|
||||
|
||||
# Extra custom output data (e.g. latent trajectories, prompt embeds)
|
||||
# passed through to OmniRequestOutput.custom_output
|
||||
custom_output: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# logged timings info, directly from Req.timings
|
||||
# timings: Optional["RequestTimings"] = None
|
||||
|
||||
|
||||
@@ -67,18 +67,19 @@ class DiffusionEngine:
|
||||
raise e
|
||||
|
||||
def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
|
||||
diffusion_engine_start_time = time.time()
|
||||
diffusion_engine_start_time = time.perf_counter()
|
||||
|
||||
# Apply pre-processing if available
|
||||
preprocess_time = 0.0
|
||||
if self.pre_process_func is not None:
|
||||
preprocess_start_time = time.time()
|
||||
preprocess_start_time = time.perf_counter()
|
||||
request = self.pre_process_func(request)
|
||||
preprocess_time = time.time() - preprocess_start_time
|
||||
preprocess_time = time.perf_counter() - preprocess_start_time
|
||||
logger.info(f"Pre-processing completed in {preprocess_time:.4f} seconds")
|
||||
|
||||
exec_start_time = time.time()
|
||||
exec_start_time = time.perf_counter()
|
||||
output = self.add_req_and_wait_for_response(request)
|
||||
exec_total_time = time.time() - exec_start_time
|
||||
exec_total_time = time.perf_counter() - exec_start_time
|
||||
|
||||
if output.error:
|
||||
raise Exception(f"{output.error}")
|
||||
@@ -97,15 +98,25 @@ class DiffusionEngine:
|
||||
for i, prompt in enumerate(request.prompts)
|
||||
]
|
||||
|
||||
postprocess_start_time = time.time()
|
||||
postprocess_start_time = time.perf_counter()
|
||||
outputs = self.post_process_func(output.output) if self.post_process_func is not None else output.output
|
||||
audio_payload = None
|
||||
if isinstance(outputs, dict):
|
||||
audio_payload = outputs.get("audio")
|
||||
outputs = outputs.get("video", outputs)
|
||||
postprocess_time = time.time() - postprocess_start_time
|
||||
postprocess_time = time.perf_counter() - postprocess_start_time
|
||||
logger.info(f"Post-processing completed in {postprocess_time:.4f} seconds")
|
||||
|
||||
step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000
|
||||
logger.info(
|
||||
"DiffusionEngine.step breakdown: preprocess=%.2f ms, "
|
||||
"add_req_and_wait=%.2f ms, postprocess=%.2f ms, total=%.2f ms",
|
||||
preprocess_time * 1000,
|
||||
exec_total_time * 1000,
|
||||
postprocess_time * 1000,
|
||||
step_total_ms,
|
||||
)
|
||||
|
||||
# Convert to OmniRequestOutput format
|
||||
# Ensure outputs is a list
|
||||
if not isinstance(outputs, list):
|
||||
@@ -113,19 +124,16 @@ class DiffusionEngine:
|
||||
|
||||
metrics = {
|
||||
"preprocess_time_ms": preprocess_time * 1000,
|
||||
"diffusion_engine_exec_time_ms": (time.time() - diffusion_engine_start_time) * 1000,
|
||||
"diffusion_engine_exec_time_ms": (time.perf_counter() - diffusion_engine_start_time) * 1000,
|
||||
"diffusion_engine_total_time_ms": exec_total_time * 1000,
|
||||
"image_num": int(request.sampling_params.num_outputs_per_prompt),
|
||||
"resolution": int(request.sampling_params.resolution),
|
||||
"postprocess_time_ms": postprocess_time * 1000,
|
||||
}
|
||||
|
||||
if self.pre_process_func is not None:
|
||||
metrics["preprocessing_time_ms"] = preprocess_time * 1000
|
||||
|
||||
# Handle single request or multiple requests
|
||||
metrics["postprocess_time_ms"] = postprocess_time * 1000
|
||||
metrics["num_inference_steps"] = int(request.sampling_params.num_inference_steps)
|
||||
|
||||
if len(request.prompts) == 1:
|
||||
# Single request: return single OmniRequestOutput
|
||||
prompt = request.prompts[0]
|
||||
@@ -155,6 +163,7 @@ class DiffusionEngine:
|
||||
prompt=prompt,
|
||||
metrics=metrics,
|
||||
latents=output.trajectory_latents,
|
||||
custom_output=output.custom_output or {},
|
||||
multimodal_output=mm_output,
|
||||
),
|
||||
]
|
||||
@@ -208,11 +217,12 @@ class DiffusionEngine:
|
||||
prompt=prompt,
|
||||
metrics=metrics,
|
||||
latents=output.trajectory_latents,
|
||||
custom_output=output.custom_output or {},
|
||||
multimodal_output=mm_output,
|
||||
),
|
||||
)
|
||||
|
||||
return results
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def make_engine(config: OmniDiffusionConfig) -> "DiffusionEngine":
|
||||
|
||||
@@ -669,6 +669,7 @@ def initialize_model_parallel(
|
||||
tensor_parallel_size: int = 1,
|
||||
pipeline_parallel_size: int = 1,
|
||||
fully_shard_degree: int = 1,
|
||||
hsdp_replicate_size: int = 1,
|
||||
backend: str | None = None,
|
||||
) -> None:
|
||||
if backend is None:
|
||||
@@ -740,6 +741,14 @@ def initialize_model_parallel(
|
||||
data_parallel_size * cfg_parallel_size * sequence_parallel_size * pipeline_parallel_size * tensor_parallel_size
|
||||
)
|
||||
|
||||
# Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
|
||||
is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1
|
||||
|
||||
# For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size) as dit_parallel_size
|
||||
# This ensures orthogonal rank generation works correctly for all HSDP workers
|
||||
if is_standalone_hsdp:
|
||||
dit_parallel_size = fully_shard_degree * hsdp_replicate_size
|
||||
|
||||
if world_size < dit_parallel_size:
|
||||
raise RuntimeError(
|
||||
f"world_size ({world_size}) is less than "
|
||||
@@ -751,12 +760,16 @@ def initialize_model_parallel(
|
||||
f"data_parallel_size ({data_parallel_size})"
|
||||
)
|
||||
|
||||
# For standalone HSDP, use (fully_shard_degree * hsdp_replicate_size) as data_parallel_size
|
||||
# so that RankGenerator.world_size matches the actual number of workers
|
||||
effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size
|
||||
|
||||
rank_generator: RankGenerator = RankGenerator(
|
||||
tensor_parallel_size,
|
||||
sequence_parallel_size,
|
||||
pipeline_parallel_size,
|
||||
cfg_parallel_size,
|
||||
data_parallel_size,
|
||||
effective_dp_size,
|
||||
fs=fully_shard_degree,
|
||||
order="tp-sp-pp-cfg-dp",
|
||||
)
|
||||
|
||||
@@ -107,7 +107,7 @@ class ModelHook:
|
||||
The output of the forward pass.
|
||||
"""
|
||||
args, kwargs = self.pre_forward(module, *args, **kwargs)
|
||||
output = module._original_forward(*args, **kwargs) # type: ignore[attr-defined]
|
||||
output = module._omni_original_forward(*args, **kwargs) # type: ignore[attr-defined]
|
||||
return self.post_forward(module, output)
|
||||
|
||||
def reset_state(self, module: nn.Module) -> nn.Module:
|
||||
@@ -131,7 +131,7 @@ class _WrappedForward:
|
||||
def __call__(self, *args: Any, **kwargs: Any):
|
||||
registry: HookRegistry | None = getattr(self.module, "_hook_registry", None)
|
||||
if registry is None or not registry._hooks:
|
||||
return self.module._original_forward(*args, **kwargs)
|
||||
return self.module._omni_original_forward(*args, **kwargs)
|
||||
return registry.dispatch(*args, **kwargs)
|
||||
|
||||
|
||||
@@ -162,8 +162,10 @@ class HookRegistry:
|
||||
setattr(module, "_hook_registry", registry)
|
||||
|
||||
# Wrap module.forward once so hooks can intercept calls.
|
||||
if not hasattr(module, "_original_forward"):
|
||||
module._original_forward = module.forward # type: ignore[attr-defined]
|
||||
# NOTE: Use `_omni_original_forward` to avoid collision with cache-dit's
|
||||
# `_original_forward` attribute on the same module.
|
||||
if not hasattr(module, "_omni_original_forward"):
|
||||
module._omni_original_forward = module.forward # type: ignore[attr-defined]
|
||||
module.forward = _WrappedForward(module) # type: ignore[assignment]
|
||||
|
||||
return registry
|
||||
@@ -212,7 +214,7 @@ class HookRegistry:
|
||||
The output of the forward pass.
|
||||
"""
|
||||
if not self._hooks:
|
||||
return self.module._original_forward(*args, **kwargs) # type: ignore[attr-defined]
|
||||
return self.module._omni_original_forward(*args, **kwargs) # type: ignore[attr-defined]
|
||||
|
||||
# For single hook case, call directly
|
||||
if len(self._hooks) == 1:
|
||||
@@ -228,7 +230,7 @@ class HookRegistry:
|
||||
args, kwargs = hook.pre_forward(self.module, *args, **kwargs)
|
||||
|
||||
# Call original forward
|
||||
output = self.module._original_forward(*args, **kwargs) # type: ignore[attr-defined]
|
||||
output = self.module._omni_original_forward(*args, **kwargs) # type: ignore[attr-defined]
|
||||
|
||||
# Apply all post_forward hooks in reverse order
|
||||
for _, hook in reversed(sorted_hooks):
|
||||
|
||||
@@ -0,0 +1,89 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""IPC utilities for transferring large tensors via POSIX shared memory.
|
||||
|
||||
Used by Hop1 (GPU worker <-> scheduler) to avoid pickling large video tensors
|
||||
through the MessageQueue. Tensors above ``_SHM_TENSOR_THRESHOLD`` are copied
|
||||
into a named shared-memory segment; only a lightweight metadata dict is
|
||||
serialised through the queue.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
|
||||
_SHM_TENSOR_THRESHOLD = 1_000_000 # 1 MB
|
||||
|
||||
|
||||
def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]:
|
||||
"""Copy a tensor into POSIX shared memory and return a metadata handle.
|
||||
|
||||
The shared memory segment remains alive after this call (the local fd is
|
||||
closed, but the segment persists until ``_tensor_from_shm`` unlinks it).
|
||||
"""
|
||||
from multiprocessing import shared_memory
|
||||
|
||||
import numpy as np
|
||||
|
||||
tensor = tensor.detach().cpu().contiguous()
|
||||
arr = tensor.numpy()
|
||||
nbytes = arr.nbytes
|
||||
shm = shared_memory.SharedMemory(create=True, size=nbytes)
|
||||
shm_arr = np.ndarray(arr.shape, dtype=arr.dtype, buffer=shm.buf[:nbytes])
|
||||
np.copyto(shm_arr, arr)
|
||||
handle = {
|
||||
"__tensor_shm__": True,
|
||||
"name": shm.name,
|
||||
"shape": list(tensor.shape),
|
||||
"torch_dtype": str(tensor.dtype),
|
||||
"numpy_dtype": str(arr.dtype),
|
||||
"nbytes": nbytes,
|
||||
}
|
||||
shm.close()
|
||||
return handle
|
||||
|
||||
|
||||
def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
|
||||
"""Reconstruct a tensor from a shared-memory handle and free the segment."""
|
||||
from multiprocessing import shared_memory
|
||||
|
||||
import numpy as np
|
||||
|
||||
shm = shared_memory.SharedMemory(name=handle["name"])
|
||||
try:
|
||||
np_dtype = np.dtype(handle["numpy_dtype"])
|
||||
arr = np.ndarray(handle["shape"], dtype=np_dtype, buffer=shm.buf[: handle["nbytes"]])
|
||||
tensor = torch.from_numpy(arr.copy())
|
||||
finally:
|
||||
shm.close()
|
||||
shm.unlink()
|
||||
return tensor
|
||||
|
||||
|
||||
def pack_diffusion_output_shm(output: DiffusionOutput) -> DiffusionOutput:
|
||||
"""Replace large tensors in *output* with shared-memory handles.
|
||||
|
||||
The DiffusionOutput is modified **in-place** so that the (now lightweight)
|
||||
object can be serialised cheaply through a MessageQueue.
|
||||
"""
|
||||
if output.output is not None and isinstance(output.output, torch.Tensor):
|
||||
if output.output.nelement() * output.output.element_size() > _SHM_TENSOR_THRESHOLD:
|
||||
output.output = _tensor_to_shm(output.output)
|
||||
if output.trajectory_latents is not None and isinstance(output.trajectory_latents, torch.Tensor):
|
||||
if output.trajectory_latents.nelement() * output.trajectory_latents.element_size() > _SHM_TENSOR_THRESHOLD:
|
||||
output.trajectory_latents = _tensor_to_shm(output.trajectory_latents)
|
||||
return output
|
||||
|
||||
|
||||
def unpack_diffusion_output_shm(output: DiffusionOutput) -> DiffusionOutput:
|
||||
"""Reconstruct tensors from shared-memory handles produced by ``pack_diffusion_output_shm``."""
|
||||
if isinstance(output.output, dict) and output.output.get("__tensor_shm__"):
|
||||
output.output = _tensor_from_shm(output.output)
|
||||
if isinstance(output.trajectory_latents, dict) and output.trajectory_latents.get("__tensor_shm__"):
|
||||
output.trajectory_latents = _tensor_from_shm(output.trajectory_latents)
|
||||
return output
|
||||
@@ -1526,6 +1526,11 @@ class Bagel(nn.Module):
|
||||
f"Use cfg_parallel_size=3 to enable image CFG in parallel mode."
|
||||
)
|
||||
|
||||
# Ensure all ranks start with the same x_t (initial noise may differ
|
||||
# across ranks when no per-request seed is set).
|
||||
x_t = x_t.contiguous()
|
||||
cfg_group.broadcast(x_t, src=0)
|
||||
|
||||
# Select this rank's branch inputs
|
||||
if cfg_rank == 0:
|
||||
# Gen branch: use main inputs directly
|
||||
@@ -1553,17 +1558,10 @@ class Bagel(nn.Module):
|
||||
|
||||
for i, t in enumerate(timesteps):
|
||||
timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
|
||||
if t > cfg_interval[0] and t <= cfg_interval[1]:
|
||||
cfg_text_scale_ = cfg_text_scale
|
||||
cfg_img_scale_ = cfg_img_scale
|
||||
else:
|
||||
cfg_text_scale_ = 1.0
|
||||
cfg_img_scale_ = 1.0
|
||||
|
||||
use_cfg_this_step = cfg_text_scale_ > 1.0
|
||||
use_cfg_this_step = t > cfg_interval[0] and t <= cfg_interval[1] and cfg_text_scale > 1.0
|
||||
|
||||
if use_cfg_this_step:
|
||||
# Each rank computes its branch's velocity
|
||||
# CFG interval: each rank computes its own branch
|
||||
local_v_t = self._forward_flow_single_branch(
|
||||
x_t=x_t,
|
||||
timestep=timestep,
|
||||
@@ -1579,46 +1577,34 @@ class Bagel(nn.Module):
|
||||
packed_key_value_indexes=branch_key_value_indexes,
|
||||
)
|
||||
|
||||
# All-gather velocities from all CFG ranks
|
||||
gathered = cfg_group.all_gather(local_v_t, separate_tensors=True)
|
||||
|
||||
# Rank 0 combines with CFG formula
|
||||
if cfg_rank == 0:
|
||||
v_t = gathered[0] # gen branch
|
||||
cfg_text_v_t = gathered[1] # text_cfg branch
|
||||
cfg_img_v_t = gathered[2] if (use_cfg_img and len(gathered) > 2) else None
|
||||
v_t = self._combine_cfg(
|
||||
v_t,
|
||||
cfg_text_v_t,
|
||||
cfg_img_v_t,
|
||||
cfg_text_scale_,
|
||||
cfg_img_scale_,
|
||||
cfg_renorm_type,
|
||||
cfg_renorm_min,
|
||||
)
|
||||
x_t = x_t - v_t.to(x_t.device) * dts[i]
|
||||
v_t = self._combine_cfg(
|
||||
gathered[0],
|
||||
gathered[1],
|
||||
gathered[2] if (use_cfg_img and len(gathered) > 2) else None,
|
||||
cfg_text_scale,
|
||||
cfg_img_scale,
|
||||
cfg_renorm_type,
|
||||
cfg_renorm_min,
|
||||
)
|
||||
else:
|
||||
# Outside cfg_interval: only rank 0 computes (no CFG needed)
|
||||
if cfg_rank == 0:
|
||||
v_t = self._forward_flow_single_branch(
|
||||
x_t=x_t,
|
||||
timestep=timestep,
|
||||
packed_vae_token_indexes=packed_vae_token_indexes,
|
||||
packed_vae_position_ids=packed_vae_position_ids,
|
||||
packed_text_ids=packed_text_ids,
|
||||
packed_text_indexes=packed_text_indexes,
|
||||
packed_indexes=packed_indexes,
|
||||
packed_position_ids=packed_position_ids,
|
||||
packed_seqlens=packed_seqlens,
|
||||
key_values_lens=key_values_lens,
|
||||
past_key_values=past_key_values,
|
||||
packed_key_value_indexes=packed_key_value_indexes,
|
||||
)
|
||||
x_t = x_t - v_t.to(x_t.device) * dts[i]
|
||||
# Outside CFG interval: all ranks compute with gen inputs, no comm
|
||||
v_t = self._forward_flow_single_branch(
|
||||
x_t=x_t,
|
||||
timestep=timestep,
|
||||
packed_vae_token_indexes=packed_vae_token_indexes,
|
||||
packed_vae_position_ids=packed_vae_position_ids,
|
||||
packed_text_ids=packed_text_ids,
|
||||
packed_text_indexes=packed_text_indexes,
|
||||
packed_indexes=packed_indexes,
|
||||
packed_position_ids=packed_position_ids,
|
||||
packed_seqlens=packed_seqlens,
|
||||
key_values_lens=key_values_lens,
|
||||
past_key_values=past_key_values,
|
||||
packed_key_value_indexes=packed_key_value_indexes,
|
||||
)
|
||||
|
||||
# Broadcast updated x_t from rank 0 to all ranks
|
||||
x_t = x_t.contiguous()
|
||||
cfg_group.broadcast(x_t, src=0)
|
||||
x_t = x_t - v_t.to(x_t.device) * dts[i]
|
||||
|
||||
unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
|
||||
return unpacked_latent
|
||||
|
||||
@@ -332,6 +332,9 @@ class BagelPipeline(nn.Module):
|
||||
else:
|
||||
gen_context["ropes"] = [seq_len]
|
||||
|
||||
if req.sampling_params.kv_metadata and "image_shape" in req.sampling_params.kv_metadata:
|
||||
image_shape = tuple(req.sampling_params.kv_metadata["image_shape"])
|
||||
|
||||
cfg_text_kv = getattr(req.sampling_params, "cfg_text_past_key_values", None)
|
||||
if cfg_text_kv is not None:
|
||||
logger.info("CFG enabled with multi-KV: using injected cfg_text KV Cache")
|
||||
|
||||
@@ -604,6 +604,7 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
use_zero_init=use_zero_init,
|
||||
zero_steps=zero_steps,
|
||||
device=device,
|
||||
generator=generator,
|
||||
)
|
||||
|
||||
if keep_first_frame and (
|
||||
@@ -752,6 +753,7 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
use_zero_init: bool = True,
|
||||
zero_steps: int = 1,
|
||||
device: torch.device | None = None,
|
||||
generator: torch.Generator | list[torch.Generator] | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Pyramid multi-stage denoising for one chunk."""
|
||||
batch_size, num_channel, num_frames_lat, height, width = latents.shape
|
||||
@@ -804,7 +806,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma)
|
||||
beta = alpha * (1 - ori_sigma) / math.sqrt(gamma)
|
||||
|
||||
noise = self.sample_block_noise(batch_size, num_channel, latents.shape[2], height, width, patch_size)
|
||||
noise = self.sample_block_noise(
|
||||
batch_size, num_channel, latents.shape[2], height, width, patch_size, generator=generator
|
||||
)
|
||||
noise = noise.to(device=device, dtype=transformer_dtype)
|
||||
latents = alpha * latents + beta * noise
|
||||
|
||||
@@ -873,18 +877,23 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
|
||||
return latents
|
||||
|
||||
def sample_block_noise(self, batch_size, channel, num_frames, height, width, patch_size=(1, 2, 2)):
|
||||
def sample_block_noise(self, batch_size, channel, num_frames, height, width, patch_size=(1, 2, 2), generator=None):
|
||||
gamma = self.scheduler.config.gamma
|
||||
_, ph, pw = patch_size
|
||||
block_size = ph * pw
|
||||
|
||||
cov = torch.eye(block_size) * (1 + gamma) - torch.ones(block_size, block_size) * gamma
|
||||
dist = torch.distributions.MultivariateNormal(torch.zeros(block_size, device=cov.device), covariance_matrix=cov)
|
||||
cov += torch.eye(block_size) * 1e-8
|
||||
cov = cov.float() # Upcast to fp32 for numerical stability — cholesky is unreliable in fp16/bf16.
|
||||
|
||||
L = torch.linalg.cholesky(cov)
|
||||
block_number = batch_size * channel * num_frames * (height // ph) * (width // pw)
|
||||
z = torch.randn(block_number, block_size, generator=generator)
|
||||
noise = z @ L.T
|
||||
|
||||
noise = dist.sample((block_number,))
|
||||
noise = noise.view(batch_size, channel, num_frames, height // ph, width // pw, ph, pw)
|
||||
noise = noise.permute(0, 1, 2, 3, 5, 4, 6).reshape(batch_size, channel, num_frames, height, width)
|
||||
|
||||
return noise
|
||||
|
||||
def predict_noise(self, **kwargs: Any) -> torch.Tensor:
|
||||
|
||||
@@ -78,7 +78,7 @@ class HeliosScheduler(SchedulerMixin, ConfigMixin):
|
||||
num_train_timesteps = self.config.num_train_timesteps
|
||||
shift = self.config.shift
|
||||
|
||||
alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps + 1)
|
||||
alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps + 1, dtype=np.float32)
|
||||
sigmas = 1.0 - alphas
|
||||
sigmas = np.flip(shift * sigmas / (1 + (shift - 1) * sigmas))[:-1].copy()
|
||||
sigmas = torch.from_numpy(sigmas)
|
||||
@@ -133,11 +133,11 @@ class HeliosScheduler(SchedulerMixin, ConfigMixin):
|
||||
timestep_ratio = self.timestep_ratios[i_s]
|
||||
timestep_max = min(self.timesteps[int(timestep_ratio[0] * training_steps)], 999)
|
||||
timestep_min = self.timesteps[min(int(timestep_ratio[1] * training_steps), training_steps - 1)]
|
||||
timesteps = np.linspace(timestep_max, timestep_min, training_steps + 1)
|
||||
timesteps = np.linspace(timestep_max, timestep_min, training_steps + 1, dtype=np.float32)
|
||||
self.timesteps_per_stage[i_s] = (
|
||||
timesteps[:-1] if isinstance(timesteps, torch.Tensor) else torch.from_numpy(timesteps[:-1])
|
||||
)
|
||||
stage_sigmas = np.linspace(0.999, 0, training_steps + 1)
|
||||
stage_sigmas = np.linspace(0.999, 0, training_steps + 1, dtype=np.float32)
|
||||
self.sigmas_per_stage[i_s] = torch.from_numpy(stage_sigmas[:-1])
|
||||
|
||||
@property
|
||||
@@ -188,10 +188,11 @@ class HeliosScheduler(SchedulerMixin, ConfigMixin):
|
||||
stage_timesteps[0].item(),
|
||||
stage_timesteps[-1].item(),
|
||||
num_inference_steps,
|
||||
dtype=np.float32,
|
||||
)
|
||||
|
||||
stage_sigmas = self.sigmas_per_stage[stage_index]
|
||||
ratios = np.linspace(stage_sigmas[0].item(), stage_sigmas[-1].item(), num_inference_steps)
|
||||
ratios = np.linspace(stage_sigmas[0].item(), stage_sigmas[-1].item(), num_inference_steps, dtype=np.float32)
|
||||
sigmas = torch.from_numpy(ratios)
|
||||
|
||||
self.timesteps = torch.from_numpy(timesteps).to(device=device)
|
||||
|
||||
@@ -6,6 +6,7 @@ from __future__ import annotations
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Iterable
|
||||
from typing import Any, cast
|
||||
|
||||
@@ -21,7 +22,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import Dist
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
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.progress_bar import ProgressBarMixin
|
||||
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
|
||||
from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
|
||||
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
@@ -29,6 +30,7 @@ from vllm_omni.inputs.data import OmniTextPrompt
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
DEBUG_PERF = False
|
||||
|
||||
|
||||
def retrieve_latents(
|
||||
@@ -424,7 +426,11 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
if generator is None and req.sampling_params.seed is not None:
|
||||
generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
|
||||
|
||||
# Encode prompts
|
||||
if DEBUG_PERF:
|
||||
# Sync GPU before timing to ensure accurate measurements
|
||||
current_omni_platform.synchronize()
|
||||
_t_pipeline_start = time.perf_counter()
|
||||
_t_text_enc_start = _t_pipeline_start
|
||||
if prompt_embeds is None:
|
||||
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
|
||||
prompt=prompt,
|
||||
@@ -443,6 +449,9 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
raise ValueError(
|
||||
"negative_prompt_embeds must be provided when prompt_embeds are given and guidance > 1."
|
||||
)
|
||||
if DEBUG_PERF:
|
||||
current_omni_platform.synchronize()
|
||||
_t_text_enc_ms = (time.perf_counter() - _t_text_enc_start) * 1000
|
||||
|
||||
# Timesteps
|
||||
self.scheduler.set_timesteps(num_steps, device=device)
|
||||
@@ -452,7 +461,8 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
if boundary_ratio is not None:
|
||||
boundary_timestep = boundary_ratio * self.scheduler.config.num_train_timesteps
|
||||
|
||||
# Handle I2V mode when expand_timesteps=True and image is provided
|
||||
if DEBUG_PERF:
|
||||
_t_latent_prep_start = time.perf_counter()
|
||||
multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
|
||||
raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
|
||||
if isinstance(raw_image, list):
|
||||
@@ -543,11 +553,15 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
generator=generator,
|
||||
latents=req.sampling_params.latents,
|
||||
)
|
||||
if DEBUG_PERF:
|
||||
current_omni_platform.synchronize()
|
||||
_t_latent_prep_ms = (time.perf_counter() - _t_latent_prep_start) * 1000
|
||||
|
||||
if attention_kwargs is None:
|
||||
attention_kwargs = {}
|
||||
|
||||
# Denoising
|
||||
if DEBUG_PERF:
|
||||
_t_denoise_start = time.perf_counter()
|
||||
with self.progress_bar(total=len(timesteps)) as pbar:
|
||||
for t in timesteps:
|
||||
self._current_timestep = t
|
||||
@@ -638,12 +652,16 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
if current_omni_platform.is_available():
|
||||
current_omni_platform.empty_cache()
|
||||
self._current_timestep = None
|
||||
if DEBUG_PERF:
|
||||
current_omni_platform.synchronize()
|
||||
_t_denoise_ms = (time.perf_counter() - _t_denoise_start) * 1000
|
||||
|
||||
# For I2V mode: blend final latents with condition
|
||||
if self.expand_timesteps and latent_condition is not None:
|
||||
latents = (1 - first_frame_mask) * latent_condition + first_frame_mask * latents
|
||||
|
||||
# Decode
|
||||
if DEBUG_PERF:
|
||||
_t_decode_start = time.perf_counter()
|
||||
if output_type == "latent":
|
||||
output = latents
|
||||
else:
|
||||
@@ -659,6 +677,28 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
latents = latents / latents_std + latents_mean
|
||||
output = self.vae.decode(latents, return_dict=False)[0]
|
||||
|
||||
if DEBUG_PERF:
|
||||
current_omni_platform.synchronize()
|
||||
_t_decode_ms = (time.perf_counter() - _t_decode_start) * 1000
|
||||
_t_pipeline_wall_ms = (time.perf_counter() - _t_pipeline_start) * 1000
|
||||
_t_stages_sum = _t_text_enc_ms + _t_latent_prep_ms + _t_denoise_ms + _t_decode_ms
|
||||
|
||||
if _is_rank_zero():
|
||||
logger.info(
|
||||
"Pipeline stage timing summary: "
|
||||
"TextEncoding=%.2f ms, LatentPreparation=%.2f ms, "
|
||||
"Denoising=%.2f ms (%d steps), Decoding=%.2f ms, "
|
||||
"StagesSum=%.2f ms, PipelineWall=%.2f ms, Unaccounted=%.2f ms",
|
||||
_t_text_enc_ms,
|
||||
_t_latent_prep_ms,
|
||||
_t_denoise_ms,
|
||||
len(timesteps),
|
||||
_t_decode_ms,
|
||||
_t_stages_sum,
|
||||
_t_pipeline_wall_ms,
|
||||
_t_pipeline_wall_ms - _t_stages_sum,
|
||||
)
|
||||
|
||||
return DiffusionOutput(output=output)
|
||||
|
||||
def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
|
||||
|
||||
@@ -5,6 +5,7 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Iterable
|
||||
from typing import Any, cast
|
||||
|
||||
@@ -22,7 +23,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.models.interface import SupportImageInput
|
||||
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
||||
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
|
||||
from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
|
||||
create_transformer_from_config,
|
||||
@@ -34,6 +35,7 @@ from vllm_omni.inputs.data import OmniTextPrompt
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
DEBUG_PERF = False
|
||||
|
||||
|
||||
def _load_model_index(model: str, local_files_only: bool) -> dict:
|
||||
@@ -388,7 +390,12 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressB
|
||||
if generator is None and req.sampling_params.seed is not None:
|
||||
generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
|
||||
|
||||
# Encode prompts
|
||||
if DEBUG_PERF:
|
||||
# Sync GPU before timing to ensure accurate measurements
|
||||
current_omni_platform.synchronize()
|
||||
_t_pipeline_start = time.perf_counter()
|
||||
_t_text_enc_start = _t_pipeline_start
|
||||
|
||||
if prompt_embeds is None:
|
||||
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
|
||||
prompt=prompt,
|
||||
@@ -404,9 +411,14 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressB
|
||||
if negative_prompt_embeds is not None:
|
||||
negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
|
||||
|
||||
if DEBUG_PERF:
|
||||
current_omni_platform.synchronize()
|
||||
_t_text_enc_ms = (time.perf_counter() - _t_text_enc_start) * 1000
|
||||
|
||||
batch_size = prompt_embeds.shape[0]
|
||||
|
||||
# Encode image embeddings (for Wan2.1-style with CLIP)
|
||||
if DEBUG_PERF:
|
||||
_t_img_enc_start = time.perf_counter()
|
||||
if self.has_image_encoder and self.transformer.config.image_dim is not None:
|
||||
if image_embeds is None:
|
||||
if last_image is None:
|
||||
@@ -418,6 +430,10 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressB
|
||||
else:
|
||||
image_embeds = None
|
||||
|
||||
if DEBUG_PERF:
|
||||
current_omni_platform.synchronize()
|
||||
_t_img_enc_ms = (time.perf_counter() - _t_img_enc_start) * 1000
|
||||
|
||||
# Timesteps
|
||||
self.scheduler.set_timesteps(num_steps, device=device)
|
||||
timesteps = self.scheduler.timesteps
|
||||
@@ -430,7 +446,8 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressB
|
||||
# Prepare latents (use out_channels=16 for VAE latent, not in_channels=36)
|
||||
num_channels_latents = self.transformer.config.out_channels
|
||||
|
||||
# Preprocess image for VAE
|
||||
if DEBUG_PERF:
|
||||
_t_latent_prep_start = time.perf_counter()
|
||||
from diffusers.video_processor import VideoProcessor
|
||||
|
||||
video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
|
||||
@@ -465,10 +482,15 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressB
|
||||
last_image=last_image_tensor,
|
||||
)
|
||||
|
||||
if DEBUG_PERF:
|
||||
current_omni_platform.synchronize()
|
||||
_t_latent_prep_ms = (time.perf_counter() - _t_latent_prep_start) * 1000
|
||||
|
||||
if attention_kwargs is None:
|
||||
attention_kwargs = {}
|
||||
|
||||
# Denoising loop
|
||||
if DEBUG_PERF:
|
||||
_t_denoise_start = time.perf_counter()
|
||||
with self.progress_bar(total=len(timesteps)) as pbar:
|
||||
for t in timesteps:
|
||||
self._current_timestep = t
|
||||
@@ -538,11 +560,17 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressB
|
||||
current_omni_platform.empty_cache()
|
||||
self._current_timestep = None
|
||||
|
||||
if DEBUG_PERF:
|
||||
current_omni_platform.synchronize()
|
||||
_t_denoise_ms = (time.perf_counter() - _t_denoise_start) * 1000
|
||||
|
||||
# For expand_timesteps mode, blend final latents with condition
|
||||
if self.expand_timesteps:
|
||||
latents = (1 - first_frame_mask) * condition + first_frame_mask * latents
|
||||
|
||||
# Decode
|
||||
if DEBUG_PERF:
|
||||
_t_decode_start = time.perf_counter()
|
||||
|
||||
if output_type == "latent":
|
||||
output = latents
|
||||
else:
|
||||
@@ -558,6 +586,30 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressB
|
||||
latents = latents / latents_std + latents_mean
|
||||
output = self.vae.decode(latents, return_dict=False)[0]
|
||||
|
||||
if DEBUG_PERF:
|
||||
current_omni_platform.synchronize()
|
||||
_t_decode_ms = (time.perf_counter() - _t_decode_start) * 1000
|
||||
_t_pipeline_wall_ms = (time.perf_counter() - _t_pipeline_start) * 1000
|
||||
_t_stages_sum = _t_text_enc_ms + _t_img_enc_ms + _t_latent_prep_ms + _t_denoise_ms + _t_decode_ms
|
||||
|
||||
if _is_rank_zero():
|
||||
logger.info(
|
||||
"Pipeline stage timing summary: "
|
||||
"TextEncoding=%.2f ms, ImageEncoding=%.2f ms, "
|
||||
"LatentPreparation=%.2f ms, Denoising=%.2f ms (%d steps), "
|
||||
"Decoding=%.2f ms, StagesSum=%.2f ms, PipelineWall=%.2f ms, "
|
||||
"Unaccounted=%.2f ms",
|
||||
_t_text_enc_ms,
|
||||
_t_img_enc_ms,
|
||||
_t_latent_prep_ms,
|
||||
_t_denoise_ms,
|
||||
len(timesteps),
|
||||
_t_decode_ms,
|
||||
_t_stages_sum,
|
||||
_t_pipeline_wall_ms,
|
||||
_t_pipeline_wall_ms - _t_stages_sum,
|
||||
)
|
||||
|
||||
return DiffusionOutput(output=output)
|
||||
|
||||
def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
|
||||
|
||||
@@ -50,6 +50,9 @@ class LayerwiseOffloadHook(ModelHook):
|
||||
# Per-block synchronization primitive: set after H2D copy completes.
|
||||
self._prefetch_done: current_omni_platform.Event | None = None
|
||||
|
||||
# Backward link to the hook that is responsible for prefetching *this* block's weights
|
||||
self._prev_hook: LayerwiseOffloadHook | None = None
|
||||
|
||||
self.next_block_parameters: dict[str, nn.Parameter] = {}
|
||||
self.next_block_buffers: dict[str, torch.Tensor] = {}
|
||||
self.dtype_cpu_flattened_weights: dict[torch.dtype, torch.Tensor] = {}
|
||||
@@ -128,6 +131,14 @@ class LayerwiseOffloadHook(ModelHook):
|
||||
|
||||
return dtype_cpu_flattened_weights, dtype_metadata
|
||||
|
||||
@property
|
||||
def is_materialized(self) -> bool:
|
||||
"""Check whether this block's parameters hold real data on device."""
|
||||
for param in self.block_parameters.values():
|
||||
return param.data.dim() > 0
|
||||
|
||||
return True
|
||||
|
||||
@torch.compiler.disable
|
||||
def prefetch_layer(self, non_blocking: bool = True) -> None:
|
||||
"""Copy layer weights from CPU -> GPU.
|
||||
@@ -185,6 +196,12 @@ class LayerwiseOffloadHook(ModelHook):
|
||||
buf.data = torch.empty((), device=self.device, dtype=buf.dtype)
|
||||
|
||||
def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
|
||||
# if the previous hook was skipped and the weights are not on device,
|
||||
# (e.g. by cache-dit block caching), ask the previous hook to
|
||||
# synchronously prefetch *this* block's weights before computation
|
||||
if not self.is_materialized and self._prev_hook is not None:
|
||||
self._prev_hook.prefetch_layer(non_blocking=False)
|
||||
|
||||
self.prefetch_layer(non_blocking=True)
|
||||
|
||||
return args, kwargs
|
||||
@@ -293,13 +310,23 @@ class LayerWiseOffloadBackend(OffloadBackend):
|
||||
# For subsequent requests, the first layer/block will be pre-fetched
|
||||
# during the last layer compute of the previous request.
|
||||
last_block, first_block = blocks[-1], blocks[0]
|
||||
hook = apply_block_hook(last_block, first_block, self.device, self.copy_stream, self.config.pin_cpu_memory)
|
||||
hook.prefetch_layer(non_blocking=False)
|
||||
last_hook = apply_block_hook(
|
||||
last_block, first_block, self.device, self.copy_stream, self.config.pin_cpu_memory
|
||||
)
|
||||
last_hook.prefetch_layer(non_blocking=False)
|
||||
|
||||
block_hooks: list[LayerwiseOffloadHook] = [last_hook]
|
||||
# Register hook for each of blocks
|
||||
for i, block in enumerate(blocks[:-1]):
|
||||
next_block = blocks[(i + 1) % num_blocks]
|
||||
apply_block_hook(block, next_block, self.device, self.copy_stream, self.config.pin_cpu_memory)
|
||||
hook = apply_block_hook(block, next_block, self.device, self.copy_stream, self.config.pin_cpu_memory)
|
||||
block_hooks.append(hook)
|
||||
|
||||
# NOTE(yuanheng-zhao): We make each hook gets a backward reference to the hook
|
||||
# that is responsible for prefetching its block's weights. This is specifically a
|
||||
# workaround for that arbitrary blocks are skipped by caching systems (e.g., cache-dit)
|
||||
for i in range(len(block_hooks)):
|
||||
block_hooks[i]._prev_hook = block_hooks[i - 1]
|
||||
|
||||
logger.info(f"Layer-wise offloading enabled on {num_blocks} layers (blocks)")
|
||||
|
||||
|
||||
@@ -35,6 +35,23 @@ class SequentialOffloadHook(ModelHook):
|
||||
self.device = device
|
||||
self.pin_memory = pin_memory
|
||||
|
||||
@staticmethod
|
||||
def _move_params(module: nn.Module, device: torch.device) -> None:
|
||||
"""Move module parameters and buffers to device.
|
||||
|
||||
This cls method specifically prevents recursion device movement,
|
||||
E.g., Cache-DiT CachedBlocks has attr `transformer` as a ref to original
|
||||
transformer blocks, thus `module.to(device)` will fail for recursion calling,
|
||||
refer to
|
||||
https://github.com/vipshop/cache-dit/blob/v1.2.3/src/cache_dit/caching/cache_blocks/__init__.py#L83
|
||||
"""
|
||||
for p in module.parameters():
|
||||
if p.data.device != device:
|
||||
p.data = p.data.to(device, non_blocking=True)
|
||||
for b in module.buffers():
|
||||
if b.device != device:
|
||||
b.data = b.data.to(device, non_blocking=True)
|
||||
|
||||
def _to_cpu(self, module: nn.Module) -> None:
|
||||
"""Move module to CPU."""
|
||||
try:
|
||||
@@ -42,13 +59,12 @@ class SequentialOffloadHook(ModelHook):
|
||||
except StopIteration:
|
||||
return
|
||||
|
||||
previous_device = param.device
|
||||
# Skip if already on CPU
|
||||
if previous_device.type == "cpu":
|
||||
if param.device.type == "cpu":
|
||||
return
|
||||
|
||||
module.to("cpu")
|
||||
torch.cuda.empty_cache()
|
||||
self._move_params(module, torch.device("cpu"))
|
||||
current_omni_platform.empty_cache()
|
||||
|
||||
if self.pin_memory:
|
||||
for p in module.parameters():
|
||||
@@ -64,7 +80,7 @@ class SequentialOffloadHook(ModelHook):
|
||||
except StopIteration:
|
||||
return
|
||||
|
||||
module.to(self.device)
|
||||
self._move_params(module, self.device)
|
||||
|
||||
def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple, dict]:
|
||||
# Offload target modules to CPU
|
||||
@@ -186,19 +202,6 @@ class ModelLevelOffloadBackend(OffloadBackend):
|
||||
except Exception as exc:
|
||||
logger.debug("Failed to move VAE to GPU: %s", exc)
|
||||
|
||||
# Initial state: keep DiT modules on CPU (encoders typically run first)
|
||||
# TODO: This part seems to be unnecessary, remove it after testing
|
||||
# for dit_mod in modules.dits:
|
||||
# dit_mod.to("cpu")
|
||||
|
||||
# torch.cuda.empty_cache()
|
||||
|
||||
# if self.config.pin_cpu_memory:
|
||||
# for dit_mod in modules.dits:
|
||||
# for p in dit_mod.parameters():
|
||||
# if p.data.device.type == "cpu" and not p.data.is_pinned():
|
||||
# p.data = p.data.pin_memory()
|
||||
|
||||
# Apply sequential offloading hooks
|
||||
apply_sequential_offload(
|
||||
dit_modules=modules.dits,
|
||||
|
||||
@@ -8,6 +8,7 @@ from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.ipc import unpack_diffusion_output_shm
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -59,12 +60,18 @@ class Scheduler:
|
||||
|
||||
# Broadcast RPC request to all workers
|
||||
self.mq.enqueue(rpc_request)
|
||||
# Wait for result from Rank 0 (or whoever sends it)
|
||||
|
||||
# Wait for result from Rank 0 (or whoever sends it)
|
||||
if self.result_mq is None:
|
||||
raise RuntimeError("Result queue not initialized")
|
||||
|
||||
output = self.result_mq.dequeue()
|
||||
|
||||
try:
|
||||
unpack_diffusion_output_shm(output)
|
||||
except Exception as e:
|
||||
logger.warning("SHM unpack failed (data may already be inline): %s", e)
|
||||
|
||||
# {"status": "error", "error": str(e)}
|
||||
if isinstance(output, dict) and output.get("status") == "error":
|
||||
raise RuntimeError("worker error")
|
||||
|
||||
@@ -34,6 +34,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm_omni.diffusion.forward_context import set_forward_context
|
||||
from vllm_omni.diffusion.ipc import pack_diffusion_output_shm
|
||||
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
|
||||
from vllm_omni.diffusion.profiler import CurrentProfiler
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
@@ -63,6 +64,7 @@ class DiffusionWorker:
|
||||
local_rank: int,
|
||||
rank: int,
|
||||
od_config: OmniDiffusionConfig,
|
||||
skip_load_model: bool = False,
|
||||
):
|
||||
self.local_rank = local_rank
|
||||
self.rank = rank
|
||||
@@ -79,8 +81,9 @@ class DiffusionWorker:
|
||||
od_config=self.od_config,
|
||||
device=self.device,
|
||||
)
|
||||
self.load_model(load_format=self.od_config.diffusion_load_format)
|
||||
self.init_lora_manager()
|
||||
if not skip_load_model:
|
||||
self.load_model(load_format=self.od_config.diffusion_load_format)
|
||||
self.init_lora_manager()
|
||||
logger.info(f"Worker {self.rank}: Initialization complete.")
|
||||
|
||||
def init_device(self) -> None:
|
||||
@@ -123,6 +126,7 @@ class DiffusionWorker:
|
||||
tensor_parallel_size=parallel_config.tensor_parallel_size,
|
||||
pipeline_parallel_size=parallel_config.pipeline_parallel_size,
|
||||
fully_shard_degree=parallel_config.hsdp_shard_size if parallel_config.use_hsdp else 1,
|
||||
hsdp_replicate_size=parallel_config.hsdp_replicate_size if parallel_config.use_hsdp else 1,
|
||||
)
|
||||
init_workspace_manager(self.device)
|
||||
|
||||
@@ -373,6 +377,10 @@ class WorkerProc:
|
||||
def return_result(self, output: DiffusionOutput):
|
||||
"""Reply to client, only on rank 0."""
|
||||
if self.result_mq is not None:
|
||||
try:
|
||||
pack_diffusion_output_shm(output)
|
||||
except Exception as e:
|
||||
logger.warning("SHM pack failed, falling back to raw enqueue: %s", e)
|
||||
self.result_mq.enqueue(output)
|
||||
|
||||
def recv_message(self):
|
||||
@@ -524,10 +532,15 @@ class WorkerWrapperBase:
|
||||
worker_class = self._prepare_worker_class()
|
||||
|
||||
# Create the actual worker instance
|
||||
# When custom_pipeline_args is provided, skip initial model loading
|
||||
# since re_init_pipeline will handle it. This avoids allocating memory
|
||||
# through CuMemAllocator twice, which causes assertion failures in
|
||||
# sleep mode.
|
||||
self.worker = worker_class(
|
||||
local_rank=gpu_id,
|
||||
rank=gpu_id,
|
||||
od_config=od_config,
|
||||
skip_load_model=(self.custom_pipeline_args is not None),
|
||||
)
|
||||
|
||||
# Re-initialize pipeline with custom pipeline if provided
|
||||
|
||||
@@ -12,6 +12,7 @@ from vllm.logger import init_logger
|
||||
|
||||
from .factory import OmniConnectorFactory
|
||||
from .utils.config import ConnectorSpec
|
||||
from .utils.kv_utils import normalize_layer_kv
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -200,8 +201,12 @@ class OmniKVTransferManager:
|
||||
logger.warning(f"Request {req_id} has no block IDs, skipping")
|
||||
continue
|
||||
|
||||
custom_metadata = data.get("custom_metadata")
|
||||
|
||||
# Extract KV cache from GPU blocks -> CPU tensors
|
||||
kv_data = self._extract_kv_cache(req_id, block_ids, seq_len, kv_caches, block_size, cache_dtype)
|
||||
kv_data = self._extract_kv_cache(
|
||||
req_id, block_ids, seq_len, kv_caches, block_size, cache_dtype, custom_metadata
|
||||
)
|
||||
if kv_data:
|
||||
# Resolve global request ID if available
|
||||
transfer_req_id = request_id_resolver(req_id) if request_id_resolver else req_id
|
||||
@@ -224,6 +229,7 @@ class OmniKVTransferManager:
|
||||
kv_caches: list[LayerKV],
|
||||
block_size: int,
|
||||
cache_dtype: str,
|
||||
custom_metadata: dict[str, Any] | None = None,
|
||||
) -> KVCacheTransferData | None:
|
||||
"""Extract KV cache from GPU blocks for a single request.
|
||||
|
||||
@@ -234,6 +240,7 @@ class OmniKVTransferManager:
|
||||
kv_caches: List of KV cache (tensor or tuple) per layer
|
||||
block_size: Size of each cache block
|
||||
cache_dtype: Data type of the cache
|
||||
custom_metadata: Optional custom metadata to include
|
||||
|
||||
Note: If key/value block counts differ, extraction uses only the overlapping
|
||||
block range. Extra key/value blocks are ignored, so returned KV may be partial.
|
||||
@@ -246,7 +253,7 @@ class OmniKVTransferManager:
|
||||
value_cache: list[torch.Tensor | None] = [None] * num_layers
|
||||
|
||||
for layer_idx, layer_kv in enumerate(kv_caches):
|
||||
kv_pair = self._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=layer_idx)
|
||||
kv_pair = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=layer_idx)
|
||||
if kv_pair is None:
|
||||
continue
|
||||
key_blocks, value_blocks = kv_pair
|
||||
@@ -289,57 +296,10 @@ class OmniKVTransferManager:
|
||||
"num_layers": num_layers,
|
||||
"dtype": str(cache_dtype),
|
||||
"seq_len": seq_len,
|
||||
**(custom_metadata or {}),
|
||||
},
|
||||
)
|
||||
|
||||
def _normalize_layer_kv(
|
||||
self,
|
||||
layer_kv: LayerKV,
|
||||
req_id: str,
|
||||
layer_idx: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor] | None:
|
||||
"""Normalize one layer KV cache to a `(key_blocks, value_blocks)` tuple.
|
||||
|
||||
Args:
|
||||
layer_kv: The raw KV cache (tensor or tuple) for the layer
|
||||
req_id: Request ID for logging
|
||||
layer_idx: Layer index for logging
|
||||
|
||||
Returns:
|
||||
Tuple of (key_blocks, value_blocks) if valid, None otherwise
|
||||
"""
|
||||
if isinstance(layer_kv, torch.Tensor):
|
||||
if layer_kv.ndim < 3 or layer_kv.shape[0] != 2:
|
||||
logger.warning(
|
||||
f"Layer {layer_idx} for request {req_id} has invalid stacked KV shape: "
|
||||
f"expected [2, blocks, block_size, ...], got {tuple(layer_kv.shape)}"
|
||||
)
|
||||
return None
|
||||
key_blocks = layer_kv[0]
|
||||
value_blocks = layer_kv[1]
|
||||
elif isinstance(layer_kv, tuple):
|
||||
if len(layer_kv) != 2:
|
||||
logger.warning(
|
||||
f"Layer {layer_idx} for request {req_id} has KV pair length {len(layer_kv)} (expected 2)"
|
||||
)
|
||||
return None
|
||||
key_blocks, value_blocks = layer_kv
|
||||
if not isinstance(key_blocks, torch.Tensor) or not isinstance(value_blocks, torch.Tensor):
|
||||
logger.warning(f"Layer {layer_idx} for request {req_id} has non-tensor KV pair entries")
|
||||
return None
|
||||
else:
|
||||
logger.warning(f"Layer {layer_idx} for request {req_id} has unsupported KV type {type(layer_kv).__name__}")
|
||||
return None
|
||||
# ensure key/value blocks are at least 2D for block indexing
|
||||
if key_blocks.ndim < 2 or value_blocks.ndim < 2:
|
||||
logger.warning(
|
||||
f"Layer {layer_idx} for request {req_id} has invalid KV block shape: "
|
||||
f"got key={tuple(key_blocks.shape)} value={tuple(value_blocks.shape)}"
|
||||
)
|
||||
return None
|
||||
|
||||
return key_blocks, value_blocks
|
||||
|
||||
def _transfer_kv_cache(self, kv_data: KVCacheTransferData, transfer_req_id: str) -> None:
|
||||
"""Transfer KV cache data to downstream stage via OmniConnector.
|
||||
|
||||
@@ -496,6 +456,8 @@ class OmniKVTransferManager:
|
||||
|
||||
if "metadata" in data:
|
||||
req.kv_metadata = data["metadata"]
|
||||
if hasattr(req, "sampling_params") and req.sampling_params is not None:
|
||||
req.sampling_params.kv_metadata = data["metadata"]
|
||||
|
||||
# Legacy compatibility method
|
||||
def receive_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool:
|
||||
|
||||
@@ -39,7 +39,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
|
||||
self.scheduler_max_num_seqs = vllm_config.scheduler_config.max_num_seqs
|
||||
self.connector = self.create_connector(model_config)
|
||||
super().__init__(model_config)
|
||||
self.model_mode = getattr(model_config, "worker_type", "ar")
|
||||
self.model_mode = getattr(model_config, "worker_type", None) or "ar"
|
||||
# State specific to Chunk management
|
||||
self.custom_process_next_stage_input_func = None
|
||||
custom_process_next_stage_input_func = getattr(model_config, "custom_process_next_stage_input_func", None)
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Utility helpers for KV cache manipulation."""
|
||||
|
||||
import torch
|
||||
from vllm.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor]
|
||||
|
||||
|
||||
def normalize_layer_kv(
|
||||
layer_kv: LayerKV,
|
||||
*,
|
||||
req_id: str = "",
|
||||
layer_idx: int = -1,
|
||||
) -> tuple[torch.Tensor, torch.Tensor] | None:
|
||||
"""Normalize one layer KV cache to a ``(key_blocks, value_blocks)`` tuple.
|
||||
|
||||
In vLLM, different attention backends return paged-attention KV blocks
|
||||
with different layouts. For example:
|
||||
|
||||
* **FlashAttention** (``vllm/v1/attention/backends/flash_attn.py``)
|
||||
returns shape ``(2, num_blocks, block_size, num_kv_heads, head_size)``
|
||||
– the key/value dimension is at **dim 0**.
|
||||
* **FlashInfer** (``vllm/v1/attention/backends/flashinfer.py``)
|
||||
returns shape ``(num_blocks, 2, block_size, num_kv_heads, head_size)``
|
||||
– the key/value dimension is at **dim 1**.
|
||||
|
||||
This utility handles both layouts (and the tuple case) so that
|
||||
downstream code can work with any backend.
|
||||
|
||||
Supported layouts:
|
||||
|
||||
* **Stacked tensor** ``[2, num_blocks, block_size, n_heads, head_dim]`` –
|
||||
dim-0 selects key / value.
|
||||
* **Stacked tensor** ``[num_blocks, 2, block_size, n_heads, head_dim]`` –
|
||||
dim-1 selects key / value.
|
||||
* **Tuple** ``(key_tensor, value_tensor)`` – returned as-is after
|
||||
validation.
|
||||
|
||||
Args:
|
||||
layer_kv: The raw KV cache (tensor or tuple) for the layer.
|
||||
req_id: Request ID used only for diagnostic log messages.
|
||||
layer_idx: Layer index used only for diagnostic log messages.
|
||||
|
||||
Returns:
|
||||
``(key_blocks, value_blocks)`` if *layer_kv* is valid, ``None``
|
||||
otherwise.
|
||||
"""
|
||||
if isinstance(layer_kv, torch.Tensor):
|
||||
if layer_kv.ndim >= 3 and layer_kv.shape[0] == 2:
|
||||
key_blocks = layer_kv[0]
|
||||
value_blocks = layer_kv[1]
|
||||
elif layer_kv.ndim >= 3 and layer_kv.shape[1] == 2:
|
||||
key_blocks = layer_kv[:, 0]
|
||||
value_blocks = layer_kv[:, 1]
|
||||
else:
|
||||
logger.warning(
|
||||
f"Layer {layer_idx} for request {req_id} has invalid stacked KV shape: "
|
||||
f"expected [2, ...] or [..., 2, ...] at dim 0/1, got {tuple(layer_kv.shape)}"
|
||||
)
|
||||
return None
|
||||
elif isinstance(layer_kv, tuple):
|
||||
if len(layer_kv) != 2:
|
||||
logger.warning(f"Layer {layer_idx} for request {req_id} has KV pair length {len(layer_kv)} (expected 2)")
|
||||
return None
|
||||
key_blocks, value_blocks = layer_kv
|
||||
if not isinstance(key_blocks, torch.Tensor) or not isinstance(value_blocks, torch.Tensor):
|
||||
logger.warning(f"Layer {layer_idx} for request {req_id} has non-tensor KV pair entries")
|
||||
return None
|
||||
else:
|
||||
logger.warning(f"Layer {layer_idx} for request {req_id} has unsupported KV type {type(layer_kv).__name__}")
|
||||
return None
|
||||
|
||||
# ensure key/value blocks are at least 2D for block indexing
|
||||
if key_blocks.ndim < 2 or value_blocks.ndim < 2:
|
||||
logger.warning(
|
||||
f"Layer {layer_idx} for request {req_id} has invalid KV block shape: "
|
||||
f"got key={tuple(key_blocks.shape)} value={tuple(value_blocks.shape)}"
|
||||
)
|
||||
return None
|
||||
|
||||
return key_blocks, value_blocks
|
||||
@@ -4,8 +4,8 @@ import asyncio
|
||||
import copy
|
||||
import time
|
||||
import weakref
|
||||
from collections.abc import AsyncGenerator, Iterable, Sequence
|
||||
from typing import Any
|
||||
from collections.abc import AsyncGenerator, Callable, Iterable, Sequence
|
||||
from typing import Any, TypeVar
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.inputs.preprocess import InputPreprocessor
|
||||
@@ -35,11 +35,20 @@ from vllm_omni.lora.request import LoRARequest
|
||||
from vllm_omni.metrics import OrchestratorAggregator
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
_R = TypeVar("_R")
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _weak_close_cleanup_async(stage_list, stage_in_queues, stage_out_queues, ray_pg, output_handler, zmq_ctx=None):
|
||||
def _weak_close_cleanup_async(
|
||||
stage_list, stage_in_queues, stage_out_queues, ray_pg, output_handler, zmq_ctx=None, inline_engine=None
|
||||
):
|
||||
"""Weak reference cleanup function for AsyncOmni instances."""
|
||||
if inline_engine is not None:
|
||||
try:
|
||||
inline_engine.close()
|
||||
except Exception as e:
|
||||
logger.warning("Failed to close inline diffusion engine: %s", e)
|
||||
if stage_list:
|
||||
for q in stage_in_queues:
|
||||
try:
|
||||
@@ -105,10 +114,17 @@ class AsyncOmni(OmniBase):
|
||||
self._pause_cond: asyncio.Condition = asyncio.Condition()
|
||||
self._paused: bool = False
|
||||
|
||||
# Sleep mode tracking
|
||||
self._is_sleeping: bool = False
|
||||
|
||||
# Request state tracking
|
||||
self.request_states: dict[str, ClientRequestState] = {}
|
||||
self.output_handler: asyncio.Task | None = None
|
||||
|
||||
# RPC results storage: {stage_id: {rpc_id: result}}
|
||||
# Used to avoid race condition between output_handler and collective_rpc
|
||||
self._rpc_results: dict[int, dict[str, dict[str, Any]]] = {}
|
||||
|
||||
super().__init__(model, **kwargs)
|
||||
|
||||
# Register weak reference cleanup (called on garbage collection)
|
||||
@@ -121,6 +137,7 @@ class AsyncOmni(OmniBase):
|
||||
self._ray_pg,
|
||||
self.output_handler,
|
||||
self._zmq_ctx,
|
||||
getattr(self, "_inline_engine", None),
|
||||
)
|
||||
|
||||
def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]:
|
||||
@@ -193,6 +210,8 @@ class AsyncOmni(OmniBase):
|
||||
"enforce_eager": kwargs.get("enforce_eager", False),
|
||||
"diffusion_load_format": kwargs.get("diffusion_load_format", "default"),
|
||||
"custom_pipeline_args": kwargs.get("custom_pipeline_args", None),
|
||||
"worker_extension_cls": kwargs.get("worker_extension_cls", None),
|
||||
"enable_sleep_mode": kwargs.get("enable_sleep_mode", False),
|
||||
"enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
|
||||
"num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
|
||||
},
|
||||
@@ -253,6 +272,26 @@ class AsyncOmni(OmniBase):
|
||||
self.io_processor = None
|
||||
self.model_config = None
|
||||
|
||||
def _setup_rpc_result_checkers(self) -> None:
|
||||
"""Override base class to use async-friendly RPC result checkers.
|
||||
|
||||
In the async path the output_handler task drains the output queues
|
||||
and stashes collective_rpc results into ``_rpc_results``. The
|
||||
checker therefore only needs to look there (no queue draining).
|
||||
"""
|
||||
for stage in self.stage_list:
|
||||
sid = stage.stage_id
|
||||
|
||||
def make_rpc_checker(stage_id: int):
|
||||
def rpc_checker(rpc_id: str) -> dict[str, Any] | None:
|
||||
if stage_id in self._rpc_results and rpc_id in self._rpc_results[stage_id]:
|
||||
return self._rpc_results[stage_id].pop(rpc_id)
|
||||
return None
|
||||
|
||||
return rpc_checker
|
||||
|
||||
stage._rpc_result_checker = make_rpc_checker(sid)
|
||||
|
||||
def shutdown(self):
|
||||
"""Shutdown, cleaning up the background proc and IPC.
|
||||
|
||||
@@ -299,6 +338,11 @@ class AsyncOmni(OmniBase):
|
||||
async with self._pause_cond:
|
||||
await self._pause_cond.wait_for(lambda: not self._paused)
|
||||
|
||||
if self._inline_diffusion:
|
||||
async for output in self._generate_inline(prompt, request_id, sampling_params_list, output_modalities):
|
||||
yield output
|
||||
return
|
||||
|
||||
logger.debug(f"[{self._name}] generate() called")
|
||||
try:
|
||||
# Start output handler on the first call to generate()
|
||||
@@ -390,6 +434,97 @@ class AsyncOmni(OmniBase):
|
||||
logger.info("[AsyncOrchestrator] Request %s aborted.", request_id)
|
||||
raise
|
||||
|
||||
async def _generate_inline(
|
||||
self,
|
||||
prompt: OmniPromptType,
|
||||
request_id: str,
|
||||
sampling_params_list: Sequence[OmniSamplingParams] | None = None,
|
||||
output_modalities: list[str] | None = None,
|
||||
) -> AsyncGenerator[OmniRequestOutput, None]:
|
||||
"""Generate using inline diffusion engine (no stage worker subprocess).
|
||||
|
||||
Eliminates Hop3 IPC overhead by running OmniDiffusion directly in the
|
||||
orchestrator process. The blocking generate() call is offloaded to a
|
||||
thread executor so the asyncio event loop remains responsive.
|
||||
"""
|
||||
_wall_start_ts = time.time()
|
||||
|
||||
if sampling_params_list is None:
|
||||
sampling_params_list = self.default_sampling_params_list
|
||||
sp0 = sampling_params_list[0]
|
||||
|
||||
stage = self.stage_list[0]
|
||||
final_stage_id_for_e2e = 0
|
||||
|
||||
metrics = OrchestratorAggregator(
|
||||
num_stages=1,
|
||||
log_stats=self.log_stats,
|
||||
wall_start_ts=_wall_start_ts,
|
||||
final_stage_id_for_e2e=final_stage_id_for_e2e,
|
||||
)
|
||||
metrics.stage_first_ts[0] = time.time()
|
||||
|
||||
logger.info(
|
||||
"[%s] Inline diffusion generate for request %s",
|
||||
self._name,
|
||||
request_id,
|
||||
)
|
||||
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
results = await loop.run_in_executor(
|
||||
None,
|
||||
self._inline_engine.generate,
|
||||
prompt,
|
||||
sp0,
|
||||
[request_id],
|
||||
)
|
||||
|
||||
for result in results:
|
||||
images = getattr(result, "images", None) or []
|
||||
finished = getattr(result, "finished", True)
|
||||
|
||||
output_to_yield = OmniRequestOutput(
|
||||
stage_id=0,
|
||||
final_output_type=stage.final_output_type,
|
||||
request_output=result,
|
||||
images=images,
|
||||
finished=finished,
|
||||
)
|
||||
|
||||
metrics.stage_last_ts[0] = time.time()
|
||||
yield output_to_yield
|
||||
|
||||
try:
|
||||
metrics.on_finalize_request(
|
||||
final_stage_id_for_e2e,
|
||||
request_id,
|
||||
_wall_start_ts,
|
||||
)
|
||||
metrics.build_and_log_summary()
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
"[%s] Failed to finalize inline metrics: %s",
|
||||
self._name,
|
||||
e,
|
||||
)
|
||||
|
||||
except (asyncio.CancelledError, GeneratorExit):
|
||||
logger.info(
|
||||
"[%s] Inline request %s cancelled.",
|
||||
self._name,
|
||||
request_id,
|
||||
)
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
"[%s] Inline diffusion failed for request %s: %s",
|
||||
self._name,
|
||||
request_id,
|
||||
e,
|
||||
)
|
||||
raise
|
||||
|
||||
async def _process_async_results(
|
||||
self,
|
||||
request_id: str,
|
||||
@@ -401,7 +536,10 @@ class AsyncOmni(OmniBase):
|
||||
) -> AsyncGenerator[OmniRequestOutput, None]:
|
||||
all_stages_finished = {stage_id: False for stage_id in range(final_stage_id_for_e2e + 1)}
|
||||
submit_flag = True
|
||||
_loop_iter = 0
|
||||
_last_progress_ts = time.time()
|
||||
while not all(all_stages_finished.values()):
|
||||
_loop_iter += 1
|
||||
for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]):
|
||||
if all_stages_finished[stage_id]:
|
||||
continue
|
||||
@@ -410,6 +548,7 @@ class AsyncOmni(OmniBase):
|
||||
except asyncio.QueueEmpty:
|
||||
await asyncio.sleep(0.001)
|
||||
continue
|
||||
_last_progress_ts = time.time()
|
||||
engine_outputs, finished, output_to_yield = self._process_single_result(
|
||||
result,
|
||||
stage,
|
||||
@@ -418,11 +557,20 @@ class AsyncOmni(OmniBase):
|
||||
)
|
||||
if submit_flag and stage_id == 0:
|
||||
submit_flag = False
|
||||
prompt_token_ids = engine_outputs.prompt_token_ids
|
||||
prompt_token_ids = getattr(engine_outputs, "prompt_token_ids", None)
|
||||
if prompt_token_ids is None:
|
||||
prompt_token_ids = []
|
||||
engine_input = copy.deepcopy(prompt)
|
||||
next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
|
||||
try:
|
||||
next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
|
||||
except Exception:
|
||||
raise
|
||||
engine_input["prompt_token_ids"] = [0] * next_prompt_len
|
||||
engine_input["multi_modal_data"] = engine_input["mm_processor_kwargs"] = None
|
||||
for _mm_key in ("mm_kwargs", "mm_hashes", "mm_placeholders", "multi_modal_uuids"):
|
||||
engine_input.pop(_mm_key, None)
|
||||
if engine_input.get("type") == "multimodal":
|
||||
engine_input["type"] = "token"
|
||||
for i in range(1, len(self.stage_list)):
|
||||
task = {
|
||||
"request_id": request_id,
|
||||
@@ -525,6 +673,7 @@ class AsyncOmni(OmniBase):
|
||||
raise RuntimeError(result)
|
||||
|
||||
engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
|
||||
|
||||
if isinstance(engine_outputs, list):
|
||||
engine_outputs = engine_outputs[0]
|
||||
|
||||
@@ -542,12 +691,17 @@ class AsyncOmni(OmniBase):
|
||||
images = engine_outputs.images
|
||||
|
||||
if stage.final_output_type == "image":
|
||||
# Propagate custom_output from inner diffusion output
|
||||
custom_output = {}
|
||||
if isinstance(engine_outputs, OmniRequestOutput):
|
||||
custom_output = engine_outputs.custom_output or {}
|
||||
output_to_yield = OmniRequestOutput(
|
||||
stage_id=stage_id,
|
||||
final_output_type=stage.final_output_type,
|
||||
request_output=engine_outputs,
|
||||
images=images,
|
||||
finished=finished,
|
||||
_custom_output=custom_output,
|
||||
)
|
||||
else:
|
||||
output_to_yield = OmniRequestOutput(
|
||||
@@ -597,6 +751,15 @@ class AsyncOmni(OmniBase):
|
||||
# so we wait for a short time and try again
|
||||
await asyncio.sleep(0.05)
|
||||
continue
|
||||
# Handle collective_rpc results separately to avoid
|
||||
# race condition with the polling in collective_rpc
|
||||
if result.get("type") == "collective_rpc_result":
|
||||
rpc_id = result.get("rpc_id")
|
||||
if rpc_id:
|
||||
if stage_id not in self._rpc_results:
|
||||
self._rpc_results[stage_id] = {}
|
||||
self._rpc_results[stage_id][rpc_id] = result
|
||||
continue
|
||||
req_id = result.get("request_id")
|
||||
req_state = request_states.get(req_id)
|
||||
if req_state is None:
|
||||
@@ -632,7 +795,8 @@ class AsyncOmni(OmniBase):
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
# Is None before the loop is started.
|
||||
if self._inline_diffusion:
|
||||
return self._inline_engine is not None
|
||||
return len(self._stage_in_queues) > 0
|
||||
|
||||
@property
|
||||
@@ -656,6 +820,10 @@ class AsyncOmni(OmniBase):
|
||||
return EngineDeadError()
|
||||
|
||||
async def abort(self, request_id: str | Iterable[str]) -> None:
|
||||
if self._inline_diffusion:
|
||||
if self._inline_engine is not None:
|
||||
self._inline_engine.engine.abort(request_id)
|
||||
return None
|
||||
abort_task = {"type": OmniStageTaskType.ABORT, "request_id": request_id}
|
||||
for stage in self.stage_list:
|
||||
stage.submit(abort_task)
|
||||
@@ -709,6 +877,8 @@ class AsyncOmni(OmniBase):
|
||||
OMNI: Required by upstream OpenAIServingModels.__init__ which
|
||||
accesses engine_client.renderer.
|
||||
"""
|
||||
if self.input_processor is None:
|
||||
return None
|
||||
return self.input_processor.renderer
|
||||
|
||||
async def do_log_stats(self) -> None:
|
||||
@@ -723,19 +893,73 @@ class AsyncOmni(OmniBase):
|
||||
async def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool:
|
||||
pass
|
||||
|
||||
async def collective_rpc(
|
||||
self,
|
||||
method: str | Callable[..., _R],
|
||||
timeout: float | None = None,
|
||||
args: tuple = (),
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
) -> list[_R]:
|
||||
"""Execute an RPC call on all stages asynchronously.
|
||||
|
||||
Args:
|
||||
method: Name of the method to execute or callable
|
||||
timeout: Optional timeout in seconds
|
||||
args: Positional arguments for the method
|
||||
kwargs: Keyword arguments for the method
|
||||
|
||||
Returns:
|
||||
List of results from each stage
|
||||
"""
|
||||
self._run_output_handler()
|
||||
# Run synchronous collective_rpc in thread pool to avoid blocking event loop
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
async def run_stage_rpc(stage: OmniStage) -> _R:
|
||||
return await loop.run_in_executor(
|
||||
None,
|
||||
stage.collective_rpc,
|
||||
method,
|
||||
timeout,
|
||||
args,
|
||||
kwargs,
|
||||
)
|
||||
|
||||
# Run all stages concurrently
|
||||
results = await asyncio.gather(*[run_stage_rpc(stage) for stage in self.stage_list])
|
||||
return list(results)
|
||||
|
||||
async def sleep(self, level: int = 1) -> None:
|
||||
pass
|
||||
self._is_sleeping = True
|
||||
await self.collective_rpc(method="sleep", args=(level,))
|
||||
|
||||
async def wake_up(self, tags: list[str] | None = None) -> None:
|
||||
pass
|
||||
self._is_sleeping = False
|
||||
await self.collective_rpc(method="wake_up", args=(tags,))
|
||||
|
||||
async def is_sleeping(self) -> bool:
|
||||
"""Check whether the engine is sleeping"""
|
||||
return False
|
||||
return getattr(self, "_is_sleeping", False)
|
||||
|
||||
async def add_lora(self, lora_request: LoRARequest) -> bool:
|
||||
"""Load a new LoRA adapter into the engine for future requests."""
|
||||
return False
|
||||
result = await self.collective_rpc(method="add_lora", args=(lora_request,))
|
||||
return result[0][0]
|
||||
|
||||
async def remove_lora(self, adapter_id: int) -> bool:
|
||||
"""Remove a LoRA adapter from the engine."""
|
||||
result = await self.collective_rpc(method="remove_lora", args=(adapter_id,))
|
||||
return result[0][0]
|
||||
|
||||
async def list_loras(self) -> list[int]:
|
||||
"""List all loaded LoRA adapter IDs."""
|
||||
result = await self.collective_rpc(method="list_loras")
|
||||
return result[0][0]
|
||||
|
||||
async def pin_lora(self, adapter_id: int) -> bool:
|
||||
"""Pin a LoRA adapter so it is not evicted from the cache."""
|
||||
result = await self.collective_rpc(method="pin_lora", args=(adapter_id,))
|
||||
return result[0][0]
|
||||
|
||||
async def encode(
|
||||
self,
|
||||
|
||||
+198
-10
@@ -9,7 +9,7 @@ import uuid
|
||||
import weakref
|
||||
from collections.abc import Callable, Generator, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import Any, Literal, overload
|
||||
from typing import Any, Literal, TypeVar, overload
|
||||
|
||||
import huggingface_hub
|
||||
import msgspec.msgpack
|
||||
@@ -40,18 +40,22 @@ from vllm_omni.entrypoints.omni_stage import OmniStage
|
||||
from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType
|
||||
from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
|
||||
from vllm_omni.entrypoints.utils import (
|
||||
filter_dataclass_kwargs,
|
||||
get_final_stage_id_for_e2e,
|
||||
inject_omni_kv_config,
|
||||
load_and_resolve_stage_configs,
|
||||
)
|
||||
from vllm_omni.entrypoints.zmq_utils import ZmqQueue
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams
|
||||
from vllm_omni.lora.request import LoRARequest
|
||||
from vllm_omni.metrics import OrchestratorAggregator, StageRequestStats
|
||||
from vllm_omni.model_executor.model_loader.weight_utils import (
|
||||
download_weights_from_hf_specific,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
_R = TypeVar("_R")
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@@ -177,6 +181,13 @@ class OmniBase:
|
||||
self._handshake_seen: set[int] = set() # Track which stage IDs have completed ZMQ handshake
|
||||
self._single_stage_id: int | None = None # Optional: deploy only a specific stage ID
|
||||
|
||||
# Sleep mode tracking
|
||||
self._is_sleeping: bool = False
|
||||
|
||||
# RPC results storage: {stage_id: {rpc_id: result}}
|
||||
# Used by collective_rpc to retrieve results collected from the output queue
|
||||
self._rpc_results: dict[int, dict[str, dict[str, Any]]] = {}
|
||||
|
||||
# Initialize stages - each stage will create appropriate instance based on stage_type
|
||||
# Stage workers will automatically create OmniLLM or OmniDiffusion instances
|
||||
# based on stage_type in YAML config (handled in omni_stage.py)
|
||||
@@ -304,6 +315,9 @@ class OmniBase:
|
||||
|
||||
def _initialize_stages(self, model: str, kwargs: dict[str, Any]) -> None:
|
||||
"""Initialize stage list management."""
|
||||
self._inline_diffusion = False
|
||||
self._inline_engine = None
|
||||
|
||||
stage_init_timeout = kwargs.get("stage_init_timeout", 20)
|
||||
shm_threshold_bytes = kwargs.get("shm_threshold_bytes", 65536)
|
||||
init_timeout = kwargs.get("init_timeout", 300)
|
||||
@@ -351,6 +365,13 @@ class OmniBase:
|
||||
self.output_modalities = [st.final_output_type for st in self.stage_list]
|
||||
logger.info(f"[{self._name}] Loaded {len(self.stage_list)} stages")
|
||||
|
||||
# Phase 1 optimization: for a single diffusion stage in async mode,
|
||||
# run the engine directly in the orchestrator process to eliminate
|
||||
# the stage worker subprocess and its IPC serialization overhead.
|
||||
if len(self.stage_list) == 1 and self.stage_list[0].stage_type == "diffusion" and self.is_async:
|
||||
self._init_inline_diffusion_engine(model, self.stage_configs[0], kwargs)
|
||||
return
|
||||
|
||||
if self.worker_backend == "ray":
|
||||
self._queue_cls = get_ray_queue_class()
|
||||
else:
|
||||
@@ -362,6 +383,71 @@ class OmniBase:
|
||||
self._start_stages(model)
|
||||
# Wait for all stages to report readiness before seeding
|
||||
self._wait_for_stages_ready(timeout=init_timeout)
|
||||
# Set up RPC result checkers so that collective_rpc works
|
||||
self._setup_rpc_result_checkers()
|
||||
|
||||
def _init_inline_diffusion_engine(
|
||||
self,
|
||||
model: str,
|
||||
stage_config: Any,
|
||||
kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
"""Initialize diffusion engine directly in the orchestrator process.
|
||||
|
||||
For single-stage diffusion pipelines, this eliminates the stage worker
|
||||
subprocess and the associated Hop3 IPC serialization overhead.
|
||||
GPU workers for tensor parallelism are still spawned by the
|
||||
DiffusionExecutor as separate processes.
|
||||
"""
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
|
||||
from vllm_omni.entrypoints.stage_utils import (
|
||||
_to_dict,
|
||||
load_func_from_config,
|
||||
set_stage_devices,
|
||||
)
|
||||
|
||||
stage_id = stage_config.stage_id
|
||||
engine_args = _to_dict(stage_config.engine_args)
|
||||
runtime_cfg = _to_dict(getattr(stage_config, "runtime", {}))
|
||||
|
||||
if os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn":
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
try:
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
device_type = current_omni_platform.device_type
|
||||
set_stage_devices(stage_id, runtime_cfg.get("devices"), device_type=device_type)
|
||||
except Exception as e:
|
||||
logger.warning("Device setup for inline diffusion failed: %s", e)
|
||||
|
||||
engine_args = filter_dataclass_kwargs(OmniDiffusionConfig, engine_args)
|
||||
engine_args.pop("model_stage", None)
|
||||
engine_args.pop("model", None)
|
||||
|
||||
cfg_kv_collect_func = load_func_from_config(getattr(stage_config, "cfg_kv_collect_func", None))
|
||||
|
||||
self._inline_engine = OmniDiffusion(
|
||||
model=model,
|
||||
stage_id=stage_id,
|
||||
engine_input_source=getattr(stage_config, "engine_input_source", []),
|
||||
cfg_kv_collect_func=cfg_kv_collect_func,
|
||||
**engine_args,
|
||||
)
|
||||
self._inline_diffusion = True
|
||||
|
||||
# These attributes are normally set by AsyncOmni._wait_for_stages_ready
|
||||
# but we skip that for inline mode. Set them to None since there is no
|
||||
# LLM stage to provide them.
|
||||
self.input_processor = None
|
||||
self.io_processor = None
|
||||
self.model_config = None
|
||||
|
||||
logger.info(
|
||||
"[%s] Inline diffusion mode active – stage worker subprocess bypassed",
|
||||
self._name,
|
||||
)
|
||||
|
||||
def _is_async_chunk_enable(self, stage_args: list) -> bool:
|
||||
"""get async chunk flag"""
|
||||
@@ -517,6 +603,117 @@ class OmniBase:
|
||||
profiler = getattr(profiler_config, "profiler", None)
|
||||
return profiler is not None
|
||||
|
||||
def _setup_rpc_result_checkers(self) -> None:
|
||||
"""Set up RPC result checkers for all stages.
|
||||
|
||||
Each checker reads from the shared ``_rpc_results`` dict so that
|
||||
``OmniStage.collective_rpc`` can retrieve results that were
|
||||
collected by the orchestrator (or, in the sync path, by the
|
||||
stage-level checker that drains the output queue).
|
||||
|
||||
Uses a weak reference to ``self`` to avoid a circular reference
|
||||
(OmniBase → stage_list → OmniStage → closure → OmniBase) that
|
||||
would prevent the instance from being freed by reference counting.
|
||||
"""
|
||||
weak_self = weakref.ref(self)
|
||||
for stage in self.stage_list:
|
||||
sid = stage.stage_id
|
||||
|
||||
def make_rpc_checker(stage_id: int):
|
||||
def rpc_checker(rpc_id: str) -> dict[str, Any] | None:
|
||||
_self = weak_self()
|
||||
if _self is None:
|
||||
return None
|
||||
# First check the shared dict
|
||||
if stage_id in _self._rpc_results and rpc_id in _self._rpc_results[stage_id]:
|
||||
return _self._rpc_results[stage_id].pop(rpc_id)
|
||||
# In the sync path there is no background output handler,
|
||||
# so drain the output queue ourselves and stash any
|
||||
# non-RPC results back.
|
||||
out_q = _self._stage_out_queues[stage_id] if stage_id < len(_self._stage_out_queues) else None
|
||||
if out_q is not None:
|
||||
import queue as _queue
|
||||
|
||||
try:
|
||||
while True:
|
||||
item = out_q.get_nowait()
|
||||
if isinstance(item, dict) and item.get("type") == "collective_rpc_result":
|
||||
item_rpc_id = item.get("rpc_id")
|
||||
if item_rpc_id == rpc_id:
|
||||
return item
|
||||
# Stash for another caller
|
||||
if stage_id not in _self._rpc_results:
|
||||
_self._rpc_results[stage_id] = {}
|
||||
_self._rpc_results[stage_id][item_rpc_id] = item
|
||||
else:
|
||||
# Non-RPC item — put it back
|
||||
out_q.put(item)
|
||||
break
|
||||
except _queue.Empty:
|
||||
pass
|
||||
return None
|
||||
|
||||
return rpc_checker
|
||||
|
||||
stage._rpc_result_checker = make_rpc_checker(sid)
|
||||
|
||||
def collective_rpc(
|
||||
self,
|
||||
method: str | Callable[..., _R],
|
||||
timeout: float | None = None,
|
||||
args: tuple = (),
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
) -> list[_R]:
|
||||
"""Execute an RPC call on all stage workers.
|
||||
|
||||
Args:
|
||||
method: Name of the worker method to execute, or a callable.
|
||||
timeout: Maximum time in seconds to wait for execution.
|
||||
args: Positional arguments to pass to the worker method.
|
||||
kwargs: Keyword arguments to pass to the worker method.
|
||||
|
||||
Returns:
|
||||
A list containing the results from each stage.
|
||||
"""
|
||||
results: list[_R] = []
|
||||
for stage in self.stage_list:
|
||||
results.append(
|
||||
stage.collective_rpc(
|
||||
method=method,
|
||||
timeout=timeout,
|
||||
args=args,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
)
|
||||
return results
|
||||
|
||||
def sleep(self, level: int = 1) -> None:
|
||||
"""Put the engine to sleep to free up resources.
|
||||
|
||||
Args:
|
||||
level: Sleep level (1 = light sleep, higher = deeper sleep).
|
||||
"""
|
||||
self._is_sleeping = True
|
||||
self.collective_rpc(method="sleep", args=(level,))
|
||||
|
||||
def wake_up(self, tags: list[str] | None = None) -> None:
|
||||
"""Wake the engine up from sleep.
|
||||
|
||||
Args:
|
||||
tags: Optional list of tags to selectively wake components.
|
||||
"""
|
||||
self._is_sleeping = False
|
||||
self.collective_rpc(method="wake_up", args=(tags,))
|
||||
|
||||
def is_sleeping(self) -> bool:
|
||||
"""Check whether the engine is sleeping."""
|
||||
return getattr(self, "_is_sleeping", False)
|
||||
|
||||
def add_lora(self, lora_request: LoRARequest) -> bool:
|
||||
"""Load a new LoRA adapter into the engine for future requests."""
|
||||
result = self.collective_rpc(method="add_lora", args=(lora_request,))
|
||||
return result[0][0]
|
||||
|
||||
def start_profile(self, stages: list[int] | None = None) -> None:
|
||||
"""Start profiling for specified stages.
|
||||
|
||||
@@ -1135,15 +1332,6 @@ class Omni(OmniBase):
|
||||
final_output_type=stage.final_output_type, # type: ignore[attr-defined]
|
||||
request_output=engine_outputs,
|
||||
)
|
||||
try:
|
||||
if stage.final_output_type == "text" or metrics.log_stats:
|
||||
output_to_yield.metrics = metrics.build_output_metrics(stage_id, req_id)
|
||||
except Exception as e:
|
||||
# Make metrics contract explicit on failure.
|
||||
output_to_yield.metrics = {}
|
||||
logger.exception(
|
||||
f"[{self._name}] Failed to attach output metrics for req {req_id} at stage {stage_id}: {e}",
|
||||
)
|
||||
|
||||
# Record audio generated frames (only when finished)
|
||||
try:
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.transformers_utils.config import get_hf_file_to_dict
|
||||
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
|
||||
@@ -12,6 +14,8 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class OmniDiffusion:
|
||||
"""
|
||||
@@ -108,6 +112,7 @@ class OmniDiffusion:
|
||||
sampling_params: OmniDiffusionSamplingParams,
|
||||
request_ids: list[str] = [],
|
||||
) -> list[OmniRequestOutput]:
|
||||
_t0 = time.perf_counter()
|
||||
if isinstance(prompts, (str, dict)):
|
||||
prompts = [prompts]
|
||||
else:
|
||||
@@ -118,7 +123,10 @@ class OmniDiffusion:
|
||||
request_ids.extend(f"{i + len(request_ids)}_{uuid.uuid4()}" for i in range(len(prompts) - len(request_ids)))
|
||||
|
||||
request = OmniDiffusionRequest(prompts, sampling_params, request_ids)
|
||||
return self._run_engine(request)
|
||||
result = self._run_engine(request)
|
||||
_t_ms = (time.perf_counter() - _t0) * 1000
|
||||
logger.info("OmniDiffusion.generate total: %.2f ms", _t_ms)
|
||||
return result
|
||||
|
||||
def _run_engine(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
|
||||
return self.engine.step(request)
|
||||
|
||||
@@ -15,10 +15,11 @@ import queue
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Sequence
|
||||
import uuid
|
||||
from collections.abc import Callable, Sequence
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import fields
|
||||
from typing import Any, Literal, cast
|
||||
from typing import Any, Literal, TypeVar, cast
|
||||
|
||||
from vllm import PromptType, RequestOutput
|
||||
from vllm.inputs import TextPrompt
|
||||
@@ -60,6 +61,8 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, O
|
||||
from vllm_omni.metrics import count_tokens_from_outputs
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
_R = TypeVar("_R")
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@@ -302,6 +305,10 @@ class OmniStage:
|
||||
self._proc: mp.Process | None = None
|
||||
self._shm_threshold_bytes: int = 65536
|
||||
self._stage_init_timeout: int = stage_init_timeout
|
||||
# Callback used by the orchestrator's output handler to stash
|
||||
# collective_rpc results so that ``collective_rpc`` can retrieve
|
||||
# them without competing for the output queue.
|
||||
self._rpc_result_checker: Callable[[str], dict | None] | None = None
|
||||
|
||||
def set_engine(self, engine: LLMEngine) -> None:
|
||||
"""Set the LLM engine for this stage.
|
||||
@@ -406,6 +413,71 @@ class OmniStage:
|
||||
logger.error(f"[Stage-{self.stage_id}] Timeout waiting for profiler results.")
|
||||
return {}
|
||||
|
||||
def collective_rpc(
|
||||
self,
|
||||
method: str | Callable[..., _R],
|
||||
timeout: float | None = None,
|
||||
args: tuple = (),
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
) -> list[_R]:
|
||||
"""Execute an RPC call on all workers via the stage engine.
|
||||
|
||||
Args:
|
||||
method: Name of the worker method to execute, or a callable that
|
||||
is serialized and sent to all workers to execute.
|
||||
|
||||
If the method is a callable, it should accept an additional
|
||||
``self`` argument, in addition to the arguments passed in
|
||||
``args`` and ``kwargs``. The ``self`` argument will be the
|
||||
worker object.
|
||||
timeout: Maximum time in seconds to wait for execution. Raises a
|
||||
:class:`TimeoutError` on timeout. ``None`` means wait
|
||||
indefinitely.
|
||||
args: Positional arguments to pass to the worker method.
|
||||
kwargs: Keyword arguments to pass to the worker method.
|
||||
|
||||
Returns:
|
||||
A list containing the results from each worker.
|
||||
|
||||
Note:
|
||||
It is recommended to use this API to only pass control messages,
|
||||
and set up data-plane communication to pass data.
|
||||
"""
|
||||
assert self._in_q is not None and self._out_q is not None, "Queues must be attached before collective_rpc"
|
||||
|
||||
# Submit collective_rpc task to worker
|
||||
rpc_id = str(uuid.uuid4())
|
||||
self._in_q.put(
|
||||
{
|
||||
"type": OmniStageTaskType.COLLECTIVE_RPC,
|
||||
"rpc_id": rpc_id,
|
||||
"method": method,
|
||||
"timeout": timeout,
|
||||
"args": args,
|
||||
"kwargs": kwargs,
|
||||
}
|
||||
)
|
||||
|
||||
start_time = time.time()
|
||||
while True:
|
||||
if timeout is not None and (time.time() - start_time) > timeout:
|
||||
raise TimeoutError(f"collective_rpc timed out after {timeout} seconds")
|
||||
|
||||
# Check if result was already collected by the orchestrator's
|
||||
# output handler (stored in a shared dict).
|
||||
result = None
|
||||
if self._rpc_result_checker is not None:
|
||||
result = self._rpc_result_checker(rpc_id)
|
||||
|
||||
if result is not None:
|
||||
if result.get("type") == "collective_rpc_result":
|
||||
if result.get("rpc_id") == rpc_id:
|
||||
if "error" in result:
|
||||
raise RuntimeError(f"collective_rpc failed: {result['error']}")
|
||||
return result["result"]
|
||||
|
||||
time.sleep(0.001) # Small sleep to avoid busy waiting
|
||||
|
||||
def init_stage_worker(
|
||||
self,
|
||||
model: str,
|
||||
@@ -538,14 +610,24 @@ class OmniStage:
|
||||
try:
|
||||
self._in_q.put_nowait(SHUTDOWN_TASK)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to send shutdown to in_q: %s", e)
|
||||
# Queue may already be closed by the weak-ref cleanup
|
||||
# (e.g. ZmqQueue socket closed) — this is expected.
|
||||
logger.debug("Failed to send shutdown to in_q: %s", e)
|
||||
close_fn = getattr(self._in_q, "close", None)
|
||||
if callable(close_fn):
|
||||
close_fn()
|
||||
try:
|
||||
close_fn()
|
||||
except Exception:
|
||||
pass
|
||||
self._in_q = None
|
||||
if self._out_q is not None:
|
||||
close_fn = getattr(self._out_q, "close", None)
|
||||
if callable(close_fn):
|
||||
close_fn()
|
||||
try:
|
||||
close_fn()
|
||||
except Exception:
|
||||
pass
|
||||
self._out_q = None
|
||||
|
||||
if hasattr(self, "_ray_actor") and self._ray_actor:
|
||||
kill_ray_actor(self._ray_actor)
|
||||
@@ -874,6 +956,40 @@ def _stage_worker(
|
||||
out_q.put({"type": "profiler_result", "data": profiler_data})
|
||||
continue
|
||||
|
||||
# Handle collective_rpc commands
|
||||
if task_type == OmniStageTaskType.COLLECTIVE_RPC:
|
||||
rpc_id = task.get("rpc_id")
|
||||
rpc_method = task.get("method")
|
||||
rpc_timeout = task.get("timeout")
|
||||
rpc_args = task.get("args", ())
|
||||
rpc_kwargs = task.get("kwargs") or {}
|
||||
try:
|
||||
rpc_result = stage_engine.collective_rpc(
|
||||
method=rpc_method,
|
||||
timeout=rpc_timeout,
|
||||
args=rpc_args,
|
||||
kwargs=rpc_kwargs,
|
||||
)
|
||||
out_q.put(
|
||||
{
|
||||
"type": "collective_rpc_result",
|
||||
"rpc_id": rpc_id,
|
||||
"stage_id": stage_id,
|
||||
"result": rpc_result,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception("[Stage-%s] collective_rpc failed: %s", stage_id, e)
|
||||
out_q.put(
|
||||
{
|
||||
"type": "collective_rpc_result",
|
||||
"rpc_id": rpc_id,
|
||||
"stage_id": stage_id,
|
||||
"error": str(e),
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
batch_tasks: list[dict[str, Any]] = [task]
|
||||
tasks_failed_to_add_to_batch: list[dict[str, Any]] = []
|
||||
start_time = _time.time()
|
||||
@@ -1003,7 +1119,6 @@ def _stage_worker(
|
||||
gen_outputs.extend(results)
|
||||
_gen_t1 = _time.time()
|
||||
_gen_ms = (_gen_t1 - _gen_t0) * 1000.0
|
||||
logger.debug(f"Generate done: batch={len(batch_tasks)}, req_ids={batch_request_ids}, gen_ms={_gen_ms:.1f}")
|
||||
|
||||
# Group outputs per request id with fallback
|
||||
req_to_outputs: dict[Any, list[Any]] = {rid: [] for rid in batch_request_ids}
|
||||
@@ -1043,6 +1158,10 @@ def _stage_worker(
|
||||
_metrics.stage_stats = None
|
||||
try:
|
||||
use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes)
|
||||
except Exception:
|
||||
use_shm, payload = False, r_outputs
|
||||
|
||||
try:
|
||||
if use_shm:
|
||||
out_q.put(
|
||||
{
|
||||
@@ -1070,10 +1189,6 @@ def _stage_worker(
|
||||
"metrics": _metrics,
|
||||
}
|
||||
)
|
||||
logger.debug(
|
||||
"Enqueued result for request %s to downstream",
|
||||
rid,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception("Failed on batch %s: %s", batch_request_ids, e)
|
||||
_tb = traceback.format_exc()
|
||||
@@ -1402,6 +1517,54 @@ async def _stage_worker_async(
|
||||
# Send result back to orchestrator for STOP command
|
||||
if task_type == OmniStageTaskType.PROFILER_STOP:
|
||||
out_q.put({"type": "profiler_result", "data": profiler_data})
|
||||
elif task_type == OmniStageTaskType.COLLECTIVE_RPC:
|
||||
rpc_id = task.get("rpc_id")
|
||||
rpc_method = task.get("method")
|
||||
rpc_timeout = task.get("timeout")
|
||||
rpc_args = task.get("args", ())
|
||||
rpc_kwargs = task.get("kwargs") or {}
|
||||
try:
|
||||
if stage_type == "diffusion":
|
||||
# DiffusionEngine.collective_rpc is synchronous
|
||||
loop = asyncio.get_event_loop()
|
||||
rpc_result = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: stage_engine.engine.collective_rpc(
|
||||
method=rpc_method,
|
||||
timeout=rpc_timeout,
|
||||
args=rpc_args,
|
||||
kwargs=rpc_kwargs,
|
||||
),
|
||||
)
|
||||
else:
|
||||
rpc_result = await cast(AsyncLLM, stage_engine).collective_rpc(
|
||||
method=rpc_method,
|
||||
timeout=rpc_timeout,
|
||||
args=rpc_args,
|
||||
kwargs=rpc_kwargs,
|
||||
)
|
||||
out_q.put(
|
||||
{
|
||||
"type": "collective_rpc_result",
|
||||
"rpc_id": rpc_id,
|
||||
"stage_id": stage_id,
|
||||
"result": rpc_result,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
"[Stage-%s] collective_rpc failed: %s",
|
||||
stage_id,
|
||||
e,
|
||||
)
|
||||
out_q.put(
|
||||
{
|
||||
"type": "collective_rpc_result",
|
||||
"rpc_id": rpc_id,
|
||||
"stage_id": stage_id,
|
||||
"error": str(e),
|
||||
}
|
||||
)
|
||||
else:
|
||||
asyncio.create_task(generation_single_request(task))
|
||||
|
||||
@@ -1555,10 +1718,16 @@ def output_strip(r_output: RequestOutput | OmniRequestOutput, final_output: bool
|
||||
if mm_output is not None:
|
||||
r_output.multimodal_output = {}
|
||||
|
||||
custom_out = getattr(r_output, "_custom_output", None)
|
||||
if custom_out is not None:
|
||||
r_output._custom_output = {}
|
||||
|
||||
outputs = getattr(r_output, "outputs", None)
|
||||
if outputs is not None:
|
||||
for out in outputs:
|
||||
if getattr(out, "multimodal_output", None):
|
||||
out.multimodal_output = {}
|
||||
if getattr(out, "_custom_output", None):
|
||||
out._custom_output = {}
|
||||
|
||||
return r_output
|
||||
@@ -882,8 +882,113 @@ async def list_voices(raw_request: Request):
|
||||
if handler is None:
|
||||
return base(raw_request).create_error_response(message="The model does not support Speech API")
|
||||
|
||||
# Get all speakers (both model built-in and uploaded)
|
||||
speakers = sorted(handler.supported_speakers) if handler.supported_speakers else []
|
||||
return JSONResponse(content={"voices": speakers})
|
||||
|
||||
# Get uploaded speakers details
|
||||
uploaded_speakers = []
|
||||
if hasattr(handler, "uploaded_speakers"):
|
||||
for voice_name, info in handler.uploaded_speakers.items():
|
||||
uploaded_speakers.append(
|
||||
{
|
||||
"name": info.get("name", voice_name),
|
||||
"consent": info.get("consent", ""),
|
||||
"created_at": info.get("created_at", 0),
|
||||
"file_size": info.get("file_size", 0),
|
||||
"mime_type": info.get("mime_type", ""),
|
||||
}
|
||||
)
|
||||
|
||||
return JSONResponse(content={"voices": speakers, "uploaded_voices": uploaded_speakers})
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v1/audio/voices",
|
||||
responses={
|
||||
HTTPStatus.OK.value: {"model": dict},
|
||||
HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
|
||||
HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
|
||||
},
|
||||
)
|
||||
async def upload_voice(
|
||||
raw_request: Request,
|
||||
audio_sample: UploadFile = File(...),
|
||||
consent: str = Form(...),
|
||||
name: str = Form(...),
|
||||
):
|
||||
"""Upload a new voice sample for voice cloning.
|
||||
|
||||
Uploads an audio file that can be used as a reference for voice cloning
|
||||
in Base task TTS requests. The voice can then be referenced by name
|
||||
in subsequent TTS requests.
|
||||
|
||||
Args:
|
||||
audio_sample: Audio file (max 10MB)
|
||||
consent: Consent recording ID
|
||||
name: Name for the new voice
|
||||
raw_request: Raw FastAPI request
|
||||
|
||||
Returns:
|
||||
JSON response with voice information
|
||||
"""
|
||||
handler = Omnispeech(raw_request)
|
||||
if handler is None:
|
||||
return base(raw_request).create_error_response(message="The model does not support Speech API")
|
||||
|
||||
try:
|
||||
# Upload the voice
|
||||
result = await handler.upload_voice(audio_sample, consent, name)
|
||||
|
||||
return JSONResponse(content={"success": True, "voice": result})
|
||||
|
||||
except ValueError as e:
|
||||
return base(raw_request).create_error_response(message=str(e))
|
||||
except Exception as e:
|
||||
logger.exception(f"Failed to upload voice: {e}")
|
||||
return base(raw_request).create_error_response(message=f"Failed to upload voice: {str(e)}")
|
||||
|
||||
|
||||
@router.delete(
|
||||
"/v1/audio/voices/{name}",
|
||||
responses={
|
||||
HTTPStatus.OK.value: {"model": dict},
|
||||
HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
|
||||
HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
|
||||
},
|
||||
)
|
||||
async def delete_voice(name: str, raw_request: Request):
|
||||
"""Delete an uploaded voice.
|
||||
|
||||
Deletes the voice sample and associated metadata. Also removes any
|
||||
cached voice clone prompts for this voice.
|
||||
|
||||
Args:
|
||||
name: Name of the voice to delete
|
||||
raw_request: Raw FastAPI request
|
||||
|
||||
Returns:
|
||||
JSON response indicating success or failure
|
||||
"""
|
||||
handler = Omnispeech(raw_request)
|
||||
if handler is None:
|
||||
return base(raw_request).create_error_response(message="The model does not support Speech API")
|
||||
|
||||
try:
|
||||
# Delete the voice
|
||||
success = await handler.delete_voice(name)
|
||||
if not success:
|
||||
return JSONResponse(
|
||||
content={"success": False, "error": f"Voice '{name}' not found"},
|
||||
status_code=HTTPStatus.NOT_FOUND.value,
|
||||
)
|
||||
|
||||
return JSONResponse(content={"success": True, "message": f"Voice '{name}' deleted successfully"})
|
||||
|
||||
except ValueError as e:
|
||||
return base(raw_request).create_error_response(message=str(e))
|
||||
except Exception as e:
|
||||
logger.exception(f"Failed to delete voice '{name}': {e}")
|
||||
return base(raw_request).create_error_response(message=f"Failed to delete voice: {str(e)}")
|
||||
|
||||
|
||||
# Health and Model endpoints for diffusion mode
|
||||
|
||||
@@ -0,0 +1,243 @@
|
||||
"""
|
||||
Metadata manager for voice samples and cache information.
|
||||
|
||||
Provides a unified interface for managing metadata.json with
|
||||
concurrency safety and data consistency across multiple processes.
|
||||
"""
|
||||
|
||||
import fcntl
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MetadataManager:
|
||||
"""
|
||||
Manages metadata for uploaded speakers and cache information.
|
||||
|
||||
Features:
|
||||
1. Single source of truth for metadata
|
||||
2. Concurrency safety with threading locks
|
||||
3. Atomic read-modify-write operations
|
||||
4. Merge updates to preserve fields from different components
|
||||
"""
|
||||
|
||||
def __init__(self, metadata_file: Path):
|
||||
"""
|
||||
Initialize the metadata manager.
|
||||
|
||||
Args:
|
||||
metadata_file: Path to metadata.json file
|
||||
"""
|
||||
self.metadata_file = metadata_file
|
||||
self._lock = threading.Lock() # For intra-process concurrency
|
||||
self._metadata = self._load_from_disk()
|
||||
|
||||
# Create lock file for cross-process synchronization
|
||||
self.lock_file = metadata_file.with_suffix(".lock")
|
||||
self.lock_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _load_from_disk(self) -> dict[str, Any]:
|
||||
"""Load metadata from disk."""
|
||||
if not self.metadata_file.exists():
|
||||
return {"uploaded_speakers": {}}
|
||||
|
||||
try:
|
||||
with open(self.metadata_file) as f:
|
||||
return json.load(f)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load metadata from {self.metadata_file}: {e}")
|
||||
return {"uploaded_speakers": {}}
|
||||
|
||||
def _save_to_disk(self, metadata: dict[str, Any]) -> bool:
|
||||
"""Save metadata to disk."""
|
||||
try:
|
||||
self.metadata_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp = self.metadata_file.with_suffix(".tmp")
|
||||
with open(tmp, "w") as f:
|
||||
json.dump(metadata, f, indent=2)
|
||||
tmp.replace(self.metadata_file)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to save metadata to {self.metadata_file}: {e}")
|
||||
return False
|
||||
|
||||
# ================================
|
||||
# Core fix: single flock overwrites RMW
|
||||
# ================================
|
||||
def _update_with_file_lock(
|
||||
self, update_fn: Callable[[dict[str, Any]], dict[str, Any] | None]
|
||||
) -> dict[str, Any] | None:
|
||||
lock_fd = os.open(self.lock_file, os.O_CREAT | os.O_RDWR)
|
||||
try:
|
||||
fcntl.flock(lock_fd, fcntl.LOCK_EX)
|
||||
|
||||
metadata = self._load_from_disk()
|
||||
result = update_fn(metadata)
|
||||
if result is None:
|
||||
return None
|
||||
|
||||
if not self._save_to_disk(metadata):
|
||||
return None
|
||||
|
||||
self._metadata = metadata
|
||||
return result
|
||||
finally:
|
||||
fcntl.flock(lock_fd, fcntl.LOCK_UN)
|
||||
os.close(lock_fd)
|
||||
|
||||
def get_uploaded_speakers(self) -> dict[str, dict[str, Any]]:
|
||||
"""Get all uploaded speakers."""
|
||||
# Read directly from disk to ensure getting the latest data
|
||||
metadata = self._load_from_disk()
|
||||
return metadata.get("uploaded_speakers", {}).copy()
|
||||
|
||||
def get_speaker(self, speaker_key: str) -> dict[str, Any] | None:
|
||||
"""Get specific speaker information."""
|
||||
# Read directly from disk to ensure getting the latest data
|
||||
metadata = self._load_from_disk()
|
||||
speakers = metadata.get("uploaded_speakers", {})
|
||||
return speakers.get(speaker_key, {}).copy() if speaker_key in speakers else None
|
||||
|
||||
def update_speaker(self, speaker_key: str, updates: dict[str, Any]) -> bool:
|
||||
"""
|
||||
Update speaker information with merge semantics.
|
||||
|
||||
Uses file locking for cross-process atomic operations.
|
||||
"""
|
||||
with self._lock:
|
||||
|
||||
def _update(metadata: dict[str, Any]):
|
||||
speakers = metadata.setdefault("uploaded_speakers", {})
|
||||
entry = speakers.get(speaker_key, {})
|
||||
entry.update(updates)
|
||||
speakers[speaker_key] = entry
|
||||
return True
|
||||
|
||||
return self._update_with_file_lock(_update) is not None
|
||||
|
||||
def create_speaker(self, speaker_key: str, speaker_data: dict[str, Any]) -> bool:
|
||||
"""
|
||||
Create a new speaker entry.
|
||||
|
||||
Uses file locking for cross-process atomic operations.
|
||||
"""
|
||||
with self._lock:
|
||||
|
||||
def _create(metadata: dict[str, Any]):
|
||||
speakers = metadata.setdefault("uploaded_speakers", {})
|
||||
if speaker_key in speakers:
|
||||
logger.warning(f"Speaker {speaker_key} already exists")
|
||||
return None
|
||||
speakers[speaker_key] = speaker_data
|
||||
return True
|
||||
|
||||
return self._update_with_file_lock(_create) is not None
|
||||
|
||||
def update_cache_info(self, speaker_key: str, cache_file_path: Path, status: str = "ready") -> bool:
|
||||
"""
|
||||
Update cache information for a speaker.
|
||||
"""
|
||||
updates = {
|
||||
"cache_status": status,
|
||||
"cache_file": str(cache_file_path),
|
||||
"cache_generated_at": time.time(),
|
||||
}
|
||||
return self.update_speaker(speaker_key, updates)
|
||||
|
||||
def delete_speaker(self, speaker_key: str) -> dict[str, Any] | None:
|
||||
"""
|
||||
Delete a speaker from metadata and clean up associated files.
|
||||
|
||||
Uses file locking for cross-process atomic operations.
|
||||
|
||||
Args:
|
||||
speaker_key: Speaker name (lowercase)
|
||||
base_dir: Base directory for file validation (optional)
|
||||
|
||||
Returns:
|
||||
dict: Deleted speaker information if successful, None if speaker doesn't exist or error
|
||||
"""
|
||||
with self._lock:
|
||||
|
||||
def _delete(metadata: dict[str, Any]):
|
||||
speakers = metadata.get("uploaded_speakers", {})
|
||||
if speaker_key not in speakers:
|
||||
logger.warning(f"Speaker {speaker_key} not found in metadata")
|
||||
return None
|
||||
|
||||
speaker_info = speakers.pop(speaker_key)
|
||||
|
||||
# Clean up associated files
|
||||
deleted_files = self._cleanup_speaker_files(speaker_info)
|
||||
if deleted_files:
|
||||
logger.info(f"Deleted {len(deleted_files)} files for speaker {speaker_key}: {deleted_files}")
|
||||
|
||||
return speaker_info
|
||||
|
||||
return self._update_with_file_lock(_delete)
|
||||
|
||||
def _cleanup_speaker_files(self, speaker_info: dict[str, Any]) -> list[str]:
|
||||
"""
|
||||
Clean up files associated with a speaker.
|
||||
|
||||
Args:
|
||||
speaker_info: Speaker information dictionary
|
||||
base_dir: Base directory for file validation (optional)
|
||||
|
||||
Returns:
|
||||
list: List of successfully deleted file paths
|
||||
"""
|
||||
deleted_files = []
|
||||
|
||||
# Helper function to safely delete a file
|
||||
def safe_delete(file_path_str: str, description: str) -> bool:
|
||||
if not file_path_str:
|
||||
return False
|
||||
|
||||
try:
|
||||
file_path = Path(file_path_str)
|
||||
|
||||
# Check if file exists
|
||||
if not file_path.exists():
|
||||
logger.debug(f"{description} not found: {file_path}")
|
||||
return False
|
||||
|
||||
# Delete the file
|
||||
file_path.unlink()
|
||||
logger.info(f"Deleted {description}: {file_path}")
|
||||
deleted_files.append(str(file_path))
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete {description} {file_path_str}: {e}")
|
||||
return False
|
||||
|
||||
# Delete audio file
|
||||
audio_file = speaker_info.get("file_path")
|
||||
if audio_file:
|
||||
safe_delete(audio_file, "audio file")
|
||||
|
||||
# Delete cache file
|
||||
cache_file = speaker_info.get("cache_file")
|
||||
if cache_file:
|
||||
safe_delete(cache_file, "cache file")
|
||||
|
||||
return deleted_files
|
||||
|
||||
def reload_from_disk(self) -> bool:
|
||||
"""Force reload metadata from disk (useful for external changes)."""
|
||||
with self._lock:
|
||||
try:
|
||||
self._metadata = self._load_from_disk()
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to reload metadata from disk: {e}")
|
||||
return False
|
||||
@@ -1,24 +1,23 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import ipaddress
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import socket
|
||||
import re
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
from urllib.request import urlopen
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
from fastapi import Request
|
||||
from fastapi import Request, UploadFile
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
from vllm.entrypoints.openai.engine.serving import OpenAIServing
|
||||
from vllm.logger import init_logger
|
||||
from vllm.multimodal.media import MediaConnector
|
||||
from vllm.utils import random_uuid
|
||||
|
||||
from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
|
||||
from vllm_omni.entrypoints.openai.metadata_manager import MetadataManager
|
||||
from vllm_omni.entrypoints.openai.protocol.audio import (
|
||||
CreateAudio,
|
||||
OpenAICreateSpeechRequest,
|
||||
@@ -27,19 +26,6 @@ from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_REF_AUDIO_TIMEOUT_S = 15
|
||||
_REF_AUDIO_MAX_BYTES = 50 * 1024 * 1024 # 50 MB
|
||||
_REF_AUDIO_BLOCKED_NETWORKS = [
|
||||
ipaddress.ip_network("127.0.0.0/8"),
|
||||
ipaddress.ip_network("10.0.0.0/8"),
|
||||
ipaddress.ip_network("172.16.0.0/12"),
|
||||
ipaddress.ip_network("192.168.0.0/16"),
|
||||
ipaddress.ip_network("169.254.0.0/16"),
|
||||
ipaddress.ip_network("::1/128"),
|
||||
ipaddress.ip_network("fc00::/7"),
|
||||
ipaddress.ip_network("fe80::/10"),
|
||||
]
|
||||
|
||||
# TTS Configuration (currently supports Qwen3-TTS)
|
||||
_TTS_MODEL_STAGES: set[str] = {"qwen3_tts"}
|
||||
_TTS_LANGUAGES: set[str] = {
|
||||
@@ -60,9 +46,52 @@ _TTS_MAX_NEW_TOKENS_MIN = 1
|
||||
_TTS_MAX_NEW_TOKENS_MAX = 4096
|
||||
|
||||
|
||||
def _sanitize_filename(filename: str) -> str:
|
||||
"""Sanitize filename to prevent path traversal attacks.
|
||||
|
||||
Only allows alphanumeric characters, underscores, hyphens, and dots.
|
||||
Replaces any other characters with underscores.
|
||||
"""
|
||||
# Remove any path components
|
||||
filename = os.path.basename(filename)
|
||||
# Replace any non-alphanumeric, underscore, hyphen, or dot with underscore
|
||||
sanitized = re.sub(r"[^a-zA-Z0-9_.\-]", "_", filename)
|
||||
# Ensure filename is not empty
|
||||
if not sanitized:
|
||||
sanitized = "file"
|
||||
# Limit length to prevent potential issues
|
||||
if len(sanitized) > 255:
|
||||
sanitized = sanitized[:255]
|
||||
return sanitized
|
||||
|
||||
|
||||
def _validate_path_within_directory(file_path: Path, directory: Path) -> bool:
|
||||
"""Validate that file_path is within the specified directory.
|
||||
|
||||
Prevents path traversal attacks by ensuring the resolved path
|
||||
is within the target directory.
|
||||
"""
|
||||
try:
|
||||
# Resolve both paths to absolute paths
|
||||
file_path_resolved = file_path.resolve()
|
||||
directory_resolved = directory.resolve()
|
||||
# Check if file_path is within directory
|
||||
return directory_resolved in file_path_resolved.parents or directory_resolved == file_path_resolved
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
# Initialize uploaded speakers storage
|
||||
speech_voice_samples_dir = os.environ.get("SPEECH_VOICE_SAMPLES", "/tmp/voice_samples")
|
||||
self.uploaded_speakers_dir = Path(speech_voice_samples_dir)
|
||||
self.uploaded_speakers_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.metadata_file = self.uploaded_speakers_dir / "metadata.json"
|
||||
|
||||
# Initialize metadata manager
|
||||
self.metadata_manager = MetadataManager(self.metadata_file)
|
||||
|
||||
# Find and cache the TTS stage (if any) during initialization
|
||||
self._tts_stage = self._find_tts_stage()
|
||||
@@ -73,9 +102,16 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
|
||||
# Load supported speakers
|
||||
self.supported_speakers = self._load_supported_speakers()
|
||||
logger.info(f"Loaded {len(self.supported_speakers)} supported speakers: {sorted(self.supported_speakers)}")
|
||||
# Load uploaded speakers
|
||||
self.uploaded_speakers = self.metadata_manager.get_uploaded_speakers()
|
||||
|
||||
# Merge supported speakers with uploaded speakers
|
||||
self.supported_speakers.update(self.uploaded_speakers.keys())
|
||||
self._tts_tokenizer = None
|
||||
|
||||
logger.info(f"Loaded {len(self.supported_speakers)} supported speakers: {sorted(self.supported_speakers)}")
|
||||
logger.info(f"Loaded {len(self.uploaded_speakers)} uploaded speakers")
|
||||
|
||||
# Load speech tokenizer codec parameters for prompt length estimation
|
||||
self._codec_frame_rate: float | None = self._load_codec_frame_rate()
|
||||
|
||||
@@ -223,6 +259,191 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
logger.warning("Failed to estimate TTS prompt length, using fallback 2048: %s", e)
|
||||
return 2048
|
||||
|
||||
def _get_uploaded_audio_data(self, voice_name: str) -> str | None:
|
||||
"""Get base64 encoded audio data for uploaded voice."""
|
||||
voice_name_lower = voice_name.lower()
|
||||
if voice_name_lower not in self.uploaded_speakers:
|
||||
return None
|
||||
|
||||
speaker_info = self.uploaded_speakers[voice_name_lower]
|
||||
file_path = Path(speaker_info["file_path"])
|
||||
|
||||
if not file_path.exists():
|
||||
logger.warning(f"Audio file not found for voice {voice_name}: {file_path}")
|
||||
return None
|
||||
|
||||
try:
|
||||
# Read audio file
|
||||
with open(file_path, "rb") as f:
|
||||
audio_bytes = f.read()
|
||||
|
||||
# Encode to base64
|
||||
audio_b64 = base64.b64encode(audio_bytes).decode("utf-8")
|
||||
|
||||
# Get MIME type from file extension
|
||||
mime_type = speaker_info.get("mime_type", "audio/wav")
|
||||
|
||||
# Return as data URL
|
||||
return f"data:{mime_type};base64,{audio_b64}"
|
||||
except Exception as e:
|
||||
logger.error(f"Could not read audio file for voice {voice_name}: {e}")
|
||||
return None
|
||||
|
||||
async def upload_voice(self, audio_file: UploadFile, consent: str, name: str) -> dict:
|
||||
"""Upload a new voice sample."""
|
||||
# Validate file size (max 10MB)
|
||||
MAX_FILE_SIZE = 10 * 1024 * 1024 # 10MB
|
||||
audio_file.file.seek(0, 2) # Seek to end
|
||||
file_size = audio_file.file.tell()
|
||||
audio_file.file.seek(0) # Reset to beginning
|
||||
|
||||
if file_size > MAX_FILE_SIZE:
|
||||
raise ValueError(f"File size exceeds maximum limit of 10MB. Got {file_size} bytes.")
|
||||
|
||||
# Detect MIME type from filename if content_type is generic
|
||||
mime_type = audio_file.content_type
|
||||
if mime_type == "application/octet-stream":
|
||||
# Simple MIME type detection based on file extension
|
||||
filename_lower = audio_file.filename.lower()
|
||||
if filename_lower.endswith(".wav"):
|
||||
mime_type = "audio/wav"
|
||||
elif filename_lower.endswith((".mp3", ".mpeg")):
|
||||
mime_type = "audio/mpeg"
|
||||
elif filename_lower.endswith(".flac"):
|
||||
mime_type = "audio/flac"
|
||||
elif filename_lower.endswith(".ogg"):
|
||||
mime_type = "audio/ogg"
|
||||
elif filename_lower.endswith(".aac"):
|
||||
mime_type = "audio/aac"
|
||||
elif filename_lower.endswith(".webm"):
|
||||
mime_type = "audio/webm"
|
||||
elif filename_lower.endswith(".mp4"):
|
||||
mime_type = "audio/mp4"
|
||||
else:
|
||||
mime_type = "audio/wav" # Default
|
||||
|
||||
# Validate MIME type
|
||||
allowed_mime_types = {
|
||||
"audio/mpeg",
|
||||
"audio/wav",
|
||||
"audio/x-wav",
|
||||
"audio/ogg",
|
||||
"audio/aac",
|
||||
"audio/flac",
|
||||
"audio/webm",
|
||||
"audio/mp4",
|
||||
}
|
||||
|
||||
if mime_type not in allowed_mime_types:
|
||||
raise ValueError(f"Unsupported MIME type: {mime_type}. Allowed: {allowed_mime_types}")
|
||||
|
||||
# Normalize voice name
|
||||
voice_name_lower = name.lower()
|
||||
|
||||
# Check if voice already exists
|
||||
if voice_name_lower in self.uploaded_speakers:
|
||||
raise ValueError(f"Voice '{name}' already exists")
|
||||
|
||||
# Sanitize name and consent to prevent path traversal
|
||||
sanitized_name = _sanitize_filename(name)
|
||||
sanitized_consent = _sanitize_filename(consent)
|
||||
|
||||
# Generate filename with sanitized inputs
|
||||
timestamp = int(time.time())
|
||||
file_suffix = Path(audio_file.filename).suffix
|
||||
file_ext = file_suffix[1:] if file_suffix and len(file_suffix) > 1 else "wav"
|
||||
# Sanitize file extension as well
|
||||
sanitized_ext = _sanitize_filename(file_ext)
|
||||
if not sanitized_ext or sanitized_ext == "file":
|
||||
sanitized_ext = "wav"
|
||||
|
||||
filename = f"{sanitized_name}_{sanitized_consent}_{timestamp}.{sanitized_ext}"
|
||||
file_path = self.uploaded_speakers_dir / filename
|
||||
|
||||
# Double-check that the path is within the upload directory
|
||||
if not _validate_path_within_directory(file_path, self.uploaded_speakers_dir):
|
||||
raise ValueError("Invalid file path: potential path traversal attack detected")
|
||||
|
||||
# Save audio file
|
||||
try:
|
||||
with open(file_path, "wb") as f:
|
||||
content = await audio_file.read()
|
||||
f.write(content)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to save audio file: {e}")
|
||||
|
||||
# Create speaker data
|
||||
speaker_data = {
|
||||
"name": name,
|
||||
"consent": consent,
|
||||
"file_path": str(file_path),
|
||||
"created_at": timestamp,
|
||||
"mime_type": mime_type,
|
||||
"original_filename": audio_file.filename,
|
||||
"file_size": file_size,
|
||||
"cache_status": "pending", # The initial cache state is pending.
|
||||
"cache_file": None, # The initial cache file is empty.
|
||||
"cache_generated_at": None, # The initial cache generation time is empty.
|
||||
}
|
||||
|
||||
# Save metadata using metadata manager (concurrency safe)
|
||||
success = self.metadata_manager.create_speaker(voice_name_lower, speaker_data)
|
||||
if not success:
|
||||
# Clean up the saved file if metadata creation failed
|
||||
try:
|
||||
file_path.unlink()
|
||||
except Exception:
|
||||
pass
|
||||
raise ValueError(f"Failed to create metadata for voice '{name}' (possibly already exists)")
|
||||
|
||||
# Update in-memory cache
|
||||
self.uploaded_speakers[voice_name_lower] = speaker_data
|
||||
self.supported_speakers.add(voice_name_lower)
|
||||
|
||||
logger.info(f"Uploaded new voice '{name}' with consent ID '{consent}'")
|
||||
|
||||
# Return voice information without exposing the server file path
|
||||
return {
|
||||
"name": name,
|
||||
"consent": consent,
|
||||
"created_at": timestamp,
|
||||
"mime_type": mime_type,
|
||||
"file_size": file_size,
|
||||
}
|
||||
|
||||
async def delete_voice(self, name: str) -> bool:
|
||||
"""
|
||||
Delete an uploaded voice.
|
||||
|
||||
Args:
|
||||
name: Voice name to delete
|
||||
|
||||
Returns:
|
||||
bool: True if successful, False if voice doesn't exist
|
||||
"""
|
||||
voice_name_lower = name.lower()
|
||||
|
||||
# Check if voice exists in memory cache
|
||||
if voice_name_lower not in self.uploaded_speakers:
|
||||
logger.warning(f"Voice '{name}' not found in memory cache")
|
||||
return False
|
||||
|
||||
# Delete from metadata manager with file cleanup
|
||||
# Pass base_dir for path validation
|
||||
deleted_info = self.metadata_manager.delete_speaker(voice_name_lower)
|
||||
if not deleted_info:
|
||||
logger.error(f"Failed to delete voice '{name}' from metadata")
|
||||
return False
|
||||
|
||||
# Update in-memory cache
|
||||
if voice_name_lower in self.uploaded_speakers:
|
||||
del self.uploaded_speakers[voice_name_lower]
|
||||
if voice_name_lower in self.supported_speakers:
|
||||
self.supported_speakers.remove(voice_name_lower)
|
||||
|
||||
logger.info(f"Deleted voice '{name}' and associated files")
|
||||
return True
|
||||
|
||||
def _is_tts_model(self) -> bool:
|
||||
"""Check if the current model is a supported TTS model."""
|
||||
stage_list = getattr(self.engine_client, "stage_list", None)
|
||||
@@ -265,19 +486,48 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
|
||||
# Validate Base task requirements
|
||||
if task_type == "Base":
|
||||
if request.ref_audio is None:
|
||||
return "Base task requires 'ref_audio' for voice cloning"
|
||||
# Validate ref_audio format
|
||||
if not (request.ref_audio.startswith(("http://", "https://")) or request.ref_audio.startswith("data:")):
|
||||
return "ref_audio must be a URL (http/https) or base64 data URL (data:...)"
|
||||
# In-context voice cloning (default) requires non-empty ref_text.
|
||||
# x_vector_only_mode skips in-context and only uses speaker embedding.
|
||||
if not request.x_vector_only_mode:
|
||||
if not request.ref_text or not request.ref_text.strip():
|
||||
return (
|
||||
"Base task requires non-empty 'ref_text' (transcript of "
|
||||
"the reference audio) unless 'x_vector_only_mode' is enabled"
|
||||
)
|
||||
if request.voice is None:
|
||||
if request.ref_audio is None:
|
||||
return "Base task requires 'ref_audio' for voice cloning"
|
||||
# Validate ref_audio format (include file:// from upstream)
|
||||
if not (
|
||||
request.ref_audio.startswith(("http://", "https://"))
|
||||
or request.ref_audio.startswith("data:")
|
||||
or request.ref_audio.startswith("file://")
|
||||
):
|
||||
return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)"
|
||||
# In-context voice cloning (default) requires non-empty ref_text.
|
||||
# x_vector_only_mode skips in-context and only uses speaker embedding.
|
||||
if not request.x_vector_only_mode:
|
||||
if not request.ref_text or not request.ref_text.strip():
|
||||
return (
|
||||
"Base task requires non-empty 'ref_text' (transcript of "
|
||||
"the reference audio) unless 'x_vector_only_mode' is enabled"
|
||||
)
|
||||
else:
|
||||
# voice is not None
|
||||
voice_lower = request.voice.lower()
|
||||
if voice_lower in self.uploaded_speakers:
|
||||
# Check if audio file exists for uploaded speaker
|
||||
speaker_info = self.uploaded_speakers[voice_lower]
|
||||
file_path = Path(speaker_info["file_path"])
|
||||
if not file_path.exists():
|
||||
return f"Audio file for uploaded speaker '{request.voice}' not found on disk"
|
||||
else:
|
||||
# need ref_audio for built-in speaker
|
||||
if request.ref_audio is None:
|
||||
return (
|
||||
f"Base task with built-in speaker '{request.voice}' requires 'ref_audio' for voice cloning"
|
||||
)
|
||||
# Validate ref_audio format for built-in speaker
|
||||
if not (
|
||||
request.ref_audio.startswith(("http://", "https://"))
|
||||
or request.ref_audio.startswith("data:")
|
||||
or request.ref_audio.startswith("file://")
|
||||
):
|
||||
return (
|
||||
"ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)"
|
||||
)
|
||||
|
||||
# Validate cross-parameter dependencies
|
||||
if task_type != "Base":
|
||||
@@ -303,44 +553,23 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
async def _resolve_ref_audio(ref_audio_str: str) -> tuple[list[float], int]:
|
||||
"""Resolve ref_audio URL/base64 to (wav_samples, sample_rate)."""
|
||||
parsed = urlparse(ref_audio_str)
|
||||
async def _resolve_ref_audio(self, ref_audio_str: str) -> tuple[list[float], int]:
|
||||
"""Resolve ref_audio to (wav_samples, sample_rate).
|
||||
|
||||
def _check_ssrf(url: str) -> None:
|
||||
host = urlparse(url).hostname
|
||||
if not host:
|
||||
raise ValueError("ref_audio URL must include a hostname")
|
||||
for info in socket.getaddrinfo(host, None):
|
||||
ip_str = str(info[4][0]).split("%", 1)[0]
|
||||
addr = ipaddress.ip_address(ip_str)
|
||||
if any(addr in net for net in _REF_AUDIO_BLOCKED_NETWORKS):
|
||||
raise ValueError(f"ref_audio URL resolves to blocked address: {addr}")
|
||||
|
||||
def _fetch_sync() -> tuple[np.ndarray, int]:
|
||||
if parsed.scheme in ("http", "https"):
|
||||
_check_ssrf(ref_audio_str)
|
||||
with urlopen(ref_audio_str, timeout=_REF_AUDIO_TIMEOUT_S) as resp:
|
||||
data = resp.read(_REF_AUDIO_MAX_BYTES + 1)
|
||||
if len(data) > _REF_AUDIO_MAX_BYTES:
|
||||
raise ValueError(f"ref_audio URL exceeds {_REF_AUDIO_MAX_BYTES} bytes")
|
||||
buf = io.BytesIO(data)
|
||||
elif ref_audio_str.startswith("data:"):
|
||||
b64 = ref_audio_str
|
||||
if "," in b64:
|
||||
b64 = b64.split(",", 1)[1]
|
||||
buf = io.BytesIO(base64.b64decode(b64))
|
||||
else:
|
||||
raise ValueError("ref_audio must be an http(s) URL or data: base64 URI")
|
||||
audio, sr = sf.read(buf, dtype="float32", always_2d=False)
|
||||
if isinstance(audio, np.ndarray) and audio.ndim > 1:
|
||||
audio = np.mean(audio, axis=-1)
|
||||
return np.asarray(audio, dtype=np.float32), int(sr)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
wav_np, sr = await loop.run_in_executor(None, _fetch_sync)
|
||||
return wav_np.tolist(), sr
|
||||
Delegates to upstream vLLM's MediaConnector which handles http(s)
|
||||
URLs, ``data:`` base64 URIs, and ``file:`` local paths (the latter
|
||||
gated by ``--allowed-local-media-path``).
|
||||
"""
|
||||
model_config = self.model_config
|
||||
connector = MediaConnector(
|
||||
allowed_local_media_path=model_config.allowed_local_media_path,
|
||||
allowed_media_domains=model_config.allowed_media_domains,
|
||||
)
|
||||
wav_np, sr = await connector.fetch_audio_async(ref_audio_str)
|
||||
wav_np = np.asarray(wav_np, dtype=np.float32)
|
||||
if wav_np.ndim > 1:
|
||||
wav_np = np.mean(wav_np, axis=-1)
|
||||
return wav_np.tolist(), int(sr)
|
||||
|
||||
async def _generate_pcm_chunks(self, generator, request_id: str):
|
||||
"""Generate PCM audio chunks for streaming response.
|
||||
@@ -447,6 +676,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
# Speaker (voice)
|
||||
if request.voice is not None:
|
||||
params["speaker"] = [request.voice]
|
||||
|
||||
# If voice is an uploaded speaker and no ref_audio provided, auto-set it
|
||||
if request.voice.lower() in self.uploaded_speakers and request.ref_audio is None:
|
||||
audio_data = self._get_uploaded_audio_data(request.voice)
|
||||
if audio_data:
|
||||
params["ref_audio"] = [audio_data]
|
||||
params["x_vector_only_mode"] = [True]
|
||||
logger.info(f"Auto-set ref_audio for uploaded voice: {request.voice}")
|
||||
else:
|
||||
raise ValueError(f"Audio file for uploaded voice '{request.voice}' is missing or corrupted")
|
||||
|
||||
elif params["task_type"][0] == "CustomVoice":
|
||||
params["speaker"] = ["Vivian"] # Default for CustomVoice
|
||||
|
||||
|
||||
@@ -141,6 +141,7 @@ class OmniOpenAIServingVideo:
|
||||
)
|
||||
|
||||
result = await self._run_generation(prompt, gen_params, request_id, raw_request)
|
||||
_t_encode_start = time.perf_counter()
|
||||
videos = self._extract_video_outputs(result)
|
||||
audios = self._extract_audio_outputs(result, expected_count=len(videos))
|
||||
output_fps = fps or 24
|
||||
@@ -161,6 +162,8 @@ class OmniOpenAIServingVideo:
|
||||
)
|
||||
for idx, video in enumerate(videos)
|
||||
]
|
||||
_t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000
|
||||
logger.info("Video response encoding (MP4+base64): %.2f ms", _t_encode_ms)
|
||||
return VideoGenerationResponse(created=int(time.time()), data=video_data)
|
||||
|
||||
def _resolve_model_name(self, raw_request: Request | None) -> str | None:
|
||||
|
||||
@@ -36,6 +36,7 @@ class OmniStageTaskType(enum.Enum):
|
||||
SHUTDOWN = "shutdown"
|
||||
PROFILER_START = "profiler_start"
|
||||
PROFILER_STOP = "profiler_stop"
|
||||
COLLECTIVE_RPC = "collective_rpc"
|
||||
|
||||
|
||||
SHUTDOWN_TASK = {"type": OmniStageTaskType.SHUTDOWN}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import copy
|
||||
import pprint
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from typing import Any, TypeAlias
|
||||
from typing import Any, TypeAlias, TypedDict
|
||||
|
||||
from vllm.inputs import PromptType
|
||||
from vllm.sampling_params import SamplingParams
|
||||
@@ -103,13 +103,34 @@ class OmniEmbedsPrompt(EmbedsPrompt):
|
||||
additional_information: NotRequired[dict[str, Any]]
|
||||
|
||||
|
||||
class OmniCustomPrompt(TypedDict, total=False):
|
||||
"""Custom prompt type for diffusion pipelines with pre-tokenized inputs.
|
||||
|
||||
Allows users to pass pre-tokenized prompt IDs, attention masks, and extra
|
||||
arguments directly, bypassing the tokenization stage in the pipeline.
|
||||
|
||||
Attributes:
|
||||
prompt_ids: Pre-tokenized prompt token IDs (single or batched)
|
||||
negative_prompt_ids: Pre-tokenized negative prompt token IDs
|
||||
prompt_mask: Attention mask tensor for the prompt
|
||||
negative_prompt_mask: Attention mask tensor for the negative prompt
|
||||
extra_args: Additional pipeline-specific arguments
|
||||
"""
|
||||
|
||||
prompt_ids: list[int] | list[list[int]]
|
||||
negative_prompt_ids: list[int] | list[list[int]]
|
||||
prompt_mask: torch.Tensor
|
||||
negative_prompt_mask: torch.Tensor
|
||||
extra_args: dict[str, Any]
|
||||
|
||||
|
||||
# Must ensure that all additional prompt types are inherited from vLLM prompt types
|
||||
# Because TypedDict doesn't support isinstance and are dict. Cannot distinguish them in runtime.
|
||||
# Inheritance ensure that there are only additional fields but not removing fields--safe to route to LLM.generate()
|
||||
OmniSingletonPrompt: TypeAlias = str | list[int] | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt
|
||||
"""Omni singleton prompt type extending vLLM's SingletonPrompt."""
|
||||
|
||||
OmniPromptType: TypeAlias = PromptType | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt
|
||||
OmniPromptType: TypeAlias = PromptType | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt | OmniCustomPrompt
|
||||
|
||||
|
||||
def token_inputs_omni(
|
||||
|
||||
+14
-50
@@ -39,7 +39,7 @@ class StageRequestStats:
|
||||
final_output_type: str | None = None
|
||||
request_id: str | None = None
|
||||
postprocess_time_ms: float = 0.0
|
||||
diffusion_metrics: dict[str, float] = None
|
||||
diffusion_metrics: dict[str, int] = None
|
||||
audio_generated_frames: int = 0
|
||||
|
||||
@property
|
||||
@@ -293,15 +293,23 @@ class OrchestratorAggregator:
|
||||
return
|
||||
|
||||
# 4. Finished with output: assign text metrics if available
|
||||
output_to_yield.metrics = self.build_output_metrics(stage_id, req_id)
|
||||
output_to_yield.metrics = {}
|
||||
stage_event = next(
|
||||
(evt for evt in reversed(self.stage_events.get(req_id, [])) if evt.stage_id == stage_id),
|
||||
None,
|
||||
)
|
||||
if stage_event is not None and stage_event.final_output_type == "text":
|
||||
output_to_yield.metrics = {
|
||||
"num_tokens_in": stage_event.num_tokens_in,
|
||||
"num_tokens_out": stage_event.num_tokens_out,
|
||||
"stage_id": stage_event.stage_id,
|
||||
"final_output_type": stage_event.final_output_type,
|
||||
}
|
||||
|
||||
# 5. Finished: record audio generated frames
|
||||
self.record_audio_generated_frames(output_to_yield, stage_id, req_id)
|
||||
|
||||
except Exception:
|
||||
if output_to_yield is not None:
|
||||
# Make metrics contract explicit on failure.
|
||||
output_to_yield.metrics = {}
|
||||
logger.exception(
|
||||
"Failed to process metrics for stage %s, req %s",
|
||||
stage_id,
|
||||
@@ -321,56 +329,12 @@ class OrchestratorAggregator:
|
||||
stats.request_id = req_id
|
||||
stats.final_output_type = final_output_type
|
||||
stats.diffusion_metrics = (
|
||||
{k: float(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()}
|
||||
{k: int(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()}
|
||||
if req_id in self.diffusion_metrics
|
||||
else None
|
||||
)
|
||||
return stats
|
||||
|
||||
def _get_stage_event(self, stage_id: int, req_id: Any) -> StageRequestStats | None:
|
||||
rid_key = str(req_id)
|
||||
for evt in reversed(self.stage_events.get(rid_key, [])):
|
||||
if evt.stage_id == stage_id:
|
||||
return evt
|
||||
return None
|
||||
|
||||
def _collect_diffusion_metrics(self, req_id: Any) -> dict[str, float]:
|
||||
"""Aggregate diffusion metrics across all stages for a request."""
|
||||
rid_key = str(req_id)
|
||||
merged: dict[str, float] = {}
|
||||
for evt in self.stage_events.get(rid_key, []):
|
||||
if not evt.diffusion_metrics:
|
||||
continue
|
||||
for key, value in evt.diffusion_metrics.items():
|
||||
merged[key] = merged.get(key, 0.0) + float(value)
|
||||
return merged
|
||||
|
||||
def build_output_metrics(self, stage_id: int, req_id: Any) -> dict[str, Any]:
|
||||
stage_event = self._get_stage_event(stage_id, req_id)
|
||||
if stage_event is None:
|
||||
return {}
|
||||
|
||||
merged: dict[str, Any] = {}
|
||||
|
||||
if stage_event.final_output_type == "text":
|
||||
merged.update(
|
||||
{
|
||||
"num_tokens_in": stage_event.num_tokens_in,
|
||||
"num_tokens_out": stage_event.num_tokens_out,
|
||||
}
|
||||
)
|
||||
|
||||
if self.log_stats:
|
||||
diffusion_metrics = self._collect_diffusion_metrics(req_id)
|
||||
if diffusion_metrics:
|
||||
merged.update(diffusion_metrics)
|
||||
|
||||
if merged:
|
||||
merged["stage_id"] = stage_event.stage_id
|
||||
merged["final_output_type"] = stage_event.final_output_type
|
||||
|
||||
return merged
|
||||
|
||||
def on_stage_metrics(
|
||||
self,
|
||||
stage_id: int,
|
||||
|
||||
@@ -1,4 +1,8 @@
|
||||
from .bagel.bagel import OmniBagelForConditionalGeneration
|
||||
from .qwen3_omni import Qwen3OmniMoeForConditionalGeneration
|
||||
from .registry import OmniModelRegistry # noqa: F401
|
||||
|
||||
__all__ = ["Qwen3OmniMoeForConditionalGeneration"]
|
||||
__all__ = [
|
||||
"Qwen3OmniMoeForConditionalGeneration",
|
||||
"OmniBagelForConditionalGeneration",
|
||||
]
|
||||
@@ -0,0 +1,3 @@
|
||||
from .bagel import OmniBagelForConditionalGeneration
|
||||
|
||||
__all__ = ["OmniBagelForConditionalGeneration"]
|
||||
File diff suppressed because it is too large.
Load diff
@@ -340,6 +340,11 @@ class Qwen3OmniMoeThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProc
|
||||
seq_len=seq_len,
|
||||
mm_counts=mm_counts,
|
||||
)
|
||||
if audio_tokens is None:
|
||||
feature_extractor = self.get_feature_extractor()
|
||||
max_audio_samples = feature_extractor.chunk_length * feature_extractor.sampling_rate
|
||||
max_audio_tokens = int(_get_feat_extract_output_lengths(torch.tensor([max_audio_samples])).item())
|
||||
audio_tokens = {"audio": max_audio_tokens}
|
||||
mm_max_tokens["audio"] = audio_tokens["audio"]
|
||||
|
||||
return mm_max_tokens
|
||||
@@ -769,8 +774,7 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix
|
||||
feature_lens=audio_feature_lengths,
|
||||
aftercnn_lens=audio_output_lengths,
|
||||
)
|
||||
# OMNI: audio_tower.forward() returns hidden_states tensor directly
|
||||
audio_features = audio_outputs
|
||||
audio_features = audio_outputs if isinstance(audio_outputs, torch.Tensor) else audio_outputs.last_hidden_state
|
||||
return audio_features.split(audio_output_lengths.tolist())
|
||||
|
||||
|
||||
|
||||
@@ -13,11 +13,9 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
import base64
|
||||
import copy
|
||||
import io
|
||||
import urllib.request
|
||||
from collections.abc import Iterable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
@@ -36,6 +34,7 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
from .configuration_qwen3_tts import Qwen3TTSConfig
|
||||
from .modeling_qwen3_tts import Qwen3TTSForConditionalGeneration
|
||||
from .processing_qwen3_tts import Qwen3TTSProcessor
|
||||
from .voice_cache_manager import VoiceCacheManager, VoiceClonePromptItem
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -60,21 +59,6 @@ AudioLike = (
|
||||
MaybeList = Any | list[Any]
|
||||
|
||||
|
||||
@dataclass
|
||||
class VoiceClonePromptItem:
|
||||
"""
|
||||
Container for one sample's voice-clone prompt information that can be fed to the model.
|
||||
|
||||
Fields are aligned with `Qwen3TTSForConditionalGeneration.generate(..., voice_clone_prompt=...)`.
|
||||
"""
|
||||
|
||||
ref_code: torch.Tensor | None # (T, Q) or (T,) depending on tokenizer 25Hz/12Hz
|
||||
ref_spk_embedding: torch.Tensor # (D,)
|
||||
x_vector_only_mode: bool
|
||||
icl_mode: bool
|
||||
ref_text: str | None = None
|
||||
|
||||
|
||||
class Qwen3TTSModelForGeneration(nn.Module):
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
@@ -130,6 +114,13 @@ class Qwen3TTSModelForGeneration(nn.Module):
|
||||
except Exception:
|
||||
logger.warning("Failed to enable CUDA Graph for decoder", exc_info=True)
|
||||
|
||||
@staticmethod
|
||||
def extract_val(d, key, default):
|
||||
val = d.get(key, default)
|
||||
if isinstance(val, list):
|
||||
return val[0] if len(val) > 0 else default
|
||||
return val
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
@@ -139,7 +130,7 @@ class Qwen3TTSModelForGeneration(nn.Module):
|
||||
**kwargs: Any,
|
||||
) -> OmniOutput:
|
||||
"""
|
||||
Forward pass for TTS generation model.
|
||||
Forward pass for TTS generation model (Patched for batched inference).
|
||||
|
||||
Args:
|
||||
input_ids: Input token IDs (required for TTS generation)
|
||||
@@ -151,46 +142,62 @@ class Qwen3TTSModelForGeneration(nn.Module):
|
||||
Returns:
|
||||
OmniOutput: Contains multimodal outputs with audio tensors
|
||||
"""
|
||||
runtime_info_list = kwargs.get("runtime_additional_information", [{}])
|
||||
if not isinstance(runtime_info_list, list):
|
||||
runtime_info_list = [runtime_info_list]
|
||||
|
||||
# Extract additional parameters from kwargs that the generation methods expect
|
||||
# Initialize lists to accumulate batched inputs
|
||||
texts = []
|
||||
task_types = []
|
||||
speakers = []
|
||||
languages = []
|
||||
instructs = []
|
||||
merged_kwargs = {}
|
||||
|
||||
runtime_additional_information = kwargs.get("model_intermediate_buffer")
|
||||
if runtime_additional_information is None:
|
||||
runtime_additional_information = kwargs.get("runtime_additional_information", [{}])
|
||||
if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs:
|
||||
logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer")
|
||||
if isinstance(runtime_additional_information, list) and len(runtime_additional_information) > 0:
|
||||
runtime_additional_information = runtime_additional_information[0]
|
||||
runtime_additional_information = copy.deepcopy(runtime_additional_information)
|
||||
text = runtime_additional_information.pop("text", [""])[0]
|
||||
# Extract task_type from kwargs, default to self.task_type
|
||||
task_type = _normalize_task_type(runtime_additional_information.pop("task_type", [self.task_type])[0])
|
||||
speaker = runtime_additional_information.pop("speaker", ["uncle_fu"])[0]
|
||||
language = runtime_additional_information.pop("language", ["Auto"])[0]
|
||||
instruct = runtime_additional_information.pop("instruct", [""])[0]
|
||||
for key, value in runtime_additional_information.items():
|
||||
if isinstance(value, list) and len(value) > 0:
|
||||
runtime_additional_information[key] = value[0]
|
||||
# Keys that the underlying model natively supports as lists for batched inference
|
||||
batched_keys = {"ref_audio", "ref_text", "x_vector_only_mode", "voice_clone_prompt"}
|
||||
|
||||
# During profile/warmup runs, text is empty and no real inputs exist.
|
||||
# Cap generation steps so the full pipeline executes (preserving
|
||||
# KV-cache profiling behaviour) but exits quickly even if the model
|
||||
# cannot converge from degenerate dummy inputs.
|
||||
if not text:
|
||||
for req_info in runtime_info_list:
|
||||
texts.append(self.extract_val(req_info, "text", ""))
|
||||
task_types.append(self.extract_val(req_info, "task_type", self.task_type))
|
||||
speakers.append(self.extract_val(req_info, "speaker", "uncle_fu"))
|
||||
languages.append(self.extract_val(req_info, "language", "Auto"))
|
||||
instructs.append(self.extract_val(req_info, "instruct", ""))
|
||||
|
||||
for k, v in req_info.items():
|
||||
if k not in ["text", "task_type", "speaker", "language", "instruct"]:
|
||||
# Extract single value from list if wrapped
|
||||
val = v[0] if isinstance(v, list) and len(v) > 0 else v
|
||||
|
||||
if k in batched_keys:
|
||||
# Accumulate as list for batched generation
|
||||
if k not in merged_kwargs:
|
||||
merged_kwargs[k] = []
|
||||
merged_kwargs[k].append(val)
|
||||
else:
|
||||
# For scalar params (e.g. max_new_tokens), take from the first request
|
||||
if k not in merged_kwargs:
|
||||
merged_kwargs[k] = val
|
||||
|
||||
# During profile/warmup runs, texts are empty.
|
||||
if all(not t for t in texts):
|
||||
logger.info("Profile run detected (empty text). Capping max_new_tokens to 2.")
|
||||
runtime_additional_information["max_new_tokens"] = 2
|
||||
merged_kwargs["max_new_tokens"] = 2
|
||||
|
||||
# Call the appropriate generation method based on task_type
|
||||
# Assume uniform task type across the batch
|
||||
if len(set(task_types)) > 1:
|
||||
raise ValueError(f"Mixed task types not supported: {set(task_types)}")
|
||||
task_type = task_types[0]
|
||||
|
||||
# Call the appropriate generation method based on task_type, passing lists
|
||||
if task_type == "CustomVoice":
|
||||
result = self.model.generate_custom_voice(
|
||||
text, speaker=speaker, language=language, instruct=instruct, **runtime_additional_information
|
||||
texts, speaker=speakers, language=languages, instruct=instructs, **merged_kwargs
|
||||
)
|
||||
elif task_type == "VoiceDesign":
|
||||
result = self.model.generate_voice_design(
|
||||
text, instruct=instruct, language=language, **runtime_additional_information
|
||||
)
|
||||
result = self.model.generate_voice_design(texts, instruct=instructs, language=languages, **merged_kwargs)
|
||||
elif task_type == "Base":
|
||||
result = self.model.generate_voice_clone(text, language=language, **runtime_additional_information)
|
||||
result = self.model.generate_voice_clone(texts, language=languages, **merged_kwargs)
|
||||
else:
|
||||
raise ValueError(f"Invalid task type: {task_type}")
|
||||
|
||||
@@ -209,17 +216,20 @@ class Qwen3TTSModelForGeneration(nn.Module):
|
||||
# Handle tuple format: (audio_tensors, sample_rate)
|
||||
if isinstance(model_outputs, tuple) and len(model_outputs) == 2:
|
||||
audio_tensors, sr = model_outputs
|
||||
# audio_tensors is a list of numpy arrays, convert first one to tensor if needed
|
||||
# audio_tensors is a list of numpy arrays, convert ALL to tensors
|
||||
if isinstance(audio_tensors, list) and len(audio_tensors) > 0:
|
||||
# Convert numpy array to tensor if needed
|
||||
audio_tensor = audio_tensors[0]
|
||||
if isinstance(audio_tensor, np.ndarray):
|
||||
audio_tensor = torch.from_numpy(audio_tensor).float()
|
||||
elif not isinstance(audio_tensor, torch.Tensor):
|
||||
audio_tensor = torch.tensor(audio_tensor, dtype=torch.float32)
|
||||
audio_tensor_list = []
|
||||
for audio_tensor in audio_tensors:
|
||||
if isinstance(audio_tensor, np.ndarray):
|
||||
audio_tensor_list.append(torch.from_numpy(audio_tensor).float())
|
||||
elif not isinstance(audio_tensor, torch.Tensor):
|
||||
audio_tensor_list.append(torch.tensor(audio_tensor, dtype=torch.float32))
|
||||
else:
|
||||
audio_tensor_list.append(audio_tensor)
|
||||
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={"model_outputs": audio_tensor, "sr": torch.tensor(sr, dtype=torch.int)},
|
||||
multimodal_outputs={"model_outputs": audio_tensor_list, "sr": torch.tensor(sr, dtype=torch.int)},
|
||||
)
|
||||
|
||||
# If it's already a tensor, wrap it
|
||||
@@ -348,6 +358,12 @@ class Qwen3TTSModel:
|
||||
self.processor = processor
|
||||
self.generate_defaults = generate_defaults or {}
|
||||
|
||||
# Initialize voice cache manager.
|
||||
# Note: this creates its own MetadataManager for the same metadata.json
|
||||
# used by serving_speech.py. Sharing is not possible across model/serving
|
||||
# layers, but file locking in MetadataManager ensures correctness.
|
||||
self.voice_cache_manager = VoiceCacheManager()
|
||||
|
||||
self.device = getattr(model, "device", None)
|
||||
if self.device is None:
|
||||
try:
|
||||
@@ -766,6 +782,7 @@ class Qwen3TTSModel:
|
||||
self,
|
||||
text: str | list[str],
|
||||
language: str | list[str] = None,
|
||||
speaker: str | None = None, # New parameter: speaker name
|
||||
ref_audio: AudioLike | list[AudioLike] | None = None,
|
||||
ref_text: str | list[str | None] | None = None,
|
||||
x_vector_only_mode: bool | list[bool] = False,
|
||||
@@ -840,7 +857,27 @@ class Qwen3TTSModel:
|
||||
|
||||
self._validate_languages(languages)
|
||||
|
||||
if voice_clone_prompt is None:
|
||||
# Cache logic: if speaker parameter is provided, try to load from cache
|
||||
cache_loaded = False
|
||||
cache_speaker = None
|
||||
cache_audio_path = None
|
||||
|
||||
if speaker:
|
||||
# Use VoiceCacheManager to load cached voice prompt, passing device parameter
|
||||
cached_items = self.voice_cache_manager.load_cached_voice_prompt(speaker, device=str(self.device))
|
||||
if cached_items is not None:
|
||||
voice_clone_prompt = cached_items
|
||||
cache_loaded = True
|
||||
|
||||
# If no cache, check if cache needs to be generated
|
||||
if not cache_loaded:
|
||||
audio_file_path = self.voice_cache_manager.get_speaker_audio_path(speaker)
|
||||
if audio_file_path:
|
||||
logger.info(f"Will generate cache for speaker: {speaker} (first use)")
|
||||
cache_speaker = speaker
|
||||
cache_audio_path = audio_file_path
|
||||
|
||||
if voice_clone_prompt is None and not cache_loaded:
|
||||
if ref_audio is None:
|
||||
# For profile run
|
||||
sample_rate = int(self.model.speaker_encoder_sample_rate)
|
||||
@@ -854,6 +891,28 @@ class Qwen3TTSModel:
|
||||
prompt_items = self.create_voice_clone_prompt(
|
||||
ref_audio=ref_audio, ref_text=ref_text, x_vector_only_mode=x_vector_only_mode
|
||||
)
|
||||
|
||||
# If cache needs to be generated, save cache file
|
||||
if cache_speaker and cache_audio_path:
|
||||
try:
|
||||
# Use VoiceCacheManager to save cache
|
||||
success = self.voice_cache_manager.save_voice_cache(cache_speaker, cache_audio_path, prompt_items)
|
||||
if success:
|
||||
logger.info(f"Cache generated and saved for speaker: {cache_speaker}")
|
||||
else:
|
||||
logger.error(f"Failed to save cache for speaker: {cache_speaker}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to save cache for speaker {cache_speaker}: {e}")
|
||||
|
||||
if len(prompt_items) == 1 and len(texts) > 1:
|
||||
prompt_items = prompt_items * len(texts)
|
||||
if len(prompt_items) != len(texts):
|
||||
raise ValueError(f"Batch size mismatch: prompt={len(prompt_items)}, text={len(texts)}")
|
||||
voice_clone_prompt_dict = self._prompt_items_to_voice_clone_prompt(prompt_items)
|
||||
ref_texts_for_ids = [it.ref_text for it in prompt_items]
|
||||
elif cache_loaded and isinstance(voice_clone_prompt, list):
|
||||
# Use cached VoiceClonePromptItem
|
||||
prompt_items = voice_clone_prompt
|
||||
if len(prompt_items) == 1 and len(texts) > 1:
|
||||
prompt_items = prompt_items * len(texts)
|
||||
if len(prompt_items) != len(texts):
|
||||
|
||||
@@ -831,7 +831,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
|
||||
def _is_url(self, s: str) -> bool:
|
||||
try:
|
||||
u = urlparse(s)
|
||||
return u.scheme in ("http", "https") and bool(u.netloc)
|
||||
if u.scheme in ("http", "https"):
|
||||
return bool(u.netloc)
|
||||
return u.scheme == "file"
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
@@ -841,12 +843,19 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
|
||||
return base64.b64decode(b64)
|
||||
|
||||
def _load_audio_to_np(self, x: str) -> tuple[np.ndarray, int]:
|
||||
"""Load audio from local path or base64 data URI (no network I/O)."""
|
||||
"""Load audio from local path, URL, or base64 data URI.
|
||||
|
||||
Uses upstream vLLM's MediaConnector for http(s) URLs and ``file:``
|
||||
URIs, with unrestricted local access (offline inference is trusted).
|
||||
"""
|
||||
import librosa
|
||||
|
||||
if self._is_url(x):
|
||||
raise ValueError("ref_audio URLs must be resolved by the serving layer before reaching the model worker.")
|
||||
if self._is_probably_base64(x):
|
||||
from vllm.multimodal.media import MediaConnector
|
||||
|
||||
connector = MediaConnector(allowed_local_media_path="/")
|
||||
audio, sr = connector.fetch_audio(x)
|
||||
elif self._is_probably_base64(x):
|
||||
wav_bytes = self._decode_base64_to_wav_bytes(x)
|
||||
with io.BytesIO(wav_bytes) as f:
|
||||
audio, sr = sf.read(f, dtype="float32", always_2d=False)
|
||||
|
||||
@@ -0,0 +1,271 @@
|
||||
# Copyright 2026 The Alibaba Qwen team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import save_file
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.entrypoints.openai.metadata_manager import MetadataManager
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class VoiceClonePromptItem:
|
||||
"""
|
||||
Container for one sample's voice-clone prompt information that can be fed to the model.
|
||||
|
||||
Fields are aligned with `Qwen3TTSForConditionalGeneration.generate(..., voice_clone_prompt=...)`.
|
||||
"""
|
||||
|
||||
ref_code: torch.Tensor | None # (T, Q) or (T,) depending on tokenizer 25Hz/12Hz
|
||||
ref_spk_embedding: torch.Tensor # (D,)
|
||||
x_vector_only_mode: bool
|
||||
icl_mode: bool
|
||||
ref_text: str | None = None
|
||||
|
||||
|
||||
class VoiceCacheManager:
|
||||
"""
|
||||
Voice cache manager, responsible for managing custom voice cache functionality.
|
||||
|
||||
Main features:
|
||||
1. Load uploaded speaker information from metadata.json
|
||||
2. Manage voice clone prompt cache
|
||||
3. Update cache status to metadata.json
|
||||
|
||||
Security properties:
|
||||
- No pickle / torch.load
|
||||
- Safetensors-only
|
||||
- Cache path confined to voice samples directory
|
||||
"""
|
||||
|
||||
def __init__(self, speech_voice_samples_dir: str | None = None, metadata_manager: MetadataManager | None = None):
|
||||
"""
|
||||
Initialize the voice cache manager.
|
||||
|
||||
Args:
|
||||
speech_voice_samples_dir: Speech voice samples directory path,
|
||||
if None, get from environment variable
|
||||
metadata_manager: Optional MetadataManager instance for shared metadata access.
|
||||
If not provided, will create its own (less efficient).
|
||||
"""
|
||||
self.speech_voice_samples_dir = speech_voice_samples_dir or os.environ.get(
|
||||
"SPEECH_VOICE_SAMPLES", "/tmp/voice_samples"
|
||||
)
|
||||
|
||||
# Initialize metadata manager
|
||||
if metadata_manager is not None:
|
||||
self.metadata_manager = metadata_manager
|
||||
else:
|
||||
metadata_file = Path(self.speech_voice_samples_dir) / "metadata.json"
|
||||
self.metadata_manager = MetadataManager(metadata_file)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Metadata helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def load_uploaded_speakers_from_metadata(self) -> dict[str, Any] | None:
|
||||
"""Load uploaded speakers from metadata manager."""
|
||||
try:
|
||||
return self.metadata_manager.get_uploaded_speakers()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load uploaded speakers from metadata: {e}")
|
||||
return None
|
||||
|
||||
def update_metadata_cache_info(self, speaker: str, cache_file_path: Path, status: str = "ready") -> bool:
|
||||
"""
|
||||
Update cache information using metadata manager.
|
||||
|
||||
Args:
|
||||
speaker: Speaker name
|
||||
cache_file_path: Cache file path
|
||||
status: Cache status, default is "ready"
|
||||
|
||||
Returns:
|
||||
bool: Whether the update was successful
|
||||
"""
|
||||
try:
|
||||
speaker_key = speaker.lower()
|
||||
return self.metadata_manager.update_cache_info(
|
||||
speaker_key=speaker_key, cache_file_path=cache_file_path, status=status
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update metadata cache info: {e}")
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Cache save (SAFE)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def save_voice_cache(
|
||||
self,
|
||||
speaker: str,
|
||||
audio_file_path: Path,
|
||||
prompt_items: list[VoiceClonePromptItem],
|
||||
) -> bool:
|
||||
"""
|
||||
Save voice cache using safetensors (no pickle, no RCE).
|
||||
"""
|
||||
try:
|
||||
cache_file_path = audio_file_path.with_suffix(".safetensors")
|
||||
|
||||
tensors: dict[str, torch.Tensor] = {}
|
||||
metadata: dict[str, str] = {}
|
||||
|
||||
tensors["__len__"] = torch.tensor(len(prompt_items), dtype=torch.int64)
|
||||
|
||||
for i, item in enumerate(prompt_items):
|
||||
prefix = f"item_{i}_"
|
||||
|
||||
tensors[prefix + "ref_spk_embedding"] = item.ref_spk_embedding.detach().cpu()
|
||||
|
||||
has_ref_code = item.ref_code is not None
|
||||
tensors[prefix + "has_ref_code"] = torch.tensor(int(has_ref_code), dtype=torch.int8)
|
||||
|
||||
if has_ref_code:
|
||||
tensors[prefix + "ref_code"] = item.ref_code.detach().cpu()
|
||||
|
||||
tensors[prefix + "x_vector_only_mode"] = torch.tensor(int(item.x_vector_only_mode), dtype=torch.int8)
|
||||
tensors[prefix + "icl_mode"] = torch.tensor(int(item.icl_mode), dtype=torch.int8)
|
||||
|
||||
if item.ref_text is not None:
|
||||
metadata[prefix + "ref_text"] = item.ref_text
|
||||
|
||||
save_file(tensors, str(cache_file_path), metadata=metadata)
|
||||
|
||||
return self.update_metadata_cache_info(
|
||||
speaker=speaker,
|
||||
cache_file_path=cache_file_path,
|
||||
status="ready",
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to save safetensors cache for speaker {speaker}: {e}")
|
||||
self.update_metadata_cache_info(speaker, Path(""), "failed")
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Cache load (SAFE)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def load_cached_voice_prompt(
|
||||
self,
|
||||
speaker: str,
|
||||
device: str | None = None,
|
||||
) -> list[VoiceClonePromptItem] | None:
|
||||
"""
|
||||
Load cached VoiceClonePromptItem list from safetensors.
|
||||
"""
|
||||
try:
|
||||
uploaded_speakers = self.load_uploaded_speakers_from_metadata()
|
||||
if not uploaded_speakers:
|
||||
return None
|
||||
|
||||
speaker_key = speaker.lower()
|
||||
if speaker_key not in uploaded_speakers:
|
||||
return None
|
||||
|
||||
speaker_info = uploaded_speakers[speaker_key]
|
||||
if speaker_info.get("cache_status") != "ready":
|
||||
return None
|
||||
|
||||
cache_file_path = Path(speaker_info.get("cache_file", "")).resolve()
|
||||
|
||||
base_dir = Path(self.speech_voice_samples_dir).resolve()
|
||||
|
||||
# ---- Path confinement (critical security check)
|
||||
if not str(cache_file_path).startswith(str(base_dir)):
|
||||
logger.error(f"Illegal cache path outside base dir: {cache_file_path}")
|
||||
return None
|
||||
|
||||
if not cache_file_path.exists():
|
||||
return None
|
||||
|
||||
if cache_file_path.suffix != ".safetensors":
|
||||
logger.error(f"Legacy or unsafe cache format rejected: {cache_file_path}")
|
||||
return None
|
||||
|
||||
with safe_open(cache_file_path, framework="pt", device="cpu") as f:
|
||||
meta = f.metadata()
|
||||
|
||||
num_items = int(f.get_tensor("__len__").item())
|
||||
result: list[VoiceClonePromptItem] = []
|
||||
|
||||
for i in range(num_items):
|
||||
prefix = f"item_{i}_"
|
||||
|
||||
has_ref_code = bool(f.get_tensor(prefix + "has_ref_code").item())
|
||||
|
||||
ref_code = f.get_tensor(prefix + "ref_code").to(device) if has_ref_code else None
|
||||
|
||||
ref_spk_embedding = f.get_tensor(prefix + "ref_spk_embedding").to(device)
|
||||
|
||||
x_vector_only_mode = bool(f.get_tensor(prefix + "x_vector_only_mode").item())
|
||||
icl_mode = bool(f.get_tensor(prefix + "icl_mode").item())
|
||||
|
||||
ref_text = meta.get(prefix + "ref_text")
|
||||
|
||||
result.append(
|
||||
VoiceClonePromptItem(
|
||||
ref_code=ref_code,
|
||||
ref_spk_embedding=ref_spk_embedding,
|
||||
x_vector_only_mode=x_vector_only_mode,
|
||||
icl_mode=icl_mode,
|
||||
ref_text=ref_text,
|
||||
)
|
||||
)
|
||||
|
||||
logger.info(f"Safetensors cache loaded for speaker: {speaker}")
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load safetensors cache for speaker {speaker}: {e}")
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Audio path helper
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_speaker_audio_path(self, speaker: str) -> Path | None:
|
||||
"""
|
||||
Get speaker's audio file path.
|
||||
|
||||
Args:
|
||||
speaker: Speaker name
|
||||
|
||||
Returns:
|
||||
Optional[Path]: Audio file path, returns None if speaker doesn't exist
|
||||
"""
|
||||
uploaded_speakers = self.load_uploaded_speakers_from_metadata()
|
||||
if not uploaded_speakers:
|
||||
return None
|
||||
|
||||
speaker_key = speaker.lower()
|
||||
if speaker_key not in uploaded_speakers:
|
||||
return None
|
||||
|
||||
audio_file_path = Path(uploaded_speakers[speaker_key]["file_path"])
|
||||
if audio_file_path.exists():
|
||||
return audio_file_path
|
||||
|
||||
logger.warning(f"Audio file not found for speaker {speaker}: {audio_file_path}")
|
||||
return None
|
||||
@@ -119,6 +119,11 @@ _OMNI_MODELS = {
|
||||
"glm_image_ar",
|
||||
"GlmImageForConditionalGeneration",
|
||||
),
|
||||
"OmniBagelForConditionalGeneration": (
|
||||
"bagel",
|
||||
"bagel",
|
||||
"OmniBagelForConditionalGeneration",
|
||||
),
|
||||
"HunyuanImage3ForCausalMM": (
|
||||
"hunyuan_image3",
|
||||
"hunyuan_image3",
|
||||
|
||||
@@ -6,15 +6,14 @@ stage_args:
|
||||
prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts
|
||||
runtime:
|
||||
devices: "0"
|
||||
# 2 = 1 user prompt + 1 CFG companion (text-unconditional).
|
||||
# For multi-prompt batches this should scale as batch_size × 2.
|
||||
max_batch_size: 2
|
||||
# 3 = 1 user prompt + 2 CFG companions (text-unconditional + image-unconditional).
|
||||
max_batch_size: 3
|
||||
engine_args:
|
||||
model_stage: thinker
|
||||
model_arch: BagelForConditionalGeneration
|
||||
model_arch: OmniBagelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.35
|
||||
gpu_memory_utilization: 0.45
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: text
|
||||
@@ -46,7 +45,7 @@ stage_args:
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: dit
|
||||
gpu_memory_utilization: 0.55
|
||||
gpu_memory_utilization: 0.45
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: image
|
||||
|
||||
@@ -11,7 +11,7 @@ stage_args:
|
||||
model_arch: BagelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.35
|
||||
gpu_memory_utilization: 0.45
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: text
|
||||
@@ -45,7 +45,7 @@ stage_args:
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: dit
|
||||
gpu_memory_utilization: 0.55
|
||||
gpu_memory_utilization: 0.45
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: image
|
||||
|
||||
@@ -4,7 +4,7 @@ stage_args:
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
max_batch_size: 10
|
||||
engine_args:
|
||||
model_stage: qwen3_tts
|
||||
model_arch: Qwen3TTSTalkerForConditionalGeneration
|
||||
|
||||
@@ -20,6 +20,7 @@ from typing import Any
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CFG_TEXT_SUFFIX = "__cfg_text"
|
||||
CFG_IMG_SUFFIX = "__cfg_img"
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -73,7 +74,37 @@ def expand_cfg_prompts(
|
||||
),
|
||||
]
|
||||
|
||||
# img2img: more complex (3 distinct caches). Reserve for future.
|
||||
if "img2img" in modalities:
|
||||
IMG2IMG_PLACEHOLDER = "<|fim_middle|>"
|
||||
|
||||
cfg_text_dict: dict[str, Any] = {
|
||||
"prompt": f"{IMG2IMG_PLACEHOLDER}{neg_prompt}",
|
||||
"modalities": ["img2img"],
|
||||
}
|
||||
mm_data = prompt.get("multi_modal_data")
|
||||
if mm_data:
|
||||
cfg_text_dict["multi_modal_data"] = mm_data
|
||||
|
||||
original_text = prompt.get("prompt", "")
|
||||
cfg_img_text = original_text.replace(IMG2IMG_PLACEHOLDER, "")
|
||||
cfg_img_dict: dict[str, Any] = {
|
||||
"prompt": cfg_img_text,
|
||||
"modalities": ["img2img"],
|
||||
}
|
||||
|
||||
return [
|
||||
ExpandedPrompt(
|
||||
prompt=cfg_text_dict,
|
||||
role="cfg_text",
|
||||
request_id_suffix=CFG_TEXT_SUFFIX,
|
||||
),
|
||||
ExpandedPrompt(
|
||||
prompt=cfg_img_dict,
|
||||
role="cfg_img",
|
||||
request_id_suffix=CFG_IMG_SUFFIX,
|
||||
),
|
||||
]
|
||||
|
||||
return []
|
||||
|
||||
|
||||
|
||||
@@ -60,6 +60,7 @@ class OmniRequestOutput:
|
||||
latents: torch.Tensor | None = None
|
||||
metrics: dict[str, Any] = field(default_factory=dict)
|
||||
_multimodal_output: dict[str, Any] = field(default_factory=dict)
|
||||
_custom_output: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def from_pipeline(
|
||||
@@ -95,6 +96,7 @@ class OmniRequestOutput:
|
||||
metrics: dict[str, Any] | None = None,
|
||||
latents: torch.Tensor | None = None,
|
||||
multimodal_output: dict[str, Any] | None = None,
|
||||
custom_output: dict[str, Any] | None = None,
|
||||
final_output_type: str = "image",
|
||||
) -> "OmniRequestOutput":
|
||||
"""Create output from diffusion model.
|
||||
@@ -105,6 +107,8 @@ class OmniRequestOutput:
|
||||
prompt: The prompt used
|
||||
metrics: Generation metrics
|
||||
latents: Optional latent tensors
|
||||
multimodal_output: Optional multimodal output dict
|
||||
custom_output: Optional custom output dict (e.g. latent trajectories, prompt embeds)
|
||||
|
||||
Returns:
|
||||
OmniRequestOutput configured for diffusion mode
|
||||
@@ -117,6 +121,7 @@ class OmniRequestOutput:
|
||||
latents=latents,
|
||||
metrics=metrics or {},
|
||||
_multimodal_output=multimodal_output or {},
|
||||
_custom_output=custom_output or {},
|
||||
finished=True,
|
||||
)
|
||||
|
||||
@@ -140,6 +145,23 @@ class OmniRequestOutput:
|
||||
return mm
|
||||
return self._multimodal_output
|
||||
|
||||
@property
|
||||
def custom_output(self) -> dict[str, Any]:
|
||||
"""Return custom output data from diffusion pipelines.
|
||||
|
||||
For diffusion outputs, returns the local _custom_output field.
|
||||
For pipeline outputs with an inner OmniRequestOutput, forwards
|
||||
the custom_output from the inner request output.
|
||||
"""
|
||||
if self.request_output is not None:
|
||||
if isinstance(self.request_output, OmniRequestOutput):
|
||||
return self.request_output._custom_output
|
||||
return self._custom_output
|
||||
|
||||
@custom_output.setter
|
||||
def custom_output(self, value: dict[str, Any]) -> None:
|
||||
self._custom_output = value
|
||||
|
||||
@property
|
||||
def num_images(self) -> int:
|
||||
"""Return the number of generated images."""
|
||||
@@ -248,7 +270,9 @@ class OmniRequestOutput:
|
||||
f"images={images_repr}",
|
||||
f"prompt={self.prompt!r}",
|
||||
f"latents={self.latents}",
|
||||
f"metrics={self.metrics}",
|
||||
f"multimodal_output={self._multimodal_output}",
|
||||
f"custom_output={self._custom_output}",
|
||||
]
|
||||
|
||||
return f"OmniRequestOutput({', '.join(parts)})"
|
||||
@@ -1,6 +1,8 @@
|
||||
import sys
|
||||
from functools import cached_property as _cached_property
|
||||
|
||||
from aenum import extend_enum
|
||||
from vllm.config import ModelConfig as _ModelConfig
|
||||
from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt
|
||||
from vllm.model_executor.layers.rotary_embedding import (
|
||||
MRotaryEmbedding as _OriginalMRotaryEmbedding,
|
||||
@@ -43,6 +45,20 @@ except ImportError:
|
||||
# GlmImageTextConfig not available, skip patching
|
||||
pass
|
||||
|
||||
|
||||
# Patch ModelConfig.is_mm_prefix_lm to include Bagel (bidirectional attention
|
||||
# for multimodal prefix positions, same as Gemma3/Molmo2/PaliGemma).
|
||||
_orig_is_mm_prefix_lm = _ModelConfig.__dict__["is_mm_prefix_lm"].func
|
||||
|
||||
|
||||
@_cached_property
|
||||
def _patched_is_mm_prefix_lm(self) -> bool:
|
||||
return _orig_is_mm_prefix_lm(self) or getattr(self.hf_config, "model_type", None) == "bagel"
|
||||
|
||||
|
||||
_patched_is_mm_prefix_lm.__set_name__(_ModelConfig, "is_mm_prefix_lm")
|
||||
_ModelConfig.is_mm_prefix_lm = _patched_is_mm_prefix_lm
|
||||
|
||||
# Extend RequestStatus enum with omni-specific statuses
|
||||
if not hasattr(RequestStatus, "WAITING_FOR_CHUNK"):
|
||||
# The value - 1 is intentionally chosen to ensure it is treated
|
||||
|
||||
@@ -98,9 +98,25 @@ class GPUARModelRunner(OmniGPUModelRunner):
|
||||
if self.execute_model_state is not None:
|
||||
raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
|
||||
|
||||
if not getattr(self, "_warmup_state_cleared", False):
|
||||
self._warmup_state_cleared = True
|
||||
if hasattr(self.model, "_clear_warmup_state"):
|
||||
self.model._clear_warmup_state()
|
||||
|
||||
# [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
|
||||
finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {})
|
||||
if finished_reqs and hasattr(self.model, "get_kv_transfer_metadata"):
|
||||
for req_id, data in finished_reqs.items():
|
||||
try:
|
||||
model_meta = self.model.get_kv_transfer_metadata(req_id)
|
||||
if model_meta:
|
||||
existing = data.get("custom_metadata") or {}
|
||||
existing.update(model_meta)
|
||||
data["custom_metadata"] = existing
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to get custom metadata from model for {req_id}: {e}")
|
||||
self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer(
|
||||
finished_reqs=getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}),
|
||||
finished_reqs=finished_reqs,
|
||||
kv_caches=self.kv_caches,
|
||||
block_size=self.cache_config.block_size,
|
||||
cache_dtype=str(self.cache_config.cache_dtype),
|
||||
@@ -290,6 +306,10 @@ class GPUARModelRunner(OmniGPUModelRunner):
|
||||
sampler=self.sampler,
|
||||
)
|
||||
|
||||
# [Omni] Map pending ropes metadata to req_ids.
|
||||
if hasattr(self.model, "flush_pending_metadata"):
|
||||
self.model.flush_pending_metadata(list(req_ids))
|
||||
|
||||
with record_function_or_nullcontext("gpu_model_runner: postprocess"):
|
||||
if self.use_aux_hidden_state_outputs:
|
||||
# True when EAGLE 3 is used.
|
||||
|
||||
Reference in new issue
Block a user