diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml new file mode 100644 index 000000000..edd98fafd --- /dev/null +++ b/.buildkite/release-pipeline.yaml @@ -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" diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 8f141576d..001c46645 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -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: diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 6c1b458b8..092d6f11a 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -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" diff --git a/README.md b/README.md index d20a90e12..df5d093b1 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone

-| Documentation | User Forum | Developer Slack | WeChat | Paper | Slides | +| Documentation | User Forum | Developer Slack | WeChat | Paper | Slides |

diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 1a6889663..cb80828eb 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -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/* diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index e9db9ee49..3e0b60650 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -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/* diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 3d59f9b04..1595165e8 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -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 \ diff --git a/docs/contributing/metrics.md b/docs/contributing/metrics.md index bc0f13efd..92dcd92cc 100644 --- a/docs/contributing/metrics.md +++ b/docs/contributing/metrics.md @@ -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** diff --git a/docs/features/custom_pipeline.md b/docs/features/custom_pipeline.md index be46b5b14..aee5cfa6b 100644 --- a/docs/features/custom_pipeline.md +++ b/docs/features/custom_pipeline.md @@ -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 diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index b22aee84c..ca5f7a2ac 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -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 | diff --git a/docs/user_guide/diffusion/quantization/gguf.md b/docs/user_guide/diffusion/quantization/gguf.md index 025180e18..9b9cb6fc7 100644 --- a/docs/user_guide/diffusion/quantization/gguf.md +++ b/docs/user_guide/diffusion/quantization/gguf.md @@ -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 ``` diff --git a/docs/user_guide/examples/offline_inference/qwen3_tts.md b/docs/user_guide/examples/offline_inference/qwen3_tts.md index 5c3b044d7..afadd39f4 100644 --- a/docs/user_guide/examples/offline_inference/qwen3_tts.md +++ b/docs/user_guide/examples/offline_inference/qwen3_tts.md @@ -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. diff --git a/docs/user_guide/examples/online_serving/qwen3_tts.md b/docs/user_guide/examples/online_serving/qwen3_tts.md index 77a3947b5..e3de00515 100644 --- a/docs/user_guide/examples/online_serving/qwen3_tts.md +++ b/docs/user_guide/examples/online_serving/qwen3_tts.md @@ -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) | diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py index 584e463e1..cffe44fc3 100644 --- a/examples/offline_inference/bagel/end2end.py +++ b/examples/offline_inference/bagel/end2end.py @@ -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) diff --git a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py index 6af06fa19..a9b3c32f8 100644 --- a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py +++ b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py @@ -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"}, diff --git a/examples/offline_inference/custom_pipeline/image_to_image/run.sh b/examples/offline_inference/custom_pipeline/image_to_image/run.sh index 1d2706e25..d62efae28 100644 --- a/examples/offline_inference/custom_pipeline/image_to_image/run.sh +++ b/examples/offline_inference/custom_pipeline/image_to_image/run.sh @@ -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 diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 305a6d54a..81fe24e21 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -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") diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 58da640d9..7b0e1bb2e 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -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, diff --git a/examples/offline_inference/qwen3_omni/README.md b/examples/offline_inference/qwen3_omni/README.md index b7eee8f74..b3e859253 100644 --- a/examples/offline_inference/qwen3_omni/README.md +++ b/examples/offline_inference/qwen3_omni/README.md @@ -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. diff --git a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py new file mode 100644 index 000000000..965595578 --- /dev/null +++ b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py @@ -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 + +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.") diff --git a/examples/offline_inference/qwen3_omni/run_multiple_prompts_async_chunk.sh b/examples/offline_inference/qwen3_omni/run_multiple_prompts_async_chunk.sh new file mode 100755 index 000000000..809054867 --- /dev/null +++ b/examples/offline_inference/qwen3_omni/run_multiple_prompts_async_chunk.sh @@ -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 \ + "$@" diff --git a/examples/offline_inference/qwen3_omni/run_single_prompt_async_chunk.sh b/examples/offline_inference/qwen3_omni/run_single_prompt_async_chunk.sh new file mode 100755 index 000000000..918c7ee4f --- /dev/null +++ b/examples/offline_inference/qwen3_omni/run_single_prompt_async_chunk.sh @@ -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 \ + "$@" diff --git a/examples/offline_inference/qwen3_tts/end2end.py b/examples/offline_inference/qwen3_tts/end2end.py index 98303533a..f60365032 100644 --- a/examples/offline_inference/qwen3_tts/end2end.py +++ b/examples/offline_inference/qwen3_tts/end2end.py @@ -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 = { diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 15cbbbf72..382971606 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -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, } diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 857530045..c1efd342e 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -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, ) diff --git a/examples/online_serving/qwen3_tts/README.md b/examples/online_serving/qwen3_tts/README.md index 2046c6508..5088a4ca6 100644 --- a/examples/online_serving/qwen3_tts/README.md +++ b/examples/online_serving/qwen3_tts/README.md @@ -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 diff --git a/pyproject.toml b/pyproject.toml index 99194c648..5bcfab7d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 = [ diff --git a/tests/conftest.py b/tests/conftest.py index c3ff10246..c775ae349 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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"]) diff --git a/tests/diffusion/distributed/test_hsdp.py b/tests/diffusion/distributed/test_hsdp.py index 3309155e3..6fcaf1e81 100644 --- a/tests/diffusion/distributed/test_hsdp.py +++ b/tests/diffusion/distributed/test_hsdp.py @@ -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.""" diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index fd0be6c98..986bad9a8 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/tests/diffusion/test_worker_wrapper_base.py @@ -96,6 +96,7 @@ class TestWorkerWrapperBaseInitialization: local_rank=0, rank=0, od_config=mock_od_config, + skip_load_model=False, ) diff --git a/tests/distributed/omni_connectors/test_kv_flow.py b/tests/distributed/omni_connectors/test_kv_flow.py index 1f3d4acab..b12fc013b 100644 --- a/tests/distributed/omni_connectors/test_kv_flow.py +++ b/tests/distributed/omni_connectors/test_kv_flow.py @@ -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 diff --git a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml index 79e825d5d..dbb93344b 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml @@ -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 diff --git a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml index eeefab236..686600f72 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml @@ -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 diff --git a/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml index d6b04258a..87707e919 100644 --- a/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml @@ -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 diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py index bf36b78ff..6e5f449c1 100644 --- a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py +++ b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py @@ -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") diff --git a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py b/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py index 998e6232e..beac14099 100644 --- a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py +++ b/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py @@ -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, ) diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni.py b/tests/e2e/offline_inference/test_qwen2_5_omni.py index c5bb65842..d3c20580c 100644 --- a/tests/e2e/offline_inference/test_qwen2_5_omni.py +++ b/tests/e2e/offline_inference/test_qwen2_5_omni.py @@ -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] diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py index 9f76f911d..8c57becef 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni.py +++ b/tests/e2e/offline_inference/test_qwen3_omni.py @@ -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] diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index 74eaa9c6b..fb9986592 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -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()) diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index e2f77af73..fd7e0704c 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -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']}" diff --git a/tests/e2e/stage_configs/qwen2_5_omni_thinker_ci.yaml b/tests/e2e/stage_configs/qwen2_5_omni_thinker_ci.yaml index 3ebf914e8..23eacafdf 100644 --- a/tests/e2e/stage_configs/qwen2_5_omni_thinker_ci.yaml +++ b/tests/e2e/stage_configs/qwen2_5_omni_thinker_ci.yaml @@ -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 diff --git a/tests/e2e/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/stage_configs/qwen3_omni_ci.yaml index 45e7fcab8..6f1293f96 100644 --- a/tests/e2e/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/stage_configs/qwen3_omni_ci.yaml @@ -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 diff --git a/tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml b/tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml index ec815ad1f..d51f7a5c8 100644 --- a/tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml +++ b/tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml @@ -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 diff --git a/tests/e2e/stage_configs/rocm/qwen3_omni_ci.yaml b/tests/e2e/stage_configs/rocm/qwen3_omni_ci.yaml index 477e6e59f..23e75aae7 100644 --- a/tests/e2e/stage_configs/rocm/qwen3_omni_ci.yaml +++ b/tests/e2e/stage_configs/rocm/qwen3_omni_ci.yaml @@ -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 diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 10dbd3a3b..8b7177365 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -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.""" diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index ce443d5a5..ed205032e 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -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, diff --git a/tests/perf/stage_configs/qwen3_omni.yaml b/tests/perf/stage_configs/qwen3_omni.yaml index 802f8dd24..01ca7d9fb 100644 --- a/tests/perf/stage_configs/qwen3_omni.yaml +++ b/tests/perf/stage_configs/qwen3_omni.yaml @@ -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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index da68adbdb..6a1e44472 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -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 diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index e5ace3b28..abd668466 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -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": diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index fbcfe0651..4c09fd7a0 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -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", ) diff --git a/vllm_omni/diffusion/hooks/base.py b/vllm_omni/diffusion/hooks/base.py index 8b330a19f..3c975408c 100644 --- a/vllm_omni/diffusion/hooks/base.py +++ b/vllm_omni/diffusion/hooks/base.py @@ -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): diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py new file mode 100644 index 000000000..cac406ef4 --- /dev/null +++ b/vllm_omni/diffusion/ipc.py @@ -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 diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index e715bdaec..09144e05a 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -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 diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 46ce4b6bd..9ed3481d7 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -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") diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index 1a33a5eea..ba15a62b0 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -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: diff --git a/vllm_omni/diffusion/models/helios/scheduling_helios.py b/vllm_omni/diffusion/models/helios/scheduling_helios.py index 847dea2d0..741e906e8 100644 --- a/vllm_omni/diffusion/models/helios/scheduling_helios.py +++ b/vllm_omni/diffusion/models/helios/scheduling_helios.py @@ -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) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 1de801c10..9a0037a6a 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -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: diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index b673cd0fa..fe15a24f5 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -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: diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index c1841dadf..5b66ae5ee 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -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)") diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 945cfc413..ebf85624f 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -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, diff --git a/vllm_omni/diffusion/scheduler.py b/vllm_omni/diffusion/scheduler.py index c30d907aa..5f1d01c12 100644 --- a/vllm_omni/diffusion/scheduler.py +++ b/vllm_omni/diffusion/scheduler.py @@ -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") diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index d4389f4d2..19ba5f31c 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -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 diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 21f09c951..2806f7315 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -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: diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index c7289d8b8..db8d5dbaa 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -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) diff --git a/vllm_omni/distributed/omni_connectors/utils/kv_utils.py b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py new file mode 100644 index 000000000..2cb48a8b3 --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py @@ -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 diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index b663430dc..6a4b88dbf 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -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, diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index dddcb12b5..c414b972b 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -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: diff --git a/vllm_omni/entrypoints/omni_diffusion.py b/vllm_omni/entrypoints/omni_diffusion.py index 9937726e5..59ddf2eeb 100644 --- a/vllm_omni/entrypoints/omni_diffusion.py +++ b/vllm_omni/entrypoints/omni_diffusion.py @@ -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) diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 4114ef3fb..b0f94e426 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -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 diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 969531b08..6db97dd2d 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -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 diff --git a/vllm_omni/entrypoints/openai/metadata_manager.py b/vllm_omni/entrypoints/openai/metadata_manager.py new file mode 100644 index 000000000..4077aa23b --- /dev/null +++ b/vllm_omni/entrypoints/openai/metadata_manager.py @@ -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 diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 8c292cc7a..de9c559dd 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -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 diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index f73df5ea4..b663eb52f 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -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: diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 110bf48cd..de4a71719 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -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} diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 16d690bf3..c0f10af2b 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -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( diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index 068b4a30b..46fec6fcc 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -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, diff --git a/vllm_omni/model_executor/models/__init__.py b/vllm_omni/model_executor/models/__init__.py index 0b2629b4a..68074bbb9 100644 --- a/vllm_omni/model_executor/models/__init__.py +++ b/vllm_omni/model_executor/models/__init__.py @@ -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", +] diff --git a/vllm_omni/model_executor/models/bagel/__init__.py b/vllm_omni/model_executor/models/bagel/__init__.py new file mode 100644 index 000000000..11e4c6667 --- /dev/null +++ b/vllm_omni/model_executor/models/bagel/__init__.py @@ -0,0 +1,3 @@ +from .bagel import OmniBagelForConditionalGeneration + +__all__ = ["OmniBagelForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py new file mode 100644 index 000000000..24a1331f9 --- /dev/null +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -0,0 +1,1016 @@ +from collections import deque +from collections.abc import Iterable, Mapping, Sequence +from math import isqrt +from typing import Any + +import torch +import torch.nn as nn +from transformers import BatchFeature +from vllm.config import VllmConfig +from vllm.config.multimodal import BaseDummyOptions +from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm +from vllm.model_executor.layers.linear import ( + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.bagel import BagelForConditionalGeneration +from vllm.model_executor.models.interfaces import MultiModalEmbeddings +from vllm.model_executor.models.qwen2 import Qwen2DecoderLayer, Qwen2MLP +from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import ( + MultiModalDataDict, + MultiModalFieldConfig, + MultiModalKwargsItems, +) +from vllm.multimodal.parse import ( + ImageEmbeddingItems, + ImageProcessorItems, + ModalityData, + ModalityDataItems, + MultiModalDataItems, + MultiModalDataParser, +) +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, + PromptUpdateDetails, +) +from vllm.transformers_utils.processors.bagel import BagelProcessor + +from vllm_omni.diffusion.models.bagel.autoencoder import ( + AutoEncoderParams, + DiagonalGaussian, + Encoder, +) +from vllm_omni.diffusion.models.bagel.bagel_transformer import ( + PositionEmbedding, + TimestepEmbedder, +) +from vllm_omni.diffusion.models.bagel.pipeline_bagel import default_ae_params + + +class OmniBagelProcessor(BagelProcessor): + def __call__(self, text=None, images=None, **kwargs): + is_img2img = kwargs.pop("is_img2img", False) + + if is_img2img and images is not None: + image_kwargs = kwargs.copy() + image_kwargs["do_resize"] = False + image_kwargs["do_rescale"] = True + if "return_tensors" not in image_kwargs: + image_kwargs["return_tensors"] = "pt" + + pixel_values = self.image_processor(images, **image_kwargs) + + text_inputs = self.tokenizer(text, **kwargs) if text is not None else None + + if pixel_values is not None and text_inputs is not None: + combined = dict(text_inputs) + combined["pixel_values"] = pixel_values["pixel_values"] + return BatchFeature(combined) + elif pixel_values is not None: + return pixel_values + elif text_inputs is not None: + return BatchFeature(dict(text_inputs)) + else: + return BatchFeature({}) + + return super().__call__(text, images, **kwargs) + + +class OmniBagelProcessingInfo(BaseProcessingInfo): + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"image": 1, "img2img": 1} + + def get_hf_processor(self, **kwargs: object): + return self.ctx.get_hf_processor(OmniBagelProcessor, **kwargs) + + def get_hf_config(self): + config = super().get_hf_config() + if not getattr(self, "_latent_size_patched", False): + self._latent_size_patched = True + self._patch_max_latent_size(config) + return config + + def _patch_max_latent_size(self, config): + """Infer correct max_latent_size from the model's latent_pos_embed + weight, since the HF config value may be stale (e.g. 32 vs 64).""" + import json + from pathlib import Path + + model_name = self.ctx.model_config.model + try: + p = Path(model_name) + if p.is_dir(): + index_path = p / "model.safetensors.index.json" + else: + from huggingface_hub import hf_hub_download + + index_path = Path(hf_hub_download(model_name, "model.safetensors.index.json")) + + if not index_path.exists(): + return + + with open(index_path) as f: + index = json.load(f) + + shard = index.get("weight_map", {}).get("latent_pos_embed.pos_embed") + if not shard: + return + + from safetensors import safe_open + + with safe_open(str(index_path.parent / shard), framework="pt") as f: + if "latent_pos_embed.pos_embed" in f.keys(): + npos = f.get_slice("latent_pos_embed.pos_embed").get_shape()[0] + side = isqrt(npos) + if side * side == npos: + old = getattr(config, "max_latent_size", 32) + if old != side: + config.max_latent_size = side + except Exception: + pass + + def get_data_parser(self) -> "OmniBagelDataParser": + return OmniBagelDataParser( + expected_hidden_size=self._get_expected_hidden_size(), + ) + + +class OmniBagelDummyInputsBuilder(BaseDummyInputsBuilder[OmniBagelProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + dummy_text = "" + if "image" in mm_counts: + dummy_text += "<|image_pad|>" * mm_counts["image"] + if "img2img" in mm_counts: + dummy_text += "<|fim_middle|>" * mm_counts["img2img"] + return dummy_text + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions] | None = None, + ) -> MultiModalDataDict: + num_images = mm_counts.get("image", 0) + hf_config = self.info.get_hf_config() + vit_config = hf_config.vit_config + + image_size = vit_config.image_size + image_overrides = mm_options.get("image") if mm_options else None + + return { + "image": self._get_dummy_images( + width=image_size, + height=image_size, + num_images=num_images, + overrides=image_overrides, + ), + "img2img": self._get_dummy_images( + width=image_size, + height=image_size, + num_images=mm_counts.get("img2img", 0), + overrides=image_overrides, + ), + } + + +class Img2ImgProcessorItems(ImageProcessorItems): + def __init__(self, data): + super().__init__(data) + self.modality = "img2img" + + def get_processor_data(self): + return {"pixel_values_img2img": self.get_all()} + + +class OmniBagelDataParser(MultiModalDataParser): + def _parse_img2img_data(self, data: ModalityData) -> ModalityDataItems | None: + items = self._parse_image_data(data) + if items is None: + return None + return Img2ImgProcessorItems(items.data) + + def _get_subparsers(self): + parsers = super()._get_subparsers() + parsers["img2img"] = self._parse_img2img_data + return parsers + + +class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingInfo]): + def _get_mm_fields_config(self, hf_inputs, hf_processor_mm_kwargs): + return { + "pixel_values": MultiModalFieldConfig.batched("image"), + "pixel_values_img2img": MultiModalFieldConfig.batched("img2img"), + } + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> "BatchFeature": + has_image = "images" in mm_data + has_img2img = "pixel_values_img2img" in mm_data + + if has_image and has_img2img: + outputs = BatchFeature() + + img_data = dict(mm_data) + if "pixel_values_img2img" in img_data: + del img_data["pixel_values_img2img"] + kwargs_img = dict(mm_kwargs) + kwargs_img["is_img2img"] = False + out_img = super()._call_hf_processor(prompt, img_data, kwargs_img, tok_kwargs) + if "pixel_values" in out_img: + outputs["pixel_values"] = out_img["pixel_values"] + for k, v in out_img.items(): + if k != "pixel_values": + outputs[k] = v + + img2img_data = dict(mm_data) + if "images" in img2img_data: + del img2img_data["images"] + img2img_data["images"] = img2img_data.pop("pixel_values_img2img") + kwargs_img2img = dict(mm_kwargs) + kwargs_img2img["is_img2img"] = True + out_img2img = super()._call_hf_processor(prompt, img2img_data, kwargs_img2img, tok_kwargs) + if "pixel_values" in out_img2img: + outputs["pixel_values_img2img"] = out_img2img["pixel_values"] + for k, v in out_img2img.items(): + if k not in outputs: + outputs[k] = v + + return outputs + + elif has_img2img: + mm_data = dict(mm_data) + mm_data["images"] = mm_data.pop("pixel_values_img2img") + mm_kwargs = dict(mm_kwargs) + mm_kwargs["is_img2img"] = True + outputs = super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + if "pixel_values" in outputs: + outputs["pixel_values_img2img"] = outputs.pop("pixel_values") + return outputs + + return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return False + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptReplacement]: + hf_config = self.info.get_hf_config() + tokenizer = self.info.get_tokenizer() + + replacements: list[PromptReplacement] = [] + + image_token_id = tokenizer.get_vocab().get("<|image_pad|>") + if image_token_id is not None: + num_patches = hf_config.vit_max_num_patch_per_side**2 + + def get_image_replacement(item_idx: int): + return [image_token_id] * num_patches + + replacements.append( + PromptReplacement( + modality="image", + target=[image_token_id], + replacement=get_image_replacement, + ) + ) + + img2img_token_id = tokenizer.get_vocab().get("<|fim_middle|>") + if img2img_token_id is not None: + vit_config = hf_config.vit_config + image_size = vit_config.image_size + num_vit_patches = (image_size // vit_config.patch_size) ** 2 + + latent_patch_size = getattr(hf_config, "latent_patch_size", 2) + downsample = hf_config.vae_config.get("downsample", 8) + latent_downsample = downsample * latent_patch_size + + def get_img2img_replacement(item_idx: int): + h, w = image_size, image_size + if "img2img" in mm_items: + item = mm_items.get_items("img2img", (Img2ImgProcessorItems, ImageEmbeddingItems)) + if hasattr(item, "get_image_size"): + size = item.get_image_size(item_idx) + h, w = size.height, size.width + + max_latent_size = getattr(hf_config, "max_latent_size", 32) + max_img_size = int(max_latent_size * latent_downsample) + stride = latent_downsample + scale = min(max_img_size / max(h, w), 1.0) + min_img_size = min(256, max_img_size) + scale = max(scale, min_img_size / min(h, w)) + new_h = max(stride, int(round(h * scale / stride) * stride)) + new_w = max(stride, int(round(w * scale / stride) * stride)) + new_h = min(new_h, max_img_size) + new_w = min(new_w, max_img_size) + + num_vae_patches = (new_h // latent_downsample) * (new_w // latent_downsample) + num_vae_total = num_vae_patches + 2 + num_vit_total = num_vit_patches + 2 + # +1 separator between VAE and ViT blocks so that + # extract_embeds_range() produces two distinct mm_prefix_range + # entries, preventing VAE tokens from attending to ViT. + total = num_vae_total + 1 + num_vit_total + tokens = [img2img_token_id] * total + + embed_mask = [True] * num_vae_total + [False] + [True] * num_vit_total + return PromptUpdateDetails( + full=tokens, + is_embed=lambda _tok, _seq, _m=embed_mask: torch.tensor(_m, dtype=torch.bool), + ) + + replacements.append( + PromptReplacement( + modality="img2img", + target=[img2img_token_id], + replacement=get_img2img_replacement, + ) + ) + + return replacements + + +class VAEEncoder(nn.Module): + """Lightweight VAE encoder (no decoder) for embedding images in the AR stage.""" + + def __init__(self, params: AutoEncoderParams): + super().__init__() + self.encoder = Encoder( + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + ch_mult=params.ch_mult, + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + self.reg = DiagonalGaussian() + self.scale_factor = params.scale_factor + self.shift_factor = params.shift_factor + + def encode(self, x: torch.Tensor) -> torch.Tensor: + z = self.reg(self.encoder(x)) + z = self.scale_factor * (z - self.shift_factor) + return z + + +@MULTIMODAL_REGISTRY.register_processor( + OmniBagelMultiModalProcessor, + info=OmniBagelProcessingInfo, + dummy_inputs=OmniBagelDummyInputsBuilder, +) +class OmniBagelForConditionalGeneration(BagelForConditionalGeneration): + """ + Omni version of BagelForConditionalGeneration. + + Extends the base model with a VAE encoder so that img2img can embed + both VAE latents and ViT features within the AR stage, producing a + combined KV cache that is then transferred to the DiT stage. + + Position IDs are adjusted so that: + - VAE tokens all share position 0 + - ViT tokens all share position 1 + - Text tokens use sequential positions starting from 2 + This matches the position scheme used by the single-stage DiT pipeline, + ensuring the transferred KV cache + ropes are directly compatible with + the DiT's denoising loop. + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + config = vllm_config.model_config.hf_config + self.latent_patch_size = getattr(config, "latent_patch_size", 2) + self.downsample = config.vae_config.get("downsample") + self.latent_downsample = self.downsample * self.latent_patch_size + self.max_latent_size = getattr(config, "max_latent_size", 32) + self.latent_channel = config.vae_config.get("z_channels") + + hidden_size = config.llm_config.hidden_size + patch_latent_dim = self.latent_patch_size**2 * self.latent_channel + self.vae = VAEEncoder(default_ae_params()) + self.vae2llm = nn.Linear(patch_latent_dim, hidden_size) + self.latent_pos_embed = PositionEmbedding(self.max_latent_size, hidden_size) + self.time_embedder = TimestepEmbedder(hidden_size) + + self._pending_img2img_info: list[tuple[int, int, int, int]] = [] + self._ropes_pending: list[dict[str, Any]] = [] + self._ropes_metadata: dict[str, dict[str, Any]] = {} + self._cfg_companion_queue: deque[tuple[tuple[int, int, int, int], int]] = deque() + + from transformers import AutoTokenizer + + tok_name = getattr(vllm_config.model_config, "tokenizer", None) or vllm_config.model_config.model + _tok = AutoTokenizer.from_pretrained(tok_name, trust_remote_code=True) + for t in ["<|vision_start|>", "<|vision_end|>"]: + if t not in _tok.get_vocab(): + _tok.add_tokens([t]) + self._start_of_image_id = int(_tok.convert_tokens_to_ids("<|vision_start|>")) + self._end_of_image_id = int(_tok.convert_tokens_to_ids("<|vision_end|>")) + + self._vae_token_mask: torch.Tensor | None = None + + self._install_mot_modules(config) + + def _install_mot_modules(self, config): + """Add generation-mode (MoT) weight modules to each Qwen2 decoder layer. + + The single-stage DiT routes VAE latent tokens through separate + ``qkv_proj_moe_gen / o_proj_moe_gen / mlp_moe_gen`` weight matrices + (``mode="gen"``). We replicate that structure here so the AR stage + produces the same KV cache values. + """ + llm_cfg = config.llm_config + hidden_size = llm_cfg.hidden_size + intermediate_size = llm_cfg.intermediate_size + num_heads = llm_cfg.num_attention_heads + num_kv_heads = llm_cfg.num_key_value_heads + head_dim = hidden_size // num_heads + rms_eps = llm_cfg.rms_norm_eps + + qwen2_model = self.language_model.model # Qwen2Model + + qwen2_model.norm_moe_gen = VllmRMSNorm(hidden_size, eps=rms_eps) + + for layer in qwen2_model.layers: + if not isinstance(layer, Qwen2DecoderLayer): + continue + attn = layer.self_attn + + attn.qkv_proj_moe_gen = QKVParallelLinear( + hidden_size, + head_dim, + num_heads, + num_kv_heads, + bias=True, + ) + attn.o_proj_moe_gen = RowParallelLinear( + num_heads * head_dim, + hidden_size, + bias=False, + ) + attn.q_norm_moe_gen = VllmRMSNorm(head_dim, eps=rms_eps) + attn.k_norm_moe_gen = VllmRMSNorm(head_dim, eps=rms_eps) + + layer.mlp_moe_gen = Qwen2MLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + hidden_act=llm_cfg.hidden_act, + ) + layer.input_layernorm_moe_gen = VllmRMSNorm(hidden_size, eps=rms_eps) + layer.post_attention_layernorm_moe_gen = VllmRMSNorm(hidden_size, eps=rms_eps) + + def _resize_to_stride(self, pixel_values: torch.Tensor) -> torch.Tensor: + """Resize pixel values to stride-aligned dimensions + (matches DiT's ``_resize_images_to_stride``).""" + H, W = pixel_values.shape[2], pixel_values.shape[3] + stride = self.latent_downsample + max_img_size = int(self.max_latent_size * stride) + + scale = min(max_img_size / max(H, W), 1.0) + min_img_size = min(256, max_img_size) + scale = max(scale, min_img_size / min(H, W)) + new_H = max(stride, int(round(H * scale / stride) * stride)) + new_W = max(stride, int(round(W * scale / stride) * stride)) + new_H = min(new_H, max_img_size) + new_W = min(new_W, max_img_size) + + if new_H != H or new_W != W: + pixel_values = torch.nn.functional.interpolate( + pixel_values, size=(new_H, new_W), mode="bicubic", align_corners=False + ) + return pixel_values + + def _clear_warmup_state(self): + """Clear stale state accumulated during warmup/profiling runs.""" + self._ropes_pending.clear() + self._ropes_metadata.clear() + self._pending_img2img_info.clear() + self._cfg_companion_queue.clear() + self._vae_token_mask = None + + def get_kv_transfer_metadata(self, req_id: str) -> dict[str, Any] | None: + return self._ropes_metadata.pop(req_id, None) + + def flush_pending_metadata(self, req_ids: list[str]) -> None: + """Map pending metadata (batch order) to req_ids after forward().""" + pending = self._ropes_pending + self._ropes_pending = [] + for i, meta in enumerate(pending): + if i < len(req_ids): + self._ropes_metadata[req_ids[i]] = meta + + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: + mm_input_by_modality = {} + + if any(k in kwargs for k in ("pixel_values", "image_embeds")): + mm_input_by_modality["img2text"] = self._parse_and_validate_image_input(**kwargs) + + img2img_keys = {"pixel_values_img2img": "pixel_values", "image_embeds_img2img": "image_embeds"} + img2img_kwargs = {img2img_keys[k]: v for k, v in kwargs.items() if k in img2img_keys} + + if img2img_kwargs: + combined_kwargs = kwargs.copy() + combined_kwargs.update(img2img_kwargs) + mm_input_by_modality["img2img"] = self._parse_and_validate_image_input(**combined_kwargs) + + return mm_input_by_modality + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None: + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) + if not mm_input_by_modality: + return None + multimodal_embeddings: tuple[torch.Tensor, ...] = () + + for modality in mm_input_by_modality: + multimodal_input = mm_input_by_modality[modality] + if modality == "img2text": + image_embeddings = self._process_img2text_input(multimodal_input) + multimodal_embeddings += tuple(image_embeddings) + if modality == "img2img": + img2img_embeddings = self._process_img2img_input(multimodal_input) + multimodal_embeddings += tuple(img2img_embeddings) + return multimodal_embeddings + + def get_flattened_position_ids(self, img_h, img_w, patch_size, max_num_patches_per_side): + num_patches_h, num_patches_w = img_h // patch_size, img_w // patch_size + coords_h = torch.arange(0, num_patches_h) + coords_w = torch.arange(0, num_patches_w) + pos_ids = (coords_h[:, None] * max_num_patches_per_side + coords_w).flatten() + return pos_ids + + def _process_img2text_input(self, multimodal_input): + return self._process_image_input(multimodal_input) + + def _process_img2img_input(self, multimodal_input): + pixel_values = multimodal_input["pixel_values"] + if pixel_values.ndim == 5: + b, n, c, h, w = pixel_values.shape + pixel_values = pixel_values.reshape(b * n, c, h, w) + + num_images = pixel_values.shape[0] + image_size = self.config.vit_config.image_size + p = self.latent_patch_size + timestep = 0 + + if self._ropes_pending: + self._ropes_pending.clear() + + vit_pixel_values = torch.nn.functional.interpolate( + pixel_values, + size=(image_size, image_size), + mode="bicubic", + align_corners=False, + ) + + vit_embeddings_tuple = self._process_image_input({"pixel_values": vit_pixel_values}) + + marker_ids = torch.tensor( + [self._start_of_image_id, self._end_of_image_id], + device=pixel_values.device, + dtype=torch.long, + ) + marker_embeds = self.language_model.model.embed_tokens(marker_ids) + start_embed = marker_embeds[0:1] + end_embed = marker_embeds[1:2] + + results = [] + + for i in range(num_images): + single_pv = pixel_values[i : i + 1] + single_pv = self._resize_to_stride(single_pv) + H, W = single_pv.shape[2:] + + padded_latent = self.vae.encode(single_pv) + h = H // self.latent_downsample + w = W // self.latent_downsample + + latent = padded_latent[0][:, : h * p, : w * p] + latent = latent.reshape(self.latent_channel, h, p, w, p) + latent = torch.einsum("chpwq->hwpqc", latent).reshape(-1, p * p * self.latent_channel) + + vae_position_ids = self.get_flattened_position_ids( + H, + W, + self.latent_downsample, + max_num_patches_per_side=self.max_latent_size, + ) + pos_embed = self.latent_pos_embed([vae_position_ids]) + packed_timesteps = torch.tensor([timestep], device=padded_latent.device) + with torch.amp.autocast("cuda", dtype=torch.bfloat16): + timestep_embeds = self.time_embedder(packed_timesteps.to(padded_latent)) + vae_embeds = self.vae2llm(latent) + timestep_embeds + pos_embed + + vit_emb = vit_embeddings_tuple[i] if i < len(vit_embeddings_tuple) else vit_embeddings_tuple[0] + + se = start_embed.to(vae_embeds.dtype) + ee = end_embed.to(vae_embeds.dtype) + combined = torch.cat([se, vae_embeds, ee, se, vit_emb, ee], dim=0) + results.append(combined) + + num_vae = h * w + 2 # +2 for start/end markers + num_vit = vit_emb.shape[0] + 2 + info = (num_vae, num_vit, int(H), int(W)) + self._pending_img2img_info.append(info) + self._cfg_companion_queue.append((info, 2)) # cfg_text + cfg_img + + return tuple(results) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors=None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor: + use_mot = False + seq_len = inputs_embeds.shape[0] if inputs_embeds is not None else positions.shape[0] + + if self._pending_img2img_info: + positions = self._adjust_positions_for_img2img(positions) + use_mot = True + + elif self._cfg_companion_queue: + cached, remaining = self._cfg_companion_queue[0] + remaining -= 1 + num_vae, num_vit, img_H, img_W = cached + num_img2img = num_vae + 1 + num_vit # +1 separator + seq_len = inputs_embeds.shape[0] if inputs_embeds is not None else positions.shape[0] + + if inputs_embeds is not None and seq_len >= num_img2img: + self._pending_img2img_info = [cached] + positions = self._adjust_positions_for_img2img(positions) + use_mot = True + else: + rope = int(positions[seq_len - 1].item()) + 1 + self._ropes_pending.append({"ropes": [rope]}) + + if remaining == 0: + self._cfg_companion_queue.popleft() + else: + self._cfg_companion_queue[0] = (cached, remaining) + + if use_mot: + return self._mot_forward(input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs) + return super().forward(input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs) + + def _adjust_positions_for_img2img(self, positions: torch.Tensor) -> torch.Tensor: + """Rewrite position IDs to match the single-stage DiT scheme: + VAE tokens -> position 0, separator -> position 0, + ViT tokens -> position 1, text -> 2, 3, ... + + Also computes ``self._vae_token_mask`` (bool tensor, True for actual + VAE latent patches that should use gen-mode weights) and pushes + per-request ropes + image_shape to the FIFO consumed by + ``get_kv_transfer_metadata``. + """ + info_list = self._pending_img2img_info + self._pending_img2img_info = [] + + if not info_list: + self._vae_token_mask = None + return positions + + boundaries = [0] + for i in range(1, len(positions)): + if positions[i] < positions[i - 1]: + boundaries.append(i) + boundaries.append(len(positions)) + + num_requests = len(boundaries) - 1 + new_positions = positions.clone() + vae_mask = torch.zeros(len(positions), dtype=torch.bool, device=positions.device) + + img2img_idx = 0 + for req_idx in range(num_requests): + start = boundaries[req_idx] + end = boundaries[req_idx + 1] + req_len = end - start + + if img2img_idx < len(info_list): + num_vae, num_vit, img_H, img_W = info_list[img2img_idx] + num_img2img = num_vae + 1 + num_vit # +1 separator + + if req_len >= num_img2img: + new_positions[start : start + num_vae] = 0 + new_positions[start + num_vae] = 0 # separator + vit_start = start + num_vae + 1 + new_positions[vit_start : vit_start + num_vit] = 1 + num_text = req_len - num_img2img + if num_text > 0: + text_start = start + num_img2img + new_positions[text_start:end] = torch.arange( + 2, 2 + num_text, device=positions.device, dtype=positions.dtype + ) + + # VAE gen-mode mask: only actual VAE patches (not markers) + vae_patches_start = start + 1 # skip start_marker + vae_patches_end = start + num_vae - 1 # before end_marker + if vae_patches_end > vae_patches_start: + vae_mask[vae_patches_start:vae_patches_end] = True + + rope = 2 + num_text + self._ropes_pending.append( + { + "ropes": [rope], + "image_shape": [img_H, img_W], + } + ) + img2img_idx += 1 + continue + + rope = int(new_positions[end - 1].item()) + 1 + self._ropes_pending.append({"ropes": [rope]}) + + self._vae_token_mask = vae_mask if vae_mask.any() else None + return new_positions + + # ------------------------------------------------------------------ + # MoT (Mixture-of-Transformers) forward path + # ------------------------------------------------------------------ + + def _mot_forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors, + inputs_embeds: torch.Tensor | None, + **kwargs, + ) -> torch.Tensor: + """Full forward pass with MoT routing for img2img requests. + + VAE latent patches are routed through ``*_moe_gen`` weight matrices + while all other tokens (markers, ViT, separator, text) use the + standard understanding-mode weights. + """ + qwen2_model = self.language_model.model # Qwen2Model + + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = qwen2_model.embed_input_ids(input_ids) + + residual = None + vae_mask = self._vae_token_mask + self._vae_token_mask = None # consumed + + for layer in qwen2_model.layers: + if not isinstance(layer, Qwen2DecoderLayer): + continue # skip PPMissingLayer (pipeline parallelism) + hidden_states, residual = self._mot_layer_forward( + layer, + positions, + hidden_states, + residual, + vae_mask, + ) + + # Final norm with MoT routing + if residual is not None: + hidden_states = hidden_states + residual + if vae_mask is not None and vae_mask.any(): + out = torch.empty_like(hidden_states) + non_vae = ~vae_mask + if non_vae.any(): + out[non_vae] = qwen2_model.norm(hidden_states[non_vae]) + out[vae_mask] = qwen2_model.norm_moe_gen(hidden_states[vae_mask]) + hidden_states = out + else: + hidden_states = qwen2_model.norm(hidden_states) + + return hidden_states + + def _mot_layer_forward( + self, + layer: Qwen2DecoderLayer, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + vae_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Single decoder-layer forward with MoT routing.""" + if vae_mask is None or not vae_mask.any(): + return layer(positions, hidden_states, residual) + + non_vae = ~vae_mask + + # ---- input layernorm (split) ---- + if residual is not None: + hidden_states = hidden_states + residual + residual = hidden_states + normed = torch.empty_like(hidden_states) + if non_vae.any(): + normed[non_vae] = layer.input_layernorm(hidden_states[non_vae]) + normed[vae_mask] = layer.input_layernorm_moe_gen(hidden_states[vae_mask]) + hidden_states = normed + + # ---- attention (split QKV / O projections) ---- + hidden_states = self._mot_attn_forward(layer.self_attn, positions, hidden_states, vae_mask) + + # ---- post-attention layernorm (split) ---- + hidden_states = hidden_states + residual + residual = hidden_states + normed = torch.empty_like(hidden_states) + if non_vae.any(): + normed[non_vae] = layer.post_attention_layernorm(hidden_states[non_vae]) + normed[vae_mask] = layer.post_attention_layernorm_moe_gen(hidden_states[vae_mask]) + hidden_states = normed + + # ---- MLP (split) ---- + mlp_out = torch.empty_like(hidden_states) + if non_vae.any(): + mlp_out[non_vae] = layer.mlp(hidden_states[non_vae]) + mlp_out[vae_mask] = layer.mlp_moe_gen(hidden_states[vae_mask]) + hidden_states = mlp_out + + return hidden_states, residual + + def _mot_attn_forward( + self, + attn, # Qwen2Attention + positions: torch.Tensor, + hidden_states: torch.Tensor, + vae_mask: torch.Tensor, + ) -> torch.Tensor: + """Attention forward with MoT routing for QKV and O projections.""" + non_vae = ~vae_mask + qkv_dim = attn.q_size + 2 * attn.kv_size + + # ---- QKV projection (split) ---- + qkv = torch.empty( + hidden_states.shape[0], + qkv_dim, + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + if non_vae.any(): + qkv_und, _ = attn.qkv_proj(hidden_states[non_vae]) + qkv[non_vae] = qkv_und + qkv_gen, _ = attn.qkv_proj_moe_gen(hidden_states[vae_mask]) + qkv[vae_mask] = qkv_gen + + q, k, v = qkv.split([attn.q_size, attn.kv_size, attn.kv_size], dim=-1) + + # ---- QK normalization (split) ---- + if attn.qk_norm: + n_tok = q.shape[0] + q = q.view(n_tok, attn.num_heads, attn.head_dim) + k = k.view(n_tok, attn.num_kv_heads, attn.head_dim) + + q_out = torch.empty_like(q) + k_out = torch.empty_like(k) + if non_vae.any(): + q_out[non_vae] = attn.q_norm(q[non_vae]) + k_out[non_vae] = attn.k_norm(k[non_vae]) + q_out[vae_mask] = attn.q_norm_moe_gen(q[vae_mask]) + k_out[vae_mask] = attn.k_norm_moe_gen(k[vae_mask]) + + q = q_out.reshape(n_tok, attn.q_size) + k = k_out.reshape(n_tok, attn.kv_size) + + # ---- RoPE + attention (same for all tokens) ---- + q, k = attn.rotary_emb(positions, q, k) + attn_output = attn.attn(q, k, v) + + # ---- O projection (split) ---- + output = torch.empty( + hidden_states.shape[0], + attn.hidden_size, + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + if non_vae.any(): + o_und, _ = attn.o_proj(attn_output[non_vae]) + output[non_vae] = o_und + o_gen, _ = attn.o_proj_moe_gen(attn_output[vae_mask]) + output[vae_mask] = o_gen + + return output + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + generation_keywords_to_skip = [ + "llm2vae", + "decoder.", + ] + + def _map_vae_weight_name(name: str) -> str: + if name.startswith("encoder."): + return "vae." + name + if name.startswith("reg."): + return "vae." + name + return name + + moe_gen_weights: list[tuple[str, torch.Tensor]] = [] + filtered_weights = [] + + for name, tensor in weights: + if any(skip in name for skip in generation_keywords_to_skip): + continue + + mapped_name = _map_vae_weight_name(name) + + if "moe_gen" in mapped_name: + moe_gen_weights.append((mapped_name, tensor)) + continue + + if "patch_embedding.weight" in mapped_name and tensor.ndim == 2: + out_channels = tensor.shape[0] + in_features = tensor.shape[1] + patch_size = self.config.vit_config.patch_size + in_channels = self.config.vit_config.num_channels + if in_features == in_channels * patch_size * patch_size: + tensor = tensor.reshape(out_channels, patch_size, patch_size, in_channels) + tensor = tensor.permute(0, 3, 1, 2).contiguous() + + if "latent_pos_embed.pos_embed" in mapped_name and tensor.ndim == 2: + npos, hdim = tensor.shape + current_param = self.latent_pos_embed.pos_embed + if current_param.shape != tensor.shape: + side = isqrt(int(npos)) + if side * side == int(npos) and hdim == current_param.shape[1]: + current_param.data = current_param.data.new_empty((npos, hdim)) + self.max_latent_size = int(side) + setattr(self.config, "max_latent_size", int(side)) + if hasattr(self.latent_pos_embed, "max_num_patch_per_side"): + self.latent_pos_embed.max_num_patch_per_side = int(side) + + filtered_weights.append((mapped_name, tensor)) + + loader = AutoWeightsLoader( + self, + skip_prefixes=["vit_pos_embed.pos_embed"], + ignore_unexpected_prefixes=["vae.", "latent_pos_embed.", "time_embedder.", "vae2llm."], + ) + loaded = loader.load_weights(filtered_weights, mapper=self.hf_to_vllm_mapper) + + loaded |= self._load_moe_gen_weights(moe_gen_weights) + + return loaded + + def _load_moe_gen_weights(self, weights: list[tuple[str, torch.Tensor]]) -> set[str]: + """Load generation-mode MoT weights with proper stacked-param mapping.""" + stacked_params = [ + ("qkv_proj_moe_gen", "q_proj_moe_gen", "q"), + ("qkv_proj_moe_gen", "k_proj_moe_gen", "k"), + ("qkv_proj_moe_gen", "v_proj_moe_gen", "v"), + ("mlp_moe_gen.gate_up_proj", "mlp_moe_gen.gate_proj", 0), + ("mlp_moe_gen.gate_up_proj", "mlp_moe_gen.up_proj", 1), + ] + + mapper = self.hf_to_vllm_mapper + prefix_map = getattr(mapper, "orig_to_new_prefix", {}) + + params_dict = dict(self.named_parameters()) + loaded: set[str] = set() + + for name, tensor in weights: + mapped = name + for orig, new in prefix_map.items(): + if mapped.startswith(orig): + mapped = new + mapped[len(orig) :] + break + + found_stacked = False + for param_name, weight_name, shard_id in stacked_params: + if weight_name not in mapped: + continue + mapped = mapped.replace(weight_name, param_name) + if mapped in params_dict: + param = params_dict[mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, tensor, shard_id) + loaded.add(mapped) + found_stacked = True + break + + if not found_stacked: + if mapped in params_dict: + param = params_dict[mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, tensor) + loaded.add(mapped) + + return loaded diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index 835dc6d6f..7dfb988eb 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -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()) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py index b70f2f3b3..b73ab33d7 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py @@ -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): diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index e80f1785f..a1f21a6a6 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -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) diff --git a/vllm_omni/model_executor/models/qwen3_tts/voice_cache_manager.py b/vllm_omni/model_executor/models/qwen3_tts/voice_cache_manager.py new file mode 100644 index 000000000..1e26a161d --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/voice_cache_manager.py @@ -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 diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 496dfe305..9bf2b4fdb 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -119,6 +119,11 @@ _OMNI_MODELS = { "glm_image_ar", "GlmImageForConditionalGeneration", ), + "OmniBagelForConditionalGeneration": ( + "bagel", + "bagel", + "OmniBagelForConditionalGeneration", + ), "HunyuanImage3ForCausalMM": ( "hunyuan_image3", "hunyuan_image3", diff --git a/vllm_omni/model_executor/stage_configs/bagel.yaml b/vllm_omni/model_executor/stage_configs/bagel.yaml index 23bbb441f..4bce816b4 100644 --- a/vllm_omni/model_executor/stage_configs/bagel.yaml +++ b/vllm_omni/model_executor/stage_configs/bagel.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml index 862e51b83..93700b6f3 100644 --- a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml +++ b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml index 564f1160c..19502add8 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_input_processors/bagel.py b/vllm_omni/model_executor/stage_input_processors/bagel.py index 80b28b230..d7055ff51 100644 --- a/vllm_omni/model_executor/stage_input_processors/bagel.py +++ b/vllm_omni/model_executor/stage_input_processors/bagel.py @@ -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 [] diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index af14dd0e6..6f035a81e 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -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)})" diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 6daef1a23..c9dea609e 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -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 diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 1ea56f724..d7d45031a 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -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.