Compare commits

..
Author SHA1 Message Date
Zhou Taichang 181eab754e Merge branch 'main' into dev/rebase-v0170 2026-03-10 09:50:31 +08:00
Taichang Zhou e942aa2640 revert changes for the docker file which will be updated after dockers are available
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-09 23:52:36 +08:00
Taichang Zhou c406622577 update installation doc
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-09 23:44:02 +08:00
Taichang Zhou b1073a6e7f Merge branch 'dev/rebase-v0170' of https://github.com/vllm-project/vllm-omni into dev/rebase-v0170 2026-03-09 23:40:54 +08:00
Taichang Zhou 0940c5a05e Merge remote-tracking branch 'remotes/origin/main' into dev/rebase-v0170
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-09 23:40:32 +08:00
tzhouam a9b4367808 Update GPU memory utilization and max number of batched tokens in qwen3_omni.yaml to optimize resource allocation during model execution.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-09 08:22:26 +00:00
tzhouam 79626586e1 Update GPU memory utilization in qwen3_omni.yaml from 0.6 to 0.45 to prevent out-of-memory errors during model execution.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-09 08:02:25 +00:00
tzhouam e6d4cdd71a Refactor model mode initialization and enhance error handling in AsyncOmni. Update worker type retrieval in omni_stage and add audio token handling in Qwen3OmniMoe classes for improved robustness.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-09 06:17:08 +00:00
Hongsheng Liu 9cdbcec44b Merge branch 'main' into dev/rebase-v0170 2026-03-08 09:18:57 +08:00
Taichang Zhou c8b695f181 smaller talker gpu util to avoid oom
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-08 01:48:42 +08:00
Taichang Zhou f405240513 Change pytest fixture scope from "module" to "class" for omni_server and omni_server_no_async_chunk in test_qwen3_tts.py to improve test isolation and resource management.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-07 14:50:59 +08:00
Taichang Zhou c630cf39c1 Enhance Dockerfile.ci by adding system dependencies installation and copying project files into the container. This improves the build process for the CI environment.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-07 14:02:48 +08:00
Taichang Zhou 3facc89dd1 Update base image in Dockerfile.rocm to v0.17.0 for improved compatibility and features.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-07 13:59:23 +08:00
Taichang Zhou 68bad871bc Fix model detection logic in is_diffusion_model function to ensure config_dict is checked for None before accessing its fields. This improves robustness against missing model_index.json files.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-07 13:57:44 +08:00
Taichang Zhou f945d2a30d Refactor Dockerfile.ci to streamline installation process and update vllm-omni installation method. Removed commented-out installation commands and ensured Python symlink is created.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-07 13:57:23 +08:00
Taichang Zhou 315f9401db Refactor imports in metrics.py and patch.py to use TokenizerLike from vllm.tokenizers instead of PreTrainedTokenizerBase from transformers.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-07 13:33:10 +08:00
Taichang Zhou fd4607f215 rebase 2026/03/07
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-07 13:30:58 +08:00
Taichang Zhou 0ecfff0694 comment out extra installation
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-07 13:29:12 +08:00
Taichang Zhou 4a43c6b017 Merge branch 'dev/rebase-v0170' of https://github.com/vllm-project/vllm-omni into dev/rebase-v0170 2026-03-07 13:26:46 +08:00
Taichang Zhou 0c2464db20 revert unnecessary changes
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-07 13:26:32 +08:00
Zhou TaichangandCyrus Leung 5f1e41970c Update vllm_omni/benchmarks/patch/patch.py
Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
2026-03-07 13:23:47 +08:00
Zhou TaichangandCyrus Leung 4ead1c8913 Update vllm_omni/benchmarks/metrics/metrics.py
Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
2026-03-07 13:23:38 +08:00
Taichang Zhou 2187db0e36 Remove obsolete RFC documents for automated rebase agent, output processor refactor, model runner clearance, long function decomposition, and model runner V2 integration. These files are no longer relevant to the current project direction.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-07 13:22:39 +08:00
Taichang Zhou 5511bc14c7 Merge origin/main into dev/rebase-v0170
- Resolved conflict in vllm_omni/engine/output_processor.py: use tracing_enabled for vLLM v0.17.0 API
- Resolved conflict in vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py: keep correct _CodePredictorDecoderLayer.forward()
- Update Dockerfile.ci vLLM wheel commit to b31e9326a (releases/v0.17.0)

Made-with: Cursor

Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-07 13:18:33 +08:00
Taichang Zhou 3a6b2e27be [Update] Bump VLLM base tag to v0.17.0 in CI Dockerfile
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-03-07 13:14:10 +08:00
tzhouam 8c41492307 Merge remote-tracking branch 'remotes/origin/main' into dev/rebase-v0169 2026-03-06 03:02:50 +00:00
tzhouam 44ac28c386 trigger RTD rebuild
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-06 03:02:07 +00:00
tzhouam 1ccffce2bc rebase to updated commit
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-06 02:32:26 +00:00
tzhouam 5be08baed3 Merge remote-tracking branch 'remotes/origin/main' into dev/rebase-v0170
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-05 08:53:26 +00:00
tzhouam 395eab8984 Merge remote-tracking branch 'remotes/origin/main' into dev/rebase-v0169 2026-03-05 02:29:14 +00:00
tzhouam 36a507e071 Merge remote-tracking branch 'remotes/origin/main' into dev/rebase-v0170
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-04 08:26:56 +00:00
tzhouam ae818f4b45 Fix Qwen3Omni VisionTransformer compatibility with updated vLLM Qwen2_5_VisionAttention
The upstream vLLM Qwen2_5_VisionAttention.forward() now requires a
`sequence_lengths` argument, but Qwen3_VisionBlock.forward() does not
pass it, causing a TypeError at init time.

Instead of patching the vLLM package directly, introduce a
Qwen3Omni_VisionTransformer subclass in vllm-omni that overrides
forward() to properly compute and propagate sequence_lengths via
MMEncoderAttention (following qwen3_vl.py's pattern). Both thinker
and talker now import the fixed subclass.

Also bump the vLLM wheel to e196028dbaa41b.

Made-with: Cursor

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-04 02:51:02 +00:00
tzhouam 4914466f0b Update Dockerfile.ci to install vLLM wheel version 3dd142adb7f7fa3b0b8345ee046ba36a12d9e530 and remove ExtractHiddenStatesProposer references from GPU model runners for cleaner code.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-03 15:09:26 +00:00
tzhouam 2398a7c1d7 Merge origin/main into dev/rebase-v0170
Resolve conflict in omni_generation_scheduler.py: keep both
num_cached_tokens tracking and additional_information caching.

Made-with: Cursor

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-03 14:31:05 +00:00
tzhouam 16824aa295 Add max_query_len parameter to _LocalPredictorKVCache for improved query handling
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-03 14:04:07 +00:00
tzhouam 2da9d28f1b Enhance reproducibility in Ulysses attention tests by replacing current_omni_platform.seed_everything with torch's manual_seed and adding support for CUDA devices. Introduce get_supported_tasks method in AsyncOmni class to dynamically determine available tasks based on stage capabilities.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-03 09:45:51 +00:00
tzhouam 4d8a76e721 Reintroduce installation of flashinfer-jit-cache in Dockerfile.ci to enhance dependency management for the CI environment, ensuring compatibility with the latest requirements.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-03 09:23:09 +00:00
tzhouam a17078974f Remove flashinfer-jit-cache installation from Dockerfile.ci to streamline dependency management and avoid potential conflicts.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-03 07:00:18 +00:00
tzhouam 8dcf1061f1 Update Dockerfile.ci to install additional dependencies including flashinfer-jit-cache, nvidia-cublas-cu12, and numpy, ensuring a more comprehensive environment setup for development.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-03 06:47:10 +00:00
tzhouam 98926707f5 rebase 2026/3/3
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-03 03:50:34 +00:00
tzhouam 3e539fcb6a Merge remote-tracking branch 'origin/main' into dev/vllm-align
Made-with: Cursor

# Conflicts:
#	vllm_omni/entrypoints/openai/serving_chat.py
#	vllm_omni/worker/gpu_model_runner.py

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-03 02:57:27 +00:00
tzhouam e0497550c4 Update Dockerfile.ci to install a specific version of flashinfer-cubin (0.6.4) directly, simplifying the installation process and ensuring compatibility with the current setup.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-03 02:51:35 +00:00
tzhouam 029bb56e3a Update Dockerfile.ci to use importlib.metadata for retrieving the current version of flashinfer, enhancing compatibility with the latest Python standards.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-03 02:28:24 +00:00
tzhouam 65169baa9d Update Dockerfile.ci to create a symbolic link for Python 3, ensuring compatibility with Python scripts during the build process.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-03 01:58:35 +00:00
tzhouam 908e3fe761 Update Dockerfile.ci to dynamically install the current version of flashinfer-cubin, ensuring compatibility with the latest changes in the flashinfer package.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-02 07:00:59 +00:00
tzhouam 8fb247a481 Update Dockerfile.ci to install specific version of flashinfer-cubin (0.6.4) and ensure vllm installation remains intact.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-02 06:48:58 +00:00
tzhouam be4226bae0 Remove unnecessary uninstallation of nvidia-cublas-cu12 and streamline vllm installation in Dockerfile.ci
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-02 06:38:24 +00:00
tzhouam 473d3ecd29 Update vLLM wheel version and adjust GPU memory utilization in CI configuration
- Updated the vLLM wheel commit hash in Dockerfile.ci to a60985b07eaf0fefa4c353ad2309b745707591ed.
- Modified the mock shared memory handle name in test_basic_connectors.py for clarity.
- Increased GPU memory utilization to 0.9 and adjusted max_num_batched_tokens to 4096 in qwen2_5_omni_ci.yaml to enable testing on H800.
- Refactored GPU model runner classes to remove deprecated activate_lora parameter and ensure compatibility with num_active_loras.
- Enhanced request handling in the OmniGPUModelRunner class to improve streaming request updates.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-02 03:08:45 +00:00
tzhouam 26744bcb35 Merge origin/main into dev/vllm-align
Resolved conflicts:
- vllm_omni/inputs/preprocess.py: took origin/main's elif mm_processor_kwargs branch, fixed undefined mm_uuids reference
- vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py: took origin/main's NPU support

Made-with: Cursor

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-02 02:10:32 +00:00
tzhouam 8d0c46487e Merge origin/main into dev/vllm-align
Resolve conflicts:
- docker/Dockerfile.ci: Keep wheel install section, update hash to 6042e66cd5
- vllm_omni/platforms/npu/worker/npu_model_runner.py: Accept main's refactoring

Made-with: Cursor

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-26 08:19:53 +00:00
tzhouam a252918008 [Update] Upgrade vLLM wheel version and enhance metrics handling
- Updated the vLLM wheel to commit 24650715105a26b41c88f4164777b5e1a0f3200b in Dockerfile.ci.
- Modified the `calculate_metrics` function to handle `tokenizer` as optional (can be None) in metrics.py.
- Adjusted the `benchmark` function to accept `tokenizer` as optional in patch.py.
- Enhanced condition checks in OmniARScheduler to ensure valid token generation in omni_ar_scheduler.py.
- Updated input processing to include supported tasks in input_processor.py.
- Cleaned up unnecessary code in output_processor.py.
- Added support for SSL ciphers in the API server configuration in api_server.py.
- Enabled speculative decoding in GPU model runners for improved performance in gpu_generation_model_runner.py and gpu_model_runner.py.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-25 10:02:23 +00:00
tzhouam 22e01c416b Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-02-25 08:15:55 +00:00
tzhouam 81b9e40bbf Enhance signature compatibility checks in OmniLLM tests
- Improved the test_overridden_method_signatures_compatible function to better handle missing parameters in method signatures.
- Added logic to account for VAR_KEYWORD parameters, ensuring more accurate detection of signature mismatches.

Signed-off-by: [Your Name] <your.email@example.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-24 07:56:47 +00:00
tzhouam 04ccb062fc Add signature compatibility tests for OmniLLM class
- Introduced tests to ensure method signatures in OmniLLM match those of the base LLM class, preventing runtime TypeErrors due to parameter mismatches.
- Updated the _run_engine method in OmniLLM to include an optional output_type parameter.

Signed-off-by: [Your Name] <your.email@example.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-24 06:58:39 +00:00
tzhouam 861698f88f Update Dockerfile for vllm wheel version and enhance SHM connector tests
- Updated Dockerfile to use a new vllm wheel version.
- Modified test_put_get_inline to clarify that it always uses SHM and updated metadata assertions.
- Added handling for speculative configuration in GPU model runners to delay clearing connector metadata.
- Refactored payload decoding in OmniGPUModelRunner for better clarity and maintainability.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-24 05:45:48 +00:00
tzhouam e22ae62396 Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-02-24 03:23:20 +00:00
tzhouam 9070f15f52 debug qwen tts
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-24 03:17:11 +00:00
tzhouam bf69cd94bb Rebase 2026/02/23
Update Dockerfile for vllm wheel version and enhance input processing in Omni models. Added decoding for byte messages in async requests, included new metrics in benchmarks, and improved pause state handling in the scheduler. Refined input and output processing methods for better multimodal support and error handling.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-23 03:36:01 +00:00
tzhouam 292da4b0d8 Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-02-23 02:09:09 +00:00
tzhouam 9869ea007e Update vLLM wheel version in Dockerfile and enhance input processing for omni models. Added functionality to reinject omni-specific fields into processed prompts and improved handling of multimodal audio extraction in serving_chat.py. Adjusted SharedMemoryConnector logic for payload size checks and refined input preprocessing methods for better compatibility with multimodal data.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-20 08:34:53 +00:00
tzhouam 663976b8ca Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-02-20 01:38:44 +00:00
tzhouam b867bacd24 rebase 2025/02/16
Qwen3 Omni offline/online passed

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-16 03:21:21 +00:00
tzhouam 93eeb1a4c8 Merge remote-tracking branch 'origin/main' into dev/rebase-v0.16.0 2026-02-15 01:24:23 +00:00
tzhouam 3051fb10b5 reverse change in SamplingParams
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-13 09:59:03 +00:00
tzhouam 15099119dd Merge remote-tracking branch 'remotes/origin/main' into dev/rebase-v0.16.0 2026-02-13 09:55:18 +00:00
tzhouam f24273dc03 Update vLLM version to 0.16.0 in quickstart and CUDA installation documentation, and adjust parameters in multimodal generation script for improved performance. Refactor request logging in serving_chat.py for better input tracking.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-13 09:50:11 +00:00
tzhouam f50cf07c74 Update Dockerfile.ci to allow the installation of 'pplx_kernels' to continue even if the compatible wheel is not found, enhancing build resilience.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-13 07:22:07 +00:00
tzhouam bc0a2598ed Update Dockerfile.ci to replace 'uv pip' with 'pip' for 'pplx_kernels' installation, ensuring consistency in package management across the CI environment.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-13 07:11:40 +00:00
tzhouam 963deaf5ad Refactor Dockerfile.ci to streamline package installations by using 'uv pip' consistently for both 'pplx_kernels' and VLLM dependencies, ensuring a more uniform CI environment setup.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-13 06:55:10 +00:00
tzhouam 3abec0173c Update Dockerfile.ci to force-reinstall the 'pplx_kernels' package, ensuring the latest version is used in the CI environment.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-13 06:41:41 +00:00
tzhouam 8a4eee9bed Enhance configuration and debugging: Added asyncio mode to pytest settings, increased timeout for Qwen3-TTS E2E test, and improved logging in GPU model runner for better traceability of input modalities. Updated Dockerfile.ci to streamline VLLM installation process.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-13 06:39:28 +00:00
tzhouam 31e9fdbb57 Refactor chat_utils to remove unused multi-modal classes and add async audio extraction from video. Update README to reflect new entry point for audio extraction. Enhance serving_chat to process audio from video URLs when specified.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-13 03:27:16 +00:00
tzhouam 39a045f439 debug Qwen 3 Omni cuda graph
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-13 02:57:05 +00:00
tzhouam 7e9cae56fc Refactor Dockerfile.ci to use precompiled VLLM installation and remove unused parse.py file. Update serving_chat.py to streamline rendering logic and enhance compatibility with new API. Modify OmniInputPreprocessor to directly handle SingletonDictPrompt, improving input processing efficiency.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-13 02:28:25 +00:00
tzhouam 31c5465e00 Update .gitignore to exclude .wav output files and remove audio_0.wav
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-13 01:52:06 +00:00
Hongsheng Liu 622c6f5f14 Merge branch 'main' into dev/rebase-v0.16.0 2026-02-13 06:05:57 +08:00
Taichang Zhou a6d633c319 debug
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-02-12 23:46:20 +08:00
tzhouam ec3814f587 debug
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-12 14:20:57 +00:00
tzhouam 59a7602d6e debug
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-12 13:55:08 +00:00
Taichang Zhou 85a7b558ab rebase to v0.16.0
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-02-12 21:30:54 +08:00
Taichang Zhou 3dc19bc305 0.16.0 docker not released yet
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-02-12 21:25:53 +08:00
Taichang Zhou 7458ba44b7 merge main
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-02-12 21:21:03 +08:00
Taichang ZhouandCursor e2ced80d59 Merge origin/main into dev/rebase-v0.16.0.
Resolve arg_utils merge conflict and align CI base image tag to v0.16.0 for the vLLM rebase track.

Co-authored-by: Cursor <cursoragent@cursor.com>

Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-02-12 21:18:03 +08:00
Taichang Zhou 1705edcce9 Refactor model path in test_qwen3_tts for improved clarity and maintainability
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-02-12 21:12:53 +08:00
tzhouam f8959bcc6e Temp fix for metrics error in test_qwen3_tts
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-12 10:02:19 +00:00
tzhouam db40563b88 Update Dockerfile.ci to pin nvidia-cublas-cu12 and add numpy installation, ensuring proper dependency management for the build environment.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-12 09:52:51 +00:00
tzhouam 86e20769fd Add SSL context support in benchmark function for improved connection handling
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-12 08:52:27 +00:00
tzhouam 8c1df5c704 update flashinfer
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-12 08:52:06 +00:00
tzhouam 3e89c99001 Merge remote-tracking branch 'remotes/origin/main' into dev/vllm-align 2026-02-12 07:16:25 +00:00
tzhouam c751f0d684 Update Dockerfile.ci to install a new version of vllm from a different extra index URL. Refactor import statement in serving_chat.py to use the correct BaseRenderer module path.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-12 06:37:17 +00:00
tzhouamandCursor 50f2b1e854 Merge origin/main into dev/vllm-align
Co-authored-by: Cursor <cursoragent@cursor.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-12 03:41:38 +00:00
tzhouam 0099eebb3f Refactor OmniEngineArgs and AsyncOmniEngineArgs to improve compatibility with partial payloads by using getattr for various attributes. Update GPUGenerationModelRunner and OmniGPUModelRunner to replace the activate_lora parameter with num_active_loras for better clarity and backward compatibility. This change enhances the handling of LoRA activation during model runs.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-12 03:33:50 +00:00
Taichang Zhou cd2ffbf2a0 update dockerfile.ci
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-02-11 23:39:15 +08:00
Taichang Zhou 9fbffc6137 Update requirements/cuda.txt to clarify the pinning of nvidia-cublas-cu12 version 12.9.1.4 due to dependency resolution issues with torch. This comment aids in understanding the force-reinstall strategy in Dockerfile.ci.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-02-11 22:44:11 +08:00
Taichang Zhou 514bd46f2e Uninstall old nvidia-cublas-cu12 version before reinstalling the pinned version to ensure consistent dependency management in Dockerfile.ci.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-02-11 22:34:05 +08:00
Taichang Zhou 169d94dab9 Refactor Dockerfile.ci to use bash for Python executable resolution during vllm-omni installation and pin nvidia-cublas-cu12 version to 12.9.1.4 for consistent dependency management.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-02-11 21:31:09 +08:00
Taichang Zhou 268a58e4e8 Update pyproject.toml to remove redundant transformers package installation and adjust Dockerfile.ci to install the current package in editable mode. Also, update fa3-fwd package version in requirements/cuda.txt to 0.0.2 for consistency.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-02-11 21:18:53 +08:00
Taichang Zhou eb05967451 Pin nvidia-cublas-cu12 to version 12.9.1.4 in Dockerfile.ci to ensure consistent dependency management in the CI environment.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-02-11 20:31:57 +08:00
tzhouam 05784d3dd0 Update pyproject.toml to include transformers package and remove redundant installation step in Dockerfile.ci for cleaner CI setup.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-11 06:27:51 +00:00
tzhouam b7f297c8c3 Add transformers package installation in Dockerfile.ci to ensure compatibility with the CI environment.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-11 06:15:00 +00:00
tzhouam 379f61d038 Enhance Dockerfile.ci by adding a debug step to print key dependency versions for CI troubleshooting, while commenting out the flashinfer version sync logic for clarity.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-11 06:01:01 +00:00
tzhouam 31b6584373 Refine flashinfer version retrieval in Dockerfile.ci to enhance reliability and maintain consistency in package installations
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-11 04:56:50 +00:00
tzhouam 50cea0b29c Refactor flashinfer version retrieval in Dockerfile.ci for improved accuracy and consistency in package installations
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-11 04:34:52 +00:00
tzhouam e897640b43 Update Dockerfile.ci to enhance flashinfer version retrieval logic and ensure package installation consistency
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-11 04:11:11 +00:00
tzhouam 3663e803d1 rebase 2026/02/11
vllm: 4a1550d22d7058e129d0e1257e726b3bf4a77025

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-11 02:43:49 +00:00
tzhouam 279b6b5050 Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-02-11 01:44:54 +00:00
tzhouam af53491cac Update sampling parameters in openai_chat_completion_client_for_multimodal_generation.py for improved model performance, adjusting temperature, top_k, and max_tokens settings.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-10 09:54:32 +00:00
tzhouam c3d35ed5fe pass local t2i test
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-10 06:20:26 +00:00
tzhouam 3c5b53358c Update Dockerfile.ci to ensure consistent formatting of flashinfer version extraction commands for improved readability and maintainability.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-10 03:53:55 +00:00
tzhouam 58d09ed656 reabse 2026/02/10
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-10 03:36:09 +00:00
tzhouam 064595a97c Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-02-10 02:22:49 +00:00
tzhouam e0f540fe98 Enhance Dockerfile.ci to synchronize flashinfer runtime package versions with installed flashinfer-python, ensuring compatibility across components. Additionally, improve error handling in ring_globals.py for flashinfer import failures, logging warnings for fallback scenarios.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-10 02:17:00 +00:00
tzhouam 670240ccd4 Enhance Dockerfile.ci to allow pre-release vLLM installations by adding --prerelease=allow flag for compatibility with nightly/commit wheels.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-09 10:01:39 +00:00
tzhouam c9fb2dc7c7 Update vLLM commit reference in Dockerfile.ci to f97ca671766c5201404e9fc812e35bf2c4e95a01
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-09 08:09:21 +00:00
tzhouam 7e4a604d4d reabse 2025/2/9
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-09 03:07:43 +00:00
tzhouamandCursor 03c6a58f69 Merge origin/main into dev/vllm-align; resolve omni_stage.py conflicts
Co-authored-by: Cursor <cursoragent@cursor.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-09 01:52:31 +00:00
tzhouam 7f54fd2f54 Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-02-06 02:26:06 +00:00
tzhouam 41e22e2770 [Refactor] Clean up code by removing unnecessary lines and adding missing imports in async_omni_llm.py, omni_stage.py, omni.py, and api_server.py
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-05 11:02:03 +00:00
tzhouam 24b04f58a0 rebase 2026/2/5 2026-02-05 10:52:28 +00:00
tzhouam 0144a424ed Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-02-05 01:44:34 +00:00
tzhouam 560a5aa57c Merge branch 'dev/vllm-align' of https://github.com/vllm-project/vllm-omni into dev/vllm-align 2026-02-03 05:59:37 +00:00
tzhouam 02c6062908 Merge branch 'main' into dev/vllm-align
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-03 05:57:38 +00:00
Taichang Zhou 87c7b7e519 Refactor comments and formatting in Qwen3 TTS configuration and modeling files for improved readability.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-02-02 21:52:10 +08:00
Taichang Zhou 22ca74432b Merge branch 'main' into dev/vllm-align 2026-02-02 21:35:45 +08:00
Taichang Zhou 354777d5e7 Merge branch 'main' of https://github.com/vllm-project/vllm-omni 2026-02-02 21:34:42 +08:00
汪志鹏 b6c846cc95 Add stage_init_timeout=300 (#1153)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
2026-02-02 17:58:48 +08:00
tzhouam 9194c1dd84 Update base image in Dockerfile.rocm from v0.14.0 to v0.15.0
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-02 08:41:35 +00:00
tzhouam ce9105267f Enhance error handling in TTS generation and improve audio output compatibility by supporting both "audio" and "model_outputs" keys in response.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-02 08:40:56 +00:00
tzhouam 39f7cade27 Fix audio sample rate access in end2end.py for correct output retrieval
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-02 08:38:24 +00:00
tzhouam 08fa85a1ad Update audio tensor access in offline inference tests and increase image fetch timeout for CI
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-02 06:10:54 +00:00
tzhouam 6568a78a26 Remove unused tokenizer argument from OmniInputProcessor class docstring
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-02 03:12:43 +00:00
tzhouam ff8b3dbbc8 update qwen 2.5 Omni
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-02 03:01:11 +00:00
tzhouam 1297a46040 Merge remote-tracking branch 'origin/main' into dev/vllm-align
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-02 01:54:12 +00:00
Taichang Zhou e033498600 Merge branch 'main' of https://github.com/vllm-project/vllm-omni 2026-01-31 10:59:36 +08:00
tzhouam 7aea8e8d91 Add OmniMRotaryEmbedding for multimodal position computation
- Introduced OmniMRotaryEmbedding to extend MRotaryEmbedding with methods for computing input positions across various multimodal scenarios, including image, video, and audio inputs.
- Updated patch.py to replace the original MRotaryEmbedding with OmniMRotaryEmbedding in the module patching logic.
- Created new files for the OmniMRotaryEmbedding implementation and its associated methods.

Signed-off-by: [Your Name] <[Your Email]>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-30 10:03:02 +00:00
tzhouam 3208afd82b Refactor OmniRequestOutput to encapsulate multimodal output
- Changed the multimodal_output field to a private attribute _multimodal_output for better encapsulation.
- Added a property method to access multimodal_output, allowing for a clearer interface and improved handling of request outputs.

Signed-off-by: [Your Name] <[Your Email]>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-30 10:02:22 +00:00
tzhouam 2b009b3cef Fix audio tensor extraction in end2end.py files across multiple examples
- Updated the audio tensor extraction logic to access the correct output structure in Qwen2.5 Omni, Qwen3 Omni, and Qwen3 TTS implementations.
- Ensured consistency in handling audio outputs for improved functionality.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-30 10:00:14 +00:00
tzhouam 4faf190e28 Merge remote-tracking branch 'origin/main' into dev/vllm-align 2026-01-30 03:49:04 +00:00
tzhouam ca57f510c7 Refactor OmniInputPreprocessor to remove unused embed processing logic
- Eliminated the _process_embeds method and related imports from preprocess.py, streamlining the input preprocessing workflow.
- Updated the main processing logic to focus solely on token inputs, enhancing clarity and maintainability.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-30 03:02:38 +00:00
tzhouam 4b895459e4 Merge branch 'main' into dev/vllm-align
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-30 02:53:03 +00:00
tzhouam 62dbf07fb6 Refactor OpenAI serving modules for improved functionality and organization
- Reintroduced the load_chat_template import in api_server.py for better template handling.
- Updated the logic in serving_chat.py to streamline the resolution of chat template content formats, enhancing clarity and functionality.
- Adjusted the application of chat templates to ensure proper handling of tools and chat templates.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-30 02:44:13 +00:00
tzhouam 104e105c93 Refactor multimodal processing imports and update configuration handling in Qwen3OmniMoeThinker
- Changed the import path for multimodal processing to improve module organization.
- Removed the unused multimodal configuration parameter from the Qwen3OmniMoeThinker initialization, streamlining the class definition.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-30 02:43:41 +00:00
tzhouam ef21879c3a Refactor GPU model runner to streamline multimodal configuration handling
- Removed the unused import of supports_mm_encoder_only from interfaces.
- Updated the dummy run check to utilize the multimodal configuration directly, enhancing clarity and future extensibility.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-30 02:43:26 +00:00
tzhouam 45880a30f0 Refactor chat_utils and OpenAI modules for improved functionality and organization
- Updated import statements in chat_utils.py to replace the internal content format type with the public one for better clarity.
- Refactored api_server.py to streamline the chat template loading process by replacing the previous method with a direct function call.
- Enhanced serving_chat.py by reintroducing safe_apply_chat_template functions for better template handling.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-29 09:00:43 +00:00
tzhouam 2ac33489b5 Update Dockerfile.ci to install vllm package using the correct URL format
- Modified the pip install command for vllm to use the correct URL format for direct installation from the specified commit hash, ensuring a more reliable installation process.
- Removed the fallback option for pre-release installation to streamline the installation command.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-29 08:25:51 +00:00
tzhouam 63fd163efb Refactor import statement in OpenAI API server for improved module organization
- Updated the import path for AnthropicServingMessages in api_server.py to enhance clarity and maintain consistency with the new module structure.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-29 08:23:17 +00:00
tzhouam fea328364a Update Dockerfile.ci to install vllm package directly from wheel file
- Changed the pip install command for vllm to download the wheel file directly instead of using an extra index URL, ensuring a more reliable installation process.
- Added a fallback option to install the pre-release version of vllm if the direct installation fails.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-29 08:05:57 +00:00
tzhouam 478c036a01 [Refactor] Update OpenAI API server and related modules for improved structure and functionality
- Refactored import statements in api_server.py to enhance module organization, including changes to the routing and server setup.
- Updated the method for removing routes to ensure compatibility with the new app structure.
- Adjusted the handling of chat completion requests to improve error responses and streamline processing.
- Modified serving modules to align with the new import paths and enhance code clarity.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-29 15:28:51 +08:00
tzhouam 8407818368 Update simple_test.sh to force reinstall vllm package with extra index URL
- Modified the pip install command for vllm to include --upgrade, --force-reinstall flags, and an extra index URL, ensuring the latest version is installed from the specified source during testing.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 17:34:22 +08:00
tzhouam 6e7cb01afe Update Dockerfile.ci to force reinstall vllm package during CI build
- Modified the pip install command for vllm to include --upgrade and --force-reinstall flags, ensuring the latest version is installed during the CI process.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 17:20:29 +08:00
tzhouam acb28f3d1d Remove commented-out import and related code for MRotaryEmbedding in patch.py to clean up the codebase.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 16:52:10 +08:00
tzhouam bf51318176 debug
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 16:40:06 +08:00
tzhouam bd8c0443cf build image from commit id
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 16:21:29 +08:00
tzhouam 554d14fd28 debug docker file
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 15:57:13 +08:00
tzhouam 49b9d0d6b1 debug docker build
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 15:47:30 +08:00
tzhouam 2c2d06d920 [Enhancement] Update GPUARWorker and GPUGenerationWorker to enforce v1 model runner
- Added a warning in GPUARWorker and GPUGenerationWorker to indicate that the v2 model runner does not include omni hooks, forcing the use of the v1 model runner.
- Updated import statements in gpu_ar_model_runner.py and gpu_generation_model_runner.py to include necessary functions from parallel_state.
- Removed unused imports from both model runner files to streamline the code.

Signed-off-by: [Your Name] <your.email@example.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 15:43:27 +08:00
tzhouam 31d127edcd [Refactor] Remove unused MRotaryEmbedding and vision position ID functions
- Deleted the MRotaryEmbedding class and its associated methods from mrope.py to streamline the codebase.
- Removed the get_llm_pos_ids_for_vision function from vision.py as it is no longer needed.
- Updated import statements in qwen2_5_omni.py to reflect the removal of the vision module.

Signed-off-by: [Your Name] <your.email@example.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 15:42:35 +08:00
tzhouam 52d0122fb5 [Enhancement] Refactor OmniInputProcessor and related classes for improved prompt handling
- Updated OmniInputProcessor to streamline input processing, including enhanced handling of prompt embeddings and additional information.
- Introduced OmniTextPrompt and OmniEmbedsInputs to support new input structures with optional metadata.
- Modified input preprocessing methods to accommodate new prompt types and ensure compatibility with multimodal inputs.
- Improved error handling and type checking for data parallel configurations.

Signed-off-by: [Your Name] <your.email@example.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 15:42:05 +08:00
tzhouam b26a9735e6 [Enhancement] Update type hint for stage_connector_spec in OmniEngineArgs and AsyncOmniEngineArgs
- Changed type hint for stage_connector_spec from 'any' to 'Any' for improved type checking and consistency.
- Ensured compatibility with typing standards in the arg_utils module.

Signed-off-by: [Your Name] <your.email@example.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 15:41:20 +08:00
tzhouam d7cfe195ea [Enhancement] Refactor OmniARScheduler and OmniGenerationScheduler to include OmniSchedulerOutput
- Updated OmniARScheduler to wrap output in OmniSchedulerOutput, including transfer metadata for finished requests needing KV cache transfer.
- Enhanced OmniGenerationScheduler to handle routed experts and finish reasons more effectively, ensuring proper output structure.
- Introduced performance statistics handling in both schedulers for improved monitoring.

Signed-off-by: [Your Name] <your.email@example.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 15:40:59 +08:00
tzhouam e1c6639f18 [Enhancement] Add mm_encoder_only parameter and deprecation warning for convert option
- Introduced mm_encoder_only parameter to OmniModelConfig.
- Added a deprecation warning for the `--convert mm_encoder_only` option, advising users to switch to `--mm-encoder-only`.

Signed-off-by: [Your Name] <your.email@example.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 15:40:29 +08:00
tzhouam ce7430b562 [Enhancement] Update OmniRequest to support torch.Tensor for prompt_embeds and add decoding method
- Modified the constructor to accept torch.Tensor as an option for prompt_embeds.
- Introduced a new static method _maybe_decode_prompt_embeds to handle the conversion from PromptEmbedsPayload to torch.Tensor.
- Preserved the original prompt_embeds payload for reference.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 15:40:17 +08:00
tzhouam ed8b03224b update the docker setup to vllm nightly
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-01-28 11:15:54 +08:00
Taichang Zhou af8501dcc7 debug Qwen TTS
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
2026-01-22 22:46:06 +08:00
249 changed files with 3697 additions and 24482 deletions

No files matched your search

@@ -1,56 +0,0 @@
#!/bin/bash
# This script build the XPU docker image and run the offline inference inside the container.
set -ex
omni_source_dir=$(git rev-parse --show-toplevel)
base_image_name="xpu/vllm-omni-ci-base:${VLLM_VERSION:?VLLM_VERSION must be set}"
image_name="xpu/vllm-omni-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set}"
container_name="xpu_${BUILDKITE_COMMIT}_$(
tr -dc A-Za-z0-9 </dev/urandom | head -c 10
echo
)"
cd "${omni_source_dir}"
if [ -z "$(docker images -q "${base_image_name}")" ]; then
docker build --target vllm-base -t "${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -f docker/Dockerfile.xpu .
fi
# Try building the docker image
docker build --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -t "${image_name}" -f docker/Dockerfile.xpu .
# Setup cleanup
remove_docker_container() {
docker rm -f "${container_name}" || true
docker image rm -f "${image_name}" || true
docker system prune -f || true
}
trap remove_docker_container EXIT
HF_CACHE="${HF_CACHE:-$(realpath ~)/.cache/huggingface}"
mkdir -p "${HF_CACHE}"
HF_MOUNT="/root/.cache/huggingface"
time timeout -k 30 30m docker run \
--device /dev/dri:/dev/dri \
--net=host \
--ipc=host \
-v /dev/dri/by-path:/dev/dri/by-path \
-v "${HF_CACHE}:${HF_MOUNT}" \
--security-opt seccomp=unconfined \
--entrypoint="" \
-e VLLM_LOGGING_LEVEL \
-e VLLM_OMNI_LOGGING_LEVEL \
-e HF_TOKEN \
-e ZE_AFFINITY_MASK \
--name "${container_name}" \
"${image_name}" \
bash -c '
set -e
echo $ZE_AFFINITY_MASK
pip install tblib==3.1.0
cd /workspace/vllm-omni
pytest -v -s -m "core_model and xpu and B60"
pytest -v -s -m "advanced_model and xpu and B60"
'
+5 -44
View File
@@ -54,6 +54,7 @@ steps:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
shm-size: "8gb"
propagate-environment: true
shm-size: "8gb"
environment:
@@ -61,42 +62,6 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: ":full_moon: Diffusion Model Test with H100"
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
# if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
timeout_in_minutes: 30
depends_on: upload-nightly-pipeline
@@ -119,7 +84,7 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: ":full_moon: Omni Model Perf Test & Test Case Statistics"
- label: ":full_moon: Omni Model Perf Test with H100"
key: nightly-performance
timeout_in_minutes: 180
depends_on: upload-nightly-pipeline
@@ -129,8 +94,6 @@ steps:
- export BENCHMARK_DIR=tests
- pytest -s -v tests/perf/scripts/run_benchmark.py
- buildkite-agent artifact upload "tests/*.json"
- python tools/nightly/buildkite_testcase_statistics.py -o tests/buildkite_testcase_statistics.html
- buildkite-agent artifact upload "tests/*.html"
agents:
queue: "mithril-h100-pool"
plugins:
@@ -160,19 +123,17 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":email: Nightly Collection & Email"
- label: ":email: Nightly Perf Collection & Email"
key: nightly-perf-distribution
depends_on:
- nightly-performance
depends_on: nightly-performance
if: build.env("NIGHTLY") == "1"
commands:
- pip install openpyxl
- export DEFAULT_INPUT_DIR=tests
- export DEFAULT_OUTPUT_DIR=tests
- buildkite-agent artifact download "tests/*.json" . --step nightly-performance
- buildkite-agent artifact download "tests/*.html" . --step nightly-performance
- python tools/nightly/generate_nightly_perf_excel.py
- python tools/nightly/send_nightly_email.py --report-file "tests/*.xlsx, tests/*.html"
- python tools/nightly/send_nightly_perf_email.py
- buildkite-agent artifact upload "tests/*.xlsx"
agents:
queue: "cpu_queue_premerge"
+9 -1
View File
@@ -34,7 +34,15 @@ steps:
- label: "Diffusion Model CPU offloading Test"
depends_on: upload-ready-pipeline
commands:
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
EXIT1=$$?
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
EXIT2=$$?
exit $$((EXIT1 | EXIT2))
'
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
+1 -2
View File
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
</h3>
<p align="center">
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://drive.google.com/file/d/12k_DOo2xoEROvd4xQxYNuEi6C2GECoRx/view?usp=drive_link"><b>Slides</b></a> |
</p>
@@ -17,7 +17,6 @@ Easy, fast, and cheap omni-modality model serving for everyone
*Latest News* 🔥
- [2026/03] Joined vLLM HK First Meetup! Video on [youtube](https://youtu.be/sgwNfsNnR9I)
- [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
- [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
- [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
+1 -16
View File
@@ -47,7 +47,6 @@ If no, check your shell running the ComfyUI process. There may be some error mes
This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
- **Generate Image** for text-to-image and image-to-image tasks
- **Generate Video** for text-to-video and image-to-video tasks
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
- **TTS** and **TTS Voice Clone** for TTS tasks
@@ -60,7 +59,7 @@ To build a simple workflow yourself,
- Drag a generation node onto the canvas.
- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
- For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
- For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
@@ -105,20 +104,6 @@ You can configure per-stage sampling parameters for multi-stage models.
> [!TIP]
> The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
### Text-to-video and image-to-video generation (e.g., Wan)
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video Generation**)
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg">
<img alt="vLLM-Omni video generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg" width=55%>
</picture>
</p>
> [!TIP]
> The node automatically choose text-to-video or image-to-video API endpoints depending on whether you connect an image input or not.
### TTS (e.g., Qwen TTS series)
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni TTS**)
-9
View File
@@ -14,21 +14,17 @@ from .comfyui_vllm_omni.nodes import (
VLLMOmniARSampling,
VLLMOmniDiffusionSampling,
VLLMOmniGenerateImage,
VLLMOmniGenerateVideo,
VLLMOmniQwenTTSParams,
VLLMOmniRemoteLoRA,
VLLMOmniSamplingParamsList,
VLLMOmniTTS,
VLLMOmniUnderstanding,
VLLMOmniVoiceClone,
VLLMOmniWanParams,
)
# A dictionary that contains all nodes you want to export with their names
NODE_CLASS_MAPPINGS = {
# === Generation ===
"VLLMOmniGenerateImage": VLLMOmniGenerateImage,
"VLLMOmniGenerateVideo": VLLMOmniGenerateVideo,
"VLLMOmniUnderstanding": VLLMOmniUnderstanding,
"VLLMOmniTTS": VLLMOmniTTS,
"VLLMOmniVoiceClone": VLLMOmniVoiceClone,
@@ -36,16 +32,13 @@ NODE_CLASS_MAPPINGS = {
"VLLMOmniARSampling": VLLMOmniARSampling,
"VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
"VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList,
"VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA,
"VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
"VLLMOmniWanParams": VLLMOmniWanParams,
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
# === Generation ===
"VLLMOmniGenerateImage": "Generate Image",
"VLLMOmniGenerateVideo": "Generate Video",
"VLLMOmniUnderstanding": "Multimodality Understanding",
"VLLMOmniTTS": "TTS (Text to Speech)",
"VLLMOmniVoiceClone": "TTS Voice Cloning",
@@ -53,9 +46,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"VLLMOmniARSampling": "AR Sampling Params",
"VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
"VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List",
"VLLMOmniRemoteLoRA": "LoRA",
"VLLMOmniQwenTTSParams": "Qwen TTS Params",
"VLLMOmniWanParams": "Wan Video Params",
}
WEB_DIRECTORY = "./web"
@@ -11,7 +11,6 @@ from .utils.types import (
AutoregressionSamplingParams,
DiffusionSamplingParams,
QwenTTSModelSpecificParams,
WanModelSpecificParams,
)
from .utils.validators import (
add_sampling_parameters_to_stage,
@@ -57,7 +56,6 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
# "video": ("VIDEO",),
# "audio": ("AUDIO",),
"sampling_params": ("SAMPLING_PARAMS",),
"lora": ("REMOTE_LORA",),
},
}
@@ -78,11 +76,9 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
audio: AudioInput | None = None, # Hidden & unused
video: VideoInput | None = None, # Hidden & unused
sampling_params: dict | list[dict] | None = None,
lora: dict | None = None,
**kwargs,
):
if kwargs:
logger.info("Uncaught kwargs: %s", kwargs)
logger.info("Uncaught kwargs: %s", kwargs)
logger.debug("Got sampling params: %s", sampling_params)
validate_model_and_sampling_params_types(model, sampling_params)
if image is None and mask is not None:
@@ -109,7 +105,6 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
height=height,
negative_prompt=negative_prompt,
sampling_params=sampling_params,
lora=lora,
)
return (output,)
elif image is not None and audio is None and video is None:
@@ -124,7 +119,6 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
negative_prompt=negative_prompt,
mask=mask,
sampling_params=sampling_params,
lora=lora,
)
return (output,)
@@ -142,91 +136,11 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
audio=audio,
video=video,
sampling_params=sampling_params,
lora=lora,
)
return (output,)
class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"url": ("STRING", {"default": "http://localhost:8000/v1"}),
"model": ("STRING", {"default": "Wan-AI/Wan2.2-T2V-A14B-Diffusers"}),
"prompt": ("STRING", {"multiline": True}),
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
"width": ("INT", {"default": 832, "min": 1}),
"height": ("INT", {"default": 480, "min": 1}),
"fps": ("INT", {"default": 16, "min": 1}),
"num_frames": ("INT", {"default": 41, "min": 1}),
},
"optional": {
"image": ("IMAGE",),
"sampling_params": ("SAMPLING_PARAMS",),
"lora": ("REMOTE_LORA",),
"model_params": ("VIDEO_PARAMS",),
},
}
RETURN_TYPES = ("VIDEO",)
RETURN_NAMES = ("video",)
FUNCTION = "generate"
async def generate(
self,
url: str,
model: str,
prompt: str,
width: int,
height: int,
fps: int,
num_frames: int,
negative_prompt: str | None = None,
image: torch.Tensor | None = None,
sampling_params: dict | list[dict] | None = None,
model_params: dict | None = None,
lora: dict | None = None,
**kwargs,
):
if kwargs:
logger.info("Uncaught kwargs: %s", kwargs)
logger.debug("Got sampling params: %s", sampling_params)
logger.debug("Got model params: %s", model_params)
validate_model_and_sampling_params_types(model, sampling_params)
# Currently, all video generation models are single-stage diffusion models
if isinstance(sampling_params, list):
if len(sampling_params) != 1:
raise ValueError(
"Video generation expects a single sampling params group. "
"Please provide one Diffusion sampling node."
)
sampling_params = sampling_params[0]
if sampling_params is not None:
sampling_params.pop("type", None) # internal fields
if model_params is not None:
model_params.pop("type", None) # internal fields
client = VLLMOmniClient(url)
output = await client.generate_video(
model=model,
prompt=prompt,
image=image, # image present => i2v, absent => t2v
width=width,
height=height,
num_frames=num_frames,
fps=fps,
negative_prompt=negative_prompt,
sampling_params=sampling_params,
lora=lora,
model_params=model_params,
)
return (output,)
class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
@classmethod
def INPUT_TYPES(cls):
@@ -274,8 +188,7 @@ class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
use_audio_in_video: bool = True,
**kwargs,
) -> tuple[str, AudioInput]:
if kwargs:
logger.info("Uncaught kwargs: %s", kwargs)
logger.info("Uncaught kwargs: %s", kwargs)
logger.debug("Got sampling params: %s", sampling_params)
validate_model_and_sampling_params_types(model, sampling_params)
@@ -634,52 +547,6 @@ class VLLMOmniSamplingParamsList:
return (params,)
class VLLMOmniRemoteLoRA:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"local_path": ("STRING", {"default": ""}),
"name": ("STRING", {"default": ""}),
"scale": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1},
),
"int_id": (
"INT",
{
"default": 0,
"min": 0,
"step": 1,
"tooltip": "0 means it is not set and the server can derive it.",
},
),
}
}
RETURN_TYPES = ("REMOTE_LORA",)
RETURN_NAMES = ("lora",)
FUNCTION = "get_lora"
CATEGORY = "vLLM-Omni"
@classmethod
def VALIDATE_INPUTS(cls, local_path, name) -> str | Literal[True]:
if not local_path.strip() or not name.strip():
return "Both local_path and name must be provided."
return True
def get_lora(self, local_path: str, name: str, scale: float, int_id: int):
local_path = local_path.strip()
name = name.strip()
lora = {
"local_path": local_path or None,
"name": name or None,
"scale": float(scale),
"int_id": int(int_id) if int_id > 0 else None,
}
return (lora,)
class VLLMOmniQwenTTSParams:
@classmethod
def INPUT_TYPES(cls):
@@ -705,32 +572,3 @@ class VLLMOmniQwenTTSParams:
def get_params(self, **kwargs):
return (QwenTTSModelSpecificParams(kwargs),)
class VLLMOmniWanParams:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"guidance_scale_2": (
"FLOAT",
{"default": 4.0, "min": 0.0, "max": 20.0, "step": 0.1},
),
"boundary_ratio": (
"FLOAT",
{"default": 0.875, "min": 0.0, "max": 1.0, "step": 0.001},
),
"flow_shift": (
"FLOAT",
{"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1},
),
}
}
RETURN_TYPES = ("VIDEO_PARAMS",)
RETURN_NAMES = ("Wan video params",)
FUNCTION = "get_params"
CATEGORY = "vLLM-Omni/Video Params"
def get_params(self, **kwargs):
return (WanModelSpecificParams(kwargs),)
@@ -6,8 +6,6 @@ The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@do
Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
"""
import asyncio
import json
from typing import Any
import aiohttp
@@ -20,7 +18,6 @@ from .format import (
base64_to_audio,
base64_to_image_tensor,
bytes_to_audio,
bytes_to_video,
image_tensor_to_base64,
image_tensor_to_png_bytes,
video_to_base64,
@@ -32,43 +29,10 @@ from .types import AudioFormat
logger = get_logger(__name__)
async def url_json(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> dict[str, Any]:
try:
async with getattr(session, verb)(url, **kwargs) as response:
if not response.ok:
error_text = await response.text()
raise (ValueError if response.status < 500 else RuntimeError)(
f"vLLM-Omni API returned status {response.status}: {error_text}"
)
try:
return await response.json()
except aiohttp.ContentTypeError as e:
raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
except aiohttp.ClientError as e:
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
async def url_bytes(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> bytes:
try:
async with getattr(session, verb)(url, **kwargs) as response:
if not response.ok:
error_text = await response.text()
raise (ValueError if response.status < 500 else RuntimeError)(
f"vLLM-Omni API returned status {response.status}: {error_text}"
)
return await response.read()
except aiohttp.ClientError as e:
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
class VLLMOmniClient:
def __init__(
self, base_url: str, timeout: float | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
):
def __init__(self, base_url: str, timeout: float = 300.0):
self.base_url = base_url
self.timeout = aiohttp.ClientTimeout(total=timeout)
self.poll_interval = poll_interval
self.max_poll_duration = max_poll_duration
async def generate_image(
self,
@@ -79,13 +43,12 @@ class VLLMOmniClient:
height: int,
negative_prompt: str | None = None,
sampling_params: dict | None = None,
lora: dict | None = None,
) -> torch.Tensor:
"""Run text-to-image generation via DALLE API"""
await self._check_model_exist(model)
size = f"{width}x{height}"
payload: dict[str, Any] = {
payload = {
"model": model,
"prompt": prompt,
"size": size,
@@ -95,8 +58,6 @@ class VLLMOmniClient:
payload["negative_prompt"] = negative_prompt
if sampling_params is not None:
payload.update(sampling_params)
if lora is not None:
payload["lora"] = lora
logger.debug("img gen payload: %s", payload)
url = self.base_url + "/images/generations"
@@ -149,7 +110,6 @@ class VLLMOmniClient:
negative_prompt: str | None = None,
mask: torch.Tensor | None = None,
sampling_params: dict | None = None,
lora: dict | None = None,
) -> torch.Tensor:
"""Run image editing via DALLE API"""
await self._check_model_exist(model)
@@ -171,8 +131,6 @@ class VLLMOmniClient:
if sampling_params is not None:
for k, v in sampling_params.items():
form.add_field(k, str(v))
if lora is not None:
form.add_field("lora", json.dumps(lora, ensure_ascii=False))
if mask is not None:
mask_filename = "mask.png"
form.add_field(
@@ -225,7 +183,6 @@ class VLLMOmniClient:
audio: AudioInput | None = None,
video: VideoInput | None = None,
sampling_params: dict | list[dict] | None = None,
lora: dict | None = None,
) -> torch.Tensor:
payload = VLLMOmniClient._prepare_chat_completion_messages(
model=model,
@@ -236,8 +193,6 @@ class VLLMOmniClient:
video=video,
sampling_params=sampling_params,
modalities=["image"],
# === below are additional `extra_body` fields, handled by **kwargs ===
lora=lora,
)
choices = await self._generate_base_chat_completion(model, payload)
@@ -251,87 +206,6 @@ class VLLMOmniClient:
return torch.stack(image_tensors, dim=0)
async def generate_video(
self,
*,
model: str,
prompt: str,
width: int,
height: int,
num_frames: int,
fps: int,
negative_prompt: str | None = None,
image: torch.Tensor | None = None,
sampling_params: dict | None = None,
model_params: dict | None = None,
lora: dict | None = None,
**extra_body,
) -> VideoInput:
form = aiohttp.FormData()
form.add_field("model", model)
form.add_field("prompt", prompt)
form.add_field("width", str(width))
form.add_field("height", str(height))
form.add_field("num_frames", str(num_frames))
form.add_field("fps", str(fps))
if negative_prompt:
form.add_field("negative_prompt", negative_prompt)
if sampling_params is not None:
for k, v in sampling_params.items():
form.add_field(k, str(v))
if model_params is not None:
for k, v in model_params.items():
form.add_field(k, str(v))
if lora is not None:
form.add_field("lora", json.dumps(lora, ensure_ascii=False))
if extra_body:
form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False))
if image is not None:
image_filename = "image.png" # Required for multipart form
form.add_field(
"input_reference",
image_tensor_to_png_bytes(image, image_filename),
filename=image_filename,
content_type="image/png",
)
async with aiohttp.ClientSession(timeout=self.timeout) as session:
# Start the video generation job
url = f"{self.base_url}/videos"
data = await url_json(session, url, "post", data=form)
if (job_id := data.get("id", None)) is None:
raise RuntimeError("API response missing job 'id' field - expected OpenAI compliant format")
if (job_status := data.get("status", None)) is None:
raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
# Poll for video generation job completion
deadline = asyncio.get_running_loop().time() + self.max_poll_duration
url = f"{self.base_url}/videos/{job_id}"
while job_status not in {"completed", "failed"}:
await asyncio.sleep(self.poll_interval)
data = await url_json(session, url)
if (job_status := data.get("status", None)) is None:
raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
if asyncio.get_running_loop().time() >= deadline:
raise RuntimeError(f"Timed out waiting for video job {job_id} to complete")
if job_status == "failed":
raise RuntimeError(f"Video job failed: {data}")
# Retrieve completed content
video_bytes = await url_bytes(session, f"{url}/content")
# Decode video and make a best effort at cleaning up server resources
try:
return bytes_to_video(video_bytes)
finally:
try:
await url_json(session, url, "delete")
except Exception as exc:
logger.warning("Failed to clean up video job %s: %s", job_id, exc)
async def generate_understanding_chat_completion(
self,
*,
@@ -6,17 +6,12 @@ Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed und
import base64
import mimetypes
from fractions import Fraction
from io import BytesIO
import av
import numpy as np
import torch
from av.audio.frame import AudioFrame
from av.audio.resampler import AudioResampler
from av.video.frame import VideoFrame
from comfy_api.input import AudioInput, VideoInput
from comfy_api.latest import InputImpl, Types
from comfy_extras import nodes_audio
from PIL import Image
@@ -143,75 +138,6 @@ def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
return f"data:{mime_type};base64,{base64_str}"
def bytes_to_video(video_bytes: bytes) -> VideoInput:
video_buffer = BytesIO(video_bytes)
try:
with av.open(video_buffer, mode="r") as container:
video_stream = next((s for s in container.streams if s.type == "video"), None)
if video_stream is None:
raise ValueError("No video stream found in decoded payload.")
frames = []
for frame in container.decode(video_stream):
if not isinstance(frame, VideoFrame):
continue
image = frame.to_ndarray(format="rgb24")
frames.append(torch.from_numpy(image).float() / 255.0)
if len(frames) == 0:
raise ValueError("No video frames found in decoded payload.")
images = torch.stack(frames, dim=0)
frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
audio: AudioInput | None = None
if len(container.streams.audio):
audio_stream = container.streams.audio[-1]
audio_frames = []
resampler = AudioResampler(format="fltp")
for frame in container.decode(audio_stream):
if not isinstance(frame, AudioFrame):
continue
resampled = resampler.resample(frame)
if not isinstance(resampled, list):
resampled = [resampled]
for audio_frame in resampled:
if audio_frame is not None:
audio_frames.append(audio_frame.to_ndarray())
if len(audio_frames) > 0:
audio_data = np.concatenate(audio_frames, axis=1)
sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1
audio = {
"waveform": torch.from_numpy(audio_data).unsqueeze(0),
"sample_rate": sample_rate,
}
components = Types.VideoComponents(
images=images,
frame_rate=frame_rate,
audio=audio,
metadata=container.metadata if container.metadata else None,
)
except Exception as e:
raise RuntimeError(f"Failed to decode video: {e}")
return InputImpl.VideoFromComponents(components)
def base64_to_video(base64_str: str) -> VideoInput:
if base64_str.startswith("data:video"):
_, base64_str = base64_str.split(",", 1)
try:
video_bytes = base64.b64decode(base64_str)
except Exception as e:
raise ValueError(f"Invalid base64 string: {e}")
return bytes_to_video(video_bytes)
def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO:
waveform = audio["waveform"][0] # Shape: (C, T)
sample_rate = audio["sample_rate"]
@@ -23,10 +23,6 @@ class QwenTTSModelSpecificParams(dict):
pass
class WanModelSpecificParams(dict):
pass
class ModelMode(Enum):
IMAGE_GENERATION = auto()
VIDEO_GENERATION = auto()
@@ -16,7 +16,7 @@ def validate_model_and_sampling_params_types(
# Skip if no spec or no sampling params
pipeline_spec, _ = lookup_model_spec(model_name)
if pipeline_spec is None:
logger.info(f"skipping sampling params check because spec for {model_name} is not found")
logger.info("skipping sampling params check because spec is not found")
return
if sampling_param_list is None:
return
@@ -70,17 +70,13 @@ def add_sampling_parameters_to_stage(
stages = pipeline_spec["stages"]
if isinstance(sampling_param_list, dict):
sampling_param_list = sampling_param_list.__class__(sampling_param_list)
sampling_param_list.update(params_to_add)
elif sampling_param_list is None:
sampling_param_list = params_to_add.copy()
else:
for i, stage in enumerate(stages):
if stage == stage_type:
stage_param = sampling_param_list[i]
stage_param = stage_param.__class__(stage_param)
stage_param.update(params_to_add)
sampling_param_list[i] = stage_param
sampling_param_list[i].update(params_to_add)
return sampling_param_list
Binary file not shown.

Before

Width:  |  Height:  |  Size: 129 KiB

@@ -1,513 +0,0 @@
{
"id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21",
"revision": 0,
"last_node_id": 14,
"last_link_id": 28,
"nodes": [
{
"id": 10,
"type": "MarkdownNote",
"pos": [
227.99306819575278,
-231.24143306069843
],
"size": [
240.3326478922474,
136.3505820791881
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [],
"title": "Note: Task and Input",
"properties": {},
"widgets_values": [
"vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload depending on whether an input image is present or not."
],
"color": "#432",
"bgcolor": "#000"
},
{
"id": 4,
"type": "LoadImage",
"pos": [
496.31859627609606,
-229.71277089860084
],
"size": [
270,
314
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"type": "COMBO",
"widget": {
"name": "image"
},
"link": null
},
{
"localized_name": "choose file to upload",
"name": "upload",
"type": "IMAGEUPLOAD",
"widget": {
"name": "upload"
},
"link": null
}
],
"outputs": [
{
"localized_name": "IMAGE",
"name": "IMAGE",
"type": "IMAGE",
"links": [
25
]
},
{
"localized_name": "MASK",
"name": "MASK",
"type": "MASK",
"links": null
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"cute-cat.jpg",
"image"
]
},
{
"id": 8,
"type": "VLLMOmniDiffusionSampling",
"pos": [
465.5140218220927,
148.80072646828987
],
"size": [
284.205078125,
226
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"localized_name": "n",
"name": "n",
"type": "INT",
"widget": {
"name": "n"
},
"link": null
},
{
"localized_name": "num_inference_steps",
"name": "num_inference_steps",
"type": "INT",
"widget": {
"name": "num_inference_steps"
},
"link": null
},
{
"localized_name": "guidance_scale",
"name": "guidance_scale",
"type": "FLOAT",
"widget": {
"name": "guidance_scale"
},
"link": null
},
{
"localized_name": "true_cfg_scale",
"name": "true_cfg_scale",
"type": "FLOAT",
"widget": {
"name": "true_cfg_scale"
},
"link": null
},
{
"localized_name": "vae_use_slicing",
"name": "vae_use_slicing",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_slicing"
},
"link": null
},
{
"localized_name": "vae_use_tiling",
"name": "vae_use_tiling",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_tiling"
},
"link": null
},
{
"localized_name": "seed",
"name": "seed",
"type": "INT",
"widget": {
"name": "seed"
},
"link": null
}
],
"outputs": [
{
"localized_name": "diffusion sampling params",
"name": "diffusion sampling params",
"type": "SAMPLING_PARAMS",
"links": [
27
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniDiffusionSampling"
},
"widgets_values": [
1,
50,
1,
1,
false,
false,
1992,
"randomize"
]
},
{
"id": 14,
"type": "VLLMOmniWanParams",
"pos": [
518.6005521845747,
428.4455663719664
],
"size": [
270,
106
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"localized_name": "guidance_scale_2",
"name": "guidance_scale_2",
"type": "FLOAT",
"widget": {
"name": "guidance_scale_2"
},
"link": null
},
{
"localized_name": "boundary_ratio",
"name": "boundary_ratio",
"type": "FLOAT",
"widget": {
"name": "boundary_ratio"
},
"link": null
},
{
"localized_name": "flow_shift",
"name": "flow_shift",
"type": "FLOAT",
"widget": {
"name": "flow_shift"
},
"link": null
}
],
"outputs": [
{
"localized_name": "Wan video params",
"name": "Wan video params",
"type": "VIDEO_PARAMS",
"links": [
28
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniWanParams"
},
"widgets_values": [
4,
0.875,
5
]
},
{
"id": 12,
"type": "VLLMOmniGenerateVideo",
"pos": [
827.2566336087867,
-73.77449831827545
],
"size": [
400,
346
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"shape": 7,
"type": "IMAGE",
"link": 25
},
{
"localized_name": "sampling_params",
"name": "sampling_params",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 27
},
{
"localized_name": "lora",
"name": "lora",
"shape": 7,
"type": "REMOTE_LORA",
"link": null
},
{
"localized_name": "model_params",
"name": "model_params",
"shape": 7,
"type": "VIDEO_PARAMS",
"link": 28
},
{
"localized_name": "url",
"name": "url",
"type": "STRING",
"widget": {
"name": "url"
},
"link": null
},
{
"localized_name": "model",
"name": "model",
"type": "STRING",
"widget": {
"name": "model"
},
"link": null
},
{
"localized_name": "prompt",
"name": "prompt",
"type": "STRING",
"widget": {
"name": "prompt"
},
"link": null
},
{
"localized_name": "negative_prompt",
"name": "negative_prompt",
"type": "STRING",
"widget": {
"name": "negative_prompt"
},
"link": null
},
{
"localized_name": "width",
"name": "width",
"type": "INT",
"widget": {
"name": "width"
},
"link": null
},
{
"localized_name": "height",
"name": "height",
"type": "INT",
"widget": {
"name": "height"
},
"link": null
},
{
"localized_name": "fps",
"name": "fps",
"type": "INT",
"widget": {
"name": "fps"
},
"link": null
},
{
"localized_name": "num_frames",
"name": "num_frames",
"type": "INT",
"widget": {
"name": "num_frames"
},
"link": null
}
],
"outputs": [
{
"localized_name": "video",
"name": "video",
"type": "VIDEO",
"links": [
26
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniGenerateVideo"
},
"widgets_values": [
"http://localhost:8000/v1",
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
"Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.",
"Cartoonish",
900,
1200,
12,
41
]
},
{
"id": 13,
"type": "SaveVideo",
"pos": [
1257.5441369501657,
-44.12956394419863
],
"size": [
374,
372.34896826553796
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"localized_name": "video",
"name": "video",
"type": "VIDEO",
"link": 26
},
{
"localized_name": "filename_prefix",
"name": "filename_prefix",
"type": "STRING",
"widget": {
"name": "filename_prefix"
},
"link": null
},
{
"localized_name": "format",
"name": "format",
"type": "COMBO",
"widget": {
"name": "format"
},
"link": null
},
{
"localized_name": "codec",
"name": "codec",
"type": "COMBO",
"widget": {
"name": "codec"
},
"link": null
}
],
"outputs": [],
"properties": {},
"widgets_values": [
"video/ComfyUI",
"auto",
"auto"
]
}
],
"links": [
[
25,
4,
0,
12,
0,
"IMAGE"
],
[
26,
12,
0,
13,
0,
"VIDEO"
],
[
27,
8,
0,
12,
1,
"SAMPLING_PARAMS"
],
[
28,
14,
0,
12,
3,
"VIDEO_PARAMS"
]
],
"groups": [
{
"id": 1,
"title": "Input",
"bounding": [
213.1706010087147,
-313.20095750667554,
560.8246471437027,
407.4246532472182
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
}
],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 1.1469075819486894,
"offset": [
101.7632950847089,
492.29121889347806
]
}
},
"version": 0.4
}
+1 -22
View File
@@ -43,27 +43,6 @@ The benchmark supports three dataset modes via `--dataset`:
### VBench dataset
`vbench` only provides prompt data (and image paths for i2v/i2i); it does not carry
per-request generation fields. In this mode, all requests share CLI values:
`--width --height --num-frames --fps --num-inference-steps`
(pass `--width` and `--height` together).
Example (`t2v`):
```bash
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--base-url http://localhost:8099 \
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--task t2v \
--dataset vbench \
--num-prompts 50 \
--width 640 --height 480 \
--num-frames 81 --fps 16 \
--num-inference-steps 40
```
Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths).
If you use i2v/i2i bench datasets and need auto-download support, you may need:
```bash
@@ -111,7 +90,7 @@ Common optional flags:
Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`.
- For `vbench` / `random`: these CLI flags act as global defaults for all generated requests.
- For `trace`: requests can carry their own fields (e.g., `width/height/num_frames/num_inference_steps`), with overrides/fallbacks as below.
- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`).
Precedence rules for `trace` (i.e., what actually gets sent):
-100
View File
@@ -1,100 +0,0 @@
# Qwen3-TTS Benchmark
Benchmarks for Qwen3-TTS text-to-speech models, comparing vLLM-Omni streaming serving against HuggingFace Transformers offline inference.
## Prerequisites
```bash
pip install matplotlib aiohttp soundfile numpy tqdm
pip install qwen_tts # for HF baseline
```
## Quick Start
Run the full benchmark (vllm-omni + HF baseline) with a single command:
```bash
cd benchmarks/qwen3-tts
bash run_benchmark.sh
```
Results (JSON + PNG plots) are saved to `results/`.
### Common options
```bash
# Only vllm-omni (skip HF baseline)
bash run_benchmark.sh --async-only
# Only HF baseline
bash run_benchmark.sh --hf-only
# Use a different model (e.g. 1.7B)
MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
# Use batch_size=4 config for higher throughput
STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
# Custom GPU, prompt count, concurrency levels
GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
```
## Manual Steps
### 1) Start the vLLM-Omni server
```bash
CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
"Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
--omni --host 127.0.0.1 --port 8000 \
--stage-configs-path benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml \
--trust-remote-code
```
### 2) Run online serving benchmark
```bash
python vllm_omni/bench_tts_serve.py \
--port 8000 \
--num-prompts 50 \
--max-concurrency 1 4 10 \
--config-name "async_chunk" \
--result-dir results/
```
### 3) Run HuggingFace baseline
```bash
python transformers/bench_tts_hf.py \
--model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
--num-prompts 50 \
--gpu-device 0 \
--result-dir results/
```
### 4) Generate comparison plots
```bash
python plot_results.py \
--results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
--labels "vllm-omni" "hf_transformers" \
--output results/comparison.png
```
## Stage Configs
| Config | Batch Size | Description |
|--------|:----------:|-------------|
| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lower latency) |
| `vllm_omni/configs/qwen3_tts_bs4.yaml` | 4 | Concurrent request processing (higher throughput) |
Both configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
The model is specified via the CLI `--model` flag (or `MODEL` env var in `run_benchmark.sh`), so the same configs work for both the 0.6B and 1.7B model variants.
## Metrics
- **TTFP (Time to First Audio Packet)**: Time from request to first audio chunk (streaming latency)
- **E2E (End-to-End Latency)**: Total time from request to complete audio response
- **RTF (Real-Time Factor)**: E2E latency / audio duration. RTF < 1.0 means faster-than-real-time synthesis
- **Throughput**: Total audio seconds generated per wall-clock second
-254
View File
@@ -1,254 +0,0 @@
"""Plot Qwen3-TTS benchmark results.
Generates comparison bar charts similar to the async_chunk design doc:
- TTFP (Time-to-First-Packet) across concurrency levels
- E2E latency across concurrency levels
- RTF (Real-Time Factor) across concurrency levels
Usage:
# Compare two configs (async_chunk vs no_async_chunk):
python plot_results.py \
--results results/bench_async_chunk_*.json results/bench_no_async_chunk_*.json \
--labels "async_chunk" "no_async_chunk" \
--output results/qwen3_tts_benchmark.png
# Single config:
python plot_results.py \
--results results/bench_async_chunk_*.json \
--labels "async_chunk" \
--output results/qwen3_tts_benchmark.png
"""
import argparse
import json
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
def load_results(result_files: list[str]) -> list[list[dict]]:
"""Load benchmark results from JSON files."""
all_results = []
for f in result_files:
with open(f) as fh:
data = json.load(fh)
all_results.append(data)
return all_results
def plot_comparison(
all_results: list[list[dict]],
labels: list[str],
output_path: str,
title_prefix: str = "Qwen3-TTS",
):
"""Generate comparison bar charts."""
n_configs = len(all_results)
# Collect concurrency levels present in ALL configs (skip missing data)
all_concurrencies = [set(r["concurrency"] for r in results) for results in all_results]
concurrencies = sorted(set.union(*all_concurrencies))
# Build data arrays, using None for missing concurrency levels
ttfp_data = {label: [] for label in labels}
e2e_data = {label: [] for label in labels}
rtf_data = {label: [] for label in labels}
throughput_data = {label: [] for label in labels}
for results, label in zip(all_results, labels):
conc_map = {r["concurrency"]: r for r in results}
for c in concurrencies:
r = conc_map.get(c)
ttfp_data[label].append(r["mean_ttfp_ms"] if r else None)
e2e_data[label].append(r["mean_e2e_ms"] if r else None)
rtf_data[label].append(r["mean_rtf"] if r else None)
throughput_data[label].append(r["audio_throughput"] if r else None)
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
fig.suptitle(f"{title_prefix} Performance Benchmark", fontsize=16, fontweight="bold")
x = np.arange(len(concurrencies))
width = 0.35 if n_configs == 2 else 0.5
if n_configs > 1:
offsets = np.linspace(-width / 2 * (n_configs - 1), width / 2 * (n_configs - 1), n_configs)
else:
offsets = [0]
colors = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107"]
def plot_metric(ax, data_dict, ylabel, title, fmt=".1f"):
bars = []
for i, (label, values) in enumerate(data_dict.items()):
# Replace None with 0 for plotting, but track which are missing
plot_values = [v if v is not None else 0 for v in values]
color = colors[i % len(colors)]
bar = ax.bar(x + offsets[i], plot_values, width, label=label, color=color, alpha=0.85)
bars.append(bar)
# Add value labels on bars (skip None/missing data)
max_val = max((v for v in values if v is not None), default=1)
for rect, val in zip(bar, values):
if val is not None and val > 0:
ax.text(
rect.get_x() + rect.get_width() / 2,
rect.get_height() + max_val * 0.02,
f"{val:{fmt}}",
ha="center",
va="bottom",
fontsize=9,
fontweight="bold",
)
ax.set_xlabel("Concurrency", fontsize=12)
ax.set_ylabel(ylabel, fontsize=12)
ax.set_title(title, fontsize=13, fontweight="bold")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.legend(fontsize=10)
ax.grid(axis="y", alpha=0.3)
ax.set_axisbelow(True)
plot_metric(axes[0, 0], ttfp_data, "TTFP (ms)", "Time to First Audio Packet (TTFP)")
plot_metric(axes[0, 1], e2e_data, "E2E Latency (ms)", "End-to-End Latency (E2E)")
plot_metric(axes[1, 0], rtf_data, "RTF", "Real-Time Factor (RTF)", fmt=".3f")
plot_metric(axes[1, 1], throughput_data, "Audio-sec / Wall-sec", "Audio Throughput", fmt=".2f")
plt.tight_layout()
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
def plot_single_summary(results: list[dict], label: str, output_path: str):
"""Generate a single-config summary with percentile breakdown."""
concurrencies = [r["concurrency"] for r in results]
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
fig.suptitle(f"Qwen3-TTS Benchmark - {label}", fontsize=15, fontweight="bold")
# TTFP breakdown
ax = axes[0]
means = [r["mean_ttfp_ms"] for r in results]
medians = [r["median_ttfp_ms"] for r in results]
p90s = [r["p90_ttfp_ms"] for r in results]
p99s = [r["p99_ttfp_ms"] for r in results]
x = np.arange(len(concurrencies))
w = 0.2
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_xlabel("Concurrency")
ax.set_ylabel("TTFP (ms)")
ax.set_title("Time to First Audio Packet")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3)
# E2E breakdown
ax = axes[1]
means = [r["mean_e2e_ms"] for r in results]
medians = [r["median_e2e_ms"] for r in results]
p90s = [r["p90_e2e_ms"] for r in results]
p99s = [r["p99_e2e_ms"] for r in results]
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_xlabel("Concurrency")
ax.set_ylabel("E2E Latency (ms)")
ax.set_title("End-to-End Latency")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3)
# RTF
ax = axes[2]
means = [r["mean_rtf"] for r in results]
medians = [r["median_rtf"] for r in results]
ax.bar(x - 0.15, means, 0.3, label="mean", color="#2196F3")
ax.bar(x + 0.15, medians, 0.3, label="median", color="#4CAF50")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_xlabel("Concurrency")
ax.set_ylabel("RTF")
ax.set_title("Real-Time Factor")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3)
plt.tight_layout()
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
def print_comparison_table(all_results: list[list[dict]], labels: list[str]):
"""Print a markdown-formatted comparison table."""
concurrencies = sorted(set(r["concurrency"] for r in all_results[0]))
print("\n## Benchmark Results\n")
header = "| Metric | Concurrency |"
sep = "| --- | --- |"
for label in labels:
header += f" {label} |"
sep += " --- |"
print(header)
print(sep)
for metric, key, fmt in [
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
("E2E (ms)", "mean_e2e_ms", ".1f"),
("RTF", "mean_rtf", ".3f"),
("Throughput (audio-s/s)", "audio_throughput", ".2f"),
]:
for c in concurrencies:
row = f"| {metric} | {c} |"
for results in all_results:
conc_map = {r["concurrency"]: r for r in results}
val = conc_map.get(c, {}).get(key, 0)
row += f" {val:{fmt}} |"
print(row)
# Improvement calculation (only if 2 configs)
if len(all_results) == 2:
print(f"\n## Improvement ({labels[0]} vs {labels[1]})\n")
print("| Metric | Concurrency | Improvement |")
print("| --- | --- | --- |")
for metric, key in [("TTFP", "mean_ttfp_ms"), ("E2E", "mean_e2e_ms"), ("RTF", "mean_rtf")]:
for c in concurrencies:
m0 = {r["concurrency"]: r for r in all_results[0]}
m1 = {r["concurrency"]: r for r in all_results[1]}
v0 = m0.get(c, {}).get(key, 0)
v1 = m1.get(c, {}).get(key, 0)
if v1 > 0:
pct = (v1 - v0) / v1 * 100
print(f"| {metric} | {c} | {pct:+.1f}% |")
def parse_args():
parser = argparse.ArgumentParser(description="Plot Qwen3-TTS benchmark results")
parser.add_argument(
"--results", type=str, nargs="+", required=True, help="Path(s) to result JSON files (one per config)"
)
parser.add_argument(
"--labels", type=str, nargs="+", required=True, help="Labels for each config (must match --results count)"
)
parser.add_argument("--output", type=str, default="results/qwen3_tts_benchmark.png", help="Output image path")
parser.add_argument("--title", type=str, default="Qwen3-TTS", help="Title prefix for the plot")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
assert len(args.results) == len(args.labels), "--results and --labels must have the same count"
all_results = load_results(args.results)
print_comparison_table(all_results, args.labels)
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
if len(all_results) == 1:
plot_single_summary(all_results[0], args.labels[0], args.output)
else:
plot_comparison(all_results, args.labels, args.output, title_prefix=args.title)
-3
View File
@@ -1,3 +0,0 @@
# Benchmark results are machine-specific - do not commit
*
!.gitignore
-272
View File
@@ -1,272 +0,0 @@
#!/bin/bash
# Qwen3-TTS Benchmark Runner
#
# Compares vllm-omni streaming serving vs HuggingFace transformers offline inference.
# Produces JSON results and comparison plots.
#
# Usage:
# # Full comparison (vllm-omni + HF):
# bash run_benchmark.sh
#
# # Only vllm-omni async_chunk config:
# bash run_benchmark.sh --async-only
#
# # Only HuggingFace baseline:
# bash run_benchmark.sh --hf-only
#
# # vllm-omni only (skip HF):
# bash run_benchmark.sh --skip-hf
#
# # Custom settings:
# GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
#
# # Use 1.7B model:
# MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
#
# # Use batch_size=4 config:
# STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
#
# Environment variables:
# GPU_DEVICE - GPU index to use (default: 0)
# NUM_PROMPTS - Number of prompts per concurrency level (default: 50)
# CONCURRENCY - Space-separated concurrency levels (default: "1 4 10")
# MODEL - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
# PORT - Server port (default: 8000)
# GPU_MEM_TALKER - gpu_memory_utilization for talker stage (default: 0.3)
# GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.2)
# STAGE_CONFIG - Path to stage config YAML (default: configs/qwen3_tts_bs1.yaml)
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
# Defaults
GPU_DEVICE="${GPU_DEVICE:-0}"
NUM_PROMPTS="${NUM_PROMPTS:-50}"
CONCURRENCY="${CONCURRENCY:-1 4 10}"
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
PORT="${PORT:-8000}"
GPU_MEM_TALKER="${GPU_MEM_TALKER:-0.3}"
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-0.2}"
NUM_WARMUPS="${NUM_WARMUPS:-3}"
STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/configs/qwen3_tts_bs1.yaml}"
RESULT_DIR="${SCRIPT_DIR}/results"
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
# Parse args
RUN_ASYNC=true
RUN_HF=true
for arg in "$@"; do
case "$arg" in
--async-only) RUN_HF=false ;;
--hf-only) RUN_ASYNC=false ;;
--skip-hf) RUN_HF=false ;;
esac
done
mkdir -p "${RESULT_DIR}"
echo "============================================================"
echo " Qwen3-TTS Benchmark"
echo "============================================================"
echo " GPU: ${GPU_DEVICE}"
echo " Model: ${MODEL}"
echo " Prompts: ${NUM_PROMPTS}"
echo " Concurrency: ${CONCURRENCY}"
echo " Port: ${PORT}"
echo " Stage config: ${STAGE_CONFIG}"
echo " Results: ${RESULT_DIR}"
echo "============================================================"
# Prepare stage config with correct GPU device and memory settings
prepare_config() {
local config_template="$1"
local config_name="$2"
local output_path="${RESULT_DIR}/${config_name}_stage_config.yaml"
# Use sed to patch GPU device and memory utilization
sed \
-e "s/devices: \"0\"/devices: \"${GPU_DEVICE}\"/g" \
-e "s/gpu_memory_utilization: 0.3/gpu_memory_utilization: ${GPU_MEM_TALKER}/g" \
-e "s/gpu_memory_utilization: 0.2/gpu_memory_utilization: ${GPU_MEM_CODE2WAV}/g" \
"${config_template}" > "${output_path}"
echo "${output_path}"
}
# Start server and wait for it to be ready
start_server() {
local stage_config="$1"
local config_name="$2"
local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log"
echo ""
echo "Starting server with config: ${config_name}"
echo " Stage config: ${stage_config}"
echo " Log file: ${log_file}"
VLLM_WORKER_MULTIPROC_METHOD=spawn \
CUDA_VISIBLE_DEVICES="${GPU_DEVICE}" \
python -m vllm_omni.entrypoints.cli.main serve "${MODEL}" \
--omni \
--host 127.0.0.1 \
--port "${PORT}" \
--stage-configs-path "${stage_config}" \
--stage-init-timeout 120 \
--trust-remote-code \
--disable-log-stats \
> "${log_file}" 2>&1 &
SERVER_PID=$!
echo " Server PID: ${SERVER_PID}"
# Wait for server to be ready
echo " Waiting for server to be ready..."
local max_wait=300
local waited=0
while [ ${waited} -lt ${max_wait} ]; do
if curl -sf "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then
echo " Server is ready! (waited ${waited}s)"
return 0
fi
# Check if process is still alive
if ! kill -0 ${SERVER_PID} 2>/dev/null; then
echo " ERROR: Server process died. Check log: ${log_file}"
tail -20 "${log_file}"
return 1
fi
sleep 2
waited=$((waited + 2))
done
echo " ERROR: Server did not start within ${max_wait}s. Check log: ${log_file}"
kill ${SERVER_PID} 2>/dev/null || true
return 1
}
# Stop the server
stop_server() {
if [ -n "${SERVER_PID:-}" ]; then
echo " Stopping server (PID: ${SERVER_PID})..."
kill ${SERVER_PID} 2>/dev/null || true
wait ${SERVER_PID} 2>/dev/null || true
# Kill any remaining child processes on the port
local pids
pids=$(lsof -ti:${PORT} 2>/dev/null || true)
if [ -n "${pids}" ]; then
echo " Cleaning up remaining processes on port ${PORT}..."
echo "${pids}" | xargs kill -9 2>/dev/null || true
fi
echo " Server stopped."
SERVER_PID=""
fi
}
# Cleanup on exit
trap 'stop_server' EXIT
# Run benchmark for a given config
run_bench() {
local config_name="$1"
local config_template="$2"
echo ""
echo "============================================================"
echo " Benchmarking: ${config_name}"
echo "============================================================"
local stage_config
stage_config=$(prepare_config "${config_template}" "${config_name}")
start_server "${stage_config}" "${config_name}"
# Convert concurrency string to args
local conc_args=""
for c in ${CONCURRENCY}; do
conc_args="${conc_args} ${c}"
done
cd "${PROJECT_ROOT}"
python "${SCRIPT_DIR}/vllm_omni/bench_tts_serve.py" \
--host 127.0.0.1 \
--port "${PORT}" \
--num-prompts "${NUM_PROMPTS}" \
--max-concurrency ${conc_args} \
--num-warmups "${NUM_WARMUPS}" \
--config-name "${config_name}" \
--result-dir "${RESULT_DIR}"
stop_server
# Allow GPU memory to settle
sleep 5
}
# Run vllm-omni benchmark
if [ "${RUN_ASYNC}" = true ]; then
run_bench "async_chunk" "${SCRIPT_DIR}/${STAGE_CONFIG}"
fi
# Run HuggingFace baseline benchmark
if [ "${RUN_HF}" = true ]; then
echo ""
echo "============================================================"
echo " Benchmarking: HuggingFace transformers (offline)"
echo "============================================================"
cd "${PROJECT_ROOT}"
python "${SCRIPT_DIR}/transformers/bench_tts_hf.py" \
--model "${MODEL}" \
--num-prompts "${NUM_PROMPTS}" \
--num-warmups "${NUM_WARMUPS}" \
--gpu-device "${GPU_DEVICE}" \
--config-name "hf_transformers" \
--result-dir "${RESULT_DIR}"
# Allow GPU memory to settle
sleep 5
fi
# Plot results
echo ""
echo "============================================================"
echo " Generating plots..."
echo "============================================================"
RESULT_FILES=""
LABELS=""
if [ "${RUN_ASYNC}" = true ]; then
ASYNC_FILE=$(ls -t "${RESULT_DIR}"/bench_async_chunk_*.json 2>/dev/null | head -1)
if [ -n "${ASYNC_FILE}" ]; then
RESULT_FILES="${ASYNC_FILE}"
LABELS="async_chunk"
fi
fi
if [ "${RUN_HF}" = true ]; then
HF_FILE=$(ls -t "${RESULT_DIR}"/bench_hf_transformers_*.json 2>/dev/null | head -1)
if [ -n "${HF_FILE}" ]; then
if [ -n "${RESULT_FILES}" ]; then
RESULT_FILES="${RESULT_FILES} ${HF_FILE}"
LABELS="${LABELS} hf_transformers"
else
RESULT_FILES="${HF_FILE}"
LABELS="hf_transformers"
fi
fi
fi
if [ -n "${RESULT_FILES}" ]; then
python "${SCRIPT_DIR}/plot_results.py" \
--results ${RESULT_FILES} \
--labels ${LABELS} \
--output "${RESULT_DIR}/qwen3_tts_benchmark_${TIMESTAMP}.png"
fi
echo ""
echo "============================================================"
echo " Benchmark complete!"
echo " Results: ${RESULT_DIR}"
echo "============================================================"
@@ -1,281 +0,0 @@
"""Benchmark Qwen3-TTS using HuggingFace transformers (qwen_tts library).
Measures E2E latency, RTF, and audio duration for offline (non-serving) inference.
Results are saved in the same JSON format as bench_tts_serve.py for unified plotting.
Usage:
python bench_tts_hf.py \
--model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
--num-prompts 50 \
--num-warmups 3 \
--gpu-device 0 \
--result-dir results/
"""
import argparse
import json
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
import numpy as np
import soundfile as sf
import torch
PROMPTS = [
"Hello, welcome to the voice synthesis benchmark test.",
"She said she would be here by noon, but nobody showed up.",
"The quick brown fox jumps over the lazy dog near the riverbank.",
"I can't believe how beautiful the sunset looks from up here on the mountain.",
"Please remember to bring your identification documents to the appointment tomorrow morning.",
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
"The train leaves at half past seven, so we need to arrive at the station before then.",
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
]
@dataclass
class BenchmarkResult:
config_name: str = ""
concurrency: int = 1 # always 1 for offline
num_prompts: int = 0
completed: int = 0
failed: int = 0
duration_s: float = 0.0
# TTFP stats - not applicable for HF offline, set to E2E for compatibility
mean_ttfp_ms: float = 0.0
median_ttfp_ms: float = 0.0
std_ttfp_ms: float = 0.0
p90_ttfp_ms: float = 0.0
p95_ttfp_ms: float = 0.0
p99_ttfp_ms: float = 0.0
# E2E stats (ms)
mean_e2e_ms: float = 0.0
median_e2e_ms: float = 0.0
std_e2e_ms: float = 0.0
p90_e2e_ms: float = 0.0
p95_e2e_ms: float = 0.0
p99_e2e_ms: float = 0.0
# RTF stats
mean_rtf: float = 0.0
median_rtf: float = 0.0
std_rtf: float = 0.0
p99_rtf: float = 0.0
# Audio stats
mean_audio_duration_s: float = 0.0
total_audio_duration_s: float = 0.0
audio_throughput: float = 0.0
request_throughput: float = 0.0
# Per-request details
per_request: list = field(default_factory=list)
def run_benchmark(args):
from qwen_tts import Qwen3TTSModel
device = f"cuda:{args.gpu_device}"
print(f"Loading model: {args.model} on {device}")
model = Qwen3TTSModel.from_pretrained(
args.model,
device_map=device,
dtype=torch.bfloat16,
)
print("Model loaded.")
# Build prompt list
prompts = [PROMPTS[i % len(PROMPTS)] for i in range(args.num_prompts)]
# Warmup
if args.num_warmups > 0:
print(f"Warming up with {args.num_warmups} requests...")
for i in range(args.num_warmups):
p = PROMPTS[i % len(PROMPTS)]
wavs, sr = model.generate_custom_voice(
text=p,
language=args.language,
speaker=args.voice,
)
# Sync GPU
torch.cuda.synchronize(device)
print("Warmup done.")
# Benchmark
print(f"Running {args.num_prompts} requests sequentially...")
e2e_times = []
rtfs = []
audio_durations = []
per_request = []
failed = 0
audio_dir = None
if args.save_audio:
audio_dir = Path(args.result_dir) / "audio_hf"
audio_dir.mkdir(parents=True, exist_ok=True)
total_start = time.perf_counter()
for i, prompt in enumerate(prompts):
try:
torch.cuda.synchronize(device)
st = time.perf_counter()
wavs, sr = model.generate_custom_voice(
text=prompt,
language=args.language,
speaker=args.voice,
)
torch.cuda.synchronize(device)
elapsed = time.perf_counter() - st
# Compute audio duration
audio_samples = wavs[0]
if isinstance(audio_samples, torch.Tensor):
audio_samples = audio_samples.cpu().numpy()
audio_dur = len(audio_samples) / sr
rtf = elapsed / audio_dur if audio_dur > 0 else 0.0
e2e_times.append(elapsed)
rtfs.append(rtf)
audio_durations.append(audio_dur)
per_request.append(
{
"e2e_ms": elapsed * 1000,
"ttfp_ms": elapsed * 1000, # no streaming, TTFP = E2E
"rtf": rtf,
"audio_duration_s": audio_dur,
"prompt": prompt,
}
)
if audio_dir:
sf.write(str(audio_dir / f"output_{i:04d}.wav"), audio_samples, sr)
if (i + 1) % 10 == 0 or i == 0:
print(
f" [{i + 1}/{args.num_prompts}] e2e={elapsed * 1000:.0f}ms rtf={rtf:.3f} audio={audio_dur:.2f}s"
)
except Exception as e:
print(f" [{i + 1}/{args.num_prompts}] FAILED: {e}")
failed += 1
total_duration = time.perf_counter() - total_start
completed = len(e2e_times)
# Compute stats
result = BenchmarkResult(
config_name=args.config_name,
concurrency=1,
num_prompts=args.num_prompts,
completed=completed,
failed=failed,
duration_s=total_duration,
)
if e2e_times:
e2e_ms = [t * 1000 for t in e2e_times]
result.mean_e2e_ms = float(np.mean(e2e_ms))
result.median_e2e_ms = float(np.median(e2e_ms))
result.std_e2e_ms = float(np.std(e2e_ms))
result.p90_e2e_ms = float(np.percentile(e2e_ms, 90))
result.p95_e2e_ms = float(np.percentile(e2e_ms, 95))
result.p99_e2e_ms = float(np.percentile(e2e_ms, 99))
# For HF offline, TTFP = E2E (no streaming)
result.mean_ttfp_ms = result.mean_e2e_ms
result.median_ttfp_ms = result.median_e2e_ms
result.std_ttfp_ms = result.std_e2e_ms
result.p90_ttfp_ms = result.p90_e2e_ms
result.p95_ttfp_ms = result.p95_e2e_ms
result.p99_ttfp_ms = result.p99_e2e_ms
result.mean_rtf = float(np.mean(rtfs))
result.median_rtf = float(np.median(rtfs))
result.std_rtf = float(np.std(rtfs))
result.p99_rtf = float(np.percentile(rtfs, 99))
result.mean_audio_duration_s = float(np.mean(audio_durations))
result.total_audio_duration_s = float(np.sum(audio_durations))
result.audio_throughput = result.total_audio_duration_s / total_duration
result.request_throughput = completed / total_duration
result.per_request = per_request
# Print summary in standardized performance template
W = 50
print("")
print(f"{'=' * W}")
print(f"{'Serving Benchmark Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Successful requests:':<40}{completed:<10}")
print(f"{'Failed requests:':<40}{failed:<10}")
print(f"{'Maximum request concurrency:':<40}{1:<10}")
print(f"{'Benchmark duration (s):':<40}{total_duration:<10.2f}")
print(f"{'Request throughput (req/s):':<40}{result.request_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'End-to-end Latency':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean E2EL (ms):':<40}{result.mean_e2e_ms:<10.2f}")
print(f"{'Median E2EL (ms):':<40}{result.median_e2e_ms:<10.2f}")
print(f"{'P99 E2EL (ms):':<40}{result.p99_e2e_ms:<10.2f}")
print(f"{'=' * W}")
print(f"{'Audio Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Total audio duration generated (s):':<40}{result.total_audio_duration_s:<10.2f}")
print(f"{'Audio throughput (audio duration/s):':<40}{result.audio_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'Time to First Packet':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_TTFP (ms):':<40}{result.mean_ttfp_ms:<10.2f}")
print(f"{'Median AUDIO_TTFP (ms):':<40}{result.median_ttfp_ms:<10.2f}")
print(f"{'P99 AUDIO_TTFP (ms):':<40}{result.p99_ttfp_ms:<10.2f}")
print(f"{'-' * W}")
print(f"{'Real Time Factor':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_RTF:':<40}{result.mean_rtf:<10.3f}")
print(f"{'Median AUDIO_RTF:':<40}{result.median_rtf:<10.3f}")
print(f"{'P99 AUDIO_RTF:':<40}{result.p99_rtf:<10.3f}")
print(f"{'=' * W}")
print("")
# Save results (as a list with single concurrency=1 entry, matching serve format)
result_dir = Path(args.result_dir)
result_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
with open(result_file, "w") as f:
json.dump([asdict(result)], f, indent=2)
print(f"Results saved to {result_file}")
return result
def parse_args():
parser = argparse.ArgumentParser(description="Qwen3-TTS HuggingFace Benchmark")
parser.add_argument(
"--model", type=str, default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", help="HuggingFace model name or path"
)
parser.add_argument("--num-prompts", type=int, default=50)
parser.add_argument("--num-warmups", type=int, default=3)
parser.add_argument("--gpu-device", type=int, default=0)
parser.add_argument("--voice", type=str, default="Vivian")
parser.add_argument("--language", type=str, default="English")
parser.add_argument(
"--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
)
parser.add_argument("--result-dir", type=str, default="results")
parser.add_argument("--save-audio", action="store_true", help="Save generated audio files")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
run_benchmark(args)
@@ -1,348 +0,0 @@
"""Benchmark client for Qwen3-TTS via /v1/audio/speech endpoint.
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF (Real-Time Factor)
across configurable concurrency levels. Saves results as JSON for plotting.
Usage:
python bench_tts_serve.py \
--host 127.0.0.1 --port 8000 \
--num-prompts 50 \
--max-concurrency 1 4 10 \
--result-dir results/
"""
import argparse
import asyncio
import json
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
import aiohttp
import numpy as np
from tqdm.asyncio import tqdm
PROMPTS = [
"Hello, welcome to the voice synthesis benchmark test.",
"She said she would be here by noon, but nobody showed up.",
"The quick brown fox jumps over the lazy dog near the riverbank.",
"I can't believe how beautiful the sunset looks from up here on the mountain.",
"Please remember to bring your identification documents to the appointment tomorrow morning.",
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
"The train leaves at half past seven, so we need to arrive at the station before then.",
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
]
@dataclass
class RequestResult:
success: bool = False
ttfp: float = 0.0 # Time to first audio packet (seconds)
e2e: float = 0.0 # End-to-end latency (seconds)
audio_bytes: int = 0 # Total audio bytes received
audio_duration: float = 0.0 # Audio duration in seconds (estimated from PCM)
rtf: float = 0.0 # Real-time factor = e2e / audio_duration
prompt: str = ""
error: str = ""
@dataclass
class BenchmarkResult:
config_name: str = ""
concurrency: int = 0
num_prompts: int = 0
completed: int = 0
failed: int = 0
duration_s: float = 0.0
# TTFP stats (ms)
mean_ttfp_ms: float = 0.0
median_ttfp_ms: float = 0.0
std_ttfp_ms: float = 0.0
p90_ttfp_ms: float = 0.0
p95_ttfp_ms: float = 0.0
p99_ttfp_ms: float = 0.0
# E2E stats (ms)
mean_e2e_ms: float = 0.0
median_e2e_ms: float = 0.0
std_e2e_ms: float = 0.0
p90_e2e_ms: float = 0.0
p95_e2e_ms: float = 0.0
p99_e2e_ms: float = 0.0
# RTF stats
mean_rtf: float = 0.0
median_rtf: float = 0.0
std_rtf: float = 0.0
p99_rtf: float = 0.0
# Audio stats
mean_audio_duration_s: float = 0.0
total_audio_duration_s: float = 0.0
audio_throughput: float = 0.0 # audio_duration / wall_time
request_throughput: float = 0.0 # requests / second
# Per-request details
per_request: list = field(default_factory=list)
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
"""Convert raw PCM byte count to duration in seconds."""
num_samples = num_bytes / sample_width
return num_samples / sample_rate
async def send_tts_request(
session: aiohttp.ClientSession,
api_url: str,
prompt: str,
voice: str = "vivian",
language: str = "English",
pbar: tqdm | None = None,
) -> RequestResult:
"""Send a streaming TTS request and measure latency metrics."""
payload = {
"input": prompt,
"voice": voice,
"language": language,
"stream": True,
"response_format": "pcm",
}
result = RequestResult(prompt=prompt)
st = time.perf_counter()
try:
async with session.post(api_url, json=payload) as response:
if response.status != 200:
result.error = f"HTTP {response.status}: {await response.text()}"
result.success = False
return result
first_chunk = True
total_bytes = 0
async for chunk in response.content.iter_any():
if first_chunk and len(chunk) > 0:
result.ttfp = time.perf_counter() - st
first_chunk = False
total_bytes += len(chunk)
result.e2e = time.perf_counter() - st
result.audio_bytes = total_bytes
result.audio_duration = pcm_bytes_to_duration(total_bytes)
if result.audio_duration > 0:
result.rtf = result.e2e / result.audio_duration
result.success = True
except Exception as e:
result.error = str(e)
result.success = False
result.e2e = time.perf_counter() - st
if pbar:
pbar.update(1)
return result
async def run_benchmark(
host: str,
port: int,
num_prompts: int,
max_concurrency: int,
num_warmups: int = 3,
voice: str = "vivian",
language: str = "English",
) -> BenchmarkResult:
"""Run benchmark at a given concurrency level."""
api_url = f"http://{host}:{port}/v1/audio/speech"
connector = aiohttp.TCPConnector(
limit=max_concurrency,
limit_per_host=max_concurrency,
keepalive_timeout=60,
)
session = aiohttp.ClientSession(
connector=connector,
timeout=aiohttp.ClientTimeout(total=600),
)
# Warmup
if num_warmups > 0:
print(f" Warming up with {num_warmups} requests...")
warmup_tasks = []
for i in range(num_warmups):
prompt = PROMPTS[i % len(PROMPTS)]
warmup_tasks.append(send_tts_request(session, api_url, prompt, voice, language))
await asyncio.gather(*warmup_tasks)
print(" Warmup done.")
# Build request list
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
# Run benchmark
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
semaphore = asyncio.Semaphore(max_concurrency)
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
async def limited_request(prompt):
async with semaphore:
return await send_tts_request(session, api_url, prompt, voice, language, pbar)
start_time = time.perf_counter()
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
results: list[RequestResult] = await asyncio.gather(*tasks)
duration = time.perf_counter() - start_time
pbar.close()
await session.close()
# Compute stats
successful = [r for r in results if r.success]
failed = [r for r in results if not r.success]
bench = BenchmarkResult(
concurrency=max_concurrency,
num_prompts=num_prompts,
completed=len(successful),
failed=len(failed),
duration_s=duration,
)
if successful:
ttfps = [r.ttfp * 1000 for r in successful] # convert to ms
e2es = [r.e2e * 1000 for r in successful]
rtfs = [r.rtf for r in successful]
audio_durs = [r.audio_duration for r in successful]
bench.mean_ttfp_ms = float(np.mean(ttfps))
bench.median_ttfp_ms = float(np.median(ttfps))
bench.std_ttfp_ms = float(np.std(ttfps))
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
bench.mean_e2e_ms = float(np.mean(e2es))
bench.median_e2e_ms = float(np.median(e2es))
bench.std_e2e_ms = float(np.std(e2es))
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
bench.mean_rtf = float(np.mean(rtfs))
bench.median_rtf = float(np.median(rtfs))
bench.std_rtf = float(np.std(rtfs))
bench.p99_rtf = float(np.percentile(rtfs, 99))
bench.mean_audio_duration_s = float(np.mean(audio_durs))
bench.total_audio_duration_s = float(np.sum(audio_durs))
bench.audio_throughput = bench.total_audio_duration_s / duration
bench.request_throughput = len(successful) / duration
bench.per_request = [
{
"ttfp_ms": r.ttfp * 1000,
"e2e_ms": r.e2e * 1000,
"rtf": r.rtf,
"audio_duration_s": r.audio_duration,
"prompt": r.prompt,
}
for r in successful
]
# Print summary in standardized performance template
W = 50
print("")
print(f"{'=' * W}")
print(f"{'Serving Benchmark Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Successful requests:':<40}{bench.completed:<10}")
print(f"{'Failed requests:':<40}{bench.failed:<10}")
print(f"{'Maximum request concurrency:':<40}{max_concurrency:<10}")
print(f"{'Benchmark duration (s):':<40}{duration:<10.2f}")
print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'End-to-end Latency':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
print(f"{'=' * W}")
print(f"{'Audio Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'Time to First Packet':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
print(f"{'-' * W}")
print(f"{'Real Time Factor':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
print(f"{'=' * W}")
print("")
if failed:
for r in failed[:3]:
print(f" [ERROR] {r.error[:200]}")
return bench
async def main(args):
all_results = []
for concurrency in args.max_concurrency:
result = await run_benchmark(
host=args.host,
port=args.port,
num_prompts=args.num_prompts,
max_concurrency=concurrency,
num_warmups=args.num_warmups,
voice=args.voice,
language=args.language,
)
result.config_name = args.config_name
all_results.append(asdict(result))
# Save results
result_dir = Path(args.result_dir)
result_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
with open(result_file, "w") as f:
json.dump(all_results, f, indent=2)
print(f"Results saved to {result_file}")
return all_results
def parse_args():
parser = argparse.ArgumentParser(description="Qwen3-TTS Benchmark Client")
parser.add_argument("--host", type=str, default="127.0.0.1")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--num-prompts", type=int, default=50, help="Number of prompts per concurrency level")
parser.add_argument( # noqa: E501
"--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
)
parser.add_argument("--num-warmups", type=int, default=3)
parser.add_argument("--voice", type=str, default="vivian")
parser.add_argument("--language", type=str, default="English")
parser.add_argument(
"--config-name", type=str, default="async_chunk", help="Label for this config (used in filenames)"
)
parser.add_argument("--result-dir", type=str, default="results")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
asyncio.run(main(args))
@@ -1,96 +0,0 @@
# Qwen3-TTS batch_size=1 config (streaming with async_chunk)
# 2-stage pipeline: Talker -> Code2Wav
async_chunk: true
stage_args:
- stage_id: 0
stage_type: llm
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: qwen3_tts
model_arch: Qwen3TTSTalkerForConditionalGeneration
hf_overrides:
architectures: [Qwen3TTSTalkerForConditionalGeneration]
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
enforce_eager: false
trust_remote_code: true
async_scheduling: false
enable_prefix_caching: false
engine_output_type: latent
gpu_memory_utilization: 0.3
distributed_executor_backend: "mp"
max_num_batched_tokens: 512
max_model_len: 4096
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
output_connectors:
to_stage_1: connector_of_shared_memory
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 4096
seed: 42
detokenize: false
repetition_penalty: 1.05
stop_token_ids: [2150]
- stage_id: 1
stage_type: llm
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen3TTSCode2Wav
hf_overrides:
architectures: [Qwen3TTSCode2Wav]
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
async_scheduling: false
enable_prefix_caching: false
engine_output_type: audio
gpu_memory_utilization: 0.2
distributed_executor_backend: "mp"
max_num_batched_tokens: 8192
max_model_len: 32768
engine_input_source: [0]
final_output: true
final_output_type: audio
input_connectors:
from_stage_0: connector_of_shared_memory
tts_args:
max_instructions_length: 500
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 65536
seed: 42
detokenize: true
repetition_penalty: 1.0
runtime:
enabled: true
defaults:
window_size: -1
max_inflight: 1
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
connector_get_max_wait: 300
codec_chunk_frames: 25
codec_left_context_frames: 25
edges:
- from: 0
to: 1
window_size: -1
@@ -1,97 +0,0 @@
# Qwen3-TTS batch_size=4 config (streaming with async_chunk)
# Enables concurrent request processing with max_inflight=4
# 2-stage pipeline: Talker -> Code2Wav
async_chunk: true
stage_args:
- stage_id: 0
stage_type: llm
runtime:
devices: "0"
max_batch_size: 4
engine_args:
model_stage: qwen3_tts
model_arch: Qwen3TTSTalkerForConditionalGeneration
hf_overrides:
architectures: [Qwen3TTSTalkerForConditionalGeneration]
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
enforce_eager: false
trust_remote_code: true
async_scheduling: false
enable_prefix_caching: false
engine_output_type: latent
gpu_memory_utilization: 0.3
distributed_executor_backend: "mp"
max_num_batched_tokens: 512
max_model_len: 4096
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
output_connectors:
to_stage_1: connector_of_shared_memory
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 4096
seed: 42
detokenize: false
repetition_penalty: 1.05
stop_token_ids: [2150]
- stage_id: 1
stage_type: llm
runtime:
devices: "0"
max_batch_size: 4
engine_args:
model_stage: code2wav
model_arch: Qwen3TTSCode2Wav
hf_overrides:
architectures: [Qwen3TTSCode2Wav]
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
async_scheduling: false
enable_prefix_caching: false
engine_output_type: audio
gpu_memory_utilization: 0.2
distributed_executor_backend: "mp"
max_num_batched_tokens: 8192
max_model_len: 32768
engine_input_source: [0]
final_output: true
final_output_type: audio
input_connectors:
from_stage_0: connector_of_shared_memory
tts_args:
max_instructions_length: 500
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 65536
seed: 42
detokenize: true
repetition_penalty: 1.0
runtime:
enabled: true
defaults:
window_size: -1
max_inflight: 4
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
connector_get_max_wait: 300
codec_chunk_frames: 25
codec_left_context_frames: 25
edges:
- from: 0
to: 1
window_size: -1
+1 -1
View File
@@ -76,7 +76,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
ENV UV_INDEX_STRATEGY="unsafe-best-match"
ENV UV_LINK_MODE="copy"
ARG VLLM_VERSION=v0.17.0
ARG VLLM_VERSION=v0.16.0
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
WORKDIR /workspace/vllm
-9
View File
@@ -13,14 +13,9 @@ nav:
- examples/README.md
- Offline Inference:
- BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
- CosyVoice3: user_guide/examples/offline_inference/cosyvoice3.md
- Fish Speech S2 Pro: user_guide/examples/offline_inference/fish_speech.md
- GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
- Helios Video Generation: user_guide/examples/offline_inference/helios.md
- HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
- Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
- MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
- MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
- Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
- Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
@@ -28,11 +23,8 @@ nav:
- Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
- Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
- X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
- Online Serving:
- BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
- vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
- Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md
- GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md
@@ -100,7 +92,6 @@ nav:
- api/vllm_omni
- CLI Reference: cli
- Community:
- Governance: community/governance.md
- community/*
- Slack: https://slack.vllm.ai
- Blog: https://blog.vllm.ai
+1 -1
View File
@@ -60,5 +60,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
For more information, checkout the following:
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1)
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
Binary file not shown.

Before

Width:  |  Height:  |  Size: 98 KiB

After

Width:  |  Height:  |  Size: 125 KiB

-52
View File
@@ -1,52 +0,0 @@
# Governance
vLLM-Omni's governance is inspired by the [vLLM governance process](https://docs.vllm.ai/en/latest/governance/process/). We share the same commitment to open source community and meritocratic norms.
## Values
vLLM-Omni aims to be the fastest and easiest-to-use omni-modality inference and serving engine. Our values are aligned with [vLLM's values](https://docs.vllm.ai/en/latest/governance/process/#values):
### Design Values
1. **Top performance**: System performance is our top priority. We monitor overheads, optimize kernels, and publish benchmarks. We never leave performance on the table.
2. **Ease of use**: vLLM-Omni must be simple to install, configure, and operate. We provide clear documentation, fast startup, clean logs, helpful error messages, and monitoring guides. Many users fork our code or study it deeply, so we keep it readable and modular.
3. **Wide coverage**: vLLM-Omni supports frontier models and high-performance accelerators. We make it easy to add new models and hardware. vLLM-Omni + PyTorch form a simple interface that avoids complexity.
4. **Production ready**: vLLM-Omni runs 24/7 in production. It must be easy to operate and monitor for health issues.
5. **Extensibility**: vLLM-Omni serves as fundamental omni-modality infrastructure. Our codebase cannot cover every use case, so we design for easy forking and customization.
### Collaboration Values
1. **Tightly Knit and Fast-Moving**: Our maintainer team is aligned on vision, philosophy, and roadmap. We work closely to unblock each other and move quickly.
2. **Individual Merit**: No one buys their way into governance. Committer status belongs to individuals, not companies. We reward contribution, maintenance, and project stewardship.
## Project Maintainers
### Lead Maintainers
Lead maintainers are responsible for the overall direction and strategy of the project:
- [@Gaohan123](https://github.com/Gaohan123)
- [@hsliuustc0106](https://github.com/hsliuustc0106)
- [@ywang96](https://github.com/ywang96)
### Active Committers
Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
- [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
- [@wtomin](https://github.com/wtomin):
- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
## Meetings
Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
## Committer Nomination Process
Every month, any active committer can nominate new committer(s) to the project. Up to **two new committers** will be admitted per month based on the quality and impact of their contributions.
+1 -1
View File
@@ -18,7 +18,7 @@ If users want to modify some part of it. The custom stage_configs file can be in
For offline (Assume necessary dependencies have ben imported):
```python
model_name = "Qwen/Qwen2.5-Omni-7B"
omni_llm = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
omni_llm = OmniLLM(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
```
For online serving:
+4 -6
View File
@@ -67,8 +67,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
<td>Online & Offline (multiple deployment scenarios):<br>real model, normal inference function, normal accuracy</td>
<td>&lt;30min</td>
<td>
/tests/e2e/online_serving/test_{model_name}.py<br>
/tests/e2e/offline_inference/test_{model_name}.py
/tests/e2e/online_serving/test_{model_name}_expansion.py<br>
/tests/e2e/offline_inference/test_{model_name}_expansion.py
</td>
<td>
<a href="#chapter-2-l3-level-testing-core-integration-performance-and-accuracy-verification">Chapter 2</a><br>
@@ -508,9 +508,7 @@ L3 level testing executes after code is merged into the main branch. Its core pu
**2.4.5 Request Execution**
```python
openai_client.send_omni_request(request_config, request_num=1) # for omni-understanding models
# or
openai_client.send_diffusion_request(request_config, request_num=1) # for diffusion models
openai_client.send_request(request_config, request_num=1)
```
**Explanation**:
@@ -519,7 +517,7 @@ L3 level testing executes after code is merged into the main branch. Its core pu
**Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n.
**Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
**Implicit Validation**: The send_request method internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
- ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
+9 -30
View File
@@ -17,7 +17,6 @@ Defined in `pyproject.toml`:
| `gpu` | Tests that run on GPU * |
| `cuda` | Tests that run on CUDA * |
| `rocm` | Tests that run on AMD/ROCm * |
| `xpu` | Tests that run on Intel XPU * |
| `npu` | Tests that run on NPU/Ascend * |
| `H100` | Tests that require H100 GPU * |
| `L4` | Tests that require L4 GPU * |
@@ -33,7 +32,7 @@ Defined in `pyproject.toml`:
| `slow` | Slow tests (may skip in quick CI) |
| `benchmark` | Benchmark tests |
\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility).
\* Means those markers are auto-added, and they will be added by the `@hardware_test` decorator.
### Example usage for markers
@@ -50,13 +49,12 @@ from tests.utils import hardware_test
def test_video_to_audio()
...
```
### Decorator: `@hardware_test`
This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions:
1. **Applies platform and resource markers**
Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`).
Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `A2`, `A3`).
```
@pytest.mark.cuda
@pytest.mark.L4
@@ -81,14 +79,14 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
- Single call for multiple platforms:
```python
@hardware_test(
res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
num_cards={"cuda": 2, "rocm": 2, "xpu": 2, "npu": 2},
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
num_cards={"cuda": 2, "rocm": 2, "npu": 2},
)
```
or
```python
@hardware_test(
res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
num_cards=2,
)
```
@@ -100,24 +98,6 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
- ROCm only: `pytest -m "distributed_rocm and MI325"`
- NPU only: `pytest -m "distributed_npu"`
### Function: `hardware_marks`
`hardware_marks` returns a list of pytest mark objects with the same signature as `@hardware_test`. Use it when you need more flexibility, such as attaching hardware marks to individual `pytest.param` entries rather than an entire test function.
```python
from tests.utils import hardware_marks
MULTI_CARD_MARKS = hardware_marks(
res={"cuda": "H100", "rocm": "MI325", "npu": "A2"}, num_cards=2
)
@pytest.mark.parametrize("omni_server", [
pytest.param(OmniServerParams(...), id="case_001", marks=MULTI_CARD_MARKS),
], indirect=True)
def test_feature(omni_server):
...
```
## Add Support for a New Platform
If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps:
@@ -152,13 +132,12 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
# Optionally: add skipif_tpu when implemented
return [test_platform, test_resource, test_distributed]
```
3. **Update `hardware_marks` to recognize your new platform**:
In the relevant place (see the `hardware_marks` implementation), add:
3. **Update `hardware_test` to recognize your new platform**:
In the relevant place (see the `hardware_test` implementation), add:
```python
if platform == "tpu":
marks = tpu_marks(res=resource, num_cards=cards)
```
(`hardware_test` calls `hardware_marks` internally, so both will pick up the change.)
4. **(Recommended) Add a test using your new markers**:
```python
@hardware_test(
@@ -172,7 +151,7 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
**Summary**:
- Add pytest markers for your new platform/resources
- Implement a marker function (`xxx_marks`)
- Plug into `hardware_marks`
- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation!
- Plug into `hardware_test`
- You're done: tests decorated with `@hardware_test` using your platform now automatically get the correct markers, distribution, and isolation!
See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).
-1
View File
@@ -158,7 +158,6 @@ vllm_omni/ tests/
├── test_diffusion_layerwise_offload.py
├── test_diffusion_lora.py
├── test_sequence_parallel.py
├── test_qwen_image_edit_expansion.py
└── stage_configs/
├── qwen2_5_omni_ci.yaml
├── qwen3_omni_ci.yaml
@@ -651,27 +651,8 @@ For a fair comparison, keep the same **prompt**, **seed**, **resolution**, **num
#### 5.2 Functionality Check in CI
To ensure project maintainability and sustainable development, please submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
For comprehensive testing guidelines and the definition of test levels (L1-L5), please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
The following tests are required to add:
- L4 test of the model's full *functionality* (i.e., all the *diffusion features* that are supported by this model), including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
- Test cases: Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`. The following test cases shall cover all features:
- 1 GPU: TeaCache & GGUF (or fallback to FP8, or disable it) & Layer-wise CPU offloading (or fallback to Module-wise)
- 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Ulysses = 2
- 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & Ring = 2
- 2 GPUs: TeaCache & FP8 (or fallback to GGUF, or disable it) & CFG Parallel = 2
- 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Tensor Parallel = 2 & VAE Patch Parallel = 2
- 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & HSDP = 2 & VAE Patch Parallel = 2
- Validation: test that the multimodal output files of your model have the correct shapes.
- Test marks: always add `advanced_model` and `diffusion`. Add `parallel` and GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/)
- To maximize code reuse, you may refer to
- `tests/conftest.py` for `omni_server` and `openai_client` fixtures, `generate_synthetic_image` and `assert_XXX_valid` helper.
- `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
- [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
- Doc: add a concise dostring for each test function.
- Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).
To ensure project maintainability and sustainable development, we encourage contributors to submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
---
+1 -8
View File
@@ -19,7 +19,7 @@ The `async_chunk` feature enables asynchronous, chunked processing of data acros
For qwen3-omni:
- **Thinker → Talker**: Per decode step (typically chunk_size=1)
- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. During the initial phase, a dynamic initial chunk size (IC) is automatically selected based on server load to reduce TTFA. Use the per-request `initial_codec_chunk_frames` API field to override.
- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. Set `initial_codec_chunk_frames` to emit smaller chunks during the initial phase for reduced TTFA
- **Code2Wav**: Streaming decode with code2wav chunk_size
With `async_chunk`:
@@ -67,13 +67,6 @@ Enabling **async_chunk** (False→True) sharply reduces time-to-first-audio (TTF
</picture>
</p>
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png">
<img alt="E2E Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png" width=100%>
</picture>
</p>
## Architecture
### Data Flow
@@ -72,241 +72,3 @@ Parameters:
For more details, refer to the
[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
---
## Design
### 1. Overview
`MooncakeStoreConnector` is the storage-oriented remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake distributed store APIs and provides a uniform `put()` / `get()` abstraction for transferring stage payloads across nodes.
Compared with `SharedMemoryConnector`, this connector is intended for multi-node deployments. Compared with `MooncakeTransferEngineConnector`, it is simpler and more object-oriented: data is written into a distributed store and fetched back by key, rather than transferred through an explicit remote-memory write protocol.
Its primary role is to provide a general-purpose remote connector for stage payloads when:
- the pipeline spans multiple hosts
- a shared-memory connector is not possible
- a simpler remote backend is preferred over the RDMA-oriented transfer engine
### 2. Relationship with the OmniConnector System
`MooncakeStoreConnector` implements `OmniConnectorBase`, so it integrates with the same connector lifecycle as the other backends:
- `OmniConnectorFactory` instantiates it from `ConnectorSpec`
- `load_omni_transfer_config()` maps edge config to the connector spec
- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
This means the rest of the pipeline does not need store-specific logic. It only relies on the generic connector contract.
### 3. Design Goals
The connector is designed around the following goals:
- **Cross-node object transfer** using a single key-based abstraction
- **Minimal connector-specific control plane**, since the store itself is the shared medium
- **Compatibility with arbitrary Python payloads** through the Omni serializer
- **Operational simplicity** compared with the transfer-engine-based remote connector
The connector is intentionally not optimized for zero-copy tensor movement or direct remote-memory access.
### 4. Core Design
#### 4.1 Object-Oriented Transfer Model
`MooncakeStoreConnector` treats the transport backend as a distributed object store:
1. Serialize the Python object into bytes.
2. Generate a unique connector key.
3. Store the bytes in Mooncake.
4. Fetch the bytes from Mooncake on the receiver side.
5. Deserialize the bytes back into the original object.
This design makes the connector conceptually close to a distributed key-value transport.
#### 4.2 Key Construction
The connector uses `OmniConnectorBase._make_key()` to derive the internal store key:
```text
{request_key}@{from_stage}_{to_stage}
```
This adds stage routing information to the logical request key and avoids collisions between different pipeline edges.
#### 4.3 No Extra Metadata Path
Unlike `SharedMemoryConnector` and `MooncakeTransferEngineConnector`, this connector returns:
```python
metadata = None
```
from `put()`.
This is a meaningful design difference:
- the store itself is the shared rendezvous point
- consumers only need the same key
- no transport handle, address, or side-channel metadata needs to be propagated
As a result, the control plane is simpler for this connector than for the SHM and RDMA variants.
### 5. Initialization
#### 5.1 Required Mooncake Components
The connector requires the Mooncake Python bindings to expose:
- `MooncakeDistributedStore`
- `ReplicateConfig`
If these symbols are unavailable, construction fails immediately with `ImportError`. This keeps startup failures explicit and avoids silent fallback behavior.
#### 5.2 Store Setup
During `_init_store()`, the connector:
1. creates `MooncakeDistributedStore()`
2. calls `store.setup(...)`
3. validates the return code
4. creates a `ReplicateConfig`
5. enables `with_soft_pin = True`
The setup step is completed during connector construction, not lazily at first use.
### 6. Put / Get Flow
#### 6.1 Producer Flow: `put()`
The producer-side flow is:
1. Validate that the store has been initialized.
2. Serialize the payload into bytes.
3. Build the internal key with stage routing information.
4. Call `store.put(key, serialized_data, self.pin)`.
5. Update metrics and return success.
The returned tuple is:
```python
(True, len(serialized_data), None)
```
#### 6.2 Consumer Flow: `get()`
The consumer-side flow is:
1. Validate that the store has been initialized.
2. Build the same internal key.
3. Poll `store.get(key)` for up to 20 retries.
4. Sleep 50 ms between retries.
5. If data is found, deserialize it and return `(data, payload_size)`.
6. If all retries are exhausted, record a timeout and return `None`.
This design gives the connector a bounded waiting model rather than an indefinite blocking get.
### 7. Integration with Stage Communication
All callers use the connector through the same `put()` / `get()` contract:
- the sender calls `put()` to serialize and store the payload
- the receiver calls `get()` to retrieve and deserialize it
- no connector-specific metadata is required, since the store key is the rendezvous point
Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, which makes the connector functional but not the highest-performance option for large payloads such as KV cache blocks.
### 8. Data Flow in the Pipeline
The end-to-end transfer model is:
```mermaid
sequenceDiagram
participant SenderStage
participant MooncakeStoreConnector
participant MooncakeStore
participant ReceiverStage
SenderStage->>MooncakeStoreConnector: put(from_stage, to_stage, put_key, data)
MooncakeStoreConnector->>MooncakeStoreConnector: serialize object
MooncakeStoreConnector->>MooncakeStore: put(internal_key, bytes)
MooncakeStoreConnector-->>SenderStage: (success, size, None)
ReceiverStage->>MooncakeStoreConnector: get(from_stage, to_stage, get_key)
MooncakeStoreConnector->>MooncakeStore: get(internal_key)
MooncakeStoreConnector->>MooncakeStoreConnector: deserialize bytes
MooncakeStoreConnector-->>ReceiverStage: (data, size)
```
This is a store-mediated remote transfer model rather than a direct peer-to-peer transport model.
### 9. Strengths and Trade-offs
#### Strengths
- Simple conceptual model: store by key, fetch by key.
- No connector-specific metadata handoff is required.
- Works naturally for remote multi-node stage transfer.
- Easy to integrate into the existing connector abstraction.
#### Trade-offs
- Full serialization and deserialization are always required.
- Large payloads are more expensive than in direct-memory transports.
- Runtime behavior depends on external Mooncake services being available.
- Cleanup semantics are weaker than request-scoped local buffer management.
### 10. Important Implementation Characteristics
#### 10.1 Cleanup Is a No-op
`cleanup()` only logs a debug message and does not actively delete remote data. The current design assumes Mooncake-side lifecycle management rather than explicit per-request removal.
This keeps the connector implementation simple but means request-level reclamation is not modeled inside the connector itself.
#### 10.2 Close Releases the Store Handle
Unlike `SharedMemoryConnector`, where `close()` is a no-op, `MooncakeStoreConnector.close()` performs a meaningful shutdown:
1. Calls `self.store.close()` to release the Mooncake store handle.
2. Sets `self.store = None` to prevent further operations.
This ensures the connector releases its connection to the external Mooncake service on shutdown. Errors during close are logged but do not propagate.
#### 10.3 Health Output Uses a Shared Schema
`health()` reports:
- `host`
- `metadata_server`
- `master`
- metrics
and also includes placeholder fields such as `protocol`, `pool_device`, and `pool_size` to keep a more uniform shape across connector health outputs. This reflects a system-level consistency choice rather than a store-specific design need.
#### 10.4 Get Is Retry-Based, Not Event-Driven
The receiver side polls the store with a short retry loop. This is simple and robust, but it means the connector is latency-sensitive to:
- store visibility delay
- network jitter
- payload size
If the payload is not visible within the retry window, the connector reports a timeout.
### 11. Summary
`MooncakeStoreConnector` is the remote, store-based connector in the OmniConnector stack. Its design is straightforward:
- serialize the payload
- store it in Mooncake under a stage-qualified key
- fetch it by the same key on the receiver side
- deserialize it back into the original object
This connector fills an important gap in the system:
1. It enables cross-node transfer without requiring shared memory.
2. It keeps the stage communication model uniform.
3. It provides a simpler operational alternative to the transfer-engine-based remote connector.
Its simplicity is also its main trade-off: it favors a clean object-transfer model over the fast-path and peer-to-peer optimizations implemented by `MooncakeTransferEngineConnector`.
@@ -205,589 +205,3 @@ for test-specific setup instructions.
For more details on the underlying engine, refer to the
[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
---
## Design
### 1. Overview
`MooncakeTransferEngineConnector` is the high-performance remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake `TransferEngine` and combines:
- a **direct data plane** for remote memory writes
- a **ZMQ side channel** for metadata lookup, handshake, and completion signaling
- a **managed local memory pool** for both send and receive buffers
Unlike `MooncakeStoreConnector`, which treats the backend as a distributed store, `MooncakeTransferEngineConnector` is designed as a peer-to-peer transport. Its goal is to move large stage payloads efficiently while still fitting the common `put()` / `get()` API defined by `OmniConnectorBase`.
It is the most performance-oriented connector in the current OmniConnector family and is intended for large remote payloads such as:
- KV cache transfer
- stage hidden-state payloads
- streaming chunk payloads
- other binary-heavy inter-stage artifacts
### 2. Relationship with the OmniConnector System
`MooncakeTransferEngineConnector` implements the same connector contract as the other backends:
- `put(from_stage, to_stage, put_key, data)`
- `get(from_stage, to_stage, get_key, metadata=None)`
- `cleanup(request_id, ...)`
- `health()`
- `close()`
It is integrated into the system through the standard connector plumbing:
- `OmniConnectorFactory` constructs the connector from `ConnectorSpec`
- `load_omni_transfer_config()` resolves the edge-level connector configuration
- `get_connectors_config_for_stage()` and `resolve_omni_kv_config_for_stage()` inject the connector role
- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
The key system-level distinction is that this connector is **role-aware**:
- sender instances expose data and listen for pull requests
- receiver instances allocate buffers and actively pull data from the sender
### 3. Design Goals
The connector is designed around four primary goals:
1. **High-throughput remote transfer**
Avoid store-mediated round trips and write directly into the receiver memory region.
2. **Fast path for raw payloads**
Support `torch.Tensor`, `bytes`, and `ManagedBuffer` without forcing all traffic through full object serialization.
3. **Unified connector abstraction**
Preserve the same `put()` / `get()` interface used by the rest of the OmniConnector stack.
4. **Safe lifecycle management**
Manage allocation, reuse, cleanup, and failure recovery for a registered memory pool.
### 4. Architecture Overview
At a high level, the connector is composed of four main subsystems:
```mermaid
classDiagram
class OmniConnectorBase {
<<abstract>>
+put(from_stage, to_stage, put_key, data)
+get(from_stage, to_stage, get_key, metadata)
+cleanup(request_id)
+health()
+close()
}
class MooncakeTransferEngineConnector {
+supports_raw_data: bool
-engine: TransferEngine
-allocator: BufferAllocator
-pool: torch.Tensor
-zmq_ctx: zmq.Context
-_local_buffers: dict
-_sender_executor: ThreadPoolExecutor
-_listener_thread: threading.Thread
+put(...)
+get(...)
+update_sender_info(sender_host, sender_zmq_port)
+get_connection_info()
+cleanup(request_id, from_stage, to_stage)
+close()
}
class BufferAllocator {
-total_size: int
-alignment: int
-free_blocks: list
+alloc(size) int
+free(offset, size)
}
class ManagedBuffer {
-allocator: BufferAllocator
-offset: int
-size: int
-pool_tensor: torch.Tensor
+tensor
+as_tensor(dtype, shape) torch.Tensor
+to_bytes() bytes
+release()
}
class TransferEngine {
+initialize(host, handshake, protocol, device_name)
+register_memory(base_ptr, size)
+batch_transfer_sync_write(remote_session, src_addrs, dst_addrs, lengths)
+unregister_memory(base_ptr)
+get_rpc_port() int
}
class QueryRequest {
+request_id: str
}
class QueryResponse {
+request_id: str
+data_size: int
+is_fast_path: bool
}
class MooncakeAgentMetadata {
+remote_hostname: str
+remote_port: int
+request_id: str
+dst_addrs: list[int]
+lengths: list[int]
}
OmniConnectorBase <|-- MooncakeTransferEngineConnector
MooncakeTransferEngineConnector *-- BufferAllocator
MooncakeTransferEngineConnector *-- TransferEngine
ManagedBuffer --> BufferAllocator : releases to
ManagedBuffer --> "1" torch.Tensor : views
MooncakeTransferEngineConnector ..> ManagedBuffer : returns / retains
MooncakeTransferEngineConnector ..> QueryRequest : decodes
MooncakeTransferEngineConnector ..> QueryResponse : encodes
MooncakeTransferEngineConnector ..> MooncakeAgentMetadata : exchanges
```
#### 4.1 Transfer Engine
Mooncake `TransferEngine` is responsible for the actual data-plane transfer. It registers local memory and performs synchronous remote writes through:
```python
batch_transfer_sync_write(...)
```
#### 4.2 Managed Memory Pool
Each connector instance owns a large pre-registered memory pool:
- CPU pinned memory when `memory_pool_device == "cpu"`
- GPU memory when `memory_pool_device == "cuda"`
This avoids repeated memory registration and allows each transfer to allocate subranges from one long-lived pool.
#### 4.3 Buffer Manager
Two helper classes control local memory ownership:
- `BufferAllocator`
Manages aligned subrange allocation and free-list merging.
- `ManagedBuffer`
Represents one live slice of the pool and exposes:
- `.tensor`
- `.as_tensor(dtype, shape)`
- `.to_bytes()`
- `.release()`
#### 4.4 ZMQ Side Channel
ZMQ is used for transport coordination, not for the data payload itself. It handles:
- metadata query from receiver to sender
- pull request submission
- completion or error signaling
- internal notification from worker threads back to the listener thread
This split makes the control plane lightweight while keeping the bulk payload on the transfer engine data plane.
### 5. Role Model
#### 5.1 Sender Role
A sender connector:
- accepts `put()` calls
- stores live transfer-ready buffers in `_local_buffers`
- starts a ZMQ listener thread
- responds to metadata queries and pull requests from receivers
#### 5.2 Receiver Role
A receiver connector:
- does not bind the sender-side ZMQ listener
- accepts `get()` calls
- allocates receive buffers from its own pool
- requests metadata or transfer service from the sender
The role is not inferred dynamically. It is injected by the stage configuration layer:
- incoming edge for a stage -> `role="receiver"`
- outgoing edge for a stage -> `role="sender"`
This is important because incorrect role assignment would break initialization semantics.
#### 5.3 Host Auto-Detection
The `host` configuration field supports the special value `"auto"`. When set, the connector auto-detects the local IP address that would be used for external communication (via a UDP socket probe to `8.8.8.8`). If that fails, it falls back to hostname resolution, and ultimately to `127.0.0.1`.
This is useful in environments where the operator does not want to hard-code IP addresses in the connector config.
#### 5.4 RDMA Device Filtering
The `device_name` configuration field allows the operator to specify which RDMA NICs to use (comma-separated, e.g. `"mlx5_0,mlx5_1"`). If not set in config, the connector also checks the `RDMA_DEVICE_NAME` environment variable.
This is important in environments with mixed InfiniBand/RoCE NICs, where not all devices are suitable for the transfer engine.
### 6. Local Memory Management
#### 6.1 Memory Pool Registration
During initialization, the connector:
1. allocates a large pool tensor
2. records its base pointer
3. registers that memory with Mooncake
4. creates a `BufferAllocator` for subrange management
This means every later transfer only allocates offsets inside the pre-registered pool rather than registering memory per request.
#### 6.2 BufferAllocator
`BufferAllocator` maintains a sorted free list of `(offset, size)` blocks and enforces alignment. Its responsibilities include:
- aligned allocation
- freeing previously allocated blocks
- adjacent block merging
- double-free detection
- overlap detection to catch corruption
This is a critical piece of the connector because both sender and receiver depend on long-lived pool reuse.
#### 6.3 ManagedBuffer
`ManagedBuffer` is the main fast-path data wrapper. It can:
- expose the pool slice as a zero-copy 1D `uint8` tensor
- reinterpret that slice as a typed tensor
- copy out the contents as Python `bytes`
- release the slice back to the allocator
The connector uses `ManagedBuffer` in two different ways:
- as a send-side holder to keep the pool slice alive
- as a receive-side return type when `is_fast_path=True`
### 7. Put Flow
#### 7.1 High-Level Behavior
`put()` is only valid in sender mode. Its job is to expose a payload for later remote pull by the receiver.
The high-level flow is:
1. validate connector state and role
2. convert the input into a pool-backed transferable representation
3. store the transfer metadata in `_local_buffers`
4. return lightweight metadata describing how the receiver can fetch the data
#### 7.2 Payload Type Handling
`put()` supports three payload classes:
**A. `ManagedBuffer`**
If the buffer belongs to the same pool, the connector can use it directly without copying. This is the most efficient path.
If the buffer comes from a different pool, the connector falls back to a copy path.
**B. `torch.Tensor` or `bytes`**
These payloads are copied into the local pool and marked as fast-path data:
- no Omni object serialization is required
- receiver can get a `ManagedBuffer` back
**C. Generic Python object**
Any other payload is serialized via `OmniSerializer.serialize(...)` and then copied into the pool.
In this case:
- `is_fast_path=False`
- the receiver will deserialize back into a Python object
#### 7.3 Sender Metadata
The sender returns:
```python
{
"source_host": self.host,
"source_port": self.zmq_port,
"data_size": size,
"is_fast_path": is_fast_path,
}
```
This metadata is intentionally lightweight. It tells the receiver:
- where the sender-side control plane lives
- how large the remote transfer will be
- whether the payload should be returned as a `ManagedBuffer` or a deserialized object
#### 7.4 Sender Buffer Table
The sender stores each live payload in `_local_buffers` under the stage-qualified key. Each entry contains:
- source addresses
- lengths
- the holder object
- ownership information (`should_release`)
- `is_fast_path`
- creation time
This table is the sender-side truth source for both metadata queries and pull requests.
### 8. Get Flow
#### 8.1 High-Level Behavior
`get()` runs on the receiver side and performs four steps:
1. resolve metadata
2. allocate a destination buffer in the local pool
3. request the sender to write into that destination buffer
4. return either a `ManagedBuffer` or a deserialized object
#### 8.2 Metadata Resolution Paths
The `metadata` parameter in `get()` is optional. The connector supports two resolution modes depending on whether the caller supplies it.
**With metadata**
When the caller passes metadata, the connector uses it directly. The metadata carries:
- `source_host` / `source_port` — sender ZMQ endpoint
- `data_size` — payload byte count
- `is_fast_path` — whether the receiver gets a `ManagedBuffer` or a deserialized object
This mode is suitable when the control plane already forwards the sender's `put()` output to the receiver.
**Without metadata**
When `get(metadata=None)` is called, the connector queries the sender over ZMQ to discover the same fields (`data_size`, `is_fast_path`). The caller must first call:
```python
update_sender_info(sender_host, sender_zmq_port)
```
so that the connector knows where to send the query.
This mode is suitable for polling-based flows (e.g. KV transfer, async chunk transfer) where the receiver does not have metadata from the control plane.
#### 8.3 Destination Allocation
Once metadata is resolved, the receiver:
1. allocates a subrange from its own local pool
2. wraps it in a `ManagedBuffer`
3. builds a `MooncakeAgentMetadata` request containing:
- receiver hostname
- receiver RPC port
- request ID
- destination addresses
- transfer lengths
This tells the sender exactly where to write the incoming data.
#### 8.4 Transfer Completion
The receiver then sends the pull request over ZMQ and waits for:
- `TRANS_DONE`
- or `TRANS_ERROR`
If the transfer succeeds:
- for `is_fast_path=True`, the receiver returns `(ManagedBuffer, size)`
- for `is_fast_path=False`, the receiver copies to bytes, deserializes, releases the buffer, and returns `(object, size)`
### 9. Sender-Side Listener Design
#### 9.1 Listener Thread
In sender mode, the connector starts `_zmq_listener_loop()` after initialization. The listener:
- binds `tcp://{host}:{zmq_port}`
- receives incoming requests
- uses a poller for socket events and internal notifications
- periodically reclaims stale buffers
If the bind fails, initialization fails immediately. The code does not silently downgrade the connector role.
#### 9.2 Worker Thread Pool
The listener hands work to `_sender_executor` so that the listener thread does not block on transfer work.
There are two request types:
- metadata query -> `_handle_query_request(...)`
- transfer request -> `_handle_pull_request(...)`
#### 9.3 Query Handling
For metadata queries, the sender looks up the request ID in `_local_buffers` and returns:
- data size
- fast-path flag
This supports consumers that only know the sender endpoint but not the original sender metadata.
#### 9.4 Pull Handling
For a transfer request, the sender:
1. locates the source addresses in `_local_buffers`
2. constructs the remote session identifier
3. calls `batch_transfer_sync_write(...)`
4. replies `TRANS_DONE` or `TRANS_ERROR`
On success, the sender immediately calls `cleanup(meta.request_id)` and frees the producer-side buffer if it owns it.
This choice is important: it makes the connector effectively a **single-consumer transfer model** for each successful put/get pair.
### 10. Fast Path Semantics
This connector explicitly advertises:
```python
supports_raw_data = True
```
That means it can move raw payloads without forcing everything through the Omni object serializer.
#### Fast Path
For `torch.Tensor`, `bytes`, or pool-local `ManagedBuffer`:
- sender returns `is_fast_path=True`
- receiver returns a `ManagedBuffer`
- caller is responsible for calling `release()`
This avoids an unnecessary copy on the receiver side.
#### Serialized Object Path
For arbitrary Python objects:
- sender serializes the payload
- receiver converts the receive buffer to bytes
- receiver deserializes the object
- receive buffer is released internally
This preserves a uniform object-oriented API while still allowing optimized raw-data transport when possible.
### 11. Failure Handling and Cleanup
#### 11.1 Timeouts and Socket Recovery
The receiver caches ZMQ REQ sockets per thread, but invalidates them after failures. This avoids reusing sockets that may be stuck in a bad state after timeout or receive errors.
Timeout is scaled based on payload size:
- a base timeout
- plus additional time for large payloads
This is intended to reduce false timeouts for large remote writes.
#### 11.2 Stale Buffer Reclamation
The sender periodically reclaims old entries from `_local_buffers` using a TTL policy. This protects the memory pool from permanent leaks if a receiver crashes or never consumes a prepared payload.
This is a practical recovery mechanism, although the code notes that TTL cleanup can still race with very long-running in-flight transfers.
#### 11.3 Connector Shutdown
`close()` is a full resource teardown routine. It:
- stops the listener thread
- shuts down the worker executor
- releases all pending buffers
- closes cached sockets
- unregisters memory from the engine when supported
- terminates the ZMQ context
- drops the pool reference
This makes `MooncakeTransferEngineConnector` the most lifecycle-aware connector in the current connector family.
### 12. Current Implementation Constraints
The current code documents several important topology constraints.
#### 12.1 One Sender to One Receiver per Successful Transfer
After a successful transfer, the sender-side buffer is cleaned up immediately. This means the same prepared payload is not retained for multiple independent receivers.
#### 12.2 One Receiver to One Active Sender Endpoint
The receiver only stores one `(sender_host, sender_zmq_port)` pair through `update_sender_info(...)`. So the metadata-query mode is currently single-sender at a time.
#### 12.3 Explicit Buffer Ownership Matters
When the connector allocates a pool slice internally, it is responsible for releasing it. When a caller passes an externally owned `ManagedBuffer`, the connector keeps it alive for transfer but does not assume ownership of its eventual release.
These constraints are consistent with the current implementation and should be treated as design assumptions rather than incidental behavior.
### 13. Data Flow in the Pipeline
The end-to-end sender/receiver interaction is:
```mermaid
sequenceDiagram
participant SenderStage
participant SenderConnector
participant ReceiverConnector
participant ReceiverStage
SenderStage->>SenderConnector: put(from_stage, to_stage, put_key, data)
SenderConnector->>SenderConnector: place payload in local memory pool
SenderConnector-->>SenderStage: metadata(source_host, source_port, data_size, is_fast_path)
ReceiverStage->>ReceiverConnector: get(..., metadata)
ReceiverConnector->>ReceiverConnector: allocate destination buffer
ReceiverConnector->>SenderConnector: ZMQ pull request with dst addr
SenderConnector->>ReceiverConnector: TransferEngine remote write
SenderConnector-->>ReceiverConnector: TRANS_DONE
ReceiverConnector-->>ReceiverStage: ManagedBuffer or deserialized object
```
For metadata-less polling, the flow simply adds a metadata query step before the pull request.
### 14. Strengths and Trade-offs
#### Strengths
- Best remote-transfer design in the current connector stack for large payloads.
- Supports raw-data fast path.
- Keeps stage communication under the same connector abstraction.
- Includes real lifecycle and memory-pool management.
- Works for both stage payload transfer and KV transfer scenarios.
#### Trade-offs
- More complex than the store-based connector.
- Correctness depends on role injection and endpoint coordination.
- Caller must release fast-path receive buffers.
- Current implementation is optimized for single-consumer transfer semantics.
### 15. Summary
`MooncakeTransferEngineConnector` is the high-performance peer-to-peer transport in the OmniConnector system. Its design combines:
- a registered memory pool
- a safe subrange allocator
- a ZMQ control plane
- a Mooncake transfer-engine data plane
This allows the connector to support both:
1. a **fast path** for raw tensors and bytes
2. a **generic object path** for arbitrary Python payloads
Within vLLM-Omni, it is the connector that most directly targets performance-sensitive remote transfer, especially for large payloads and KV cache movement. Its additional complexity is deliberate: it is the connector that turns the generic OmniConnector abstraction into a transport capable of efficient remote memory movement rather than simple object storage.
@@ -7,10 +7,9 @@ auto-configured when no explicit connector is specified for an edge.
## How It Works
All payloads are serialized and stored in shared memory (`/dev/shm`); the SHM
segment name is returned in metadata. The configuration exposes a
`shm_threshold_bytes` field for a future inline-vs-SHM split, but the current
implementation always uses shared memory regardless of payload size.
- Small payloads (< threshold): serialized and passed inline in metadata.
- Large payloads (>= threshold): stored in shared memory; the SHM name is
returned in metadata.
## Configuration
@@ -26,234 +25,3 @@ runtime:
## Notes
- Auto-mode uses SharedMemoryConnector if no connector is declared for an edge.
---
## Design
### 1. Overview
`SharedMemoryConnector` is the default same-node connector in `vllm_omni/distributed/omni_connectors`. It is designed for stage-to-stage transfer when producer and consumer processes run on the same host and can share `/dev/shm`.
The connector provides a unified `put()` / `get()` API for arbitrary Python objects while keeping the control plane lightweight:
- The payload is serialized by the connector.
- The serialized bytes are placed in shared memory.
- The queue/control plane only carries a small metadata handle.
This makes `SharedMemoryConnector` the simplest connector in the OmniConnector family and the default fallback when an edge does not explicitly configure another backend.
### 2. Relationship with the OmniConnector System
`SharedMemoryConnector` implements `OmniConnectorBase`, so it follows the same lifecycle and API contract as the other connectors:
- `put(from_stage, to_stage, put_key, data)`
- `get(from_stage, to_stage, get_key, metadata=None)`
- `cleanup(request_id)`
- `health()`
- `close()`
Within the larger system:
- `load_omni_transfer_config()` automatically fills missing edges with `SharedMemoryConnector`.
- Callers interact with the connector exclusively through the `put()` / `get()` / `cleanup()` contract — the connector does not require caller-specific logic.
Compared with the remote Mooncake-based connectors, `SharedMemoryConnector` is intentionally minimal and local-only.
### 3. Design Goals
The connector is built around the following goals:
- **Low-friction local transfer** for single-node multi-process pipelines.
- **Unified object semantics** for arbitrary Python payloads.
- **Small control-plane overhead** by passing only metadata through queues.
- **Zero external dependencies** beyond Python shared memory and the existing stage utilities.
It is not intended to provide cross-node transfer, RDMA, or raw tensor zero-copy semantics across processes.
### 4. Core Design
#### 4.1 Serialization Model
`SharedMemoryConnector` always starts from a Python object and serializes it through the shared Omni serializer:
```python
payload = self.serialize_obj(data)
```
This keeps the connector behavior consistent with the rest of the connector stack:
- producer code does not need connector-specific serialization logic
- consumer code always receives the original object after deserialization
- the connector can reuse the same serializer used by other backends
#### 4.2 Shared Memory as the Data Plane
The actual data plane is a shared-memory segment created by:
- `shm_write_bytes(...)`
- `shm_read_bytes(...)`
The connector stores a small metadata object such as:
```python
{
"shm": {"name": ..., "size": ...},
"size": ...
}
```
This metadata is passed over the control plane and allows the downstream stage to locate the shared-memory segment.
#### 4.3 Locking Model
To avoid races between the producer and consumer, the connector uses a lock file per request:
```text
/dev/shm/shm_{put_key}_lockfile.lock
```
Locking is done with `fcntl.flock`:
- producer uses `LOCK_EX`
- consumer uses `LOCK_EX`
Both sides acquire an exclusive lock. This ensures that the shared-memory segment is not read while it is still being written and makes the handoff safer in a multi-process environment.
### 5. Put / Get Flow
#### 5.1 Producer Flow: `put()`
The producer-side flow is:
1. Serialize the input object to bytes.
2. Compute the payload size.
3. Acquire the per-request lock file.
4. Write the bytes into shared memory.
5. Return lightweight metadata to the caller.
The returned tuple is:
```python
(success, serialized_size, metadata)
```
where `metadata` contains the shared-memory handle needed by the consumer.
#### 5.2 Consumer Flow: `get(metadata=...)`
The primary consumer path is metadata-driven:
1. Extract the shared-memory handle from `metadata`.
2. Acquire the exclusive lock.
3. Read the raw bytes from shared memory.
4. Deserialize the bytes back into the original Python object.
5. Remove the lock file if it still exists.
This is the path used by the current stage-to-stage connector flow.
#### 5.3 Compatibility Flow: `get(metadata=None)`
The connector also keeps a compatibility path for callers that only know the key:
1. Attempt to open the shared-memory segment by name via `SharedMemory(name=get_key)`.
2. If the segment exists and has non-zero size, acquire the exclusive lock and read the bytes.
3. Deserialize the bytes and return the object.
If the segment does not exist or any exception occurs, the call returns `None` immediately. There is no retry loop in this path -- it is a single-attempt open.
This path is mainly for older code paths and is not the preferred mode for the current connector pipeline.
### 6. Key Implementation Characteristics
#### 6.1 Threshold Exists, but the Current Code Always Uses SHM
The class keeps a `shm_threshold_bytes` field and still exposes metrics for inline writes. However, the current implementation uses:
```python
if True:
...
```
inside `put()`, which means the current code path always writes to shared memory.
So the design still suggests a future split between:
- small payloads inline
- large payloads in shared memory
but the current behavior is effectively:
- all payloads go through shared memory
This should be documented because it affects real runtime behavior.
#### 6.2 Cleanup Is Currently Passive
`cleanup()` is currently a no-op. The intended assumption is:
- the consumer reads the segment
- the underlying shared-memory helpers unlink it
If the consumer never executes `get()`, the shared-memory segment may remain allocated. This means the connector relies on the normal success path for resource reclamation.
#### 6.3 Close Is Currently Minimal
`close()` is also a no-op. There is no connector-owned background thread, socket, or memory pool to tear down, so the lifecycle is simple. The trade-off is that `close()` does not scan or recover leaked shared-memory resources.
### 7. Data Flow in the Pipeline
The typical flow with `SharedMemoryConnector` is:
```mermaid
sequenceDiagram
participant SenderStage
participant SharedMemoryConnector
participant QueueOrControlPlane
participant ReceiverStage
SenderStage->>SharedMemoryConnector: put(from_stage, to_stage, put_key, data)
SharedMemoryConnector->>SharedMemoryConnector: serialize object
SharedMemoryConnector->>SharedMemoryConnector: write bytes to /dev/shm
SharedMemoryConnector-->>SenderStage: metadata {shm, size}
SenderStage->>QueueOrControlPlane: forward connector metadata
QueueOrControlPlane->>ReceiverStage: task + connector metadata
ReceiverStage->>SharedMemoryConnector: get(from_stage, to_stage, get_key, metadata)
SharedMemoryConnector->>SharedMemoryConnector: read bytes from /dev/shm
SharedMemoryConnector->>SharedMemoryConnector: deserialize object
SharedMemoryConnector-->>ReceiverStage: (data, size)
```
This is a classic split-control-plane / data-plane design, but constrained to a single host.
### 8. Strengths and Trade-offs
#### Strengths
- Very simple deployment model.
- No external service dependency.
- Fits naturally into the existing queue-driven orchestration flow.
- Good default for local multi-process pipelines.
#### Trade-offs
- Same-node only.
- Full object serialization and deserialization are still required.
- Resource cleanup depends on the normal consumer path.
- Shared memory capacity is limited by host configuration.
### 9. Summary
`SharedMemoryConnector` is the baseline local transport for the OmniConnector system. Its design is intentionally straightforward:
- serialize object
- place bytes in shared memory
- pass metadata through the control plane
- deserialize on the receiving side
It plays two important roles in vLLM-Omni:
1. It is the simplest production-ready connector for same-node stage pipelines.
2. It serves as the automatic fallback connector when no explicit edge transport is configured.
Although the current implementation is deliberately minimal, it provides the foundation for reliable local connector semantics and keeps the stage communication model uniform across the system.
@@ -94,265 +94,8 @@ stage_args:
Parameters:
- host: datasystem worker host.
- port: datasystem worker port (default: `35001` if omitted; the example above uses `31501` to match the worker startup command).
- port: datasystem worker port.
- get_sub_timeout_ms: get timeout in milliseconds (0 for no timeout).
For more details, refer to the
[Yuanrong Datasystem repository](https://atomgit.com/openeuler/yuanrong-datasystem).
---
## Design
### 1. Overview
`YuanrongConnector` is the Datasystem-based remote connector in `vllm_omni/distributed/omni_connectors`. It uses Yuanrong Datasystem's distributed key-value client as the transport backend and exposes the same `put()` / `get()` interface as the other OmniConnectors.
Like `MooncakeStoreConnector`, it is a store-oriented remote connector rather than a direct peer-to-peer transport. Its role is to let stage payloads move across nodes through a deterministic key-based storage abstraction while keeping the rest of the pipeline on the common connector API.
It is intended for deployments that already use Yuanrong Datasystem and want a remote connector that integrates with the existing OmniConnector configuration and orchestration model.
### 2. Relationship with the OmniConnector System
`YuanrongConnector` implements `OmniConnectorBase`, so it participates in the same connector lifecycle as the other backends:
- `OmniConnectorFactory` constructs it from a `ConnectorSpec`
- stage edge configuration is resolved by `load_omni_transfer_config()`
- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
This means the connector is not exposed directly to stage logic. Stages only interact with the generic connector contract, and the backend choice remains a configuration concern.
### 3. Design Goals
The connector is built around the following goals:
1. **Cross-node payload transfer through Datasystem**
Reuse Yuanrong Datasystem as the remote exchange medium for stage data.
2. **Uniform object transfer semantics**
Allow arbitrary Python objects to be transmitted through the shared Omni serializer.
3. **Minimal connector-specific control plane**
Use deterministic keys so that consumers can retrieve data without an extra transport metadata handoff.
4. **Operational reuse of existing infrastructure**
Fit into environments that already deploy Yuanrong Datasystem workers and etcd.
The connector is not designed for direct remote-memory writes or tensor-specific fast-path transfer.
### 4. Core Design
#### 4.1 Store-Oriented Transfer Model
`YuanrongConnector` treats the transport backend as a distributed object store:
1. serialize the Python payload
2. build a deterministic connector key
3. write the serialized bytes into Datasystem
4. read the bytes back on the receiver side
5. deserialize them into the original object
This is the same broad architectural class as `MooncakeStoreConnector`, but implemented on top of Yuanrong Datasystem APIs instead of Mooncake store APIs.
#### 4.2 Deterministic Keying
Unlike the default `_make_key()` in `OmniConnectorBase`, `YuanrongConnector` defines its own key format:
```text
{request_id}:{from_stage}_{to_stage}
```
This has two design implications:
- the request identifier remains the primary lookup handle
- stage routing information is embedded in the key so that the same logical request ID can safely appear on different edges
The explicit override also makes the key format easier to align with Datasystem-side debugging and operational inspection.
#### 4.3 No Extra Metadata Hand-off
`put()` returns:
```python
(success, serialized_size, None)
```
and does not generate connector-specific metadata.
This design works because the receiver can reconstruct the exact same key from:
- `get_key`
- `from_stage`
- `to_stage`
As a result, the connector does not require a separate side-channel metadata handoff.
### 5. Initialization
#### 5.1 Datasystem Client Dependency
The connector requires the Datasystem Python bindings to expose:
- `KVClient`
- `SetParam`
- `WriteMode`
If any of these symbols are unavailable, connector construction fails immediately with `ImportError`. This keeps configuration errors explicit and avoids a partially initialized runtime.
#### 5.2 Client Setup
During `_init_client()`, the connector:
1. reads `host` and `port`
2. creates `KVClient(host, port)`
3. calls `client.init()`
At construction time it also creates a `SetParam` and fixes:
```python
self.set_param.write_mode = WriteMode.NONE_L2_CACHE_EVICT
```
This means the connector has a stable write policy for all writes and does not currently expose write-mode selection as a higher-level connector option.
### 6. Put / Get Flow
#### 6.1 Producer Flow: `put()`
The producer-side flow is:
1. verify that the Datasystem client has been initialized
2. serialize the input object with the shared Omni serializer
3. build the Datasystem key using the connector-specific `_make_key()`
4. call `client.set(key, serialized_data, self.set_param.write_mode)`
5. update metrics and return success
The returned metadata is always `None`, because the Datasystem key itself is the lookup contract between producer and consumer.
#### 6.2 Consumer Flow: `get()`
The consumer-side flow is:
1. verify that the Datasystem client has been initialized
2. rebuild the same key with `from_stage`, `to_stage`, and `get_key`
3. call:
```python
client.get([key], False, self.get_sub_timeout_ms)
```
4. take the first returned element if present
5. deserialize it and return `(data, payload_size)`
If the returned list is empty or contains no data for the key, `get()` returns `None`.
### 7. Timeout and Retrieval Semantics
The connector uses `get_sub_timeout_ms` as its read timeout. Unlike `MooncakeStoreConnector`, which performs an explicit retry loop in Python, `YuanrongConnector` delegates waiting behavior more directly to the Datasystem client call.
This leads to a slightly different retrieval model:
- `MooncakeStoreConnector`: retry-oriented polling in connector code
- `YuanrongConnector`: single client call with backend-managed timeout behavior
From the connector API perspective the result is the same, but operationally the waiting behavior is more dependent on Datasystem client semantics.
### 8. Integration with Stage Communication
All callers use the connector through the same `put()` / `get()` contract:
- the sender calls `put()` to serialize and store the payload
- the receiver calls `get()` to retrieve and deserialize it
- no connector-specific metadata is required, since the Datasystem key is the rendezvous point
Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, and there is no raw tensor fast path, which makes the connector functional but not optimized for the largest payloads.
### 9. Data Flow in the Pipeline
The end-to-end transfer model is:
```mermaid
sequenceDiagram
participant SenderStage
participant YuanrongConnector
participant Datasystem
participant ReceiverStage
SenderStage->>YuanrongConnector: put(from_stage, to_stage, put_key, data)
YuanrongConnector->>YuanrongConnector: serialize object
YuanrongConnector->>Datasystem: set(key, bytes)
YuanrongConnector-->>SenderStage: (success, size, None)
ReceiverStage->>YuanrongConnector: get(from_stage, to_stage, get_key)
YuanrongConnector->>Datasystem: get([key], timeout)
YuanrongConnector->>YuanrongConnector: deserialize bytes
YuanrongConnector-->>ReceiverStage: (data, size)
```
This is a store-mediated remote connector design with deterministic key lookup and no explicit side-channel metadata exchange.
### 10. Strengths and Trade-offs
#### Strengths
- Reuses existing Yuanrong Datasystem infrastructure.
- Keeps the connector contract simple and uniform.
- Requires no connector-specific metadata handoff.
- Suitable for remote stage transfer in Datasystem-based deployments.
#### Trade-offs
- Always pays full serialization and deserialization cost.
- Does not support raw bytes / tensor fast-path semantics.
- Depends on external Datasystem worker availability.
- Retrieval and timeout behavior are partly delegated to the backend client.
### 11. Important Implementation Characteristics
#### 11.1 Cleanup Is a No-op
`cleanup()` only logs a debug message and does not explicitly remove data from Datasystem. The current design assumes backend-side lifecycle or garbage collection rather than request-scoped delete semantics inside the connector.
This mirrors the same design trade-off seen in other store-based connectors: simplicity over explicit per-request data reclamation.
#### 11.2 Close Only Releases the Local Client Handle
`close()` does not perform a remote shutdown. It simply clears the local `client` reference and marks the connector as closed from the process perspective:
```python
self.client = None
```
This is appropriate for a client-based store connector, but it also means that backend resource lifecycle remains outside the connector's control.
#### 11.3 Error and Timeout Metrics Are Coarse-Grained
The connector tracks:
- `puts`
- `gets`
- `bytes_transferred`
- `errors`
- `timeouts`
In the current code, failed `put()` increments `errors`, while `get()` exceptions increment `timeouts`. This is operationally useful, but it does not distinguish between:
- backend timeout
- not-found result
- transport failure
- deserialization failure
So the metrics should be read as high-level health indicators, not detailed root-cause diagnostics.
### 12. Summary
`YuanrongConnector` is the Datasystem-backed remote connector in the OmniConnector family. Its design is straightforward:
- serialize payloads
- store them under a deterministic stage-qualified key
- retrieve them by the same key
- deserialize them on the receiving side
Within vLLM-Omni, it provides a clean integration point for Yuanrong Datasystem-based deployments while preserving the same connector abstraction used by the rest of the pipeline.
Its design priorities are simplicity, infrastructure reuse, and API consistency rather than direct-memory transport optimization.
+5 -35
View File
@@ -55,8 +55,6 @@ The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The
- Tensor operations happen at `nn.Module` boundaries
- Predictable sharding/gathering patterns
This is the ideal approach for integrating sequence parallelism into new models, as it is easier to maintain and ensure compatibility with other types of acceleration.
**How it works:**
1. Declare `_sp_plan` dict in your transformer class
2. Framework automatically applies hooks when `sequence_parallel_size > 1`
@@ -203,36 +201,6 @@ class TransformerWithRoPE(nn.Module):
}
```
**Pattern 3: Shard RoPE for Dual Stream Attention**
In some cases, different streams in attention may need to handle sequence parallelism differently. For example, we may want to shard the image embeddings, while replicating the text embeddings to correctly configure joint attention.
```python
class DualStreamTransformer(nn.Module):
"""
Dual-stream model where we need to replicate the text components, but shard
the image components to correctly handle sequence parallelism.
"""
_sp_plan = {
# In this case, the rope_preparer returns a tuple of len 4, where the
# first 2 items correspond to the text, and the second 2 correspond to
# visual inputs, so we only shard the second.
"rope_preparer": {
# Outputs 0, 1 (text) - NOT sharded (replicated)
# Outputs 2, 3 (image) - sharded
2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), # img_cos
3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), # img_sin
},
# Shard transformer block INPUT
"transformer_blocks.0": {
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
},
# Gather at output
"proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
}
```
NOTE: be careful to test adequately when refactoring classes that take this style of plan, as changing the order of the return values will break sequence parallelism.
### API Reference
**SequenceParallelInput Parameters:**
@@ -272,7 +240,7 @@ NOTE: be careful to test adequately when refactoring classes that take this styl
## Approach 2: Intrusive Modification (For Complex Cases)
For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls.
For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls. Importantly, when taking this approach, be careful to ensure that you correctly manage the `_sp_shard_depth`; if the sequence parallel shard depth is 0, Ulysses will not be used.
**When to use:**
@@ -285,15 +253,18 @@ from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
def forward(self, hidden_states, ...):
if self.parallel_config.sequence_parallel_size > 1:
# <Increment the _sp_shard_depth on the fwd context>
hidden_states = sp_shard(hidden_states, dim=1)
# ... computation ...
if self.parallel_config.sequence_parallel_size > 1:
output = sp_gather(output, dim=1)
# <Decrement the _sp_shard_depth on the fwd context>
return output
```
Note that currently, `sp_shard` / `sp_gather` do *not* automatically manage the `_sp_shard_depth`; you need to be careful to manage it yourself.
---
@@ -468,7 +439,6 @@ Complete examples in the codebase:
| Model | Path | Pattern | Notes |
|-------|------|---------|-------|
| **LongCat** | `vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py` | Dual-stream | Text components replicated, image components sharded |
| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Dual-stream + preprocessing | auto_pad, separate RoPE |
| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Dual-Transformer + RoPE | Video transformer |
| **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Unified sequence | Concatenated input |
@@ -483,7 +453,7 @@ Complete examples in the codebase:
Adding Sequence Parallel support to a transformer:
1. ✅ **Choose approach** - Use `_sp_plan` for standard cases, intrusive modification for complex cases
2. ✅ **Identify sharding boundaries** - Where should tensors be split/gathered? And which module boundaries need to be moved to facilitate this?
2. ✅ **Identify sharding boundaries** - Where should tensors be split/gathered?
3. ✅ **Extract inline operations** - Move `torch.cat`, `pad_sequence`, etc. to submodules
4. ✅ **Define `_sp_plan`** - Declare shard/gather points as class attribute
5. ✅ **Use `auto_pad` for variable lengths** - Support non-uniform sequences
+2 -3
View File
@@ -47,8 +47,7 @@ If no, check your shell running the ComfyUI process. There may be some error mes
This extension offers the following nodes based on the output modalities:
- **Generate Image** for text-to-image and image-to-image tasks
- **Generate Video** for text-to-video and image-to-video tasks
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
- **TTS** and **TTS Voice Clone** for TTS tasks
This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
@@ -60,7 +59,7 @@ To build a simple workflow yourself,
- Drag a generation node onto the canvas.
- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
- For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
- For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
-10
View File
@@ -132,19 +132,9 @@ class Example:
".dat",
".db",
".sqlite", # Data files
".pyc",
".pyo",
".pyd", # Python compiled
".npy",
".npz",
".pkl",
".pickle", # Serialized data
}
excluded_dirs = {"__pycache__", ".git", "node_modules", ".tox", ".mypy_cache"}
def is_other_file(file: Path) -> bool:
if any(part in excluded_dirs for part in file.parts):
return False
return file.is_file() and file != self.main_file and file.suffix.lower() not in binary_extensions
return [file for file in self.path.rglob("*") if is_other_file(file)]
-4
View File
@@ -49,9 +49,6 @@ th {
|`GlmImageForConditionalGeneration` | GLM-Image | `zai-org/GLM-Image` |
|`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` |
|`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` |
|`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
|`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` |
|`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` |
## List of Supported Models for NPU
@@ -67,7 +64,6 @@ th {
|--------------|--------|-------------------|
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
+338
View File
@@ -0,0 +1,338 @@
# Audio Generate API
vLLM-Omni provides an API for text-to-audio generation using diffusion-based models such as Stable Audio.
Unlike the [Speech API](speech_api.md) which targets text-to-speech synthesis, the Audio Generate API is designed for general-purpose audio generation from text descriptions (sound effects, music, ambient soundscapes, etc.).
Each server instance runs a single model (specified at startup via `vllm-omni serve <model> --omni`).
## Quick Start
### Start the Server
```bash
vllm-omni serve stabilityai/stable-audio-open-1.0 \
--host 0.0.0.0 \
--port 8000 \
--gpu-memory-utilization 0.9 \
--trust-remote-code \
--enforce-eager \
--omni
```
### Generate Audio
**Using curl:**
```bash
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "The sound of a cat purring",
"audio_length": 10.0
}' --output cat.wav
```
**Using Python:**
```python
import httpx
response = httpx.post(
"http://localhost:8000/v1/audio/generate",
json={
"input": "The sound of a cat purring",
"audio_length": 10.0,
},
timeout=300.0,
)
with open("cat.wav", "wb") as f:
f.write(response.content)
```
## API Reference
### Endpoint
```
POST /v1/audio/generate
Content-Type: application/json
```
### Request Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input` | string | **required** | Text prompt describing the audio to generate |
| `model` | string | server's model | Model to use (optional, should match server if specified) |
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
#### Diffusion Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `audio_length` | float | null | Audio duration in seconds (default value is the max ~47s for `stable-audio-open-1.0`) |
| `audio_start` | float | 0.0 | Audio start time in seconds |
| `negative_prompt` | string | null | Text describing what to avoid in generation |
| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
| `seed` | int | null | Random seed for reproducible generation |
### Response Format
Returns binary audio data with the appropriate `Content-Type` header:
| `response_format` | Content-Type |
|--------------------|--------------|
| `wav` | `audio/wav` |
| `mp3` | `audio/mpeg` |
| `flac` | `audio/flac` |
| `pcm` | `audio/pcm` |
| `aac` | `audio/aac` |
| `opus` | `audio/opus` |
## Examples
### Basic Generation
Generate audio with only a text prompt (model defaults for all other parameters):
```bash
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "The sound of ocean waves crashing on a beach"
}' --output ocean.wav
```
### Custom Duration
Specify an explicit audio length in seconds:
```bash
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "A dog barking",
"audio_length": 5.0
}' --output dog_5s.wav
```
### High Quality with Negative Prompt
Use a negative prompt to steer generation away from undesired characteristics, and increase inference steps for higher quality:
```bash
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "A piano playing a gentle melody",
"audio_length": 10.0,
"negative_prompt": "Low quality, distorted, noisy",
"guidance_scale": 8.0,
"num_inference_steps": 150
}' --output piano_hq.wav
```
### Reproducible Generation
Set a `seed` to get deterministic results across runs:
```bash
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "Thunder and rain sounds",
"audio_length": 15.0,
"seed": 42
}' --output thunder.wav
```
### Full Control
Combine all parameters for precise control over generation:
```bash
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "Thunder and rain sounds",
"audio_length": 15.0,
"negative_prompt": "Low quality",
"guidance_scale": 7.0,
"num_inference_steps": 100,
"seed": 42
}' --output thunder_rain.wav
```
### Quick Generation (Fewer Steps)
For faster generation with slightly lower quality:
```bash
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "Birds chirping in a forest",
"audio_length": 8.0,
"num_inference_steps": 50
}' --output birds_quick.wav
```
### Python Client
```python
import httpx
response = httpx.post(
"http://localhost:8000/v1/audio/generate",
json={
"input": "Thunder and rain",
"audio_length": 15.0,
"negative_prompt": "Low quality",
"guidance_scale": 7.0,
"num_inference_steps": 100,
"seed": 42,
"response_format": "wav",
},
timeout=300.0,
)
with open("thunder.wav", "wb") as f:
f.write(response.content)
```
## Parameter Tuning Guide
### `guidance_scale`
Controls how closely the generated audio follows the text prompt.
| Range | Behaviour |
|-------|-----------|
| 3 - 5 | More creative / varied output |
| 7 (default) | Balanced adherence |
| 10+ | Strict adherence to the prompt |
### `num_inference_steps`
Controls the number of denoising steps in the diffusion process.
| Steps | Quality | Speed | Use Case |
|-------|---------|-------|----------|
| 50 | Good | Fast | Quick previews |
| 100 | Very Good | Medium | General purpose |
| 150+ | Excellent | Slow | Final / critical audio |
### `audio_length`
Duration of the generated audio clip. For `stable-audio-open-1.0`, the maximum is approximately 47 seconds. If omitted, the model uses its own default length.
### `negative_prompt`
Describes characteristics to avoid. Common negative prompts include:
- `"Low quality, distorted, noisy"`
- `"Silence, static"`
- `"Music"` (when generating sound effects only)
## Supported Models
| Model | Description |
|-------|-------------|
| `stabilityai/stable-audio-open-1.0` | Open-source audio generation model, up to ~47 seconds, 44.1 kHz stereo |
## Error Responses
### 400 Bad Request
Invalid or missing parameters:
```json
{
"error": {
"message": "Audio generation model did not produce audio output.",
"type": "BadRequestError",
"param": null,
"code": 400
}
}
```
### 404 Not Found
Model mismatch:
```json
{
"error": {
"message": "The model `xxx` does not exist.",
"type": "NotFoundError",
"param": "model",
"code": 404
}
}
```
### 422 Unprocessable Entity
Pydantic validation failure (e.g. invalid `response_format`, `speed` out of range):
```json
{
"detail": [
{
"type": "literal_error",
"msg": "Input should be 'wav', 'pcm', 'flac', 'mp3', 'aac' or 'opus'",
...
}
]
}
```
## Troubleshooting
### "Audio generation model did not produce audio output"
The model finished but returned no audio data. Verify the server started successfully and the model loaded without errors.
### Server Not Responding
```bash
# Check if the server is healthy
curl http://localhost:8000/health
```
### Audio Quality Issues
- Increase `num_inference_steps` (e.g. 150).
- Add a negative prompt: `"Low quality, distorted, noisy"`.
- Increase `guidance_scale` for stronger prompt adherence.
### Generation Timeout
- Reduce `num_inference_steps`.
- Reduce `audio_length`.
- Check GPU memory with `nvidia-smi`.
### Out of Memory
- Lower `--gpu-memory-utilization` (e.g. 0.8).
- Reduce `audio_length`.
## Development
Enable debug logging:
```bash
vllm-omni serve stabilityai/stable-audio-open-1.0 \
--host 0.0.0.0 \
--port 8000 \
--gpu-memory-utilization 0.9 \
--trust-remote-code \
--enforce-eager \
--omni \
--uvicorn-log-level debug
```
+3 -61
View File
@@ -1,9 +1,6 @@
# Speech API
vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. Supported TTS models include:
- **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
- **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation using Qwen3-TTS models.
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
@@ -12,22 +9,13 @@ Each server instance runs a single model (specified at startup via `vllm serve <
### Start the Server
```bash
# Qwen3-TTS: CustomVoice model (predefined speakers)
# CustomVoice model (predefined speakers)
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
--omni \
--port 8091 \
--trust-remote-code \
--enforce-eager
# Fish Speech S2 Pro
vllm-omni serve fishaudio/s2-pro \
--stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
--omni \
--port 8091 \
--trust-remote-code \
--enforce-eager \
--gpu-memory-utilization 0.9
```
### Generate Speech
@@ -108,7 +96,7 @@ Content-Type: application/json
| `language` | string | "Auto" | Language (see supported languages below) |
| `instructions` | string | "" | Voice style/emotion instructions |
| `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
| `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
| `initial_codec_chunk_frames` | integer | null | Initial chunk size for reduced TTFA (overrides stage config) |
**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
@@ -176,42 +164,6 @@ curl -X POST http://localhost:8091/v1/audio/voices \
-F "name=custom_voice_1"
```
## Streaming Text Input (WebSocket)
The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally and generates audio per sentence as boundaries are detected.
> Note: text input is always streamed incrementally. Audio output remains sentence-scoped:
> use `stream_audio=false` for one binary frame per sentence, or `stream_audio=true` for one or more PCM chunks per sentence.
### WebSocket Protocol
Client -> Server:
| Message | Description |
|---------|-------------|
| `{"type": "session.config", ...}` | Session configuration (sent once, first message) |
| `{"type": "input.text", "text": "..."}` | Text chunk |
| `{"type": "input.done"}` | End of input, flushes remaining buffer |
Server -> Client:
| Message | Description |
|---------|-------------|
| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for a sentence |
| Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for a sentence |
| `{"type": "session.done", "total_sentences": N}` | Session complete |
| `{"type": "error", "message": "..."}` | Non-fatal error |
### Session Config Parameters
All REST API parameters are supported, plus:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `stream_audio` | bool | false | Stream one or more PCM chunks per sentence over WebSocket |
| `split_granularity` | string | "sentence" | Text splitting granularity |
```bash
DELETE /v1/audio/voices/{name}
@@ -329,8 +281,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \
## Supported Models
### Qwen3-TTS
| Model | Task Type | Description |
|-------|-----------|-------------|
| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
@@ -339,14 +289,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \
| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
### Fish Speech S2 Pro
| Model | Description |
|-------|-------------|
| `fishaudio/s2-pro` | 4B dual-AR TTS with DAC codec (44.1 kHz). Supports text-to-speech and voice cloning. |
Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech online serving example](../user_guide/examples/online_serving/fish_speech.md) for details.
## Error Responses
### 400 Bad Request
Binary file not shown.

Before

Width:  |  Height:  |  Size: 34 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 37 KiB

After

Width:  |  Height:  |  Size: 36 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 45 KiB

After

Width:  |  Height:  |  Size: 33 KiB

@@ -18,28 +18,23 @@ The following parallelism methods are currently supported in vLLM-Omni:
6. [HSDP](#hsdp): Hybrid Sharded Data Parallel shards model weights across GPUs using PyTorch FSDP2. This reduces per-GPU memory usage, enabling inference of large models on GPUs with limited memory.
7. [Expert Parallel](#expert-parallelism): Expert Parallelism shards the Experts of a Mixture-of-Experts (MoE) layer across multiple devices. During the forward, a gating mechanism routes tokens to their designated experts, necessitating cross-cards communication(all-to-all) to dispatch tokens to the correct ranks and combine the results. This parallelism allows for massive scaling of model parameters without a proportional increase in the computational load per device.
The following table shows which models are currently supported by parallelism method:
### ImageGen
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel |
|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A |
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A |
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | N/A |
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | N/A |
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | N/A |
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ | N/A |
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A |
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | N/A |
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A |
| **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ |
| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ✅ | ❌ | ❌ | N/A |
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel |
|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ |
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ |
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ |
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ |
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ |
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ |
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ |
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ |
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ |
!!! note "TP Limitations for Diffusion Models"
We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
@@ -419,28 +414,3 @@ vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --usp
#### Adding HSDP Support to New Models
For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../design/feature/hsdp.md).
### Expert Parallelism
Unlike Tensor Parallelism which shards every layer's weights, EP only shards the MoE expert MLP blocks. This significantly reduces the memory footprint of MoE models (e.g., HunyuanImage3.0) while maintaining constant dense-equivalent compute efficiency. Expert Parallelism is enabled via `DiffusionParallelConfig.enable_expert_parallel`. And `self.ep = tp * sp * cfg * dp` for now, so at least one of TP/SP/CFG/DP should set when EP enabled.
#### Offline Inference
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="tencent/HunyuanImage-3.0",
parallel_config=DiffusionParallelConfig(tensor_parallel_size=8, enable_expert_parallel=True),
)
outputs = omni.generate(
"A brown and white dog is running on the grass",
OmniDiffusionSamplingParams(
num_inference_steps=50,
width=1024,
height=1024,
),
)
```
@@ -69,7 +69,6 @@ The following table shows which models are currently supported by each accelerat
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
### VideoGen
@@ -77,7 +76,6 @@ The following table shows which models are currently supported by each accelerat
|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ |
### Quantization
@@ -1,74 +0,0 @@
# CosyVoice3
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/cosyvoice3>.
## Setup
Install dependencies:
```
uv pip install -e .
```
> **Note:** This includes required libraries such as `librosa`, `soundfile`,
> `onnxruntime`, `x-transformers`, and `einops` via
> `requirements/common.txt` and platform-specific requirements files.
Download the model snapshot:
```
from huggingface_hub import snapshot_download
snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', local_dir='pretrained_models/Fun-CosyVoice3-0.5B')
```
Add `config.json` in `pretrained_models/Fun-CosyVoice3-0.5B/`:
```json
{
"model_type": "cosyvoice3",
"architectures": [
"CosyVoice3Model"
]
}
```
> **Why `config.json` is required:**
> `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers a class mapping.
> The loader still needs `model_type: "cosyvoice3"` from `config.json` to select that class.
> If no `config.json` is present, model type cannot be inferred automatically.
> If your downloaded checkpoint already includes a valid `config.json` with
> `model_type: "cosyvoice3"`, this manual step can be skipped.
Run the offline verification script:
```
python examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py \
--model pretrained_models/Fun-CosyVoice3-0.5B \
--tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN
```
## Implementation Overview
CosyVoice3 runs as a **2-stage Omni pipeline**.
- **Stage 0** (`talker`) converts text + prompt audio to speech tokens.
- **Stage 1** (`code2wav`) converts speech tokens + prompt features to waveform via flow matching and HiFiGAN.
Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py`.
- `CosyVoice3MultiModalProcessor` builds the multimodal inputs:
- Tokenizes `prompt` and `prompt_text`.
- Extracts speech tokens and mel features from the prompt audio.
- Extracts a speaker embedding.
- `CosyVoice3Model` implements both stages:
- Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features.
- Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform.
Stage wiring is configured in `vllm_omni/model_executor/stage_configs/cosyvoice3.yaml`.
- Stage 0 emits latent speech tokens.
- Stage 1 consumes them via `custom_process_input_func` and outputs audio.
## Example materials
??? abstract "verify_e2e_cosyvoice.py"
``````py
--8<-- "examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py"
``````
@@ -1,97 +0,0 @@
# Fish Speech S2 Pro
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/fish_speech>.
This directory contains an offline demo for running Fish Speech S2 Pro (`fishaudio/s2-pro`) with vLLM Omni. It supports text-only synthesis and voice cloning with reference audio.
## Model Overview
[Fish Speech S2 Pro](https://huggingface.co/fishaudio/s2-pro) is a 4B dual-AR text-to-speech model by FishAudio. It uses:
- **Slow AR**: Generates semantic tokens autoregressively (Qwen3-based backbone)
- **Fast AR**: Predicts residual codebook tokens from semantic tokens
- **DAC Decoder**: Converts codec codes to 44.1 kHz audio waveform
The model produces high-quality speech with voice cloning capabilities.
## Prerequisites
Install the `fish-speech` package for the DAC codec:
```bash
pip install fish-speech
```
## 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.
## Quick Start
### Text-Only Synthesis
```bash
python end2end.py --text "Hello, this is a test of the Fish Speech text to speech system."
```
Generated audio files are saved to `output_audio/` by default.
### Voice Cloning
Provide a reference audio file and its transcript to clone a voice:
```bash
python end2end.py \
--text "Hello, this is a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "Transcript of the reference audio."
```
## Streaming Mode
Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
```bash
python end2end.py --text "Hello, this is a test." --streaming
```
Each DAC decoder chunk is logged as it arrives. The final WAV file is written once generation completes.
## Configuration
| Argument | Default | Description |
|----------|---------|-------------|
| `--text` | `"Hello, this is a test..."` | Text to synthesize |
| `--ref-audio` | None | Path to reference audio for voice cloning |
| `--ref-text` | None | Transcript of the reference audio |
| `--model` | `fishaudio/s2-pro` | HuggingFace model path |
| `--stage-configs-path` | `fish_speech_s2_pro.yaml` | Path to stage configs YAML |
| `--output-dir` | `output_audio` | Output directory for WAV files |
| `--streaming` | False | Enable streaming via AsyncOmni |
| `--stage-init-timeout` | 600 | Stage initialization timeout (seconds) |
## Notes
- Output audio is 44.1 kHz mono WAV.
- The DAC codec weights (`codec.pth`) are loaded lazily from the model directory.
- Voice cloning encodes the reference audio using the DAC codec to extract semantic codes, then prepends them as a system message.
## Example materials
??? abstract "diag_codebook.py"
``````py
--8<-- "examples/offline_inference/fish_speech/diag_codebook.py"
``````
??? abstract "diag_reference_compare.py"
``````py
--8<-- "examples/offline_inference/fish_speech/diag_reference_compare.py"
``````
??? abstract "end2end.py"
``````py
--8<-- "examples/offline_inference/fish_speech/end2end.py"
``````
??? abstract "standalone_test.py"
``````py
--8<-- "examples/offline_inference/fish_speech/standalone_test.py"
``````
@@ -99,8 +99,7 @@ python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_str
```
Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames`
in the stage config). The initial chunk size is dynamically selected based on server load for reduced
TTFA, and can be overridden per-request via the `initial_codec_chunk_frames` API field. The final WAV file is written once generation
and `initial_codec_chunk_frames` in the stage config). The final WAV file is written once generation
completes. This demonstrates that audio data is available progressively rather than only at the end.
> **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.
@@ -1,58 +0,0 @@
# X-To-Video-Audio
The `DreamID-Omni` pipeline generates short videos from text, image and video.
## Local CLI Usage
### Download the Model locally
Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
```bash
python download_dreamid_omni.py --output-dir ./dreamid_omni
```
After download, the model directory will look like this:
```
dreamid_omni/
├── DreamID-Omni/
│ ├── dreamid_omni.safetensors
├── MMAudio/
│ ├── ext_weights/
│ │ ├── best_netG.pt
│ │ ├── v1-16.pth
├── Wan2.2-TI2V-5B/
│ ├── google/*
│ ├── models_t5_umt5-xxl-enc-bf16.pth
│ ├── Wan2.2_VAE.pth
│
├── model_index.json # create by download_dreamid_omni.py
```
### Run the Inference
```
python x_to_video_audio.py \
--model /xx/dreamid_omni \
--prompt "Two people walking together and singing happily" \
--image-path ./example0.png ./example1.png \
--audio-path ./example0.wav ./example1.wav \
--video-negative-prompt "jitter, bad hands, blur, distortion" \
--audio-negative-prompt "robotic, muffled, echo, distorted" \
--cfg-parallel-size 2 \
--num-inference-steps 45 \
--height 704 \
--width 1280 \
--output dreamid_omni.mp4
```
In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
Key arguments:
- `--prompt`: text description (string).
- `--model`: path to the model local directory.
- `--height/--width`: output resolution (defaults 704 * 1024).
- `--image-path`: path to the input image list.
- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
- `--num-inference-steps`: number of denoising steps (defaults 45).
- `--video-negative-prompt`: negative prompt for video generation.
- `--audio-negative-prompt`: negative prompt for audio generation.
- `--enable-cpu-offload`: enable CPU offload (defaults False).
@@ -1,194 +0,0 @@
# Fish Speech S2 Pro
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/fish_speech>.
## Model
| Model | Description |
|-------|-------------|
| `fishaudio/s2-pro` | Fish Speech S2 Pro -- 4B dual-AR TTS model with DAC codec (44.1 kHz) |
## Gradio Demo
An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support.
```bash
# Option 1: Launch server + Gradio together
./run_gradio_demo.sh
# Option 2: If server is already running
python gradio_demo.py --api-base http://localhost:8091
```
Then open http://localhost:7860 in your browser.
Features:
- Text-to-speech synthesis
- Voice cloning from uploaded audio or URL
- Streaming mode (progressive PCM playback)
## Launch the Server
```bash
vllm-omni serve fishaudio/s2-pro \
--stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
--omni \
--port 8091 \
--trust-remote-code \
--enforce-eager \
--gpu-memory-utilization 0.9
```
Or use the convenience script:
```bash
./run_server.sh
```
## Send TTS Request
### Using curl
```bash
# Basic TTS
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, how are you?",
"voice": "default",
"response_format": "wav"
}' --output output.wav
```
### Voice Cloning
Provide a reference audio (URL or base64 data URL) and its transcript:
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, this is a cloned voice.",
"voice": "default",
"ref_audio": "https://example.com/reference.wav",
"ref_text": "Transcript of the reference audio."
}' --output cloned.wav
```
### Using Python
```python
import httpx
# Basic TTS
response = httpx.post(
"http://localhost:8091/v1/audio/speech",
json={
"input": "Hello, how are you?",
"voice": "default",
"response_format": "wav",
},
timeout=300.0,
)
with open("output.wav", "wb") as f:
f.write(response.content)
```
### Using the CLI Client
```bash
cd examples/online_serving/fish_speech
# Basic TTS
python speech_client.py --text "Hello, how are you?"
# Voice cloning
python speech_client.py \
--text "Hello, this is a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "Transcript of the reference audio."
# Streaming PCM output
python speech_client.py --text "Hello world" --stream --output output.pcm
```
The CLI client supports:
- `--api-base`: API base URL (default: `http://localhost:8091`)
- `--model` (or `-m`): Model name (default: `fishaudio/s2-pro`)
- `--text`: Text to synthesize (required)
- `--ref-audio`: Reference audio for voice cloning (local path or URL)
- `--ref-text`: Transcript of the reference audio
- `--stream`: Enable streaming (PCM output)
- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav)
- `--output` (or `-o`): Output file path
## Streaming
Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded:
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, how are you?",
"voice": "default",
"stream": true,
"response_format": "pcm"
}' --no-buffer | play -t raw -r 44100 -e signed -b 16 -c 1 -
```
**Note:** Fish Speech outputs at 44.1 kHz (unlike Qwen3-TTS which outputs at 24 kHz).
## API Parameters
Fish Speech uses the same `/v1/audio/speech` endpoint as Qwen3-TTS. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full parameter documentation.
Key parameters for Fish Speech:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input` | string | **required** | Text to synthesize |
| `voice` | string | "default" | Voice name (use "default" for Fish Speech) |
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
| `ref_audio` | string | null | Reference audio URL or base64 data URL for voice cloning |
| `ref_text` | string | null | Transcript of reference audio (required for voice cloning) |
| `max_new_tokens` | int | 4096 | Maximum tokens to generate |
| `stream` | bool | false | Stream raw PCM chunks |
## Prerequisites
Install the `fish-speech` package for the DAC codec:
```bash
pip install fish-speech
```
## Troubleshooting
1. **No audio output**: Make sure the `fish-speech` package is installed for the DAC decoder
2. **Connection refused**: Ensure the server is running on the correct port
3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings
4. **Out of memory**: Reduce `--gpu-memory-utilization` or use a GPU with more VRAM
## Example materials
??? abstract "gradio_demo.py"
``````py
--8<-- "examples/online_serving/fish_speech/gradio_demo.py"
``````
??? abstract "run_gradio_demo.sh"
``````sh
--8<-- "examples/online_serving/fish_speech/run_gradio_demo.sh"
``````
??? abstract "run_server.sh"
``````sh
--8<-- "examples/online_serving/fish_speech/run_server.sh"
``````
??? abstract "speech_client.py"
``````py
--8<-- "examples/online_serving/fish_speech/speech_client.py"
``````
@@ -29,35 +29,6 @@ The script allows overriding:
- `CACHE_BACKEND` (default: `none`)
- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
## Async Job Behavior
`POST /v1/videos` is asynchronous. It creates a video job and immediately
returns metadata like the job ID and initial `queued` status. To get the final
artifact, poll the job status and then download the completed file from the
content endpoint.
The main endpoints are:
- `POST /v1/videos`: create a video generation job
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
- `GET /v1/videos`: list stored video jobs
- `GET /v1/videos/{video_id}/content`: download the generated video file
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
## Storage
Generated video files are stored on local disk by the async video API.
Local file storage behavior can be controlled via the following environment variables:
- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
Example:
```bash
export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
```
## API Calls
### Method 1: Using curl
@@ -67,7 +38,7 @@ export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
bash run_curl_image_to_video.sh
# Or execute directly (OpenAI-style multipart)
create_response=$(curl -s http://localhost:8091/v1/videos \
curl -X POST http://localhost:8091/v1/videos \
-H "Accept: application/json" \
-F "prompt=A bear playing with yarn, smooth motion" \
-F "negative_prompt=low quality, blurry, static" \
@@ -81,23 +52,7 @@ create_response=$(curl -s http://localhost:8091/v1/videos \
-F "guidance_scale_2=1.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=12.0" \
-F "seed=42")
video_id=$(echo "$create_response" | jq -r '.id')
while true; do
status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
if [ "$status" = "completed" ]; then
break
fi
if [ "$status" = "failed" ]; then
echo "Video generation failed"
exit 1
fi
sleep 2
done
curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o wan22_i2v_output.mp4
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_i2v_output.mp4
```
## Request Format
@@ -111,18 +66,6 @@ curl -X POST http://localhost:8091/v1/videos \
-F "input_reference=@/path/to/qwen-bear.png"
```
### Alternative JSON-Safe Reference Input
Use `image_reference` when you want to pass a URL or JSON-safe image reference
instead of uploading a file. Do not send `input_reference` and
`image_reference` together.
```bash
curl -X POST http://localhost:8091/v1/videos \
-F "prompt=A bear playing with yarn, smooth motion" \
-F 'image_reference={"image_url":"https://example.com/qwen-bear.png"}'
```
### Generation with Parameters
```bash
@@ -142,63 +85,6 @@ curl -X POST http://localhost:8091/v1/videos \
-F "seed=42"
```
## Create Response Format
`POST /v1/videos` returns a job record.
```json
{
"id": "video_gen_123",
"object": "video",
"status": "queued",
"model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
"prompt": "A bear playing with yarn, smooth motion",
"created_at": 1234567890
}
```
## Retrieve, List, Download, and Delete
### Retrieve a job
```bash
curl -s http://localhost:8091/v1/videos/${video_id} | jq .
```
### List jobs
```bash
curl -s http://localhost:8091/v1/videos | jq .
```
### Download the completed video
```bash
curl -L http://localhost:8091/v1/videos/${video_id}/content -o wan22_i2v_output.mp4
```
### Delete a job and its stored file
```bash
curl -X DELETE http://localhost:8091/v1/videos/${video_id} | jq .
```
## Poll Until Complete
```bash
while true; do
status=$(curl -s http://localhost:8091/v1/videos/${video_id} | jq -r '.status')
if [ "$status" = "completed" ]; then
break
fi
if [ "$status" = "failed" ]; then
echo "Video generation failed"
exit 1
fi
sleep 2
done
```
## Example materials
??? abstract "run_curl_image_to_video.sh"
@@ -17,31 +17,6 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
## Gradio Demo
Two interactive Gradio demos are available, both supporting all 3 task types:
| Demo | File | Transport | Streaming Quality |
| ---- | ---- | --------- | ----------------- |
| Standard | `gradio_demo.py` | HTTP chunked | May have small gaps between chunks |
| FastRTC | `gradio_fastrtc_demo.py` | WebRTC | Gapless streaming (requires `pip install fastrtc`) |
```bash
# Option 1: Launch server + Standard Gradio together
./run_gradio_demo.sh # CustomVoice (default)
./run_gradio_demo.sh --task-type VoiceDesign # VoiceDesign
./run_gradio_demo.sh --task-type Base # Voice cloning
# Option 2: If server is already running
python gradio_demo.py --api-base http://localhost:8000
# Option 3: FastRTC demo (gapless streaming)
pip install fastrtc
python gradio_fastrtc_demo.py --api-base http://localhost:8000
```
Then open http://localhost:7860 in your browser.
## Run examples (Qwen3-TTS)
### Launch the Server
@@ -137,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`)
@@ -212,59 +187,6 @@ 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
```
@@ -274,6 +196,20 @@ 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
@@ -284,7 +220,7 @@ This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/
"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
@@ -317,7 +253,6 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
| `language` | string | "Auto" | Language (see supported languages below) |
| `instructions` | string | "" | Voice style/emotion instructions |
| `max_new_tokens` | int | 2048 | Maximum tokens to generate |
| `initial_codec_chunk_frames` | int | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
@@ -326,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) |
@@ -366,27 +301,11 @@ curl -X POST http://localhost:8091/v1/audio/speech \
## Example materials
??? abstract "gradio_demo.py"
``````py
--8<-- "examples/online_serving/qwen3_tts/gradio_demo.py"
``````
??? abstract "gradio_fastrtc_demo.py"
``````py
--8<-- "examples/online_serving/qwen3_tts/gradio_fastrtc_demo.py"
``````
??? abstract "openai_speech_client.py"
``````py
--8<-- "examples/online_serving/qwen3_tts/openai_speech_client.py"
``````
??? abstract "run_gradio_demo.sh"
``````sh
--8<-- "examples/online_serving/qwen3_tts/run_gradio_demo.sh"
``````
??? abstract "run_server.sh"
``````sh
--8<-- "examples/online_serving/qwen3_tts/run_server.sh"
``````
??? abstract "tts_common.py"
``````py
--8<-- "examples/online_serving/qwen3_tts/tts_common.py"
``````
@@ -0,0 +1,193 @@
# Text-To-Audio
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/stable_audio>.
This example demonstrates how to deploy Stable Audio models for online text-to-audio generation using vLLM-Omni.
## Supported Models
| Model | Description |
|-------|-------------|
| `stabilityai/stable-audio-open-1.0` | Open-source audio generation, up to ~47 seconds, 44.1 kHz stereo |
## Start Server
### Basic Start
```bash
vllm-omni serve stabilityai/stable-audio-open-1.0 \
--host 0.0.0.0 \
--port 8000 \
--gpu-memory-utilization 0.9 \
--trust-remote-code \
--enforce-eager \
--omni
```
## API Calls
### Method 1: Using curl
```bash
# Run all curl examples
bash curl_examples.sh
# Or execute directly
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "The sound of a cat purring",
"audio_length": 10.0
}' --output cat.wav
```
### Method 2: Using Python Client
```bash
cd examples/online_serving/stable_audio
# Simple generation
python stable_audio_client.py \
--text "The sound of a cat purring"
# With custom duration
python stable_audio_client.py \
--text "A dog barking" \
--audio_length 5.0
# With all parameters
python stable_audio_client.py \
--text "Thunder and rain" \
--audio_length 15.0 \
--negative_prompt "Low quality" \
--guidance_scale 7.0 \
--num_inference_steps 100 \
--seed 42 \
--output thunder.wav
```
The Python client supports the following command-line arguments:
- `--api_url`: API endpoint URL (default: `http://localhost:8000/v1/audio/generate`)
- `--text`: Text prompt for audio generation (default: `"The sound of a cat purring"`)
- `--audio_length`: Audio length in seconds (default: `10.0`, max ~47s for `stable-audio-open-1.0`)
- `--audio_start`: Audio start time in seconds (default: `0.0`)
- `--negative_prompt`: Negative prompt for classifier-free guidance (default: `"Low quality"`)
- `--guidance_scale`: Guidance scale for diffusion (default: `7.0`)
- `--num_inference_steps`: Number of inference steps (default: `100`)
- `--seed`: Random seed for reproducibility (default: `None`)
- `--response_format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`
- `--output`: Output file path (default: `stable_audio_output.wav`)
### Method 3: Using Python httpx
```python
import httpx
response = httpx.post(
"http://localhost:8000/v1/audio/generate",
json={
"input": "The sound of ocean waves crashing on a beach",
"audio_length": 10.0,
"negative_prompt": "Low quality, distorted",
"guidance_scale": 7.0,
"num_inference_steps": 100,
},
timeout=300.0,
)
with open("ocean.wav", "wb") as f:
f.write(response.content)
```
## Request Format
### Simple Generation
```json
{
"input": "The sound of ocean waves"
}
```
### Generation with Parameters
```json
{
"input": "A piano playing a gentle melody",
"audio_length": 10.0,
"negative_prompt": "Low quality, distorted, noisy",
"guidance_scale": 8.0,
"num_inference_steps": 150,
"seed": 42,
"response_format": "wav"
}
```
## API Reference
### Endpoint
```
POST /v1/audio/generate
Content-Type: application/json
```
### Generation Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input` | string | **required** | Text prompt describing the audio to generate |
| `model` | string | server's model | Model to use (optional, should match server if specified) |
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
| `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) |
| `audio_start` | float | 0.0 | Audio start time in seconds |
| `negative_prompt` | string | null | Text describing what to avoid in generation |
| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
| `seed` | int | null | Random seed for reproducible generation |
### Response Format
Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
## Tuning Tips
1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`.
2. **Quality vs Speed**:
- 50 steps: Fast, decent quality (quick previews)
- 100 steps: Good balance (general purpose)
- 150+ steps: High quality, slower (final / critical audio)
3. **Guidance Scale**:
- Lower (3 - 5): More creative / varied output
- Default (7): Good balance
- Higher (10+): Strict adherence to the prompt
4. **Negative Prompts**: Use to avoid unwanted characteristics such as `"Low quality"`, `"distorted"`, `"noisy"`.
5. **Seeds**: Set a fixed seed to get deterministic, reproducible results.
## File Description
| File | Description |
|------|-------------|
| `curl_examples.sh` | Curl examples covering common use cases |
| `stable_audio_client.py` | Python client with full CLI argument support |
## Troubleshooting
1. **Audio generation model did not produce audio output**: Verify the server started successfully and the model loaded without errors.
2. **Connection refused**: Make sure the server is running on the correct port.
3. **Generation timeout**: Reduce `num_inference_steps` or `audio_length`, and check GPU memory with `nvidia-smi`.
4. **Out of memory**: Lower `--gpu-memory-utilization` or reduce `audio_length`.
5. **Audio quality issues**: Increase `num_inference_steps`, add a negative prompt, or raise `guidance_scale`.
## Example materials
??? abstract "stable_audio_client.py"
``````py
--8<-- "examples/online_serving/stable_audio/stable_audio_client.py"
``````
??? abstract "curl_examples.sh"
``````sh
--8<-- "examples/online_serving/stable_audio/curl_examples.sh"
``````
@@ -29,34 +29,6 @@ The script allows overriding:
- `CACHE_BACKEND` (default: `none`)
- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
## Async Job Behavior
`POST /v1/videos` is asynchronous. It creates a video job and immediately
returns metadata like the job ID and initial `queued` status. To get the final
artifact, poll the job status and then download the completed file from the
content endpoint.
The main endpoints are:
- `POST /v1/videos`: create a video generation job
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
- `GET /v1/videos`: list stored video jobs
- `GET /v1/videos/{video_id}/content`: download the generated video file
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
## Storage
Generated video files are stored on local disk by the async video API.
- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
Example:
```bash
export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
```
## API Calls
### Method 1: Using curl
@@ -66,7 +38,7 @@ export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
bash run_curl_text_to_video.sh
# Or execute directly (OpenAI-style multipart)
create_response=$(curl -s http://localhost:8091/v1/videos \
curl -s http://localhost:8091/v1/videos \
-H "Accept: application/json" \
-F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
-F "width=832" \
@@ -78,24 +50,7 @@ create_response=$(curl -s http://localhost:8091/v1/videos \
-F "guidance_scale=4.0" \
-F "guidance_scale_2=4.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=5.0" \
-F "seed=42")
video_id=$(echo "$create_response" | jq -r '.id')
while true; do
status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
if [ "$status" = "completed" ]; then
break
fi
if [ "$status" = "failed" ]; then
echo "Video generation failed"
exit 1
fi
sleep 2
done
curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o wan22_output.mp4
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
```
## Request Format
@@ -130,9 +85,8 @@ curl -X POST http://localhost:8091/v1/videos \
| Parameter | Type | Default | Description |
| --------------------- | ------ | ------- | ------------------------------------------------ |
| `prompt` | str | - | Text description of the desired video |
| `seconds` | str | None | Clip duration in seconds |
| `size` | str | None | Output size in `WIDTHxHEIGHT` format |
| `negative_prompt` | str | None | Negative prompt |
| `n` | int | 1 | Number of videos to generate |
| `width` | int | None | Video width in pixels |
| `height` | int | None | Video height in pixels |
| `num_frames` | int | None | Number of frames to generate |
@@ -144,62 +98,24 @@ curl -X POST http://localhost:8091/v1/videos \
| `flow_shift` | float | None | Scheduler flow shift (Wan2.2) |
| `seed` | int | None | Random seed (reproducible) |
| `lora` | object | None | LoRA configuration |
| `extra_body` | object | None | Model-specific extra parameters |
## Create Response Format
`POST /v1/videos` returns a job record.
## Response Format
```json
{
"id": "video_gen_123",
"object": "video",
"status": "queued",
"model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
"prompt": "A cinematic view of a futuristic city at sunset",
"created_at": 1234567890
"created": 1234567890,
"data": [
{ "b64_json": "<base64-mp4>" }
]
}
```
## Retrieve, List, Download, and Delete
### Retrieve a job
## Extract Video
```bash
curl -s http://localhost:8091/v1/videos/${video_id} | jq .
```
### List jobs
```bash
curl -s http://localhost:8091/v1/videos | jq .
```
### Download the completed video
```bash
curl -L http://localhost:8091/v1/videos/${video_id}/content -o wan22_output.mp4
```
### Delete a job and its stored file
```bash
curl -X DELETE http://localhost:8091/v1/videos/${video_id} | jq .
```
## Poll Until Complete
```bash
while true; do
status=$(curl -s http://localhost:8091/v1/videos/${video_id} | jq -r '.status')
if [ "$status" = "completed" ]; then
break
fi
if [ "$status" = "failed" ]; then
echo "Video generation failed"
exit 1
fi
sleep 2
done
# Extract base64 from response and decode to video
cat response.json | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
```
## Example materials
@@ -1,64 +0,0 @@
# CosyVoice3
## Setup
Install dependencies:
```
uv pip install -e .
```
> **Note:** This includes required libraries such as `librosa`, `soundfile`,
> `onnxruntime`, `x-transformers`, and `einops` via
> `requirements/common.txt` and platform-specific requirements files.
Download the model snapshot:
```
from huggingface_hub import snapshot_download
snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', local_dir='pretrained_models/Fun-CosyVoice3-0.5B')
```
Add `config.json` in `pretrained_models/Fun-CosyVoice3-0.5B/`:
```json
{
"model_type": "cosyvoice3",
"architectures": [
"CosyVoice3Model"
]
}
```
> **Why `config.json` is required:**
> `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers a class mapping.
> The loader still needs `model_type: "cosyvoice3"` from `config.json` to select that class.
> If no `config.json` is present, model type cannot be inferred automatically.
> If your downloaded checkpoint already includes a valid `config.json` with
> `model_type: "cosyvoice3"`, this manual step can be skipped.
Run the offline verification script:
```
python examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py \
--model pretrained_models/Fun-CosyVoice3-0.5B \
--tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN
```
## Implementation Overview
CosyVoice3 runs as a **2-stage Omni pipeline**.
- **Stage 0** (`talker`) converts text + prompt audio to speech tokens.
- **Stage 1** (`code2wav`) converts speech tokens + prompt features to waveform via flow matching and HiFiGAN.
Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py`.
- `CosyVoice3MultiModalProcessor` builds the multimodal inputs:
- Tokenizes `prompt` and `prompt_text`.
- Extracts speech tokens and mel features from the prompt audio.
- Extracts a speaker embedding.
- `CosyVoice3Model` implements both stages:
- Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features.
- Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform.
Stage wiring is configured in `vllm_omni/model_executor/stage_configs/cosyvoice3.yaml`.
- Stage 0 emits latent speech tokens.
- Stage 1 consumes them via `custom_process_input_func` and outputs audio.
@@ -1,197 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import os
from pathlib import Path
import librosa
import numpy as np
import soundfile as sf
from vllm import SamplingParams
from vllm.assets.audio import AudioAsset
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
def _ensure_mel_filters_asset() -> None:
repo_root = Path(__file__).resolve().parents[3]
filters_path = repo_root / "vllm_omni" / "model_executor" / "models" / "cosyvoice3" / "assets" / "mel_filters.npz"
if filters_path.exists():
return
source_url = "https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz"
raise FileNotFoundError(
"Missing CosyVoice3 mel filter asset:\n"
f" {filters_path}\n"
"Download it with:\n"
f" mkdir -p {filters_path.parent} && "
f"curl -L {source_url} -o {filters_path}"
)
def run_e2e():
parser = argparse.ArgumentParser()
# ""FunAudioLLM/Fun-CosyVoice3-0.5B-2512
parser.add_argument(
"--model",
type=str,
required=True,
help="Path to CosyVoice3 model directory (e.g., pretrained_models/Fun-CosyVoice3-0.5B/).",
)
parser.add_argument("--stage-config", type=str, default="vllm_omni/model_executor/stage_configs/cosyvoice3.yaml")
parser.add_argument("--prompt", type=str, default="Hello, this is a test of the CosyVoice system capability.")
parser.add_argument(
"--prompt-text",
type=str,
default="You are a helpful assistant.<|endofprompt|>Testing my voices. Why should I not?",
)
parser.add_argument("--audio-path", type=str, default="prompt.wav")
parser.add_argument(
"--tokenizer",
type=str,
required=True,
help="Path to tokenizer directory (e.g., <model_path>/CosyVoice-BlankEN).",
)
args = parser.parse_args()
_ensure_mel_filters_asset()
# Ensure tokenizer directory exists
if not os.path.exists(args.tokenizer):
raise FileNotFoundError(f"{args.tokenizer} does not exist!")
# Ensure stage config exists
if not os.path.exists(args.stage_config):
raise FileNotFoundError(f"{args.stage_config} does not exist!")
print(f"Initializing cosyvoice E2E with model={args.model}")
# Initialize Omni
# This spins up the engine(s) based on the stage config
# We pass trust_remote_code=True same as Qwen examples
omni = Omni(
model=args.model,
stage_configs_path=args.stage_config,
trust_remote_code=True,
tokenizer=args.tokenizer,
log_stats=True,
)
sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
print("Model initialized. Preparing inputs...")
if args.audio_path:
if not os.path.exists(args.audio_path):
raise FileNotFoundError(f"Audio file not found: {args.audio_path}")
# Load at native sample rate
audio_signal, sr = librosa.load(args.audio_path, sr=None)
# Validate sample rate before processing (similar to original CosyVoice)
min_sr = 16000
if sr < min_sr:
raise ValueError(
f"Audio sample rate {sr} Hz is too low. "
f"Minimum required: {min_sr} Hz. "
f"Please provide audio with sample rate >= {min_sr} Hz."
)
audio_data = (audio_signal.astype(np.float32), sr)
else:
audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
prompts = {
"prompt": args.prompt,
"multi_modal_data": {
"audio": audio_data,
},
"mm_processor_kwargs": {
"prompt_text": args.prompt_text,
"sample_rate": audio_data[1],
},
}
print(f"Generating for prompt: {args.prompt}")
config = CosyVoice3Config()
tokenizer = get_qwen_tokenizer(
token_path=args.tokenizer,
skip_special_tokens=config.skip_special_tokens,
version=config.version,
)
_, text_token_len = extract_text_token(args.prompt, tokenizer, config.allowed_special)
base_len = int(text_token_len)
min_len = int(base_len * config.min_token_text_ratio)
max_len = int(base_len * config.max_token_text_ratio)
# Build SamplingParams for each stage (GPT, S2Mel, Vocoder)
gpt_sampling = SamplingParams(
temperature=1.0,
top_p=sampling_cfg["top_p"],
top_k=sampling_cfg["top_k"],
repetition_penalty=2.0,
min_tokens=min_len,
max_tokens=max_len,
stop_token_ids=[sampling_cfg["eos_token_id"]],
# allowed_token_ids=list(range(6561+3)),
detokenize=False,
)
# Not used
s2mel_sampling = SamplingParams(
temperature=1.0,
top_p=1.0,
top_k=-1,
repetition_penalty=2.0,
max_tokens=256,
detokenize=False,
)
sampling_params_list = [gpt_sampling, s2mel_sampling]
# Start profiling (requires VLLM_TORCH_PROFILER_DIR env var)
if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
print("Starting profiler...")
omni.start_profile()
# Generate (Omni orchestrator requires a per-stage SamplingParams list)
outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list[:2]))
# Stop profiling and get results
if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
print("Stopping profiler...")
profile_results = omni.stop_profile()
print(f"Profile traces saved to: {profile_results}")
print(outputs)
# Verify outputs
print(f"Received {len(outputs)} outputs.")
for i, output in enumerate(outputs):
try:
ro_list = output.request_output or []
if not ro_list:
print("No request_output found.")
continue
for ro in ro_list:
# Multimodal output may be attached to RequestOutput or CompletionOutput.
mm = getattr(ro, "multimodal_output", None)
if not mm and ro.outputs:
mm = getattr(ro.outputs[0], "multimodal_output", None)
if mm:
print(f"Multimodal output keys: {mm.keys()}")
if "audio" in mm:
audio_out = mm["audio"]
print(f"Generated Audio Shape: {audio_out.shape}")
out_path = f"output_{i}.wav"
sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
print(f"Saved audio to {out_path}")
else:
print("No multimodal output found.")
except Exception as e:
print(f"Error inspecting output: {e}")
omni.close()
if __name__ == "__main__":
run_e2e()
@@ -1,75 +0,0 @@
# Fish Speech S2 Pro
This directory contains an offline demo for running Fish Speech S2 Pro (`fishaudio/s2-pro`) with vLLM Omni. It supports text-only synthesis and voice cloning with reference audio.
## Model Overview
[Fish Speech S2 Pro](https://huggingface.co/fishaudio/s2-pro) is a 4B dual-AR text-to-speech model by FishAudio. It uses:
- **Slow AR**: Generates semantic tokens autoregressively (Qwen3-based backbone)
- **Fast AR**: Predicts residual codebook tokens from semantic tokens
- **DAC Decoder**: Converts codec codes to 44.1 kHz audio waveform
The model produces high-quality speech with voice cloning capabilities.
## Prerequisites
Install the `fish-speech` package for the DAC codec:
```bash
pip install fish-speech
```
## 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.
## Quick Start
### Text-Only Synthesis
```bash
python end2end.py --text "Hello, this is a test of the Fish Speech text to speech system."
```
Generated audio files are saved to `output_audio/` by default.
### Voice Cloning
Provide a reference audio file and its transcript to clone a voice:
```bash
python end2end.py \
--text "Hello, this is a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "Transcript of the reference audio."
```
## Streaming Mode
Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
```bash
python end2end.py --text "Hello, this is a test." --streaming
```
Each DAC decoder chunk is logged as it arrives. The final WAV file is written once generation completes.
## Configuration
| Argument | Default | Description |
|----------|---------|-------------|
| `--text` | `"Hello, this is a test..."` | Text to synthesize |
| `--ref-audio` | None | Path to reference audio for voice cloning |
| `--ref-text` | None | Transcript of the reference audio |
| `--model` | `fishaudio/s2-pro` | HuggingFace model path |
| `--stage-configs-path` | `fish_speech_s2_pro.yaml` | Path to stage configs YAML |
| `--output-dir` | `output_audio` | Output directory for WAV files |
| `--streaming` | False | Enable streaming via AsyncOmni |
| `--stage-init-timeout` | 600 | Stage initialization timeout (seconds) |
## Notes
- Output audio is 44.1 kHz mono WAV.
- The DAC codec weights (`codec.pth`) are loaded lazily from the model directory.
- Voice cloning encodes the reference audio using the DAC codec to extract semantic codes, then prepends them as a system message.
@@ -1,252 +0,0 @@
"""Offline inference demo for Fish Speech S2 Pro via vLLM Omni.
Generates speech from text using the fishaudio/s2-pro model.
Supports both text-only synthesis and voice cloning with reference audio.
Usage:
# Text-only synthesis
python end2end.py --text "Hello, this is a test."
# Voice cloning with reference audio
python end2end.py --text "Hello, this is a test." --ref-audio ref.wav --ref-text "Reference text."
# Streaming mode
python end2end.py --text "Hello, this is a test." --streaming
"""
import asyncio
import logging
import os
import time
import soundfile as sf
import torch
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import AsyncOmni, Omni
logger = logging.getLogger(__name__)
DEFAULT_MODEL = "fishaudio/s2-pro"
DEFAULT_STAGE_CONFIG = os.path.join(
os.path.dirname(__file__),
"..",
"..",
"..",
"vllm_omni",
"model_executor",
"stage_configs",
"fish_speech_s2_pro.yaml",
)
def build_prompt(
text: str,
ref_audio_path: str | None = None,
ref_text: str | None = None,
model_name: str = DEFAULT_MODEL,
) -> dict:
"""Build a prompt for Fish Speech S2 Pro.
Uses Qwen3 chat template format with <|voice|> modality marker:
<|im_start|>user\n<|speaker:0|>text<|im_end|>\n<|im_start|>assistant\n<|voice|>
The model then generates semantic tokens autoregressively.
"""
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
# Build user message with speaker tag prefix.
user_text = f"<|speaker:0|>{text}"
# Use chat template for user turn + assistant generation prompt.
messages = [{"role": "user", "content": user_text}]
prompt_ids = tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
)
# Append <|voice|> token (151673) to signal voice generation mode.
voice_token_id = tokenizer.encode("<|voice|>", add_special_tokens=False)
prompt_ids = prompt_ids + voice_token_id
additional_information = {
"text": [text],
"max_new_tokens": [4096],
}
if ref_audio_path:
additional_information["ref_audio"] = [ref_audio_path]
if ref_text:
additional_information["ref_text"] = [ref_text]
return {
"prompt_token_ids": prompt_ids,
"additional_information": additional_information,
}
def _save_wav(output_dir: str, request_id: str, mm: dict) -> None:
"""Concatenate audio chunks and write to a wav file."""
audio_data = mm["audio"]
sr_raw = mm["sr"]
sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
out_wav = os.path.join(output_dir, f"output_{request_id}.wav")
sf.write(out_wav, audio_tensor.float().cpu().numpy().flatten(), samplerate=sr, format="WAV")
logger.info("Request %s: saved audio to %s (sr=%d)", request_id, out_wav, sr)
def main(args):
"""Run offline inference with Omni."""
output_dir = args.output_dir
os.makedirs(output_dir, exist_ok=True)
stage_configs_path = args.stage_configs_path or DEFAULT_STAGE_CONFIG
model_name = args.model or DEFAULT_MODEL
inputs = [
build_prompt(
text=args.text,
ref_audio_path=args.ref_audio,
ref_text=args.ref_text,
model_name=model_name,
)
]
omni = Omni(
model=model_name,
stage_configs_path=stage_configs_path,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
t_start = time.perf_counter()
for stage_outputs in omni.generate(inputs):
for output in stage_outputs.request_output:
_save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
t_end = time.perf_counter()
logger.info("Total inference time: %.1f ms", (t_end - t_start) * 1000)
async def main_streaming(args):
"""Run offline inference with AsyncOmni for streaming audio chunks."""
output_dir = args.output_dir
os.makedirs(output_dir, exist_ok=True)
stage_configs_path = args.stage_configs_path or DEFAULT_STAGE_CONFIG
model_name = args.model or DEFAULT_MODEL
prompt = build_prompt(
text=args.text,
ref_audio_path=args.ref_audio,
ref_text=args.ref_text,
model_name=model_name,
)
omni = AsyncOmni(
model=model_name,
stage_configs_path=stage_configs_path,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
request_id = "0"
t_start = time.perf_counter()
t_prev = t_start
chunk_idx = 0
async for stage_output in omni.generate(prompt, request_id=request_id):
mm = stage_output.request_output.outputs[0].multimodal_output
if not stage_output.finished:
t_now = time.perf_counter()
audio = mm.get("audio")
n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1)
dt_ms = (t_now - t_prev) * 1000
ttfa_ms = (t_now - t_start) * 1000
if chunk_idx == 0:
logger.info("Request %s: chunk %d samples=%d TTFA=%.1fms", request_id, chunk_idx, n, ttfa_ms)
else:
logger.info("Request %s: chunk %d samples=%d inter_chunk=%.1fms", request_id, chunk_idx, n, dt_ms)
t_prev = t_now
chunk_idx += 1
else:
t_end = time.perf_counter()
total_ms = (t_end - t_start) * 1000
logger.info("Request %s: done total=%.1fms chunks=%d", request_id, total_ms, chunk_idx)
_save_wav(output_dir, request_id, mm)
def parse_args():
parser = FlexibleArgumentParser(
description="Fish Speech S2 Pro offline inference with vLLM Omni",
)
parser.add_argument(
"--text",
type=str,
default="Hello, this is a test of the Fish Speech text to speech system.",
help="Text to synthesize.",
)
parser.add_argument(
"--ref-audio",
type=str,
default=None,
help="Path to reference audio for voice cloning.",
)
parser.add_argument(
"--ref-text",
type=str,
default=None,
help="Reference text corresponding to the reference audio.",
)
parser.add_argument(
"--model",
type=str,
default=None,
help=f"HuggingFace model path (default: {DEFAULT_MODEL}).",
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="Path to stage configs YAML.",
)
parser.add_argument(
"--output-dir",
default="output_audio",
help="Output directory for generated wav files.",
)
parser.add_argument(
"--log-stats",
action="store_true",
default=False,
help="Enable writing detailed statistics.",
)
parser.add_argument(
"--stage-init-timeout",
type=int,
default=600,
help="Timeout for initializing stages in seconds.",
)
parser.add_argument(
"--streaming",
action="store_true",
default=False,
help="Stream audio chunks as they arrive via AsyncOmni.",
)
return parser.parse_args()
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
args = parse_args()
if args.streaming:
asyncio.run(main_streaming(args))
else:
main(args)
@@ -208,11 +208,6 @@ def parse_args() -> argparse.Namespace:
default=1,
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
)
parser.add_argument(
"--enable-expert-parallel",
action="store_true",
help="Enable expert parallelism for MoE layers.",
)
parser.add_argument("--layers", type=int, default=4, help="Number of layers to decompose the input image into.")
parser.add_argument(
"--resolution",
@@ -353,7 +348,6 @@ def main():
ring_degree=args.ring_degree,
cfg_parallel_size=args.cfg_parallel_size,
tensor_parallel_size=args.tensor_parallel_size,
enable_expert_parallel=args.enable_expert_parallel,
)
# Configure cache based on backend type
@@ -408,7 +402,7 @@ def main():
else:
print(f" Input image size: {input_image.size}")
print(
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}"
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}"
)
print(f"{'=' * 60}\n")
@@ -96,8 +96,7 @@ python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_str
```
Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames`
in the stage config). The initial chunk size is dynamically selected based on server load for reduced
TTFA, and can be overridden per-request via the `initial_codec_chunk_frames` API field. The final WAV file is written once generation
and `initial_codec_chunk_frames` in the stage config). The final WAV file is written once generation
completes. This demonstrates that audio data is available progressively rather than only at the end.
> **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.
@@ -106,23 +106,6 @@ python text_to_image.py \
--seed 42
```
### Flux.2-dev Models
To start Flux.2-dev with a single GPU, cpu-offload must be enabled because the total size of its weights exceeds the 80GB memory capacity of the GPU.
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model black-forest-labs/FLUX.2-dev \
--prompt "a lovely bunny holding a sign that says 'vllm-omni'" \
--seed 42 \
--tensor-parallel-size 1 \
--num-images-per-prompt 1 \
--num-inference-steps 50 \
--guidance-scale 4.0 \
--height 1024 \
--width 1024 \
--enable-cpu-offload \
--output flux2-dev.png
```
### Key Arguments
**Common arguments:**
@@ -170,11 +170,6 @@ def parse_args() -> argparse.Namespace:
default=1,
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
)
parser.add_argument(
"--enable-expert-parallel",
action="store_true",
help="Enable expert parallelism for MoE layers.",
)
parser.add_argument(
"--lora-path",
type=str,
@@ -263,7 +258,6 @@ def main():
cfg_parallel_size=args.cfg_parallel_size,
tensor_parallel_size=args.tensor_parallel_size,
vae_patch_parallel_size=args.vae_patch_parallel_size,
enable_expert_parallel=args.enable_expert_parallel,
)
# Check if profiling is requested via environment variable
@@ -329,7 +323,7 @@ def main():
print(
f" Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, "
f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, "
f"vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}."
f"vae_patch_parallel_size={args.vae_patch_parallel_size}"
)
print(f" CPU offload: {args.enable_cpu_offload}")
print(f" Image size: {args.width}x{args.height}")
@@ -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(
"--enable-expert-parallel",
action="store_true",
help="Enable expert parallelism for MoE layers.",
)
return parser.parse_args()
@@ -165,7 +160,6 @@ def main():
cfg_parallel_size=args.cfg_parallel_size,
tensor_parallel_size=args.tensor_parallel_size,
vae_patch_parallel_size=args.vae_patch_parallel_size,
enable_expert_parallel=args.enable_expert_parallel,
)
# Check if profiling is requested via environment variable
@@ -199,7 +193,7 @@ def main():
print(
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree},"
f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size},"
f" vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}"
f" vae_patch_parallel_size={args.vae_patch_parallel_size}"
)
print(f" Video size: {args.width}x{args.height}")
print(f"{'=' * 60}\n")
@@ -1,103 +0,0 @@
import argparse
import fcntl
import os
import site
import subprocess
import tempfile
import time
from pathlib import Path
from huggingface_hub import snapshot_download
DEPENDENCY_REPO = "https://github.com/bytedance/DreamID-V.git"
DEPENDENCY_BRANCH = "omni"
CACHE_DIR = Path(tempfile.gettempdir()) / "vllm-omni-dependency"
LOCK_FILE = CACHE_DIR / ".install.lock"
DEPENDENCY_DIR = CACHE_DIR / "DreamID-Omni"
def download_dependency():
CACHE_DIR.mkdir(parents=True, exist_ok=True)
with open(LOCK_FILE, "w") as f:
fcntl.flock(f, fcntl.LOCK_EX)
if not DEPENDENCY_DIR.exists():
print(f"Downloading DreamID-Omni to {DEPENDENCY_DIR} ...")
subprocess.run(
["git", "clone", "--depth", "1", DEPENDENCY_REPO, "--branch", DEPENDENCY_BRANCH, str(DEPENDENCY_DIR)],
check=True,
)
print("Download finished.")
fcntl.flock(f, fcntl.LOCK_UN)
# write .pth to site-packages
site_packages = Path(site.getsitepackages()[0])
pth_file = site_packages / "vllm_omni_dependency.pth"
pth_file.write_text(str(DEPENDENCY_DIR))
print(f"Added {DEPENDENCY_DIR} to site-packages via {pth_file}")
def timed_download(repo_id: str, local_dir: str, allow_patterns: list | None = None):
"""Download files from HF repo and log time + destination."""
if os.path.exists(local_dir):
print(f"Directory {local_dir} already exists. Skipping download.")
return
print(f"Starting download from {repo_id} into {local_dir}")
start_time = time.time()
snapshot_download(
repo_id=repo_id,
local_dir=local_dir,
local_dir_use_symlinks=False,
allow_patterns=allow_patterns,
)
elapsed = time.time() - start_time
print(f"✅ Finished downloading {repo_id} in {elapsed:.2f} seconds. Files saved at: {local_dir}")
def main(output_dir: str):
# Wan2.2
wan_dir = os.path.join(output_dir, "Wan2.2-TI2V-5B")
timed_download(
repo_id="Wan-AI/Wan2.2-TI2V-5B",
local_dir=wan_dir,
allow_patterns=["google/*", "models_t5_umt5-xxl-enc-bf16.pth", "Wan2.2_VAE.pth"],
)
# MMAudio
mm_audio_dir = os.path.join(output_dir, "MMAudio")
timed_download(
repo_id="hkchengrex/MMAudio",
local_dir=mm_audio_dir,
allow_patterns=["ext_weights/best_netG.pt", "ext_weights/v1-16.pth"],
)
dreamid_dir = os.path.join(output_dir, "DreamID-Omni")
timed_download(repo_id="XuGuo699/DreamID-Omni", local_dir=dreamid_dir)
# Now we construct the config file
import json
data = {
"_class_name": "DreamIDOmniPipeline",
"fusion": "DreamID-Omni/dreamid_omni.safetensors",
}
with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
json.dump(data, f, indent=2)
print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}")
# now we download the dependency code
download_dependency()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Download models from Hugging Face")
parser.add_argument(
"--output-dir", type=str, default="./dreamid_omni", help="Base directory to save downloaded models"
)
args = parser.parse_args()
main(args.output_dir)
@@ -1,58 +0,0 @@
# X-To-Video-Audio
The `DreamID-Omni` pipeline generates short videos from text, image and video.
## Local CLI Usage
### Download the Model locally
Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
```bash
python download_dreamid_omni.py --output-dir ./dreamid_omni
```
After download, the model directory will look like this:
```
dreamid_omni/
├── DreamID-Omni/
│ ├── dreamid_omni.safetensors
├── MMAudio/
│ ├── ext_weights/
│ │ ├── best_netG.pt
│ │ ├── v1-16.pth
├── Wan2.2-TI2V-5B/
│ ├── google/*
│ ├── models_t5_umt5-xxl-enc-bf16.pth
│ ├── Wan2.2_VAE.pth
│
├── model_index.json # create by download_dreamid_omni.py
```
### Run the Inference
```
python x_to_video_audio.py \
--model /xx/dreamid_omni \
--prompt "Two people walking together and singing happily" \
--image-path ./example0.png ./example1.png \
--audio-path ./example0.wav ./example1.wav \
--video-negative-prompt "jitter, bad hands, blur, distortion" \
--audio-negative-prompt "robotic, muffled, echo, distorted" \
--cfg-parallel-size 2 \
--num-inference-steps 45 \
--height 704 \
--width 1280 \
--output dreamid_omni.mp4
```
In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
Key arguments:
- `--prompt`: text description (string).
- `--model`: path to the model local directory.
- `--height/--width`: output resolution (defaults 704 * 1024).
- `--image-path`: path to the input image list.
- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
- `--num-inference-steps`: number of denoising steps (defaults 45).
- `--video-negative-prompt`: negative prompt for video generation.
- `--audio-negative-prompt`: negative prompt for audio generation.
- `--enable-cpu-offload`: enable CPU offload (defaults False).
@@ -1,148 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import re
import time
import librosa
from PIL import Image
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"}
AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Offline inference for DreamID-Omni (video + audio).")
parser.add_argument("--model", required=True, help="DreamID ckpt root directory.")
parser.add_argument("--model-type", default="dreamid-omni", help="Model type.")
parser.add_argument("--prompt", default=None, help="Text prompt.")
parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path")
parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
parser.add_argument("--height", type=int, default=704, help="Video height.")
parser.add_argument("--width", type=int, default=1280, help="Video width.")
parser.add_argument("--num-inference-steps", type=int, default=45, help="Sampling steps.")
parser.add_argument("--solver-name", default="unipc", help="Solver name: unipc|dpm++|euler.")
parser.add_argument("--shift", type=float, default=5.0, help="Scheduler shift.")
parser.add_argument("--seed", type=int, default=103, help="Random seed for reproducible generation.")
parser.add_argument(
"--cfg-parallel-size",
type=int,
default=1,
choices=[1, 2],
help="Number of GPUs used for classifier free guidance parallel size.",
)
parser.add_argument(
"--video-negative-prompt",
default="jitter, bad hands, blur, distortion",
help="Negative prompt for video.",
)
parser.add_argument(
"--audio-negative-prompt",
default="robotic, muffled, echo, distorted",
help="Negative prompt for audio.",
)
parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
parser.add_argument(
"--enable-cpu-offload",
action="store_true",
default=False,
help="Enable CPU offloading for diffusion models.",
)
return parser.parse_args()
def load_image_and_audio(image_paths, audio_paths):
image = []
audio = []
for path in image_paths:
with Image.open(path) as img:
img = img.convert("RGB")
image.append(img)
for path in audio_paths:
audio_array, sr = librosa.load(path, sr=16000)
audio_array = audio_array[int(sr * 1) : int(sr * 3)]
audio.append(audio_array)
return image, audio
def main() -> None:
args = parse_args()
if args.prompt is None and args.prompt_file is None:
raise ValueError("Either --prompt or --prompt-file must be provided.")
text_prompt = args.prompt
if args.prompt_file:
import json
with open(args.prompt_file) as f:
text_prompt = json.load(f)
text_prompt = re.sub(
r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", "", text_prompt, flags=re.DOTALL
).strip()
text_prompt = re.sub(
r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", "", text_prompt, flags=re.DOTALL
).strip()
text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
image, audio = load_image_and_audio(args.image_path, args.audio_path)
prompt = {
"prompt": text_prompt,
"video_negative_prompt": args.video_negative_prompt,
"audio_negative_prompt": args.audio_negative_prompt,
"multi_modal_data": {"image": image, "audio": audio},
}
sampling_params = OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
num_inference_steps=args.num_inference_steps,
seed=args.seed,
extra_args={
"solver_name": args.solver_name,
"shift": args.shift,
},
)
parallel_config = DiffusionParallelConfig(
cfg_parallel_size=args.cfg_parallel_size,
)
omni = Omni(
model=args.model,
parallel_config=parallel_config,
model_type=args.model_type,
enable_cpu_offload=args.enable_cpu_offload,
)
start = time.perf_counter()
outputs = omni.generate(prompt, sampling_params)
elapsed = time.perf_counter() - start
if not outputs:
raise RuntimeError("No output returned from DreamID-Omni.")
output = outputs[0].request_output
generated_video = output[0].images[0][0]
generated_audio = output[0].images[0][1]
try:
from dreamid_omni.utils.io_utils import save_video
except Exception as e:
raise RuntimeError(f"Failed to extract video and audio from DreamID-Omni output. Error: {e}")
output_path = args.output
save_video(output_path, generated_video, generated_audio, fps=24, sample_rate=16000)
print(f"Saved generated video to {output_path}")
print(f"Total time: {elapsed:.2f}s")
if __name__ == "__main__":
main()
@@ -1,172 +0,0 @@
# Fish Speech S2 Pro
## Model
| Model | Description |
|-------|-------------|
| `fishaudio/s2-pro` | Fish Speech S2 Pro -- 4B dual-AR TTS model with DAC codec (44.1 kHz) |
## Gradio Demo
An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support.
```bash
# Option 1: Launch server + Gradio together
./run_gradio_demo.sh
# Option 2: If server is already running
python gradio_demo.py --api-base http://localhost:8091
```
Then open http://localhost:7860 in your browser.
Features:
- Text-to-speech synthesis
- Voice cloning from uploaded audio or URL
- Streaming mode (progressive PCM playback)
## Launch the Server
```bash
vllm-omni serve fishaudio/s2-pro \
--stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
--omni \
--port 8091 \
--trust-remote-code \
--enforce-eager \
--gpu-memory-utilization 0.9
```
Or use the convenience script:
```bash
./run_server.sh
```
## Send TTS Request
### Using curl
```bash
# Basic TTS
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, how are you?",
"voice": "default",
"response_format": "wav"
}' --output output.wav
```
### Voice Cloning
Provide a reference audio (URL or base64 data URL) and its transcript:
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, this is a cloned voice.",
"voice": "default",
"ref_audio": "https://example.com/reference.wav",
"ref_text": "Transcript of the reference audio."
}' --output cloned.wav
```
### Using Python
```python
import httpx
# Basic TTS
response = httpx.post(
"http://localhost:8091/v1/audio/speech",
json={
"input": "Hello, how are you?",
"voice": "default",
"response_format": "wav",
},
timeout=300.0,
)
with open("output.wav", "wb") as f:
f.write(response.content)
```
### Using the CLI Client
```bash
cd examples/online_serving/fish_speech
# Basic TTS
python speech_client.py --text "Hello, how are you?"
# Voice cloning
python speech_client.py \
--text "Hello, this is a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "Transcript of the reference audio."
# Streaming PCM output
python speech_client.py --text "Hello world" --stream --output output.pcm
```
The CLI client supports:
- `--api-base`: API base URL (default: `http://localhost:8091`)
- `--model` (or `-m`): Model name (default: `fishaudio/s2-pro`)
- `--text`: Text to synthesize (required)
- `--ref-audio`: Reference audio for voice cloning (local path or URL)
- `--ref-text`: Transcript of the reference audio
- `--stream`: Enable streaming (PCM output)
- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav)
- `--output` (or `-o`): Output file path
## Streaming
Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded:
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, how are you?",
"voice": "default",
"stream": true,
"response_format": "pcm"
}' --no-buffer | play -t raw -r 44100 -e signed -b 16 -c 1 -
```
**Note:** Fish Speech outputs at 44.1 kHz (unlike Qwen3-TTS which outputs at 24 kHz).
## API Parameters
Fish Speech uses the same `/v1/audio/speech` endpoint as Qwen3-TTS. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full parameter documentation.
Key parameters for Fish Speech:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input` | string | **required** | Text to synthesize |
| `voice` | string | "default" | Voice name (use "default" for Fish Speech) |
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
| `ref_audio` | string | null | Reference audio URL or base64 data URL for voice cloning |
| `ref_text` | string | null | Transcript of reference audio (required for voice cloning) |
| `max_new_tokens` | int | 4096 | Maximum tokens to generate |
| `stream` | bool | false | Stream raw PCM chunks |
## Prerequisites
Install the `fish-speech` package for the DAC codec:
```bash
pip install fish-speech
```
## Troubleshooting
1. **No audio output**: Make sure the `fish-speech` package is installed for the DAC decoder
2. **Connection refused**: Ensure the server is running on the correct port
3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings
4. **Out of memory**: Reduce `--gpu-memory-utilization` or use a GPU with more VRAM
@@ -1,272 +0,0 @@
"""Gradio demo for Fish Speech S2 Pro online serving via /v1/audio/speech API.
Supports:
- Text-to-speech synthesis
- Voice cloning from reference audio (upload or URL)
- Streaming (progressive PCM chunks) and non-streaming modes
Usage:
# Start the server first (see run_server.sh), then:
python gradio_demo.py --api-base http://localhost:8091
# Or use run_gradio_demo.sh to start both server and demo together.
"""
import argparse
import base64
import io
import gradio as gr
import httpx
import numpy as np
import soundfile as sf
PCM_SAMPLE_RATE = 44100
DEFAULT_API_BASE = "http://localhost:8091"
def encode_audio_to_base64(audio_data: tuple) -> str:
"""Encode Gradio audio input (sample_rate, numpy_array) to base64 data URL."""
sample_rate, audio_np = audio_data
if audio_np.dtype != np.int16:
if audio_np.dtype in (np.float32, np.float64):
audio_np = np.clip(audio_np, -1.0, 1.0)
audio_np = (audio_np * 32767).astype(np.int16)
else:
audio_np = audio_np.astype(np.int16)
buf = io.BytesIO()
sf.write(buf, audio_np, sample_rate, format="WAV")
wav_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
return f"data:audio/wav;base64,{wav_b64}"
def build_payload(
text: str,
ref_audio: tuple | None,
ref_audio_url: str,
ref_text: str,
response_format: str = "wav",
stream: bool = False,
) -> dict:
"""Build the /v1/audio/speech request payload."""
if not text or not text.strip():
raise gr.Error("Please enter text to synthesize.")
payload: dict = {
"input": text.strip(),
"voice": "default",
"response_format": "pcm" if stream else response_format,
"stream": stream,
}
# Voice cloning: ref_audio takes priority over URL.
ref_url_stripped = ref_audio_url.strip() if ref_audio_url else ""
if ref_audio is not None:
payload["ref_audio"] = encode_audio_to_base64(ref_audio)
elif ref_url_stripped:
payload["ref_audio"] = ref_url_stripped
if "ref_audio" in payload:
if not ref_text or not ref_text.strip():
raise gr.Error("Voice cloning requires a transcript of the reference audio.")
payload["ref_text"] = ref_text.strip()
return payload
def generate_speech(api_base: str, text: str, ref_audio, ref_audio_url, ref_text, response_format):
"""Non-streaming: call /v1/audio/speech and return full audio."""
payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=False)
try:
with httpx.Client(timeout=300.0) as client:
resp = client.post(
f"{api_base}/v1/audio/speech",
json=payload,
headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
)
except httpx.TimeoutException:
raise gr.Error("Request timed out. The server may be busy.")
except httpx.ConnectError:
raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?")
if resp.status_code != 200:
raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
content_type = resp.headers.get("content-type", "")
if "application/json" in content_type:
try:
raise gr.Error(f"Server error: {resp.json()}")
except ValueError:
pass
try:
if response_format == "pcm":
audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0
return (PCM_SAMPLE_RATE, audio_np)
audio_np, sample_rate = sf.read(io.BytesIO(resp.content))
if audio_np.ndim > 1:
audio_np = audio_np[:, 0]
return (sample_rate, audio_np.astype(np.float32))
except Exception as e:
raise gr.Error(f"Failed to decode audio: {e}")
def generate_speech_stream(api_base, text, ref_audio, ref_audio_url, ref_text, response_format, chunk_seconds=0.5):
"""Streaming: yield PCM chunks for progressive playback."""
payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=True)
min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds)
prebuffer_chunks = 2
pending: list[np.ndarray] = []
pending_len = 0
chunks_yielded = 0
prebuffer: list[np.ndarray] = []
try:
leftover = b""
with httpx.Client(timeout=300.0) as client:
with client.stream(
"POST",
f"{api_base}/v1/audio/speech",
json=payload,
headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
) as resp:
if resp.status_code != 200:
resp.read()
raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
for chunk in resp.iter_bytes():
if not chunk:
continue
raw = leftover + chunk
usable = len(raw) - (len(raw) % 2)
leftover = raw[usable:]
if usable == 0:
continue
float_samples = np.frombuffer(raw[:usable], dtype=np.int16).copy().astype(np.float32) / 32767.0
pending.append(float_samples)
pending_len += len(float_samples)
if pending_len >= min_chunk_samples:
audio_chunk = np.concatenate(pending)
pending.clear()
pending_len = 0
if chunks_yielded < prebuffer_chunks:
prebuffer.append(audio_chunk)
chunks_yielded += 1
if chunks_yielded == prebuffer_chunks:
yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer))
prebuffer.clear()
else:
yield (PCM_SAMPLE_RATE, audio_chunk)
remaining = prebuffer + pending
if remaining:
yield (PCM_SAMPLE_RATE, np.concatenate(remaining))
except httpx.TimeoutException:
raise gr.Error("Request timed out. The server may be busy.")
except httpx.ConnectError:
raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?")
def on_stream_change(stream: bool):
"""Lock format to PCM when streaming."""
if stream:
return gr.update(value="pcm", interactive=False)
return gr.update(value="wav", interactive=True)
def build_interface(api_base: str, stream_chunk_seconds: float = 0.5):
"""Build the Gradio interface."""
with gr.Blocks(title="Fish Speech S2 Pro Demo") as demo:
gr.Markdown("# Fish Speech S2 Pro - Text to Speech")
gr.Markdown(f"**Server:** `{api_base}` | **Model:** fishaudio/s2-pro | **Output:** 44.1kHz")
with gr.Row():
with gr.Column(scale=3):
text_input = gr.Textbox(
label="Text to Synthesize",
placeholder="Enter text here...",
lines=4,
)
with gr.Accordion("Voice Cloning (optional)", open=False):
gr.Markdown(
"Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice."
)
ref_audio = gr.Audio(
label="Reference Audio",
type="numpy",
sources=["upload", "microphone"],
)
ref_audio_url = gr.Textbox(
label="Reference Audio URL (alternative to upload)",
placeholder="https://example.com/reference.wav",
lines=1,
)
ref_text = gr.Textbox(
label="Reference Audio Transcript (required for cloning)",
placeholder="Exact transcript of the reference audio...",
lines=2,
)
with gr.Row():
response_format = gr.Dropdown(
choices=["wav", "mp3", "flac", "pcm"],
value="wav",
label="Audio Format",
scale=1,
)
stream_checkbox = gr.Checkbox(
label="Stream output",
value=False,
info="Progressive PCM streaming",
scale=1,
)
generate_btn = gr.Button("Generate Speech", variant="primary", size="lg")
with gr.Column(scale=2):
audio_output = gr.Audio(
label="Generated Audio",
interactive=False,
streaming=True,
autoplay=True,
)
gr.Markdown(
"### About\n"
"- **Fish Speech S2 Pro** by FishAudio: 4B dual-AR model\n"
"- **Voice cloning**: upload 10-30s reference + transcript\n"
"- **Streaming**: real-time PCM output\n"
"- **44.1kHz** output via DAC codec"
)
stream_checkbox.change(fn=on_stream_change, inputs=[stream_checkbox], outputs=[response_format])
all_inputs = [text_input, ref_audio, ref_audio_url, ref_text, response_format]
def dispatch(stream_enabled, *args):
if stream_enabled:
yield from generate_speech_stream(api_base, *args, chunk_seconds=stream_chunk_seconds)
else:
yield generate_speech(api_base, *args)
generate_btn.click(fn=dispatch, inputs=[stream_checkbox] + all_inputs, outputs=[audio_output])
demo.queue()
return demo
def main():
parser = argparse.ArgumentParser(description="Gradio demo for Fish Speech S2 Pro")
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})")
parser.add_argument("--host", default="0.0.0.0", help="Gradio host (default: 0.0.0.0)")
parser.add_argument("--port", type=int, default=7860, help="Gradio port (default: 7860)")
parser.add_argument("--share", action="store_true", help="Share publicly")
parser.add_argument("--stream-chunk-seconds", type=float, default=0.5, help="Seconds per streaming chunk")
args = parser.parse_args()
print(f"Connecting to vLLM server at: {args.api_base}")
demo = build_interface(args.api_base, stream_chunk_seconds=args.stream_chunk_seconds)
demo.launch(server_name=args.host, server_port=args.port, share=args.share)
if __name__ == "__main__":
main()
@@ -1,48 +0,0 @@
#!/bin/bash
# Launch Fish Speech S2 Pro server + Gradio demo together.
#
# Usage:
# ./run_gradio_demo.sh
# CUDA_VISIBLE_DEVICES=0 PORT=8091 GRADIO_PORT=7860 ./run_gradio_demo.sh
set -e
MODEL="${MODEL:-fishaudio/s2-pro}"
PORT="${PORT:-8091}"
GRADIO_PORT="${GRADIO_PORT:-7860}"
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
echo "Starting Fish Speech S2 Pro server (port $PORT)..."
FLASHINFER_DISABLE_VERSION_CHECK=1 \
vllm-omni serve "$MODEL" \
--stage-configs-path "$REPO_ROOT/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml" \
--host 0.0.0.0 \
--port "$PORT" \
--gpu-memory-utilization 0.9 \
--trust-remote-code \
--enforce-eager \
--omni &
SERVER_PID=$!
cleanup() {
echo "Stopping server (PID $SERVER_PID)..."
kill $SERVER_PID 2>/dev/null
wait $SERVER_PID 2>/dev/null
}
trap cleanup EXIT
# Wait for server to be ready.
echo "Waiting for server to start..."
for i in $(seq 1 120); do
if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then
echo "Server ready."
break
fi
sleep 2
done
echo "Starting Gradio demo (port $GRADIO_PORT)..."
python "$SCRIPT_DIR/gradio_demo.py" \
--api-base "http://localhost:$PORT" \
--port "$GRADIO_PORT"
@@ -1,23 +0,0 @@
#!/bin/bash
# Launch vLLM-Omni server for Fish Speech S2 Pro
#
# Usage:
# ./run_server.sh
# CUDA_VISIBLE_DEVICES=0 ./run_server.sh
set -e
MODEL="${MODEL:-fishaudio/s2-pro}"
PORT="${PORT:-8091}"
echo "Starting Fish Speech S2 Pro server with model: $MODEL"
FLASHINFER_DISABLE_VERSION_CHECK=1 \
vllm-omni serve "$MODEL" \
--stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
--host 0.0.0.0 \
--port "$PORT" \
--gpu-memory-utilization 0.9 \
--trust-remote-code \
--enforce-eager \
--omni
@@ -1,129 +0,0 @@
"""Client for Fish Speech S2 Pro via /v1/audio/speech endpoint.
Examples:
# Basic TTS
python speech_client.py --text "Hello, how are you?"
# Voice cloning
python speech_client.py --text "Hello, how are you?" \
--ref-audio ref.wav --ref-text "This is the reference transcript."
# Streaming PCM output
python speech_client.py --text "Hello world" --stream --output output.pcm
"""
import argparse
import base64
import os
import httpx
DEFAULT_API_BASE = "http://localhost:8091"
DEFAULT_API_KEY = "EMPTY"
def encode_audio_to_base64(audio_path: str) -> str:
"""Encode a local audio file to base64 data URL."""
if not os.path.exists(audio_path):
raise FileNotFoundError(f"Audio file not found: {audio_path}")
ext = audio_path.lower().rsplit(".", 1)[-1]
mime_map = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac", "ogg": "audio/ogg"}
mime_type = mime_map.get(ext, "audio/wav")
with open(audio_path, "rb") as f:
audio_b64 = base64.b64encode(f.read()).decode("utf-8")
return f"data:{mime_type};base64,{audio_b64}"
def run_tts(args) -> None:
"""Generate speech via /v1/audio/speech API."""
payload = {
"model": args.model,
"input": args.text,
"voice": "default",
"response_format": args.response_format,
}
# Voice cloning parameters.
if args.ref_audio:
if args.ref_audio.startswith(("http://", "https://")):
payload["ref_audio"] = args.ref_audio
else:
payload["ref_audio"] = encode_audio_to_base64(args.ref_audio)
if args.ref_text:
payload["ref_text"] = args.ref_text
if args.stream:
payload["stream"] = True
payload["response_format"] = "pcm"
print(f"Model: {args.model}")
print(f"Text: {args.text}")
if args.ref_audio:
print(f"Voice cloning: ref_audio={args.ref_audio}, ref_text={args.ref_text}")
print("Generating audio...")
api_url = f"{args.api_base}/v1/audio/speech"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {args.api_key}",
}
if args.stream:
output_path = args.output or "output.pcm"
with httpx.Client(timeout=300.0) as client:
with client.stream("POST", api_url, json=payload, headers=headers) as resp:
if resp.status_code != 200:
print(f"Error: {resp.status_code}")
print(resp.read().decode())
return
total_bytes = 0
with open(output_path, "wb") as f:
for chunk in resp.iter_bytes():
f.write(chunk)
total_bytes += len(chunk)
print(f"Streamed {total_bytes} bytes to: {output_path}")
else:
with httpx.Client(timeout=300.0) as client:
response = client.post(api_url, json=payload, headers=headers)
if response.status_code != 200:
print(f"Error: {response.status_code}")
print(response.text)
return
try:
text = response.content.decode("utf-8")
if text.startswith('{"error"'):
print(f"Error: {text}")
return
except UnicodeDecodeError:
pass
output_path = args.output or "output.wav"
with open(output_path, "wb") as f:
f.write(response.content)
print(f"Audio saved to: {output_path}")
def main():
parser = argparse.ArgumentParser(description="Fish Speech S2 Pro TTS client")
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
parser.add_argument("--model", "-m", default="fishaudio/s2-pro", help="Model name")
parser.add_argument("--text", required=True, help="Text to synthesize")
parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning (path or URL)")
parser.add_argument("--ref-text", default=None, help="Transcript of reference audio")
parser.add_argument("--stream", action="store_true", help="Enable streaming (PCM output)")
parser.add_argument(
"--response-format",
default="wav",
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
help="Audio format (default: wav)",
)
parser.add_argument("--output", "-o", default=None, help="Output file path")
args = parser.parse_args()
run_tts(args)
if __name__ == "__main__":
main()
@@ -26,35 +26,6 @@ The script allows overriding:
- `CACHE_BACKEND` (default: `none`)
- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
## Async Job Behavior
`POST /v1/videos` is asynchronous. It creates a video job and immediately
returns metadata like the job ID and initial `queued` status. To get the final
artifact, poll the job status and then download the completed file from the
content endpoint.
The main endpoints are:
- `POST /v1/videos`: create a video generation job
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
- `GET /v1/videos`: list stored video jobs
- `GET /v1/videos/{video_id}/content`: download the generated video file
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
## Storage
Generated video files are stored on local disk by the async video API.
Local file storage behavior can be controlled via the following environment variables:
- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
Example:
```bash
export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
```
## API Calls
### Method 1: Using curl
@@ -64,7 +35,7 @@ export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
bash run_curl_image_to_video.sh
# Or execute directly (OpenAI-style multipart)
create_response=$(curl -s http://localhost:8091/v1/videos \
curl -X POST http://localhost:8091/v1/videos \
-H "Accept: application/json" \
-F "prompt=A bear playing with yarn, smooth motion" \
-F "negative_prompt=low quality, blurry, static" \
@@ -78,23 +49,7 @@ create_response=$(curl -s http://localhost:8091/v1/videos \
-F "guidance_scale_2=1.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=12.0" \
-F "seed=42")
video_id=$(echo "$create_response" | jq -r '.id')
while true; do
status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
if [ "$status" = "completed" ]; then
break
fi
if [ "$status" = "failed" ]; then
echo "Video generation failed"
exit 1
fi
sleep 2
done
curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o wan22_i2v_output.mp4
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_i2v_output.mp4
```
## Request Format
@@ -108,18 +63,6 @@ curl -X POST http://localhost:8091/v1/videos \
-F "input_reference=@/path/to/qwen-bear.png"
```
### Alternative JSON-Safe Reference Input
Use `image_reference` when you want to pass a URL or JSON-safe image reference
instead of uploading a file. Do not send `input_reference` and
`image_reference` together.
```bash
curl -X POST http://localhost:8091/v1/videos \
-F "prompt=A bear playing with yarn, smooth motion" \
-F 'image_reference={"image_url":"https://example.com/qwen-bear.png"}'
```
### Generation with Parameters
```bash
@@ -138,60 +81,3 @@ curl -X POST http://localhost:8091/v1/videos \
-F "flow_shift=12.0" \
-F "seed=42"
```
## Create Response Format
`POST /v1/videos` returns a job record, not inline base64 video data.
```json
{
"id": "video_gen_123",
"object": "video",
"status": "queued",
"model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
"prompt": "A bear playing with yarn, smooth motion",
"created_at": 1234567890
}
```
## Retrieve, List, Download, and Delete
### Retrieve a job
```bash
curl -s http://localhost:8091/v1/videos/${video_id} | jq .
```
### List jobs
```bash
curl -s http://localhost:8091/v1/videos | jq .
```
### Download the completed video
```bash
curl -L http://localhost:8091/v1/videos/${video_id}/content -o wan22_i2v_output.mp4
```
### Delete a job and its stored file
```bash
curl -X DELETE http://localhost:8091/v1/videos/${video_id} | jq .
```
## Poll Until Complete
```bash
while true; do
status=$(curl -s http://localhost:8091/v1/videos/${video_id} | jq -r '.status')
if [ "$status" = "completed" ]; then
break
fi
if [ "$status" = "failed" ]; then
echo "Video generation failed"
exit 1
fi
sleep 2
done
```
@@ -1,75 +1,33 @@
#!/bin/bash
# Wan2.2 image-to-video curl example using the async video job API.
set -euo pipefail
# Wan2.2 image-to-video curl example (OpenAI-style multipart)
INPUT_IMAGE="${INPUT_IMAGE:-../../offline_inference/image_to_video/qwen-bear.png}"
BASE_URL="${BASE_URL:-http://localhost:8099}"
OUTPUT_PATH="${OUTPUT_PATH:-wan22_i2v_output.mp4}"
NEGATIVE_PROMPT="${NEGATIVE_PROMPT:-}"
POLL_INTERVAL="${POLL_INTERVAL:-2}"
if [ ! -f "$INPUT_IMAGE" ]; then
echo "Input image not found: $INPUT_IMAGE"
exit 1
fi
create_cmd=(
curl -sS -X POST "${BASE_URL}/v1/videos"
-H "Accept: application/json"
-F "prompt=A bear playing with yarn, smooth motion"
-F "input_reference=@${INPUT_IMAGE}"
-F "seconds=2"
-F "size=832x480"
-F "fps=16"
-F "num_inference_steps=40"
-F "guidance_scale=1.0"
-F "guidance_scale_2=1.0"
-F "boundary_ratio=0.875"
-F "flow_shift=12.0"
-F "seed=42"
)
if [ -n "${NEGATIVE_PROMPT}" ]; then
create_cmd+=(-F "negative_prompt=${NEGATIVE_PROMPT}")
NEGATIVE_PROMPT_FLAG=""
if [ -n "$NEGATIVE_PROMPT" ]; then
NEGATIVE_PROMPT_FLAG="-F negative_prompt=${NEGATIVE_PROMPT}"
fi
create_response="$("${create_cmd[@]}")"
video_id="$(echo "${create_response}" | jq -r '.id')"
if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then
echo "Failed to create video job:"
echo "${create_response}" | jq .
exit 1
fi
curl -X POST http://localhost:8099/v1/videos \
-H "Accept: application/json" \
-F "prompt=A bear playing with yarn, smooth motion" \
$NEGATIVE_PROMPT_FLAG \
-F "input_reference=@${INPUT_IMAGE}" \
-F "seconds=2" \
-F "size=832x480" \
-F "fps=16" \
-F "num_inference_steps=40" \
-F "guidance_scale=1.0" \
-F "guidance_scale_2=1.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=12.0" \
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > "${OUTPUT_PATH}"
echo "Created video job ${video_id}"
echo "${create_response}" | jq .
while true; do
status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
status="$(echo "${status_response}" | jq -r '.status')"
case "${status}" in
queued|in_progress)
echo "Video job ${video_id} status: ${status}"
sleep "${POLL_INTERVAL}"
;;
completed)
echo "${status_response}" | jq .
break
;;
failed)
echo "Video generation failed:"
echo "${status_response}" | jq .
exit 1
;;
*)
echo "Unexpected status response:"
echo "${status_response}" | jq .
exit 1
;;
esac
done
curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
echo "Saved video to ${OUTPUT_PATH}"
+1 -62
View File
@@ -14,31 +14,6 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
## Gradio Demo
Two interactive Gradio demos are available, both supporting all 3 task types:
| Demo | File | Transport | Streaming Quality |
| ---- | ---- | --------- | ----------------- |
| Standard | `gradio_demo.py` | HTTP chunked | May have small gaps between chunks |
| FastRTC | `gradio_fastrtc_demo.py` | WebRTC | Gapless streaming (requires `pip install fastrtc`) |
```bash
# Option 1: Launch server + Standard Gradio together
./run_gradio_demo.sh # CustomVoice (default)
./run_gradio_demo.sh --task-type VoiceDesign # VoiceDesign
./run_gradio_demo.sh --task-type Base # Voice cloning
# Option 2: If server is already running
python gradio_demo.py --api-base http://localhost:8000
# Option 3: FastRTC demo (gapless streaming)
pip install fastrtc
python gradio_fastrtc_demo.py --api-base http://localhost:8000
```
Then open http://localhost:7860 in your browser.
## Run examples (Qwen3-TTS)
### Launch the Server
@@ -314,7 +289,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
| `language` | string | "Auto" | Language (see supported languages below) |
| `instructions` | string | "" | Voice style/emotion instructions |
| `max_new_tokens` | int | 2048 | Maximum tokens to generate |
| `initial_codec_chunk_frames` | int | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
| `initial_codec_chunk_frames` | int | null | Initial chunk size for reduced TTFA (overrides stage config) |
| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
@@ -349,42 +324,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \
- `speed` adjustment is not supported when streaming.
- Requires the server stage config to have `async_chunk: true` (default in `qwen3_tts.yaml`).
## Streaming Text Input (WebSocket)
The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally, buffers and splits it at sentence boundaries, and generates audio per sentence.
When `stream_audio=true`, each sentence is emitted as `audio.start`, one or more binary PCM frames, and `audio.done`.
### Quick Start
```bash
python streaming_speech_client.py \
--text "Hello world. How are you? I am fine."
python streaming_speech_client.py \
--text "Hello world. How are you? I am fine." \
--simulate-stt --stt-delay 0.1
```
### WebSocket Protocol
Client -> Server:
```jsonc
{"type": "session.config", "voice": "Vivian", "task_type": "CustomVoice", "language": "Auto", "split_granularity": "sentence", "stream_audio": true, "response_format": "pcm"}
{"type": "input.text", "text": "Hello, how are you? "}
{"type": "input.done"}
```
Server -> Client:
```jsonc
{"type": "audio.start", "sentence_index": 0, "sentence_text": "Hello, how are you?", "format": "pcm", "sample_rate": 24000}
// binary PCM frame(s)
{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}
{"type": "session.done", "total_sentences": 1}
```
## Limitations
- **Single request**: Batch processing is not yet optimized for online serving.
@@ -1,404 +0,0 @@
"""Gradio demo for Qwen3-TTS online serving via /v1/audio/speech API.
Uses Gradio's built-in audio component for playback. Supports streaming
(progressive PCM chunks) and non-streaming (full audio download) modes.
For gapless streaming playback, see gradio_fastrtc_demo.py which uses WebRTC.
Supports all 3 task types:
- CustomVoice: Predefined speaker with optional style instructions
- VoiceDesign: Natural language voice description
- Base: Voice cloning from reference audio (upload or URL)
Usage:
# Start the server first (see run_server.sh), then:
python gradio_demo.py --api-base http://localhost:8000
# Or use run_gradio_demo.sh to start both server and demo together.
"""
import argparse
import io
import gradio as gr
import httpx
import numpy as np
import soundfile as sf
from tts_common import (
PCM_SAMPLE_RATE,
SUPPORTED_LANGUAGES,
TASK_TYPES,
add_common_args,
build_payload,
fetch_voices,
on_task_type_change,
stream_pcm_chunks,
)
def generate_speech(
api_base: str,
text: str,
task_type: str,
voice: str,
language: str,
instructions: str,
ref_audio: tuple | None,
ref_audio_url: str,
ref_text: str,
x_vector_only: bool,
response_format: str,
speed: float,
):
"""Non-streaming: call /v1/audio/speech and return audio."""
payload = build_payload(
text,
task_type,
voice,
language,
instructions,
ref_audio,
ref_audio_url,
ref_text,
x_vector_only,
response_format,
speed,
stream=False,
)
try:
with httpx.Client(timeout=300.0) as client:
resp = client.post(
f"{api_base}/v1/audio/speech",
json=payload,
headers={
"Content-Type": "application/json",
"Authorization": "Bearer EMPTY",
},
)
except httpx.TimeoutException:
raise gr.Error("Request timed out. The server may be busy.")
except httpx.ConnectError:
raise gr.Error(f"Cannot connect to server at {api_base}. Make sure the vLLM server is running.")
if resp.status_code != 200:
raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
content_type = resp.headers.get("content-type", "")
if "application/json" in content_type:
try:
error_data = resp.json()
raise gr.Error(f"Server error: {error_data}")
except ValueError:
pass
try:
if response_format == "pcm":
audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0
return (PCM_SAMPLE_RATE, audio_np)
audio_np, sample_rate = sf.read(io.BytesIO(resp.content))
if audio_np.ndim > 1:
audio_np = audio_np[:, 0]
return (sample_rate, audio_np.astype(np.float32))
except Exception as e:
raise gr.Error(f"Failed to decode audio response: {e}")
def generate_speech_stream(
api_base: str,
text: str,
task_type: str,
voice: str,
language: str,
instructions: str,
ref_audio: tuple | None,
ref_audio_url: str,
ref_text: str,
x_vector_only: bool,
response_format: str,
speed: float,
chunk_seconds: float = 0.25,
):
"""Streaming: yield individual PCM chunks for progressive playback."""
payload = build_payload(
text,
task_type,
voice,
language,
instructions,
ref_audio,
ref_audio_url,
ref_text,
x_vector_only,
response_format,
speed,
stream=True,
)
min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds)
# Pre-buffer 2 chunks before first playback to avoid the gap
# between chunk 1 and chunk 2 during progressive streaming.
prebuffer_chunks = 2
pending: list[np.ndarray] = []
pending_len = 0
chunks_yielded = 0
prebuffer: list[np.ndarray] = []
try:
for samples in stream_pcm_chunks(api_base, payload):
float_samples = samples.astype(np.float32) / 32767.0
pending.append(float_samples)
pending_len += len(float_samples)
if pending_len >= min_chunk_samples:
audio_chunk = np.concatenate(pending)
pending.clear()
pending_len = 0
if chunks_yielded < prebuffer_chunks:
prebuffer.append(audio_chunk)
chunks_yielded += 1
if chunks_yielded == prebuffer_chunks:
yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer))
prebuffer.clear()
else:
yield (PCM_SAMPLE_RATE, audio_chunk)
# Flush: prebuffer (if stream was shorter than 2 chunks) + remaining pending
remaining = prebuffer + pending
if remaining:
yield (PCM_SAMPLE_RATE, np.concatenate(remaining))
except httpx.TimeoutException:
raise gr.Error("Request timed out. The server may be busy.")
except httpx.ConnectError:
raise gr.Error(f"Cannot connect to server at {api_base}. Make sure the vLLM server is running.")
def on_stream_change(stream: bool):
"""When streaming is enabled, lock format to PCM and disable speed."""
if stream:
return (
gr.update(value="pcm", interactive=False),
gr.update(interactive=False),
)
return (
gr.update(value="wav", interactive=True),
gr.update(interactive=True),
)
def build_interface(api_base: str, stream_chunk_seconds: float = 0.25):
"""Build the Gradio interface."""
voices = fetch_voices(api_base)
css = """
#generate-btn button { width: 100%; }
.task-info { padding: 8px 12px; border-radius: 6px;
background: #f0f4ff; margin-bottom: 8px; }
"""
with gr.Blocks(css=css, title="Qwen3-TTS Demo") as demo:
gr.Markdown("# Qwen3-TTS Online Serving Demo")
gr.Markdown(f"**Server:** `{api_base}`")
with gr.Row():
# Left column: inputs
with gr.Column(scale=3):
text_input = gr.Textbox(
label="Text to Synthesize",
placeholder="Enter text here, e.g., Hello, how are you?",
lines=4,
)
with gr.Row():
task_type = gr.Radio(
choices=TASK_TYPES,
value="CustomVoice",
label="Task Type",
scale=2,
)
language = gr.Dropdown(
choices=SUPPORTED_LANGUAGES,
value="Auto",
label="Language",
scale=1,
)
voice = gr.Dropdown(
choices=voices,
value=voices[0] if voices else None,
label="Speaker",
visible=True,
)
instructions = gr.Textbox(
label="Instructions",
placeholder=("e.g., Speak with excitement / A warm, friendly female voice"),
lines=2,
visible=True,
info="Optional style/emotion instructions",
)
ref_audio = gr.Audio(
label="Reference Audio (upload for voice cloning)",
type="numpy",
sources=["upload", "microphone"],
visible=False,
)
ref_audio_url = gr.Textbox(
label="Reference Audio URL",
placeholder=("https://example.com/reference.wav (alternative to uploading)"),
lines=1,
visible=False,
)
ref_text = gr.Textbox(
label="Reference Audio Transcript",
placeholder=("Transcript of the reference audio (optional, improves quality)"),
lines=2,
visible=False,
)
x_vector_only = gr.Checkbox(
label="Use x-vector only",
value=False,
visible=False,
info=("Skip reference transcript, use speaker embedding only (lower quality)"),
)
with gr.Row():
response_format = gr.Dropdown(
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
value="wav",
label="Audio Format",
scale=1,
)
speed = gr.Slider(
minimum=0.25,
maximum=4.0,
value=1.0,
step=0.05,
label="Speed",
scale=1,
)
stream_checkbox = gr.Checkbox(
label="Stream output",
value=False,
info=(
"Enable streaming (PCM format, speed disabled). "
"For gapless streaming, use gradio_fastrtc_demo.py"
),
scale=1,
)
generate_btn = gr.Button(
"Generate Speech",
variant="primary",
size="lg",
elem_id="generate-btn",
)
# Right column: output
with gr.Column(scale=2):
audio_output = gr.Audio(
label="Generated Audio",
interactive=False,
streaming=True,
autoplay=True,
)
gr.Markdown(
"### Task Types\n"
"- **CustomVoice**: Use a predefined speaker "
"(Vivian, Ryan, etc.) with optional style instructions\n"
"- **VoiceDesign**: Describe the desired voice in natural "
"language (instructions required)\n"
"- **Base**: Clone a voice from reference audio "
"(upload a file or provide a URL)"
)
# Dynamic UI updates
task_type.change(
fn=on_task_type_change,
inputs=[task_type],
outputs=[
voice,
instructions,
ref_audio,
ref_audio_url,
ref_text,
x_vector_only,
],
)
stream_checkbox.change(
fn=on_stream_change,
inputs=[stream_checkbox],
outputs=[response_format, speed],
)
all_inputs = [
text_input,
task_type,
voice,
language,
instructions,
ref_audio,
ref_audio_url,
ref_text,
x_vector_only,
response_format,
speed,
]
def dispatch(stream_enabled, *args):
if stream_enabled:
yield from generate_speech_stream(
api_base,
*args,
chunk_seconds=stream_chunk_seconds,
)
else:
yield generate_speech(api_base, *args)
generate_btn.click(
fn=dispatch,
inputs=[stream_checkbox] + all_inputs,
outputs=[audio_output],
)
demo.queue()
return demo
def parse_args():
parser = argparse.ArgumentParser(description="Gradio demo for Qwen3-TTS online serving.")
add_common_args(parser)
parser.add_argument(
"--stream-chunk-seconds",
type=float,
default=0.25,
help=(
"Seconds of audio to buffer per streaming chunk (default: 0.25). "
"Lower = smoother but more overhead, higher = chunkier."
),
)
return parser.parse_args()
def main():
args = parse_args()
print(f"Connecting to vLLM server at: {args.api_base}")
demo = build_interface(
args.api_base,
stream_chunk_seconds=args.stream_chunk_seconds,
)
try:
demo.launch(
server_name=args.host,
server_port=args.port,
share=args.share,
)
except KeyboardInterrupt:
print("\nShutting down...")
if __name__ == "__main__":
main()
@@ -1,274 +0,0 @@
"""Gradio demo for Qwen3-TTS using FastRTC for gapless streaming audio.
Uses WebRTC via FastRTC to deliver audio chunks as a continuous real-time
stream, avoiding the inter-chunk gaps present in Gradio's built-in streaming.
Requires: pip install fastrtc
Supports all 3 task types:
- CustomVoice: Predefined speaker with optional style instructions
- VoiceDesign: Natural language voice description
- Base: Voice cloning from reference audio (URL only in this demo)
Usage:
# Start the server first (see run_server.sh), then:
python gradio_fastrtc_demo.py --api-base http://localhost:8000
# With public sharing:
python gradio_fastrtc_demo.py --api-base http://localhost:8000 --share
"""
import argparse
import logging
from collections.abc import Generator
import gradio as gr
import numpy as np
from fastrtc import WebRTC
from tts_common import (
PCM_SAMPLE_RATE,
SUPPORTED_LANGUAGES,
TASK_TYPES,
add_common_args,
build_payload,
fetch_voices,
stream_pcm_chunks,
)
logger = logging.getLogger(__name__)
def make_tts_generator(api_base: str):
"""Create a generator function for FastRTC's mode='receive'.
FastRTC calls this with the additional_inputs values whenever the user
clicks "Start Stream". It must be a synchronous generator yielding
(sample_rate, int16_ndarray) tuples.
"""
def tts_generate(
text: str,
task_type: str,
voice: str,
language: str,
instructions: str,
ref_audio_url: str,
ref_text: str,
x_vector_only: bool,
) -> Generator[tuple[int, np.ndarray], None, None]:
if not text or not text.strip():
return
payload = build_payload(
text,
task_type,
voice,
language,
instructions,
ref_audio=None, # WebRTC UI doesn't support file upload
ref_audio_url=ref_audio_url,
ref_text=ref_text,
x_vector_only=x_vector_only,
stream=True,
)
try:
for samples in stream_pcm_chunks(api_base, payload):
yield (PCM_SAMPLE_RATE, samples)
except Exception:
logger.exception("TTS streaming error")
return tts_generate
def build_interface(api_base: str):
"""Build a custom Gradio Blocks UI matching gradio_demo.py layout."""
voices = fetch_voices(api_base)
handler = make_tts_generator(api_base)
css = """
#generate-btn button { width: 100%; }
.task-info { padding: 8px 12px; border-radius: 6px;
background: #f0f4ff; margin-bottom: 8px; }
"""
with gr.Blocks(css=css, title="Qwen3-TTS (FastRTC)") as demo:
gr.Markdown("# Qwen3-TTS Online Serving Demo (FastRTC)")
gr.Markdown(f"**Server:** `{api_base}` | **Transport:** WebRTC (gapless streaming)")
with gr.Row():
# Left column: inputs
with gr.Column(scale=3):
text_input = gr.Textbox(
label="Text to Synthesize",
placeholder="Enter text here, e.g., Hello, how are you?",
lines=4,
)
with gr.Row():
task_type = gr.Radio(
choices=TASK_TYPES,
value="CustomVoice",
label="Task Type",
scale=2,
)
language = gr.Dropdown(
choices=SUPPORTED_LANGUAGES,
value="Auto",
label="Language",
scale=1,
)
voice = gr.Dropdown(
choices=voices,
value=voices[0] if voices else None,
label="Speaker",
visible=True,
)
instructions = gr.Textbox(
label="Instructions",
placeholder=("e.g., Speak with excitement / A warm, friendly female voice"),
lines=2,
visible=True,
info="Optional style/emotion instructions",
)
# Base (voice clone) controls
# Note: file upload is not supported with WebRTC,
# use a URL instead.
ref_audio_url = gr.Textbox(
label="Reference Audio URL",
placeholder=("https://example.com/reference.wav (alternative to uploading)"),
lines=1,
visible=False,
)
ref_text = gr.Textbox(
label="Reference Audio Transcript",
placeholder=("Transcript of the reference audio (optional, improves quality)"),
lines=2,
visible=False,
)
x_vector_only = gr.Checkbox(
label="Use x-vector only",
value=False,
visible=False,
info=("Skip reference transcript, use speaker embedding only (lower quality)"),
)
generate_btn = gr.Button(
"Generate Speech",
variant="primary",
size="lg",
elem_id="generate-btn",
)
# Right column: output
with gr.Column(scale=2):
webrtc_output = WebRTC(
label="Generated Audio (WebRTC)",
mode="receive",
modality="audio",
)
gr.Markdown(
"### Task Types\n"
"- **CustomVoice**: Use a predefined speaker "
"(Vivian, Ryan, etc.) with optional style instructions\n"
"- **VoiceDesign**: Describe the desired voice in natural "
"language (instructions required)\n"
"- **Base**: Clone a voice from reference audio "
"(provide a URL)"
)
# Dynamic UI updates for task type
# FastRTC doesn't support file upload, so we only toggle
# 5 components (no ref_audio)
def on_task_type_change_fastrtc(task_type: str):
if task_type == "CustomVoice":
return (
gr.update(visible=True), # voice
gr.update(visible=True, info="Optional style/emotion instructions"),
gr.update(visible=False), # ref_audio_url
gr.update(visible=False), # ref_text
gr.update(visible=False), # x_vector_only
)
elif task_type == "VoiceDesign":
return (
gr.update(visible=False),
gr.update(visible=True, info="Required: describe the voice style"),
gr.update(visible=False),
gr.update(visible=False),
gr.update(visible=False),
)
elif task_type == "Base":
return (
gr.update(visible=False),
gr.update(visible=False),
gr.update(visible=True),
gr.update(visible=True),
gr.update(visible=True),
)
return (
gr.update(visible=True),
gr.update(visible=True),
gr.update(visible=False),
gr.update(visible=False),
gr.update(visible=False),
)
task_type.change(
fn=on_task_type_change_fastrtc,
inputs=[task_type],
outputs=[
voice,
instructions,
ref_audio_url,
ref_text,
x_vector_only,
],
)
# Wire up WebRTC streaming
all_inputs = [
text_input,
task_type,
voice,
language,
instructions,
ref_audio_url,
ref_text,
x_vector_only,
]
webrtc_output.stream(
fn=handler,
inputs=all_inputs,
outputs=[webrtc_output],
trigger=generate_btn.click,
)
demo.queue()
return demo
def parse_args():
parser = argparse.ArgumentParser(description="Gradio FastRTC demo for Qwen3-TTS (gapless streaming).")
add_common_args(parser)
return parser.parse_args()
def main():
args = parse_args()
logging.basicConfig(level=logging.INFO)
print(f"Connecting to vLLM server at: {args.api_base}")
demo = build_interface(args.api_base)
demo.launch(
server_name=args.host,
server_port=args.port,
share=args.share,
)
if __name__ == "__main__":
main()
@@ -1,203 +0,0 @@
#!/bin/bash
# Launch both vLLM server and Gradio demo for Qwen3-TTS
#
# Usage:
# ./run_gradio_demo.sh # Default: CustomVoice
# ./run_gradio_demo.sh --task-type VoiceDesign # VoiceDesign model
# ./run_gradio_demo.sh --task-type Base --gradio-port 7861
#
# Options:
# --task-type TYPE Task type: CustomVoice, VoiceDesign, Base (default: CustomVoice)
# --server-port PORT Port for vLLM server (default: 8000)
# --gradio-port PORT Port for Gradio demo (default: 7860)
# --server-host HOST Host for vLLM server (default: 0.0.0.0)
# --gradio-ip IP IP for Gradio demo (default: 127.0.0.1)
# --share Share Gradio demo publicly
set -e
# Default values
TASK_TYPE="CustomVoice"
SERVER_PORT=8000
GRADIO_PORT=7860
SERVER_HOST="0.0.0.0"
GRADIO_IP="127.0.0.1"
GRADIO_SHARE=false
# Parse command line arguments
while [[ $# -gt 0 ]]; do
case $1 in
--task-type)
TASK_TYPE="$2"
shift 2
;;
--server-port)
SERVER_PORT="$2"
shift 2
;;
--gradio-port)
GRADIO_PORT="$2"
shift 2
;;
--server-host)
SERVER_HOST="$2"
shift 2
;;
--gradio-ip)
GRADIO_IP="$2"
shift 2
;;
--share)
GRADIO_SHARE=true
shift
;;
--help)
echo "Usage: $0 [OPTIONS]"
echo ""
echo "Options:"
echo " --task-type TYPE Task type: CustomVoice, VoiceDesign, Base (default: CustomVoice)"
echo " --server-port PORT Port for vLLM server (default: 8000)"
echo " --gradio-port PORT Port for Gradio demo (default: 7860)"
echo " --server-host HOST Host for vLLM server (default: 0.0.0.0)"
echo " --gradio-ip IP IP for Gradio demo (default: 127.0.0.1)"
echo " --share Share Gradio demo publicly"
echo ""
exit 0
;;
*)
echo "Unknown option: $1"
echo "Use --help for usage information"
exit 1
;;
esac
done
# Map task type to model
case "$TASK_TYPE" in
CustomVoice)
MODEL="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
;;
VoiceDesign)
MODEL="Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
;;
Base)
MODEL="Qwen/Qwen3-TTS-12Hz-1.7B-Base"
;;
*)
echo "Unknown task type: $TASK_TYPE"
echo "Supported: CustomVoice, VoiceDesign, Base"
exit 1
;;
esac
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
API_BASE="http://localhost:${SERVER_PORT}"
echo "=========================================="
echo "Qwen3-TTS Gradio Demo"
echo "=========================================="
echo "Task Type : $TASK_TYPE"
echo "Model : $MODEL"
echo "Server : http://${SERVER_HOST}:${SERVER_PORT}"
echo "Gradio : http://${GRADIO_IP}:${GRADIO_PORT}"
echo "=========================================="
# Cleanup on exit
cleanup() {
echo ""
echo "Shutting down..."
if [ -n "$SERVER_PID" ]; then
echo "Stopping vLLM server (PID: $SERVER_PID)..."
kill "$SERVER_PID" 2>/dev/null || true
wait "$SERVER_PID" 2>/dev/null || true
fi
if [ -n "$GRADIO_PID" ]; then
echo "Stopping Gradio demo (PID: $GRADIO_PID)..."
kill "$GRADIO_PID" 2>/dev/null || true
wait "$GRADIO_PID" 2>/dev/null || true
fi
echo "Cleanup complete"
exit 0
}
trap cleanup SIGINT SIGTERM
# Start vLLM server
echo ""
echo "Starting vLLM server..."
LOG_FILE="/tmp/vllm_tts_server_${SERVER_PORT}.log"
vllm-omni serve "$MODEL" \
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
--host "$SERVER_HOST" \
--port "$SERVER_PORT" \
--gpu-memory-utilization 0.9 \
--trust-remote-code \
--enforce-eager \
--omni 2>&1 | tee "$LOG_FILE" &
SERVER_PID=$!
# Wait for server startup
echo ""
echo "Waiting for vLLM server to be ready..."
STARTUP_FLAG="/tmp/vllm_tts_startup_flag_${SERVER_PORT}.tmp"
rm -f "$STARTUP_FLAG"
(
tail -f "$LOG_FILE" 2>/dev/null | grep -m 1 "Application startup complete" > /dev/null && touch "$STARTUP_FLAG"
) &
TAIL_PID=$!
MAX_WAIT=300
ELAPSED=0
while [ $ELAPSED -lt $MAX_WAIT ]; do
if [ -f "$STARTUP_FLAG" ]; then
kill "$TAIL_PID" 2>/dev/null || true
wait "$TAIL_PID" 2>/dev/null || true
echo ""
echo "vLLM server is ready!"
break
fi
if ! kill -0 "$SERVER_PID" 2>/dev/null; then
kill "$TAIL_PID" 2>/dev/null || true
echo ""
echo "Error: vLLM server failed to start"
exit 1
fi
sleep 1
ELAPSED=$((ELAPSED + 1))
done
rm -f "$STARTUP_FLAG"
if [ $ELAPSED -ge $MAX_WAIT ]; then
kill "$TAIL_PID" 2>/dev/null || true
echo "Error: Server startup timed out after ${MAX_WAIT}s"
kill "$SERVER_PID" 2>/dev/null || true
exit 1
fi
# Start Gradio demo
echo ""
echo "Starting Gradio demo..."
cd "$SCRIPT_DIR"
GRADIO_CMD=("python" "gradio_demo.py" "--api-base" "$API_BASE" "--ip" "$GRADIO_IP" "--port" "$GRADIO_PORT")
if [ "$GRADIO_SHARE" = true ]; then
GRADIO_CMD+=("--share")
fi
"${GRADIO_CMD[@]}" &
GRADIO_PID=$!
echo ""
echo "=========================================="
echo "Both services are running!"
echo "=========================================="
echo "vLLM Server : http://${SERVER_HOST}:${SERVER_PORT}"
echo "Gradio Demo : http://${GRADIO_IP}:${GRADIO_PORT}"
echo ""
echo "Press Ctrl+C to stop both services"
echo "=========================================="
echo ""
wait $SERVER_PID $GRADIO_PID || true
cleanup
@@ -1,249 +0,0 @@
"""WebSocket client for streaming text-input TTS.
Connects to the /v1/audio/speech/stream endpoint, sends text incrementally
(simulating real-time STT output), and saves per-sentence audio files.
Usage:
# Send full text at once
python streaming_speech_client.py --text "Hello world. How are you? I am fine."
# Simulate STT: send text word-by-word with delay
python streaming_speech_client.py \
--text "Hello world. How are you? I am fine." \
--simulate-stt --stt-delay 0.1
# VoiceDesign task
python streaming_speech_client.py \
--text "Today is a great day. The weather is nice." \
--task-type VoiceDesign \
--instructions "A cheerful young female voice"
# Base task (voice cloning)
python streaming_speech_client.py \
--text "Hello world. How are you?" \
--task-type Base \
--ref-audio /path/to/reference.wav \
--ref-text "Transcript of reference audio"
Requirements:
pip install websockets
"""
import argparse
import asyncio
import json
import os
try:
import websockets
except ImportError:
print("Please install websockets: pip install websockets")
raise SystemExit(1)
async def stream_tts(
url: str,
text: str,
config: dict,
output_dir: str,
simulate_stt: bool = False,
stt_delay: float = 0.1,
) -> None:
"""Connect to the streaming TTS endpoint and process audio responses."""
os.makedirs(output_dir, exist_ok=True)
async with websockets.connect(url) as ws:
# 1. Send session config
config_msg = {"type": "session.config", **config}
await ws.send(json.dumps(config_msg))
print(f"Sent session config: {config}")
# 2. Send text (either all at once or word-by-word)
async def send_text():
if simulate_stt:
words = text.split(" ")
for i, word in enumerate(words):
chunk = word + (" " if i < len(words) - 1 else "")
await ws.send(
json.dumps(
{
"type": "input.text",
"text": chunk,
}
)
)
print(f" Sent: {chunk!r}")
await asyncio.sleep(stt_delay)
else:
await ws.send(
json.dumps(
{
"type": "input.text",
"text": text,
}
)
)
print(f"Sent full text: {text!r}")
# 3. Signal end of input
await ws.send(json.dumps({"type": "input.done"}))
print("Sent input.done")
# Run sender and receiver concurrently
sender_task = asyncio.create_task(send_text())
response_format = config.get("response_format", "wav")
current_sentence_index = 0
current_chunks: list[bytes] = []
try:
while True:
message = await ws.recv()
if isinstance(message, bytes):
current_chunks.append(message)
print(f" Received audio chunk for sentence {current_sentence_index}: {len(message)} bytes")
else:
# JSON frame
msg = json.loads(message)
msg_type = msg.get("type")
if msg_type == "audio.start":
current_sentence_index = msg["sentence_index"]
current_chunks = []
print(f" [sentence {msg['sentence_index']}] Generating: {msg['sentence_text']!r}")
elif msg_type == "audio.done":
filename = os.path.join(
output_dir,
f"sentence_{msg['sentence_index']:03d}.{response_format}",
)
with open(filename, "wb") as f:
f.write(b"".join(current_chunks))
print(
f" [sentence {msg['sentence_index']}] Done"
f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}"
f" error={msg.get('error', False)}"
f" -> {filename}"
)
current_chunks = []
elif msg_type == "session.done":
print(f"\nSession complete: {msg['total_sentences']} sentence(s) generated")
break
elif msg_type == "error":
print(f" ERROR: {msg['message']}")
else:
print(f" Unknown message: {msg}")
finally:
sender_task.cancel()
try:
await sender_task
except asyncio.CancelledError:
pass # Task cancellation is expected during shutdown
print(f"\nAudio files saved to: {output_dir}/")
def main():
parser = argparse.ArgumentParser(description="Streaming text-input TTS client")
parser.add_argument(
"--url",
default="ws://localhost:8000/v1/audio/speech/stream",
help="WebSocket endpoint URL",
)
parser.add_argument(
"--text",
required=True,
help="Text to synthesize",
)
parser.add_argument(
"--output-dir",
default="streaming_tts_output",
help="Directory to save audio files (default: streaming_tts_output)",
)
# Session config options
parser.add_argument("--model", default=None, help="Model name")
parser.add_argument("--voice", default="Vivian", help="Speaker voice")
parser.add_argument(
"--task-type",
default="CustomVoice",
choices=["CustomVoice", "VoiceDesign", "Base"],
help="TTS task type",
)
parser.add_argument("--language", default="Auto", help="Language")
parser.add_argument("--instructions", default=None, help="Voice style instructions")
parser.add_argument(
"--response-format",
default="wav",
choices=["wav", "pcm", "flac", "mp3", "aac", "opus"],
help="Audio format",
)
parser.add_argument(
"--stream-audio",
action="store_true",
help="Receive one or more PCM chunks per sentence (requires --response-format pcm)",
)
parser.add_argument("--speed", type=float, default=1.0, help="Playback speed (0.25-4.0)")
parser.add_argument("--max-new-tokens", type=int, default=None, help="Max tokens")
# Base task options
parser.add_argument("--ref-audio", default=None, help="Reference audio")
parser.add_argument("--ref-text", default=None, help="Reference text")
parser.add_argument(
"--x-vector-only-mode",
action="store_true",
default=False,
help="Speaker embedding only mode",
)
# STT simulation
parser.add_argument(
"--simulate-stt",
action="store_true",
help="Simulate STT by sending text word-by-word",
)
parser.add_argument(
"--stt-delay",
type=float,
default=0.1,
help="Delay between words in STT simulation (seconds)",
)
args = parser.parse_args()
# Build session config (only include non-None values)
config = {}
for key in [
"model",
"voice",
"task_type",
"language",
"instructions",
"response_format",
"speed",
"max_new_tokens",
"ref_audio",
"ref_text",
]:
val = getattr(args, key.replace("-", "_"), None)
if val is not None:
config[key] = val
if args.stream_audio:
config["stream_audio"] = True
if args.x_vector_only_mode:
config["x_vector_only_mode"] = True
asyncio.run(
stream_tts(
url=args.url,
text=args.text,
config=config,
output_dir=args.output_dir,
simulate_stt=args.simulate_stt,
stt_delay=args.stt_delay,
)
)
if __name__ == "__main__":
main()
@@ -1,218 +0,0 @@
"""Shared constants, helpers, and payload building for Qwen3-TTS Gradio demos."""
import base64
import io
import gradio as gr
import httpx
import numpy as np
import soundfile as sf
SUPPORTED_LANGUAGES = [
"Auto",
"Chinese",
"English",
"Japanese",
"Korean",
"German",
"French",
"Russian",
"Portuguese",
"Spanish",
"Italian",
]
TASK_TYPES = ["CustomVoice", "VoiceDesign", "Base"]
PCM_SAMPLE_RATE = 24000
DEFAULT_API_BASE = "http://localhost:8000"
def fetch_voices(api_base: str) -> list[str]:
"""Fetch available voices from the server."""
try:
with httpx.Client(timeout=10.0) as client:
resp = client.get(
f"{api_base}/v1/audio/voices",
headers={"Authorization": "Bearer EMPTY"},
)
if resp.status_code == 200:
data = resp.json()
return data.get("voices", ["Vivian", "Ryan"])
except Exception:
pass
return ["Vivian", "Ryan"]
def encode_audio_to_base64(audio_data: tuple) -> str:
"""Encode Gradio audio input (sample_rate, numpy_array) to base64 data URL."""
sample_rate, audio_np = audio_data
if audio_np.dtype != np.int16:
if audio_np.dtype in (np.float32, np.float64):
audio_np = np.clip(audio_np, -1.0, 1.0)
audio_np = (audio_np * 32767).astype(np.int16)
else:
audio_np = audio_np.astype(np.int16)
buf = io.BytesIO()
sf.write(buf, audio_np, sample_rate, format="WAV")
wav_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
return f"data:audio/wav;base64,{wav_b64}"
def build_payload(
text: str,
task_type: str,
voice: str,
language: str,
instructions: str,
ref_audio: tuple | None,
ref_audio_url: str,
ref_text: str,
x_vector_only: bool,
response_format: str = "pcm",
speed: float = 1.0,
stream: bool = True,
) -> dict:
"""Build the /v1/audio/speech request payload.
Raises gr.Error for invalid input so callers don't need to validate.
"""
if not text or not text.strip():
raise gr.Error("Please enter text to synthesize.")
payload: dict = {
"input": text.strip(),
"response_format": "pcm" if stream else response_format,
"stream": stream,
}
if not stream:
payload["speed"] = speed
if task_type:
payload["task_type"] = task_type
if language:
payload["language"] = language
if task_type == "CustomVoice":
if voice:
payload["voice"] = voice
if instructions and instructions.strip():
payload["instructions"] = instructions.strip()
elif task_type == "VoiceDesign":
if not instructions or not instructions.strip():
raise gr.Error("VoiceDesign task requires voice style instructions.")
payload["instructions"] = instructions.strip()
elif task_type == "Base":
ref_audio_url_stripped = ref_audio_url.strip() if ref_audio_url else ""
if ref_audio_url_stripped:
payload["ref_audio"] = ref_audio_url_stripped
elif ref_audio is not None:
payload["ref_audio"] = encode_audio_to_base64(ref_audio)
else:
raise gr.Error("Base (voice clone) task requires reference audio. Upload a file or provide a URL.")
if ref_text and ref_text.strip():
payload["ref_text"] = ref_text.strip()
if x_vector_only:
payload["x_vector_only_mode"] = True
return payload
def on_task_type_change(task_type: str):
"""Update UI visibility based on selected task type."""
if task_type == "CustomVoice":
return (
gr.update(visible=True), # voice dropdown
gr.update(visible=True, info="Optional style/emotion instructions"),
gr.update(visible=False), # ref_audio
gr.update(visible=False), # ref_audio_url
gr.update(visible=False), # ref_text
gr.update(visible=False), # x_vector_only
)
elif task_type == "VoiceDesign":
return (
gr.update(visible=False),
gr.update(visible=True, info="Required: describe the voice style"),
gr.update(visible=False),
gr.update(visible=False),
gr.update(visible=False),
gr.update(visible=False),
)
elif task_type == "Base":
return (
gr.update(visible=False),
gr.update(visible=False),
gr.update(visible=True),
gr.update(visible=True),
gr.update(visible=True),
gr.update(visible=True),
)
return (
gr.update(visible=True),
gr.update(visible=True),
gr.update(visible=False),
gr.update(visible=False),
gr.update(visible=False),
gr.update(visible=False),
)
def stream_pcm_chunks(api_base: str, payload: dict):
"""Stream raw PCM bytes from the server, yielding int16 numpy arrays.
Handles odd-byte boundaries between network chunks.
"""
leftover = b""
with httpx.Client(timeout=300.0) as client:
with client.stream(
"POST",
f"{api_base}/v1/audio/speech",
json=payload,
headers={
"Content-Type": "application/json",
"Authorization": "Bearer EMPTY",
},
) as resp:
if resp.status_code != 200:
resp.read()
raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
for chunk in resp.iter_bytes():
if not chunk:
continue
raw = leftover + chunk
usable = len(raw) - (len(raw) % 2)
leftover = raw[usable:]
if usable == 0:
continue
yield np.frombuffer(raw[:usable], dtype=np.int16).copy()
def add_common_args(parser):
"""Add CLI arguments shared by both demos."""
parser.add_argument(
"--api-base",
default=DEFAULT_API_BASE,
help=f"Base URL for the vLLM API server (default: {DEFAULT_API_BASE}).",
)
parser.add_argument(
"--host",
default="0.0.0.0",
help="Host/IP for Gradio server (default: 0.0.0.0).",
)
parser.add_argument(
"--port",
type=int,
default=7860,
help="Port for Gradio server (default: 7860).",
)
parser.add_argument(
"--share",
action="store_true",
help="Share the Gradio demo publicly.",
)
return parser
@@ -0,0 +1,234 @@
# Stable Audio Online Serving
Generate audio from text prompts using Stable Audio models via an OpenAI-compatible API endpoint.
## Features
- **OpenAI-compatible API**: Use `/v1/audio/generate` endpoint
- **Flexible control**: Adjust audio length, guidance scale, inference steps
- **Quality control**: Use negative prompts to avoid unwanted characteristics
- **Reproducible**: Set random seed for deterministic generation
## Quick Start
### 1. Start the Server
```bash
vllm-omni serve stabilityai/stable-audio-open-1.0 \
--host 0.0.0.0 \
--port 8000 \
--gpu-memory-utilization 0.9 \
--trust-remote-code \
--enforce-eager \
--omni
```
### 2. Generate Audio
#### Using curl
```bash
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "The sound of a cat purring",
"audio_length": 10.0
}' --output cat.wav
```
#### Using Python Client
```bash
python stable_audio_client.py \
--text "The sound of a cat purring" \
--audio_length 10.0 \
--output cat.wav
```
#### Using Bash Script
```bash
bash curl_examples.sh
```
## API Reference
### Endpoint
```
POST /v1/audio/generate
```
### Request Body
```json
{
"input": "Text description of the audio",
"audio_length": 10.0,
"audio_start": 0.0,
"negative_prompt": "Low quality",
"guidance_scale": 7.0,
"num_inference_steps": 100,
"seed": 42,
"response_format": "wav"
}
```
### Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input` | string | **required** | Text prompt describing the audio to generate |
| `audio_length` | float | ~47s | Audio duration in seconds (max ~47s for stable-audio-open-1.0) |
| `audio_start` | float | 0.0 | Audio start time in seconds |
| `negative_prompt` | string | null | Text describing what to avoid in generation |
| `guidance_scale` | float | 7.0 | Classifier-free guidance scale (higher = more adherence to prompt) |
| `num_inference_steps` | int | 50 | Number of denoising steps (higher = better quality, slower) |
| `seed` | int | null | Random seed for reproducibility |
| `response_format` | string | "wav" | Output format: wav, mp3, flac, pcm |
### Response
Returns audio data in the requested format (default: WAV).
## Usage Examples
### Basic Generation
```bash
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "The sound of ocean waves"
}' --output ocean.wav
```
### Custom Duration
```bash
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "A dog barking",
"audio_length": 5.0
}' --output dog_5s.wav
```
### High Quality with Negative Prompt
```bash
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "A piano playing a gentle melody",
"audio_length": 10.0,
"negative_prompt": "Low quality, distorted, noisy",
"guidance_scale": 8.0,
"num_inference_steps": 150
}' --output piano_hq.wav
```
### Reproducible Generation
```bash
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "Thunder and rain sounds",
"audio_length": 15.0,
"seed": 42
}' --output thunder.wav
```
### Quick Generation (Fewer Steps)
For faster generation with slightly lower quality:
```bash
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "Birds chirping in a forest",
"audio_length": 8.0,
"num_inference_steps": 50
}' --output birds_quick.wav
```
## Python Client Examples
### Simple Generation
```bash
python stable_audio_client.py \
--text "The sound of a cat purring"
```
### Custom Parameters
```bash
python stable_audio_client.py \
--text "Thunder and rain" \
--audio_length 15.0 \
--negative_prompt "Low quality" \
--guidance_scale 7.0 \
--num_inference_steps 100 \
--seed 42 \
--output thunder.wav
```
### Different Output Format
```bash
python stable_audio_client.py \
--text "Guitar playing" \
--response_format mp3 \
--output guitar.mp3
```
## Tips
1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`
2. **Quality vs Speed**:
- 50 steps: Fast, decent quality
- 100 steps: Good balance (default)
- 150+ steps: High quality, slower
3. **Guidance Scale**:
- Lower (3-5): More creative/varied
- Default (7): Good balance
- Higher (10+): More literal to prompt
4. **Negative Prompts**: Use to avoid "Low quality", "distorted", "noisy", etc.
5. **Seeds**: Use same seed for reproducible results
## Performance
| Inference Steps | Quality | Speed | Use Case |
|----------------|---------|-------|----------|
| 50 | Good | Fast | Quick previews |
| 100 (default) | Very Good | Medium | Production |
| 150+ | Excellent | Slow | Final/critical audio |
## Troubleshooting
### Server not responding
- Check if server is running: `curl http://localhost:8000/health`
- Check server logs for errors
### Audio quality issues
- Increase `num_inference_steps` (e.g., 150)
- Add negative prompts: `"Low quality, distorted, noisy"`
- Increase `guidance_scale` for more prompt adherence
### Generation timeout
- Reduce `num_inference_steps`
- Reduce `audio_length`
- Check GPU memory with `nvidia-smi`
### Wrong audio length
- Ensure `audio_length` is within model limits (~47s max)
- Adjust `audio_start` if trimming is needed
## See Also
- [Offline Inference Example](../../offline_inference/text_to_audio/README.md)
- [Stable Audio Model Card](https://huggingface.co/stabilityai/stable-audio-open-1.0)
- [vLLM-Omni Documentation](https://github.com/vllm-project/vllm-omni)
+54
View File
@@ -0,0 +1,54 @@
#!/bin/bash
# Examples for using Stable Audio with curl via /v1/audio/generate endpoint
# Example 1: Simple request with default parameters
echo "Example 1: Simple request with default parameters"
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "The sound audience clapping and cheering in a stadium"
}' --output stadium.wav
# Example 2: Request with custom audio_length
echo "Example 2: Custom audio length (5 seconds)"
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "The sound of a dog barking",
"audio_length": 5.0
}' --output dog_5s.wav
# Example 3: Request with negative prompt for quality control
echo "Example 3: With negative prompt"
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "A piano playing a gentle melody",
"audio_length": 10.0,
"negative_prompt": "Low quality, distorted, noisy"
}' --output piano.wav
# Example 4: Full control with all parameters
echo "Example 4: Full control (custom length, guidance, steps, seed)"
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "Thunder and rain sounds",
"audio_length": 15.0,
"negative_prompt": "Low quality",
"guidance_scale": 7.0,
"num_inference_steps": 100,
"seed": 42
}' --output thunder_rain.wav
# Example 5: Quick generation with fewer steps (faster but lower quality)
echo "Example 5: Quick generation (fewer steps)"
curl -X POST http://localhost:8000/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "Ocean waves crashing on a beach",
"audio_length": 8.0,
"num_inference_steps": 50
}' --output ocean.wav
echo "All examples completed!"
+170
View File
@@ -0,0 +1,170 @@
#!/usr/bin/env python3
"""
OpenAI-compatible client for Stable Audio via /v1/audio/generate endpoint.
This script demonstrates how to use the OpenAI-compatible speech API
to generate audio from text using Stable Audio models.
Examples:
# Simple generation
python stable_audio_client.py --text "The sound of a cat purring"
# With custom duration
python stable_audio_client.py --text "A dog barking" --audio_length 5.0
# With all parameters
python stable_audio_client.py --text "Thunder and rain" \
--audio_length 15.0 \
--negative_prompt "Low quality" \
--guidance_scale 7.0 \
--num_inference_steps 100 \
--seed 42 \
--output thunder.wav
"""
import argparse
import sys
import requests
def parse_args():
parser = argparse.ArgumentParser(description="Generate audio with Stable Audio via OpenAI-compatible API")
parser.add_argument(
"--api_url",
default="http://localhost:8000/v1/audio/generate",
help="API endpoint URL",
)
parser.add_argument(
"--text",
default="The sound of a cat purring",
help="Text prompt for audio generation",
)
parser.add_argument(
"--audio_length",
type=float,
default=10.0,
help="Audio length in seconds (max ~47s for stable-audio-open-1.0)",
)
parser.add_argument(
"--audio_start",
type=float,
default=0.0,
help="Audio start time in seconds",
)
parser.add_argument(
"--negative_prompt",
default="Low quality",
help="Negative prompt for classifier-free guidance",
)
parser.add_argument(
"--guidance_scale",
type=float,
default=7.0,
help="Guidance scale for diffusion (higher = more adherence to prompt)",
)
parser.add_argument(
"--num_inference_steps",
type=int,
default=100,
help="Number of inference steps (higher = better quality, slower)",
)
parser.add_argument(
"--seed",
type=int,
default=None,
help="Random seed for reproducibility",
)
parser.add_argument(
"--output",
default="stable_audio_output.wav",
help="Output file path",
)
parser.add_argument(
"--response_format",
default="wav",
choices=["wav", "mp3", "flac", "pcm"],
help="Audio output format",
)
return parser.parse_args()
def generate_audio(args):
"""Generate audio using the API."""
# Build request payload
payload = {
"input": args.text,
"audio_length": args.audio_length,
"audio_start": args.audio_start,
"response_format": args.response_format,
}
# Add optional parameters
if args.negative_prompt:
payload["negative_prompt"] = args.negative_prompt
if args.guidance_scale:
payload["guidance_scale"] = args.guidance_scale
if args.num_inference_steps:
payload["num_inference_steps"] = args.num_inference_steps
if args.seed is not None:
payload["seed"] = args.seed
print(f"\n{'=' * 60}")
print("Stable Audio - Text-to-Audio Generation")
print(f"{'=' * 60}")
print(f"API URL: {args.api_url}")
print(f"Prompt: {args.text}")
print(f"Audio length: {args.audio_length}s")
print(f"Negative prompt: {args.negative_prompt}")
print(f"Guidance scale: {args.guidance_scale}")
print(f"Inference steps: {args.num_inference_steps}")
if args.seed is not None:
print(f"Seed: {args.seed}")
print(f"Output: {args.output}")
print(f"{'=' * 60}\n")
try:
# Make the API request
print("Generating audio...")
response = requests.post(
args.api_url,
json=payload,
headers={"Content-Type": "application/json"},
timeout=300, # 5 minute timeout for long generations
)
# Check for errors
if response.status_code != 200:
print(f"Error: API returned status code {response.status_code}")
print(f"Response: {response.text}")
return False
# Save the audio
with open(args.output, "wb") as f:
f.write(response.content)
print(f"✓ Audio saved to {args.output}")
print(f" File size: {len(response.content) / 1024:.1f} KB")
return True
except requests.exceptions.Timeout:
print("Error: Request timed out. Try reducing inference steps or audio length.")
return False
except requests.exceptions.ConnectionError:
print(f"Error: Could not connect to {args.api_url}")
print("Make sure the server is running.")
return False
except Exception as e:
print(f"Error: {e}")
return False
def main():
args = parse_args()
success = generate_audio(args)
sys.exit(0 if success else 1)
if __name__ == "__main__":
main()
+12 -97
View File
@@ -26,35 +26,6 @@ The script allows overriding:
- `CACHE_BACKEND` (default: `none`)
- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
## Async Job Behavior
`POST /v1/videos` is asynchronous. It creates a video job and immediately
returns metadata like the job ID and initial `queued` status. To get the final
artifact, poll the job status and then download the completed file from the
content endpoint.
The main endpoints are:
- `POST /v1/videos`: create a video generation job
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
- `GET /v1/videos`: list stored video jobs
- `GET /v1/videos/{video_id}/content`: download the generated video file
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
## Storage
Generated video files are stored on local disk by the async video API.
Local file storage behavior can be controlled via the following environment variables:
- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
Example:
```bash
export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
```
## API Calls
### Method 1: Using curl
@@ -64,7 +35,7 @@ export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
bash run_curl_text_to_video.sh
# Or execute directly (OpenAI-style multipart)
create_response=$(curl -s http://localhost:8091/v1/videos \
curl -s http://localhost:8091/v1/videos \
-H "Accept: application/json" \
-F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
-F "width=832" \
@@ -76,24 +47,7 @@ create_response=$(curl -s http://localhost:8091/v1/videos \
-F "guidance_scale=4.0" \
-F "guidance_scale_2=4.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=5.0" \
-F "seed=42")
video_id=$(echo "$create_response" | jq -r '.id')
while true; do
status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
if [ "$status" = "completed" ]; then
break
fi
if [ "$status" = "failed" ]; then
echo "Video generation failed"
exit 1
fi
sleep 2
done
curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o wan22_output.mp4
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
```
## Request Format
@@ -128,9 +82,8 @@ curl -X POST http://localhost:8091/v1/videos \
| Parameter | Type | Default | Description |
| --------------------- | ------ | ------- | ------------------------------------------------ |
| `prompt` | str | - | Text description of the desired video |
| `seconds` | str | None | Clip duration in seconds |
| `size` | str | None | Output size in `WIDTHxHEIGHT` format |
| `negative_prompt` | str | None | Negative prompt |
| `n` | int | 1 | Number of videos to generate |
| `width` | int | None | Video width in pixels |
| `height` | int | None | Video height in pixels |
| `num_frames` | int | None | Number of frames to generate |
@@ -142,60 +95,22 @@ curl -X POST http://localhost:8091/v1/videos \
| `flow_shift` | float | None | Scheduler flow shift (Wan2.2) |
| `seed` | int | None | Random seed (reproducible) |
| `lora` | object | None | LoRA configuration |
| `extra_body` | object | None | Model-specific extra parameters |
## Create Response Format
`POST /v1/videos` returns a job record, not inline base64 video data.
## Response Format
```json
{
"id": "video_gen_123",
"object": "video",
"status": "queued",
"model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
"prompt": "A cinematic view of a futuristic city at sunset",
"created_at": 1234567890
"created": 1234567890,
"data": [
{ "b64_json": "<base64-mp4>" }
]
}
```
## Retrieve, List, Download, and Delete
### Retrieve a job
## Extract Video
```bash
curl -s http://localhost:8091/v1/videos/${video_id} | jq .
```
### List jobs
```bash
curl -s http://localhost:8091/v1/videos | jq .
```
### Download the completed video
```bash
curl -L http://localhost:8091/v1/videos/${video_id}/content -o wan22_output.mp4
```
### Delete a job and its stored file
```bash
curl -X DELETE http://localhost:8091/v1/videos/${video_id} | jq .
```
## Poll Until Complete
```bash
while true; do
status=$(curl -s http://localhost:8091/v1/videos/${video_id} | jq -r '.status')
if [ "$status" = "completed" ]; then
break
fi
if [ "$status" = "failed" ]; then
echo "Video generation failed"
exit 1
fi
sleep 2
done
# Extract base64 from response and decode to video
cat response.json | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
```
@@ -1,63 +1,20 @@
#!/bin/bash
# Wan2.2 text-to-video curl example using the async video job API.
# Wan2.2 text-to-video curl example
set -euo pipefail
OUTPUT_PATH="wan22_output.mp4"
BASE_URL="${BASE_URL:-http://localhost:8098}"
OUTPUT_PATH="${OUTPUT_PATH:-wan22_output.mp4}"
POLL_INTERVAL="${POLL_INTERVAL:-2}"
curl -X POST http://localhost:8098/v1/videos \
-H "Accept: application/json" \
-F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
-F "seconds=2" \
-F "size=832x480" \
-F "negative_prompt=色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \
-F "fps=16" \
-F "num_inference_steps=40" \
-F "guidance_scale=4.0" \
-F "guidance_scale_2=4.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=5.0" \
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > "${OUTPUT_PATH}"
create_response=$(
curl -sS -X POST "${BASE_URL}/v1/videos" \
-H "Accept: application/json" \
-F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
-F "seconds=2" \
-F "size=832x480" \
-F "negative_prompt=色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \
-F "fps=16" \
-F "num_inference_steps=40" \
-F "guidance_scale=4.0" \
-F "guidance_scale_2=4.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=5.0" \
-F "seed=42"
)
video_id="$(echo "${create_response}" | jq -r '.id')"
if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then
echo "Failed to create video job:"
echo "${create_response}" | jq .
exit 1
fi
echo "Created video job ${video_id}"
echo "${create_response}" | jq .
while true; do
status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
status="$(echo "${status_response}" | jq -r '.status')"
case "${status}" in
queued|in_progress)
echo "Video job ${video_id} status: ${status}"
sleep "${POLL_INTERVAL}"
;;
completed)
echo "${status_response}" | jq .
break
;;
failed)
echo "Video generation failed:"
echo "${status_response}" | jq .
exit 1
;;
*)
echo "Unexpected status response:"
echo "${status_response}" | jq .
exit 1
;;
esac
done
curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
echo "Saved video to ${OUTPUT_PATH}"
-3
View File
@@ -167,7 +167,6 @@ markers = [
"gpu: Tests that run on GPU (auto-added)",
"cuda: Tests that run on CUDA (auto-added)",
"rocm: Tests that run on AMD/ROCm (auto-added)",
"xpu: Tests that run on XPU (auto-added)",
"npu: Tests that run on NPU/Ascend (auto-added)",
# specified computation resources marks (auto-added)
"H100: Tests that require H100 GPU",
@@ -177,11 +176,9 @@ markers = [
"A3: Tests that require A3 NPU",
"distributed_cuda: Tests that require multi cards on CUDA platform",
"distributed_rocm: Tests that require multi cards on ROCm platform",
"distributed_xpu: Tests that require multi cards on XPU platform",
"distributed_npu: Tests that require multi cards on NPU platform",
"skipif_cuda: Skip if the num of CUDA cards is less than the required",
"skipif_rocm: Skip if the num of ROCm cards is less than the required",
"skipif_xpu: Skip if the num of XPU cards is less than the required",
"skipif_npu: Skip if the num of NPU cards is less than the required",
# more detailed markers
"slow: Slow tests (may skip in quick CI)",
+1 -1
View File
@@ -6,7 +6,7 @@ diffusers>=0.36.0
accelerate==1.12.0
gradio==5.50
soundfile>=0.13.1
cache-dit==1.3.0
cache-dit==1.2.0
tqdm>=4.66.0
torchsde>=0.2.6
openai-whisper>=20250625
+2 -9
View File
@@ -3,25 +3,18 @@ from pathlib import Path
import pytest
from tests.conftest import OmniServerParams
from tests.utils import hardware_test
from vllm_omni.platforms import current_omni_platform
models = ["Qwen/Qwen2.5-Omni-7B"]
stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_ci.yaml")]
if current_omni_platform.is_xpu():
stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "xpu" / "qwen2_5_omni_ci.yaml")]
# Create parameter combinations for model and stage config
test_params = [
OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
]
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
@pytest.mark.core_model
@pytest.mark.benchmark
@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=3)
@hardware_test(res={"cuda": "L4"}, num_cards=3)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_bench_serve_chat(omni_server):
command = [
-7
View File
@@ -63,12 +63,6 @@ def _setup_comfyui_test_environment():
mock_comfy_api_input.AudioInput = AudioInput
mock_comfy_api_input.VideoInput = VideoInput
mock_comfy_api.input = mock_comfy_api_input
mock_comfy_api_latest = MagicMock()
mock_comfy_api_latest.Types.VideoComponents = MagicMock(side_effect=lambda **kwargs: kwargs)
mock_comfy_api_latest.InputImpl.VideoFromComponents = MagicMock(
side_effect=lambda _: VideoInput(b"mock_video_from_components")
)
mock_comfy_api.latest = mock_comfy_api_latest
def mock_load(_: str | BinaryIO):
"""Mock nodes_audio.load that returns a waveform tensor (channels, samples) and sample rate."""
@@ -84,6 +78,5 @@ def _setup_comfyui_test_environment():
# Install mock modules BEFORE importing any comfyui_vllm_omni code
sys.modules["comfy_api"] = mock_comfy_api
sys.modules["comfy_api.input"] = mock_comfy_api_input
sys.modules["comfy_api.latest"] = mock_comfy_api_latest
sys.modules["comfy_extras"] = mock_comfy_extras
sys.modules["comfy_extras.nodes_audio"] = mock_nodes_audio
+6 -174
View File
@@ -20,12 +20,11 @@ import torch
from comfy_api.input import AudioInput, VideoInput
from comfyui_vllm_omni.nodes import (
VLLMOmniGenerateImage,
VLLMOmniGenerateVideo,
VLLMOmniTTS,
VLLMOmniUnderstanding,
VLLMOmniVoiceClone,
)
from comfyui_vllm_omni.utils.types import AutoregressionSamplingParams, DiffusionSamplingParams, WanModelSpecificParams
from comfyui_vllm_omni.utils.types import AutoregressionSamplingParams, DiffusionSamplingParams
from PIL import Image
from vllm import SamplingParams
from vllm.outputs import CompletionOutput, RequestOutput
@@ -35,8 +34,6 @@ from vllm_omni.entrypoints.cli.serve import OmniServeCommand
from vllm_omni.inputs.data import OmniSamplingParams
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class ServerCase(NamedTuple):
"""Parametrizing the model to serve."""
@@ -52,7 +49,6 @@ class SamplingCase(NamedTuple):
kind: "SamplingKind"
sampling_params: dict | list[dict] | None
lora: dict | None = None
class SamplingKind(StrEnum):
@@ -62,17 +58,11 @@ class SamplingKind(StrEnum):
UNDERSTANDING_AR_LIST = auto()
TTS_NONE = auto()
TTS_DIFFUSION_SINGLE = auto()
VIDEO_NONE = auto()
VIDEO_DIFFUSION_SINGLE = auto()
# Pre-defined arguments to be used in function calls during the tests
IMAGE_WIDTH = 64
IMAGE_HEIGHT = 64
VIDEO_WIDTH = 32
VIDEO_HEIGHT = 32
VIDEO_FPS = 8
VIDEO_NUM_FRAMES = 5
DIFFUSION_SINGLE_SAMPLING_PARAMS = DiffusionSamplingParams(
{
"n": 2,
@@ -82,15 +72,6 @@ DIFFUSION_SINGLE_SAMPLING_PARAMS = DiffusionSamplingParams(
}
)
DIFFUSION_VIDEO_SINGLE_SAMPLING_PARAMS = DiffusionSamplingParams(
{
"num_inference_steps": 30,
"guidance_scale": 6.0,
"true_cfg_scale": 1.5,
}
)
AR_LIST_SAMPLING_PARAMS = [
AutoregressionSamplingParams(
{
@@ -121,10 +102,6 @@ AR_LIST_SAMPLING_PARAMS = [
),
]
VIDEO_MODEL_PARAMS = WanModelSpecificParams({"guidance_scale_2": 5.0, "boundary_ratio": 0.98, "flow_shift": 12.0})
LORA_PARAMS = {"local_path": "test_lora_path", "name": "test_name", "scale": 0.7, "int_id": 10}
def _build_image_output(size: tuple[int, int] = (IMAGE_WIDTH, IMAGE_HEIGHT), color: str = "red") -> Image.Image:
return Image.new("RGB", size, color=color)
@@ -204,16 +181,6 @@ def _build_diffusion_image_output_for_images_endpoint() -> OmniRequestOutput:
)
def _build_diffusion_video_output() -> OmniRequestOutput:
# Small video: VIDEO_NUM_FRAMES frames of (VIDEO_HEIGHT x VIDEO_WIDTH) RGB, shape (F, H, W, C)
video_frames = torch.zeros((VIDEO_NUM_FRAMES, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32)
return OmniRequestOutput.from_diffusion(
request_id="test_req_video",
images=[video_frames],
final_output_type="video",
)
def _build_diffusion_image_output_for_chat_endpoint() -> OmniRequestOutput:
request_output = MagicMock()
request_output.images = [_build_image_output(color="blue")]
@@ -226,40 +193,12 @@ def _build_diffusion_image_output_for_chat_endpoint() -> OmniRequestOutput:
)
def _assert_sampling_param_values(
received: OmniSamplingParams, expected: dict[str, Any], expected_lora: dict | None = None
):
def _assert_sampling_param_values(received: OmniSamplingParams, expected: dict[str, Any]):
for key, expected_value in expected.items():
actual_value = getattr(received, key)
actual_value = getattr(received, key, None)
assert actual_value == expected_value, (
f"Expected sampling param '{key}'={expected_value}, got {actual_value}. The received sampling params: {received}"
)
if expected_lora:
assert received.lora_request.lora_name == expected_lora["name"], (
f"Expected lora name={(expected_lora['name'])}, got {received.lora_request.lora_name}. The received sampling params: {received}"
)
assert received.lora_request.lora_int_id == expected_lora["int_id"], (
f"Expected lora int_id={expected_lora['int_id']}, got {received.lora_request.lora_int_id}. The received sampling params: {received}"
)
assert received.lora_request.lora_path == expected_lora["local_path"], (
f"Expected lora path={expected_lora['local_path']}, got {received.lora_request.lora_path}. The received sampling params: {received}"
)
assert received.lora_scale == expected_lora["scale"], (
f"Expected lora scale={expected_lora['scale']}, got {received.lora_scale}. The received sampling params: {received}"
)
def _assert_model_param_values(received: OmniSamplingParams, expected: dict):
for key, expected_value in expected.items():
try:
actual_value = getattr(received, key)
expected_param_name = key
except AttributeError:
actual_value = received.extra_args.get(key, None)
expected_param_name = f'extra_args["{key}"]'
assert actual_value == expected_value, (
f"Expected model param '{expected_param_name}'={expected_value}, got {actual_value}. The received sampling params: {received}"
)
def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: SamplingCase, server_case: ServerCase):
@@ -296,7 +235,6 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
"height": IMAGE_HEIGHT,
**expected,
},
LORA_PARAMS,
)
elif sampling_case.kind is SamplingKind.UNDERSTANDING_NONE:
assert len(received_sampling_params_list) == 3
@@ -306,33 +244,6 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
_assert_sampling_param_values(received_sampling_params_list[i], expected)
elif sampling_case.kind in {SamplingKind.TTS_NONE, SamplingKind.TTS_DIFFUSION_SINGLE}:
assert len(received_sampling_params_list) == 1
elif sampling_case.kind is SamplingKind.VIDEO_NONE:
assert len(received_sampling_params_list) == 1
_assert_sampling_param_values(
received_sampling_params_list[0],
{
"width": VIDEO_WIDTH,
"height": VIDEO_HEIGHT,
"num_frames": VIDEO_NUM_FRAMES,
"fps": VIDEO_FPS,
},
)
elif sampling_case.kind is SamplingKind.VIDEO_DIFFUSION_SINGLE:
assert len(received_sampling_params_list) == 1
expected = DIFFUSION_VIDEO_SINGLE_SAMPLING_PARAMS.copy()
# expected["num_outputs_per_prompt"] = expected.pop("n") # convert from n to num_outputs_per_prompt
_assert_sampling_param_values(
received_sampling_params_list[0],
{
"width": VIDEO_WIDTH,
"height": VIDEO_HEIGHT,
"num_frames": VIDEO_NUM_FRAMES,
"fps": VIDEO_FPS,
**expected,
},
LORA_PARAMS,
)
_assert_model_param_values(received_sampling_params_list[0], VIDEO_MODEL_PARAMS)
else:
raise AssertionError(f"Unknown sampling case: {sampling_case.kind}")
@@ -490,15 +401,9 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
@pytest.mark.parametrize(
"sampling_case",
[
pytest.param(SamplingCase(kind=SamplingKind.IMAGE_NONE, sampling_params=None), id="no-sampling-params"),
pytest.param(
SamplingCase(kind=SamplingKind.IMAGE_NONE, sampling_params=None, lora=None), id="no-sampling-params"
),
pytest.param(
SamplingCase(
kind=SamplingKind.IMAGE_DIFFUSION_SINGLE,
sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS,
lora=LORA_PARAMS,
),
SamplingCase(kind=SamplingKind.IMAGE_DIFFUSION_SINGLE, sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS),
id="single-diffusion-sampling-params",
),
],
@@ -518,8 +423,7 @@ async def test_image_generation_node(api_server: str, model: str, image_input: b
kwargs["image"] = torch.zeros((1, IMAGE_WIDTH, IMAGE_HEIGHT, 3), dtype=torch.float32)
if sampling_case.sampling_params is not None:
kwargs["sampling_params"] = sampling_case.sampling_params
if sampling_case.lora:
kwargs["lora"] = sampling_case.lora
print(f"!!!!!! Calling {model} node.generate with kwargs: {sampling_case.sampling_params}")
result = await node.generate(**kwargs)
@@ -661,75 +565,3 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
assert result[0]["sample_rate"] == 24000
assert isinstance(result[0]["waveform"], torch.Tensor)
assert result[0]["waveform"].shape == (1, 1, 24000)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"server_case,model,image_input",
[
pytest.param(
ServerCase(
served_model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
stage_list=["diffusion"],
stage_configs=[{"stage_type": "diffusion"}],
outputs=[_build_diffusion_video_output()],
),
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
False,
id="text-to-video",
),
pytest.param(
ServerCase(
served_model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
stage_list=["diffusion"],
stage_configs=[{"stage_type": "diffusion"}],
outputs=[_build_diffusion_video_output()],
),
"Wan-AI/Wan2.2-I2V-A14B-Diffusers",
True,
id="image-to-video",
),
],
indirect=["server_case"],
)
@pytest.mark.parametrize(
"sampling_case",
[
pytest.param(SamplingCase(kind=SamplingKind.VIDEO_NONE, sampling_params=None), id="no-sampling-params"),
pytest.param(
SamplingCase(
kind=SamplingKind.VIDEO_DIFFUSION_SINGLE,
sampling_params=DIFFUSION_VIDEO_SINGLE_SAMPLING_PARAMS,
lora=LORA_PARAMS,
),
id="single-diffusion-sampling-params",
),
],
indirect=["sampling_case"],
)
async def test_video_generation_node(api_server: str, model: str, image_input: bool, sampling_case: SamplingCase):
node = VLLMOmniGenerateVideo()
kwargs = {
"url": api_server,
"model": model,
"prompt": "A beautiful sunset timelapse",
"negative_prompt": "",
"width": VIDEO_WIDTH,
"height": VIDEO_HEIGHT,
"fps": VIDEO_FPS,
"num_frames": VIDEO_NUM_FRAMES,
"model_params": VIDEO_MODEL_PARAMS,
}
if image_input:
kwargs["image"] = torch.zeros((1, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32)
if sampling_case.sampling_params is not None:
kwargs["sampling_params"] = sampling_case.sampling_params
if sampling_case.lora:
kwargs["lora"] = sampling_case.lora
result = await node.generate(**kwargs)
assert isinstance(result, tuple)
assert len(result) == 1
assert isinstance(result[0], VideoInput)
+28 -254
View File
@@ -18,21 +18,16 @@ import subprocess
import sys
import threading
import time
from collections.abc import Generator
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path
from typing import Any, NamedTuple
from typing import Any
import imageio.v3 as iio
import numpy as np
import psutil
import pytest
import soundfile as sf
import torch
import yaml
from openai import OpenAI, omit
from PIL import Image
from openai import OpenAI
from vllm import TextPrompt
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
from vllm.logger import init_logger
@@ -49,65 +44,6 @@ PromptImageInput = list[Any] | Any | None
PromptVideoInput = list[Any] | Any | None
class OmniServerParams(NamedTuple):
model: str
port: int | None = None
stage_config_path: str | None = None
server_args: list[str] | None = None
def assert_image_valid(image: Path | Image.Image, *, width: int | None = None, height: int | None = None):
"""Assert the file is a loadable image with optional exact dimensions."""
if isinstance(image, Path):
assert image.exists(), f"Image not found: {image}"
image = Image.open(image)
image.load()
assert image.width > 0 and image.height > 0
if width is not None:
assert image.width == width, f"Expected width={width}, got {image.width} in {image.name}"
if height is not None:
assert image.height == height, f"Expected height={height}, got {image.height} in {image.name}"
return image
def assert_video_valid(frames: Path | np.ndarray, *, width: int, height: int, num_frames: int) -> None:
"""Assert the MP4 has the expected resolution and exact frame count."""
if isinstance(frames, Path):
assert frames.exists(), f"Video not found: {frames}"
frames = iio.imread(str(frames), plugin="pyav", index=None)
assert frames.shape[0] == num_frames, f"Expected {num_frames} frames, got {frames.shape[0]}"
assert frames.shape[1] == height, f"Expected height={height}, got {frames.shape[1]}"
assert frames.shape[2] == width, f"Expected width={width}, got {frames.shape[2]}"
def assert_audio_valid(path: Path, *, sample_rate: int, channels: int, duration_s: float) -> None:
"""Assert the WAV has the expected sample rate, channel count, and duration."""
assert path.exists(), f"Audio not found: {path}"
info = sf.info(str(path))
assert info.samplerate == sample_rate, f"Expected sample_rate={sample_rate}, got {info.samplerate}"
assert info.channels == channels, f"Expected {channels} channel(s), got {info.channels}"
expected_frames = int(duration_s * sample_rate)
assert info.frames == expected_frames, (
f"Expected {expected_frames} frames ({duration_s}s @ {sample_rate} Hz), got {info.frames}"
)
def decode_b64_image(b64: str):
img = Image.open(BytesIO(base64.b64decode(b64)))
img.load()
return img
@pytest.fixture(scope="session")
def model_prefix() -> str:
"""Optional model-path prefix from MODEL_PREFIX env var.
Useful if models are downloaded to non-default local directories.
"""
prefix = os.environ.get("MODEL_PREFIX", "")
return f"{prefix.rstrip('/')}/" if prefix else ""
@pytest.fixture(autouse=True)
def default_vllm_config():
"""Set a default VllmConfig for all tests.
@@ -1121,16 +1057,15 @@ class OmniServer:
max_wait = 1200 # 20 minutes
start_time = time.time()
while time.time() - start_time < max_wait:
# Check for process status
ret = self.proc.poll()
if ret is not None:
raise RuntimeError(f"Server processes exited with code {ret} before becoming ready.")
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.settimeout(1)
result = sock.connect_ex((self.host, self.port))
if result == 0:
print(f"Server ready on {self.host}:{self.port}")
return
try:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.settimeout(1)
result = sock.connect_ex((self.host, self.port))
if result == 0:
print(f"Server ready on {self.host}:{self.port}")
return
except Exception:
pass
time.sleep(2)
raise RuntimeError(f"Server failed to start within {max_wait} seconds")
@@ -1215,9 +1150,7 @@ def pytest_addoption(parser):
@pytest.fixture(scope="session")
def run_level(request) -> str:
"""A command-line argument that specifies the level of tests to run in this session.
See https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/CI_5levels/"""
def run_level(request):
return request.config.getoption("--run-level")
@@ -1225,17 +1158,15 @@ _omni_server_lock = threading.Lock()
@pytest.fixture(scope="module")
def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: str) -> Generator[OmniServer, Any, None]:
def omni_server(request, run_level):
"""Start vLLM-Omni server as a subprocess with actual model weights.
Uses session scope so the server starts only once for the entire test session.
Multi-stage initialization can take 10-20+ minutes.
"""
with _omni_server_lock:
params: OmniServerParams = request.param
model = model_prefix + params.model
port = params.port
stage_config_path = params.stage_config_path
if run_level == "advanced_model" and stage_config_path is not None:
*port, model, stage_config_path = request.param
port = port[0] if port else None
if run_level == "advanced_model":
stage_config_path = modify_stage_config(
stage_config_path,
deletes={
@@ -1247,11 +1178,7 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st
},
)
server_args = params.server_args or []
server_args = ["--stage-init-timeout", "120", *server_args]
if stage_config_path is not None:
server_args += ["--stage-configs-path", stage_config_path]
server_args = ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]
with OmniServer(model, server_args, port=port) if port else OmniServer(model, server_args) as server:
print("OmniServer started successfully")
yield server
@@ -1271,17 +1198,6 @@ class OmniResponse:
error_message: str | None = None
@dataclass
class DiffusionResponse:
text_content: str | None = None
images: list[Image.Image] | None = None
audios: list[Any] | None = None
videos: list[Any] | None = None
e2e_latency: float | None = None
success: bool = False
error_message: str | None = None
def assert_omni_response(response: OmniResponse, request_config: dict[str, Any], run_level):
"""
Validate response results.
@@ -1330,51 +1246,6 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any],
print(f"similarity is: {response.similarity}")
def assert_diffusion_response(response: DiffusionResponse, request_config: dict[str, Any], run_level: str = None):
"""
Validate diffusion response results.
Args:
response: DiffusionResponse object. Any not-None content will be validated based on the request_config.
request_config: Request configuration dictionary containing parameters like model, messages, extra_body.
When validating a certain modality, the corresponding params in request_config['extra_body'] must present.
It will be used to check against the multimedia file in the response.
run_level: Test run level (e.g., "core_model", "advanced_model")
Raises:
AssertionError: When the response does not meet validation criteria
KeyError: When the request_config does not contain necessary parameters for validation
"""
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}")
extra_body = request_config.get("extra_body", {})
num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt", 1)
if response.images is not None:
assert len(response.images) > 0, "No images in response"
assert len(response.images) == num_outputs_per_prompt, (
f"Expected {num_outputs_per_prompt} images, got {len(response.images)}"
)
if run_level == "advanced_model":
expected_width = extra_body["width"] # intend to raise KeyError
expected_height = extra_body["height"] # intend to raise KeyError
for img in response.images:
assert_image_valid(img, width=expected_width, height=expected_height)
if response.videos is not None:
raise NotImplementedError(
"Video validation is not implemented yet"
) # consider using assert_video_valid defined above
if response.audios is not None:
raise NotImplementedError(
"Audio validation is not implemented yet"
) # consider using assert_audio_valid defined above
class OpenAIClientHandler:
"""
OpenAI client handler class, encapsulating both streaming and non-streaming response processing logic.
@@ -1397,7 +1268,7 @@ class OpenAIClientHandler:
self.client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key=api_key)
self.run_level = run_level
def _process_stream_omni_response(self, chat_completion) -> OmniResponse:
def _process_stream_response(self, chat_completion) -> OmniResponse:
"""
Process streaming responses.
@@ -1458,7 +1329,7 @@ class OpenAIClientHandler:
return result
def _process_non_stream_omni_response(self, chat_completion) -> OmniResponse:
def _process_non_stream_response(self, chat_completion) -> OmniResponse:
"""
Process non-streaming responses.
@@ -1512,46 +1383,7 @@ class OpenAIClientHandler:
return result
def _process_diffusion_response(self, chat_completion) -> DiffusionResponse:
"""
Process diffusion responses (image generation/editing).
Args:
chat_completion: OpenAI response object
Returns:
DiffusionResponse: Processed response object
"""
result = DiffusionResponse()
start_time = time.perf_counter()
try:
images = []
# [TODO] reading video and audio output from API response for later validation
for choice in chat_completion.choices:
if hasattr(choice.message, "content") and choice.message.content is not None:
content = choice.message.content
if isinstance(content, list):
for item in content:
if hasattr(item, "image_url") and item.image_url is not None:
image_url = item.image_url.url
if image_url.startswith("data:image"):
b64_data = image_url.split(",", 1)[1]
img = decode_b64_image(b64_data)
images.append(img)
result.e2e_latency = time.perf_counter() - start_time
result.images = images if images else None
result.success = True
except Exception as e:
result.error_message = f"Diffusion response processing error: {str(e)}"
print(f"Error: {result.error_message}")
return result
def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
def send_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
"""
Send OpenAI requests.
@@ -1577,9 +1409,9 @@ class OpenAIClientHandler:
)
if stream:
response = self._process_stream_omni_response(chat_completion)
response = self._process_stream_response(chat_completion)
else:
response = self._process_non_stream_omni_response(chat_completion)
response = self._process_non_stream_response(chat_completion)
assert_omni_response(response, request_config, run_level=self.run_level)
responses.append(response)
@@ -1605,82 +1437,25 @@ class OpenAIClientHandler:
chat_completion = future.result()
if stream:
response = self._process_stream_omni_response(chat_completion)
response = self._process_stream_response(chat_completion)
else:
response = self._process_non_stream_omni_response(chat_completion)
response = self._process_non_stream_response(chat_completion)
assert_omni_response(response, request_config, run_level=self.run_level)
responses.append(response)
return responses
def send_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
"""
Send OpenAI requests for diffusion models.
Args:
request_config: Request configuration dictionary containing parameters like model, messages
request_num: Number of requests to send concurrently, defaults to 1 (single request)
Returns:
List[OmniResponse]: List of response objects
"""
responses = []
stream = request_config.get("stream", False)
modalities = request_config.get("modalities", omit) # Most diffusion models don't require modalities param
extra_body = request_config.get("extra_body", None)
if stream:
raise NotImplementedError("Streaming is not currently implemented for diffusion model e2e test")
if request_num == 1:
# Send single request
chat_completion = self.client.chat.completions.create(
model=request_config.get("model"),
messages=request_config.get("messages"),
extra_body=extra_body,
modalities=modalities,
)
response = self._process_diffusion_response(chat_completion)
assert_diffusion_response(response, request_config, run_level=self.run_level)
responses.append(response)
else:
# Send concurrent requests
with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
futures = []
# Submit all request tasks
for _ in range(request_num):
future = executor.submit(
self.client.chat.completions.create,
model=request_config.get("model"),
messages=request_config.get("messages"),
modalities=modalities,
extra_body=extra_body,
)
futures.append(future)
# Process completed tasks
for future in concurrent.futures.as_completed(futures):
chat_completion = future.result()
response = self._process_diffusion_response(chat_completion)
assert_diffusion_response(response, request_config, run_level=self.run_level)
responses.append(response)
return responses
@pytest.fixture
def openai_client(omni_server: OmniServer, run_level: str):
"""Create OpenAIClientHandler fixture to facilitate communication with OmniServer
with encapsulated request sending, concurrent requests, response handling, and validation."""
def openai_client(omni_server, run_level):
"""Create OpenAIClientHandler fixture"""
return OpenAIClientHandler(host=omni_server.host, port=omni_server.port, api_key="EMPTY", run_level=run_level)
class OmniRunner:
"""
Offline test runner for Omni models.
Test runner for Omni models.
"""
def __init__(
@@ -1951,10 +1726,9 @@ class OmniRunner:
@pytest.fixture(scope="module")
def omni_runner(request, model_prefix):
def omni_runner(request):
with _omni_server_lock:
model, stage_config_path = request.param
model = model_prefix + model
with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, stage_init_timeout=300) as runner:
print("OmniRunner started successfully")
yield runner
+13 -31
View File
@@ -1,29 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for Ulysses + Ring sequence-parallel attention correctness.
What is tested
--------------
* ``test_sequence_parallel`` verifies that the ``Attention`` layer produces
numerically equivalent results whether the sequence is processed on a single
rank (baseline, SP=1) or sharded across multiple ranks via Ulysses/Ring SP.
The test spawns two separate multi-process runs with ``torch.multiprocessing.spawn``:
1. Baseline – world_size=1, ulysses_degree=1, ring_degree=1.
2. SP run – world_size=ulysses_degree*ring_degree, each rank holds a
contiguous slice of the full sequence.
After both runs, rank-0 output tensors are compared element-wise with a
tolerance appropriate for the dtype (bfloat16).
SP-plan hooks are NOT applied in this test
------------------------------------------
``ForwardContext.sp_plan_hooks_applied`` remains ``False``. As a result,
``ForwardContext.sp_active`` falls back to the "no hooks" branch: SP is
considered active whenever ``parallel_config.sequence_parallel_size > 1``.
This makes the test self-contained and suitable for standalone CI runs that
do not exercise the full model-registry pipeline.
"""
import os
import pickle
import tempfile
@@ -51,11 +27,6 @@ def update_environment_variables(envs_dict: dict[str, str]):
os.environ[k] = v
def seed_everything(seed: int):
torch.manual_seed(seed)
current_omni_platform.manual_seed(seed)
class TestAttentionModel(torch.nn.Module):
"""Test model using Attention layer."""
@@ -167,10 +138,11 @@ class TestMultiLayerAttentionModel(torch.nn.Module):
@pytest.mark.parametrize("num_heads", [8])
@pytest.mark.parametrize("head_size", [8])
@pytest.mark.parametrize("causal", [False])
@pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) # [torch.float16, torch.bfloat16]
@pytest.mark.parametrize("use_sync", [False])
@pytest.mark.parametrize("dynamic", [False])
@pytest.mark.parametrize("use_compile", [False])
@pytest.mark.parametrize("attn_backend", ["sdpa", "flash_attn"])
def test_sequence_parallel(
ulysses_degree: int,
ring_degree: int,
@@ -184,6 +156,7 @@ def test_sequence_parallel(
seq_len: int,
num_heads: int,
head_size: int,
attn_backend: str,
):
"""Test Ulysses attention by comparing with and without SP enabled."""
sequence_parallel_size = ulysses_degree * ring_degree
@@ -227,6 +200,7 @@ def test_sequence_parallel(
model_state_file,
input_data_file,
True, # is_baseline
attn_backend,
),
nprocs=1,
)
@@ -254,6 +228,7 @@ def test_sequence_parallel(
model_state_file,
input_data_file,
False, # is_baseline
attn_backend,
),
nprocs=sequence_parallel_size,
)
@@ -349,11 +324,12 @@ def ulysses_attention_on_test_model(
model_state_file: str,
input_data_file: str,
is_baseline: bool,
attn_backend: str,
):
"""Run Ulysses attention test on a test model and save results for comparison."""
# Use fixed seed for reproducibility across baseline and SP runs
RANDOM_SEED = 42
seed_everything(RANDOM_SEED)
current_omni_platform.seed_everything(RANDOM_SEED)
mode_str = "Baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})"
print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}")
@@ -390,6 +366,7 @@ def ulysses_attention_on_test_model(
model="test_model",
dtype=dtype,
parallel_config=parallel_config,
attention_backend=attn_backend, # Set the attention backend here
)
# Initialize model parallel
@@ -434,6 +411,11 @@ def ulysses_attention_on_test_model(
with open(model_state_file, "wb") as f:
pickle.dump(model_state, f)
# Generate and save full input data with fixed seed
# Reinitialize RNG to ensure reproducibility
torch.manual_seed(42)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(42)
full_hidden_states = torch.randn(
(batch_size, seq_len, hidden_size),
dtype=dtype,
+1 -1
View File
@@ -11,7 +11,7 @@ import torch.nn as nn
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model]
pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu]
class TestHSDPInferenceConfig:
@@ -11,7 +11,10 @@ import torch
from vllm_omni.diffusion.distributed.parallel_state import RankGenerator, set_seq_parallel_pg
pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.core_model, pytest.mark.cpu]
pytestmark = [
pytest.mark.diffusion,
pytest.mark.parallel,
]
def _fake_new_group_factory(created_groups: list[SimpleNamespace]):
@@ -47,8 +47,6 @@ requires_distributed = pytest.mark.skipif(
pytestmark = [
pytest.mark.diffusion,
pytest.mark.parallel,
pytest.mark.core_model,
pytest.mark.cpu,
]
# =============================================================================
@@ -8,8 +8,6 @@ import torch
from vllm_omni.diffusion.distributed import vae_patch_parallel as vae_patch_parallel
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _DummyConfig:
def __init__(self, **attrs):
-47
View File
@@ -508,50 +508,3 @@ def test_scale_keys_are_rounded():
# Currently we round keys to 3 decimal places
manager._update_adapter_scale(adapter_id, 0.0031)
assert manager._adapter_scales[adapter_id] == 0.003
def test_lora_manager_uses_valid_max_rank(monkeypatch):
"""Ensure that the LoRA manager uses a valid max rank for vLLM."""
manager = DiffusionLoRAManager(
pipeline=torch.nn.Module(),
device=torch.device("cpu"),
dtype=torch.bfloat16,
)
# Ensure that the rank is correctly adjusted to the smallest valid max rank
supported_max_rank = 64
unsupported_max_rank = 63
assert supported_max_rank in DiffusionLoRAManager._VALID_MAX_RANKS
assert unsupported_max_rank not in DiffusionLoRAManager._VALID_MAX_RANKS
def _fake_load(_req: LoRARequest):
lora_model = type("LM", (), {"id": _req.lora_int_id})()
peft_helper = type("PH", (), {"r": unsupported_max_rank})()
return lora_model, peft_helper
monkeypatch.setattr(manager, "_load_adapter", _fake_load)
req1 = _dummy_lora_request(1)
manager.add_adapter(req1)
assert manager._max_lora_rank == supported_max_rank
@pytest.mark.parametrize("rank", [-1, 0, DiffusionLoRAManager._VALID_MAX_RANKS[-1] + 1])
def test_lora_manager_max_rank_validation(monkeypatch, rank):
"""Check that invalid max ranks are handled correctly."""
manager = DiffusionLoRAManager(
pipeline=torch.nn.Module(),
device=torch.device("cpu"),
dtype=torch.bfloat16,
)
lora_rank = rank
def _fake_load(_req: LoRARequest):
lora_model = type("LM", (), {"id": _req.lora_int_id})()
peft_helper = type("PH", (), {"r": lora_rank})()
return lora_model, peft_helper
monkeypatch.setattr(manager, "_load_adapter", _fake_load)
req1 = _dummy_lora_request(1)
with pytest.raises(ValueError):
manager.add_adapter(req1)
@@ -1,244 +0,0 @@
from unittest.mock import MagicMock, patch
import pytest
import torch
from vllm_omni.diffusion.models.flux2.flux2_transformer import (
Flux2PosEmbed,
Flux2Transformer2DModel,
)
# Initialize TP group before tests
@pytest.fixture(scope="function", autouse=True)
def setup_tp_group():
"""Set up TP group for each test function"""
with patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=2):
with patch("vllm.distributed.parallel_state.get_tp_group") as mock_get_tp_group:
mock_tp_group = MagicMock()
mock_tp_group.world_size = 2
mock_get_tp_group.return_value = mock_tp_group
yield
class TestFlux2TransformerWeightLoading:
"""Test Flux2Transformer weight loading functionality"""
def test_weight_loading_tp2(self, setup_tp_group):
"""Verify weights load correctly with TP=2"""
# Prepare test data
model = Flux2Transformer2DModel(
num_layers=2,
num_single_layers=1,
num_attention_heads=48,
attention_head_dim=128,
joint_attention_dim=15360,
)
# Mock TP=2 weight loading
mock_weights = []
# 1. Test regular weight loading
mock_weights.append(("x_embedder.weight", torch.randn(6144, 128)))
mock_weights.append(("context_embedder.weight", torch.randn(6144, 15360)))
mock_weights.append(("proj_out.weight", torch.randn(128, 6144)))
# 2. Test stacked_params_mapping weight loading
# Full weights - load_weights handles sharding internally
to_qkv_weight = torch.randn(18432, 6144) # Full weights
mock_weights.append(("transformer_blocks.0.attn.to_qkv.weight", to_qkv_weight))
# Add_kv_proj weights
add_kv_proj_weight = torch.randn(18432, 6144) # Full weights
mock_weights.append(("transformer_blocks.0.attn.add_kv_proj.weight", add_kv_proj_weight))
# 3. Test single block weight loading
single_block_qkv_mlp_weight = torch.randn(18432 + 18432 * 3, 6144) # Full weights
mock_weights.append(("single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight", single_block_qkv_mlp_weight))
# Execute weight loading
loaded_params = model.load_weights(mock_weights)
# Verify
assert len(loaded_params) > 0, "Parameters should be loaded"
# Verify stacked_params_mapping is correctly set
assert model.stacked_params_mapping is not None
# Should have 6 mappings: 3 for to_qkv + 3 for add_kv_proj
assert len(model.stacked_params_mapping) == 6, "Should have 6 mappings"
# Verify weights are correctly loaded to corresponding layers
assert hasattr(model.transformer_blocks[0].attn.to_qkv, "weight")
# With TP=2, weight dimension on each GPU should be 18432/2 = 9216
assert model.transformer_blocks[0].attn.to_qkv.weight.shape == (9216, 6144), (
f"With TP=2, to_qkv weight dimension should be (9216, 6144), got {model.transformer_blocks[0].attn.to_qkv.weight.shape}"
)
class TestFlux2RopePositionEmbedding:
"""Test Flux2 RoPE position embedding functionality"""
def test_rope_position_embedding(self):
"""Verify RoPE produces correct embeddings for 4D coordinates"""
# Prepare test data - use model default configuration
# axes_dims_rope default is (32, 32, 32, 32)
# get_1d_rotary_pos_embed outputs half the dimension for real/imag parts
# So actual output dimension should be (16+16+16+16) = 64
axes_dims = (32, 32, 32, 32) # Use model default
rope_theta = 2000 # Model default is 2000
pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims)
# Create test IDs
seq_len = 10
ids = torch.randint(0, 100, (seq_len, 4)) # [S, 4]
# Forward pass
freqs_cos, freqs_sin = pos_embed(ids)
# Verify output shape - based on model config, expected dimension is 64
# Each axes_dim=32 outputs 16-dim real part, sum of 4 dimensions = 64
expected_dim = sum(axes_dims) // 2 # 128/2 = 64
assert freqs_cos.shape == (seq_len, expected_dim), (
f"Expected shape {(seq_len, expected_dim)}, got {freqs_cos.shape}"
)
assert freqs_sin.shape == (seq_len, expected_dim), (
f"Expected shape {(seq_len, expected_dim)}, got {freqs_sin.shape}"
)
# Verify output type - NPU may use float32, other devices use float64
assert freqs_cos.dtype in [torch.float32, torch.float64], "Should be float type"
assert freqs_sin.dtype in [torch.float32, torch.float64], "Should be float type"
# Verify value range
assert torch.all(freqs_cos >= -1) and torch.all(freqs_cos <= 1), "cos values should be in [-1, 1]"
assert torch.all(freqs_sin >= -1) and torch.all(freqs_sin <= 1), "sin values should be in [-1, 1]"
# Verify trigonometric relationship: cos^2 + sin^2 ≈ 1
cos_sq_sin_sq = freqs_cos**2 + freqs_sin**2
assert torch.allclose(cos_sq_sin_sq, torch.ones_like(cos_sq_sin_sq), atol=1e-6), "cos^2 + sin^2 should ≈ 1"
# Verify different positions produce different embeddings
ids_diff = torch.randint(0, 100, (seq_len, 4))
freqs_cos_diff, freqs_sin_diff = pos_embed(ids_diff)
assert not torch.allclose(freqs_cos, freqs_cos_diff), "Different positions should produce different embeddings"
# Verify same positions produce same embeddings
ids_same = ids.clone()
freqs_cos_same, freqs_sin_same = pos_embed(ids_same)
assert torch.allclose(freqs_cos, freqs_cos_same), "Same positions should produce same embeddings"
assert torch.allclose(freqs_sin, freqs_sin_same), "Same positions should produce same embeddings"
class TestFlux2PackedModuleMapping:
"""Test Flux2 packed module mapping functionality"""
def test_packed_module_mapping(self, setup_tp_group):
"""Verify to_qkv packing matches HF checkpoint"""
model = Flux2Transformer2DModel(
num_layers=1,
num_single_layers=0,
num_attention_heads=48,
attention_head_dim=128,
joint_attention_dim=15360,
)
# Verify stacked_params_mapping is correctly initialized
model.load_weights([]) # Trigger stacked_params_mapping initialization
assert model.stacked_params_mapping is not None
# Verify mapping configuration
expected_mappings = [
(".to_qkv.", ".to_q.", "q"),
(".to_qkv.", ".to_k.", "k"),
(".to_qkv.", ".to_v.", "v"),
(".add_kv_proj", ".add_q_proj", "q"),
(".add_kv_proj", ".add_k_proj", "k"),
(".add_kv_proj", ".add_v_proj", "v"),
]
assert model.stacked_params_mapping == expected_mappings
# Create mock HF checkpoint weights
hf_weights = []
# Mock HF separated Q/K/V weights
attn_block = model.transformer_blocks[0].attn
head_dim = 128
num_heads = 48
hidden_size = 6144
# Full weight dimension
full_shard_size = num_heads * head_dim # 6144
# Q projection weights (full weights)
q_weight = torch.randn(full_shard_size, hidden_size)
hf_weights.append(("transformer_blocks.0.attn.to_q.weight", q_weight))
# K projection weights (full weights)
k_weight = torch.randn(full_shard_size, hidden_size)
hf_weights.append(("transformer_blocks.0.attn.to_k.weight", k_weight))
# V projection weights (full weights)
v_weight = torch.randn(full_shard_size, hidden_size)
hf_weights.append(("transformer_blocks.0.attn.to_v.weight", v_weight))
# Mock HF separated add_q/k/v projection weights
add_q_weight = torch.randn(full_shard_size, hidden_size)
hf_weights.append(("transformer_blocks.0.attn.add_q_proj.weight", add_q_weight))
add_k_weight = torch.randn(full_shard_size, hidden_size)
hf_weights.append(("transformer_blocks.0.attn.add_k_proj.weight", add_k_weight))
add_v_weight = torch.randn(full_shard_size, hidden_size)
hf_weights.append(("transformer_blocks.0.attn.add_v_proj.weight", add_v_weight))
# Execute weight loading
loaded_params = model.load_weights(hf_weights)
# Verify weights are loaded
assert len(loaded_params) > 0
# Verify final QKV weights are correctly combined (considering TP sharding)
# With TP=2, dimension on each GPU should be half of full dimension
expected_qkv_shape = (full_shard_size * 3 // 2, hidden_size) # 18432/2 = 9216
assert attn_block.to_qkv.weight.shape == expected_qkv_shape, (
f"to_qkv weight dimension should be {expected_qkv_shape}, got {attn_block.to_qkv.weight.shape}"
)
expected_add_kv_shape = (full_shard_size * 3 // 2, hidden_size)
assert attn_block.add_kv_proj.weight.shape == expected_add_kv_shape, (
f"add_kv_proj weight dimension should be {expected_add_kv_shape}, got {attn_block.add_kv_proj.weight.shape}"
)
def test_packed_mapping_edge_cases(self, setup_tp_group):
"""Test edge cases for packed mapping"""
model = Flux2Transformer2DModel(
num_layers=1,
num_single_layers=1,
num_attention_heads=48,
attention_head_dim=128,
joint_attention_dim=15360,
)
model.load_weights([])
# Test invalid weight names
invalid_weights = [("invalid.weight", torch.randn(10, 10))]
loaded_params = model.load_weights(invalid_weights)
assert len(loaded_params) == 0, "Should not load invalid weights"
# Test to_out weight renaming
to_out_weight = torch.randn(6144, 6144)
weights = [("transformer_blocks.0.attn.to_out.0.weight", to_out_weight)]
loaded_params = model.load_weights(weights)
# Check if to_out related weights are loaded
to_out_loaded = any("to_out" in p for p in loaded_params)
assert to_out_loaded, "to_out weights should be correctly renamed and loaded"
# Test partial weight loading
partial_weights = [
("x_embedder.weight", torch.randn(6144, 128)),
("transformer_blocks.0.attn.to_q.weight", torch.randn(6144, 6144)), # Full weights
]
loaded_params = model.load_weights(partial_weights)
assert len(loaded_params) == 2, "Should load two weights"
assert "x_embedder.weight" in loaded_params
assert any("to_qkv" in p for p in loaded_params), "to_q should be mapped to to_qkv"
Loaded 100 of 249 files, more files were not shown because too many files have changed in this diff. Show more