Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bfdf9534e8 | ||
|
|
67bd02f767 | ||
|
|
0814d845dc | ||
|
|
bb973cb10c | ||
|
|
fb3abb38fb | ||
|
|
eb81914c64 | ||
|
|
1951dca786 | ||
|
|
d71cacc74c | ||
|
|
465f1464d2 | ||
|
|
f6894f4308 | ||
|
|
a2d1dafc65 | ||
|
|
6dc4638a7e | ||
|
|
d35cdd2dab | ||
|
|
0958e62ecc | ||
|
|
3901ad8165 | ||
|
|
b84c99c122 | ||
|
|
733f4f6a36 | ||
|
|
67d14b5ac4 | ||
|
|
64efa66869 | ||
|
|
6646a66d41 | ||
|
|
378359fc11 | ||
|
|
9c260b3774 | ||
|
|
9177b37e2f | ||
|
|
50bc51ffdf | ||
|
|
044240d8eb | ||
|
|
f4d922d8f6 | ||
|
|
eaa16e1941 | ||
|
|
8a6786731e |
No files matched your search
@@ -14,7 +14,7 @@ steps:
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -23,7 +23,7 @@ steps:
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -33,6 +33,6 @@ steps:
|
||||
key: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-nightly.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
@@ -13,13 +13,16 @@ Test pipeline mode (e.g. test-merge.yml):
|
||||
recognize this uploader-only key).
|
||||
|
||||
Usage:
|
||||
python3 upload_pipeline.py [--upload] <pipeline.yml>
|
||||
python3 upload_pipeline.py [--upload] [--all] <pipeline.yml>
|
||||
|
||||
# Bootstrap (replaces upload_pipeline_with_skip_ci.sh):
|
||||
python3 upload_pipeline.py --upload .buildkite/pipeline.yml
|
||||
|
||||
# Test pipeline (replaces upload_test_pipeline_with_diff_skip.py):
|
||||
python3 upload_pipeline.py --upload .buildkite/test-merge.yml
|
||||
|
||||
# Full suite (rebase pipeline): keep all steps, still strip the uploader-only key:
|
||||
python3 upload_pipeline.py --upload --all .buildkite/test-merge.yml
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -437,11 +440,17 @@ def resolve_pipeline_path(arg: str) -> Path:
|
||||
return ROOT / path
|
||||
|
||||
|
||||
def render_pipeline(path: Path) -> str:
|
||||
def render_pipeline(path: Path, *, force_all: bool = False) -> str:
|
||||
text = path.read_text(encoding="utf-8")
|
||||
diff_range = resolve_diff_range()
|
||||
changed_files = resolve_changed_files()
|
||||
|
||||
# ``--all`` forces the keep-all-steps path (no diff-aware skipping) while still
|
||||
# stripping ``source_file_dependencies``. Used by the rebase pipeline so main builds
|
||||
# run the full e2e suite (see .buildkite/rebase-pipeline.yaml).
|
||||
if force_all:
|
||||
changed_files = None
|
||||
|
||||
if BOOTSTRAP_MARKER in text:
|
||||
rendered = render_bootstrap_pipeline(
|
||||
text,
|
||||
@@ -480,6 +489,11 @@ def main() -> int:
|
||||
action="store_true",
|
||||
help="Pipe rendered YAML to buildkite-agent pipeline upload",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--all",
|
||||
action="store_true",
|
||||
help="Keep all steps (disable diff-aware skipping); still strips source_file_dependencies",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
path = resolve_pipeline_path(args.pipeline)
|
||||
@@ -487,7 +501,7 @@ def main() -> int:
|
||||
_log(f"missing pipeline file: {path}")
|
||||
return 1
|
||||
|
||||
rendered = render_pipeline(path)
|
||||
rendered = render_pipeline(path, force_all=args.all)
|
||||
if args.upload:
|
||||
upload_to_buildkite(rendered)
|
||||
else:
|
||||
|
||||
@@ -363,6 +363,42 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "TTS · IndexTTS2 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_indextts2.py
|
||||
- tests/e2e/offline_inference/test_indextts2.py
|
||||
- tests/model_executor/models/indextts2/
|
||||
- tests/model_executor/stage_input_processors/test_indextts2.py
|
||||
- tests/assets/indextts2/
|
||||
- vllm_omni/model_executor/models/indextts2/
|
||||
- vllm_omni/model_executor/stage_input_processors/indextts2.py
|
||||
- vllm_omni/deploy/indextts2.yaml
|
||||
- vllm_omni/deploy/indextts2_low_latency.yaml
|
||||
- vllm_omni/entrypoints/openai/serving_speech.py
|
||||
- vllm_omni/entrypoints/openai/tts_adapters/indextts2.py
|
||||
- vllm_omni/config/pipeline_registry.py
|
||||
- vllm_omni/model_executor/models/registry.py
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
pytest -s -v tests/e2e/online_serving/test_indextts2.py tests/e2e/offline_inference/test_indextts2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
gpus: all
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni · Qwen3-Omni Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
@@ -509,6 +545,34 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · SoulX-Singer Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_soulxsinger.py
|
||||
- vllm_omni/diffusion/models/soulx_singer/
|
||||
- vllm_omni/deploy/soulxsinger_svc.yaml
|
||||
- vllm_omni/deploy/soulxsinger_svs.yaml
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/offline_inference/test_soulxsinger.py -k 'svc or precomputed' -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion · Qwen Image Edit Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_edit.py
|
||||
|
||||
@@ -1279,7 +1279,7 @@ steps:
|
||||
- label: ":full_moon: Diffusion X2V · Function Test"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py tests/e2e/online_serving/test_ltx2_expansion.py -m "full_model and cuda" --run-level "full_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -1322,6 +1322,7 @@ steps:
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
|
||||
+10
-10
@@ -22,30 +22,30 @@
|
||||
/.buildkite/ @yenuo26 @congw729
|
||||
|
||||
# Engine and model_executor
|
||||
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan
|
||||
/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan
|
||||
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan
|
||||
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
|
||||
/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
|
||||
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks
|
||||
/vllm_omni/model_executor/ @tzhouam @gcanlin
|
||||
/vllm_omni/worker/ @tzhouam @gcanlin
|
||||
|
||||
# Diffusion
|
||||
/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666
|
||||
/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
|
||||
/vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
|
||||
|
||||
# Diffusion CPU offloading
|
||||
/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
|
||||
/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao @xuechendi
|
||||
|
||||
# Omni and TTS support
|
||||
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
|
||||
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter
|
||||
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin
|
||||
|
||||
# Benchmarks
|
||||
/benchmarks/accuracy/ @david6666666
|
||||
/benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
|
||||
/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa
|
||||
/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi
|
||||
/benchmarks/distributed/ @princepride
|
||||
/benchmarks/glm_image/ @princepride @JaredforReal
|
||||
/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride
|
||||
/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin
|
||||
|
||||
# Usability and Observability
|
||||
/vllm_omni/deploy/ @tzhouam @lishunyang12
|
||||
@@ -53,13 +53,13 @@
|
||||
/vllm_omni/profiler/ @gcanlin
|
||||
|
||||
# ComfyUI
|
||||
/apps/ @fhfuih
|
||||
/apps/ @fhfuih @alex-jw-brooks
|
||||
|
||||
# Large Scale Serving
|
||||
/vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
|
||||
|
||||
# Quantization and configuration
|
||||
/vllm_omni/config/ @lishunyang12
|
||||
/vllm_omni/config/ @lishunyang12 @alex-jw-brooks
|
||||
/vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
|
||||
|
||||
# Hardware plugins and accelerators
|
||||
|
||||
@@ -19,6 +19,7 @@ body:
|
||||
|
||||
Consider redacting or replacing sensitive values with placeholders like `<YOUR_TOKEN_HERE>` when sharing configuration or code examples.
|
||||
- type: textarea
|
||||
id: current-environment
|
||||
attributes:
|
||||
label: Your current environment
|
||||
description: |
|
||||
@@ -42,6 +43,7 @@ body:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: code-version
|
||||
attributes:
|
||||
label: Your code version
|
||||
description: |
|
||||
@@ -66,6 +68,7 @@ body:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: bug-description
|
||||
attributes:
|
||||
label: 🐛 Describe the bug
|
||||
description: |
|
||||
|
||||
@@ -241,6 +241,7 @@ datasets/
|
||||
*.json
|
||||
!tests/dfx/perf/tests/*.json
|
||||
!apps/ComfyUI-vLLM-Omni/example_workflows/*.json
|
||||
!tests/assets/soulxsinger/*.json
|
||||
*.jsonl
|
||||
*.parquet
|
||||
|
||||
|
||||
@@ -101,7 +101,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
|
||||
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
|
||||
|
||||
ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
|
||||
ENV LD_LIBRARY_PATH="/opt/intel/oneapi/ccl/2021.15/lib:$LD_LIBRARY_PATH:/usr/local/lib/"
|
||||
|
||||
ENV VLLM_TARGET_DEVICE=xpu
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
@@ -33,16 +33,20 @@ Lead maintainers are responsible for the overall direction and strategy of the p
|
||||
|
||||
Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
|
||||
|
||||
- [@alex-jw-brooks](https://github.com/alex-jw-brooks): Frontend, Configuration, and Entrypoint Support
|
||||
- [@amy-why-3459](https://github.com/amy-why-3459): Omni Modality Model and World Model Support
|
||||
- [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
|
||||
- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
|
||||
- [@gcanlin](https://github.com/gcanlin): Hardware plugin, NPU integration, and TTS Model Support
|
||||
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
|
||||
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
|
||||
- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
|
||||
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
|
||||
- [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs
|
||||
- [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache
|
||||
- [@Sy0307](https://github.com/Sy0307): TTS Model and AR Engine Support
|
||||
- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
|
||||
- [@wtomin](https://github.com/wtomin): Diffusion models, parallel, and docs
|
||||
- [@xuechendi](https://github.com/xuechendi): Diffusion and Intel GPU Support
|
||||
- [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
|
||||
- [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
|
||||
- [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
|
||||
@@ -53,7 +57,7 @@ This section breaks down diffusion-related responsibilities. If you have PRs tou
|
||||
|
||||
### Runtime-related
|
||||
|
||||
- Diffusion models: @RuixiangMa, @wtomin
|
||||
- Diffusion models: @RuixiangMa, @wtomin, @xuechendi
|
||||
|
||||
### Optimization
|
||||
|
||||
@@ -61,6 +65,7 @@ This section breaks down diffusion-related responsibilities. If you have PRs tou
|
||||
- Attention backend: @ZJY0516
|
||||
- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
|
||||
- Offload: @yuanheng-zhao
|
||||
- Intel GPU support: @xuechendi
|
||||
- Quantization: @lishunyang12, @david6666666
|
||||
- Kernel fusion / communication-computation: @ZJY0516
|
||||
|
||||
|
||||
@@ -37,7 +37,7 @@ th {
|
||||
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
@@ -63,6 +63,8 @@ th {
|
||||
| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | |
|
||||
| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | |
|
||||
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
@@ -74,6 +76,7 @@ th {
|
||||
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
|
||||
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
|
||||
| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | |
|
||||
| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | |
|
||||
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
|
||||
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
|
||||
|
||||
@@ -32,6 +32,20 @@ curl -X POST http://localhost:8000/v1/images/generations \
|
||||
}' | jq -r '.data[0].b64_json' | base64 -d > dragon.png
|
||||
```
|
||||
|
||||
**Using curl save to file:**
|
||||
|
||||
```bash
|
||||
curl -o dragon.png -X POST http://localhost:8000/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"prompt": "a dragon laying over the spine of the Green Mountains of Vermont",
|
||||
"size": "1024x1024",
|
||||
"seed": 42,
|
||||
"response_format":"file"
|
||||
}'
|
||||
```
|
||||
|
||||
|
||||
**Using Python:**
|
||||
|
||||
```python
|
||||
@@ -57,6 +71,36 @@ img = Image.open(io.BytesIO(img_bytes))
|
||||
img.save("cat.png")
|
||||
```
|
||||
|
||||
**Using Python save to file:**
|
||||
|
||||
```python
|
||||
import requests
|
||||
import base64
|
||||
from PIL import Image
|
||||
import io
|
||||
import re
|
||||
|
||||
response = requests.post(
|
||||
"http://localhost:8000/v1/images/generations",
|
||||
json={
|
||||
"prompt": "a black and white cat wearing a princess tiara",
|
||||
"size": "1024x1024",
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42,
|
||||
"response_format":"file"
|
||||
}
|
||||
)
|
||||
|
||||
# save to file
|
||||
content_disposition = response.headers.get("Content-Disposition", "")
|
||||
match = re.search(r'filename="?(.+)"?', content_disposition)
|
||||
filename = match.group(1) if match else "save.png"
|
||||
with open(filename, "wb") as f:
|
||||
for chunk in response.iter_content(8192):
|
||||
f.write(chunk)
|
||||
print("saved:", filename)
|
||||
```
|
||||
|
||||
**Using OpenAI SDK:**
|
||||
|
||||
```python
|
||||
|
||||
@@ -197,6 +197,7 @@ Factory function `get_offload_backend()` selects the appropriate backend based o
|
||||
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` |
|
||||
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
|
||||
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
|
||||
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | `"layers"` |
|
||||
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | `"layers"`, `"customized modules"` |
|
||||
|
||||
**Notes:**
|
||||
|
||||
@@ -146,7 +146,7 @@ The following tables show which models support each feature:
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.2-S2V** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.2-S2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
|
||||
@@ -1,14 +0,0 @@
|
||||
{
|
||||
"prompt": "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. The robotic arm is in the process of pouring a transparent liquid into the glass. To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. The lighting is bright, casting soft shadows on the table. The robotic arm's movements are smooth and controlled, demonstrating precision in its task. As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. The other robotic arm on the right side also stays stationary throughout the video. The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper.",
|
||||
"negative_prompt": "",
|
||||
"vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4",
|
||||
"height": 720,
|
||||
"width": 1280,
|
||||
"num_frames": 189,
|
||||
"num_inference_steps": 35,
|
||||
"guidance_scale": 6.0,
|
||||
"flow_shift": 10.0,
|
||||
"fps": 24,
|
||||
"condition_frame_indexes_vision": [0, 1],
|
||||
"condition_video_keep": "first"
|
||||
}
|
||||
@@ -74,6 +74,12 @@ def parse_args():
|
||||
parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
|
||||
parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
|
||||
parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
|
||||
parser.add_argument(
|
||||
"--stream",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable streaming CoT display; print AR text token-by-token in real time.",
|
||||
)
|
||||
parser.add_argument("--log-stats", action="store_true", default=False)
|
||||
parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.")
|
||||
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
|
||||
@@ -222,6 +228,8 @@ def main():
|
||||
prompt_dict["modalities"] = ["text"]
|
||||
formatted_prompts.append(prompt_dict)
|
||||
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
|
||||
params_list = list(omni.default_sampling_params_list)
|
||||
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
@@ -253,6 +261,10 @@ def main():
|
||||
sp.width = args.width
|
||||
elif hasattr(sp, "stop_token_ids"):
|
||||
sp.stop_token_ids = ar_stop_token_ids
|
||||
# When --stream is set, request DELTA output from the AR stage
|
||||
# so we can display CoT text token-by-token in real time.
|
||||
if args.stream and hasattr(sp, "output_kind"):
|
||||
sp.output_kind = RequestOutputKind.DELTA
|
||||
|
||||
print(f"\n{'=' * 60}")
|
||||
print("HunyuanImage-3.0 Generation Configuration:")
|
||||
@@ -281,22 +293,32 @@ def main():
|
||||
print(f" Prompts: {prompts}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
|
||||
# When --stream is set, print AR CoT text token-by-token in real time.
|
||||
# Otherwise, collect and print the full AR text once when stage 0 finishes.
|
||||
omni_outputs = omni.generate(
|
||||
prompts=formatted_prompts,
|
||||
sampling_params_list=params_list,
|
||||
py_generator=True,
|
||||
use_tqdm=False,
|
||||
)
|
||||
img_idx = 0
|
||||
for req_output in omni_outputs:
|
||||
ro = getattr(req_output, "request_output", None)
|
||||
txt = ""
|
||||
if ro and getattr(ro, "outputs", None):
|
||||
txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
|
||||
if not txt:
|
||||
ar_text = getattr(req_output, "custom_output", {}).get("ar_generated_text")
|
||||
if isinstance(ar_text, list):
|
||||
txt = "\n".join(text for text in ar_text if text)
|
||||
else:
|
||||
txt = ar_text or ""
|
||||
if txt:
|
||||
print(f"[Output] Text:\n{txt}")
|
||||
stage_id = getattr(req_output, "stage_id", None)
|
||||
|
||||
# AR stage text — each CompletionOutput.text is already a delta when
|
||||
# output_kind=DELTA, so we can print it directly (matching the pattern
|
||||
# in serving_chat.py).
|
||||
if stage_id == 0 and ro and getattr(ro, "outputs", None):
|
||||
for o in ro.outputs:
|
||||
text = getattr(o, "text", "") or ""
|
||||
if text:
|
||||
print(text, end="", flush=True)
|
||||
# Non-streaming: one shot with full text — emit a trailing newline.
|
||||
if not args.stream:
|
||||
print(flush=True)
|
||||
|
||||
# Collect images from diffusion stage
|
||||
images = getattr(req_output, "images", None)
|
||||
if not images and ro and hasattr(ro, "images"):
|
||||
images = ro.images
|
||||
@@ -304,7 +326,7 @@ def main():
|
||||
for j, img in enumerate(images):
|
||||
save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
|
||||
img.save(save_path)
|
||||
print(f"[Output] Saved image to {save_path}")
|
||||
print(f"\n[Output] Saved image to {save_path}")
|
||||
img_idx += 1
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# Image-To-Video
|
||||
|
||||
This example demonstrates how to generate videos from images using Wan2.2 Image-to-Video models with vLLM-Omni's offline inference API.
|
||||
This shared example generates videos from images with VACE, Wan2.2, LTX2,
|
||||
HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
@@ -48,19 +49,38 @@ python image_to_video.py \
|
||||
--output i2v_output.mp4
|
||||
```
|
||||
|
||||
### Cosmos3
|
||||
|
||||
```bash
|
||||
# Cosmos3 bundles example frames under assets/ (any RGB image works too):
|
||||
python image_to_video.py \
|
||||
--model nvidia/Cosmos3-Nano \
|
||||
--image /path/to/Cosmos3-Nano/assets/example_i2v_input.jpg \
|
||||
--prompt "The scene comes to life with smooth, natural motion." \
|
||||
--negative-prompt "blurry, distorted, low quality" \
|
||||
--height 720 --width 1280 --num-frames 189 --fps 24 \
|
||||
--num-inference-steps 35 --guidance-scale 6.0 \
|
||||
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}' \
|
||||
--output cosmos3_i2v.mp4
|
||||
```
|
||||
|
||||
Key arguments:
|
||||
|
||||
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V).
|
||||
- `--image`: Path to input image (required).
|
||||
- `--image`: Path to the first-frame or source image.
|
||||
- `--last-image`: Optional last-frame condition for models such as VACE.
|
||||
- `--mask-image`: Optional inpainting mask. White pixels are regenerated and black pixels are preserved.
|
||||
- `--reference-image`: Optional reference image; repeat it to provide multiple references.
|
||||
- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3.
|
||||
- `--prompt`: Text description of desired motion/animation.
|
||||
- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16.
|
||||
- `--num-frames`: Number of frames (default 81).
|
||||
- `--num-frames`: Number of frames (default: model-specific — Wan 81, LTX2 121, Cosmos3 189).
|
||||
- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
|
||||
- `--negative-prompt`: Optional list of artifacts to suppress.
|
||||
- `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
|
||||
- `--flow-shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
|
||||
- `--flow-shift`: Scheduler flow shift (default: model-specific — Wan/LTX2 5.0, Cosmos3 10.0).
|
||||
- `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints.
|
||||
- `--num-inference-steps`: Number of denoising steps (default 50).
|
||||
- `--num-inference-steps`: Number of denoising steps (default: model-specific — Wan 50, LTX2 40, Cosmos3 35).
|
||||
- `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
|
||||
- `--output`: Path to save the generated video.
|
||||
- `--vae-use-slicing`: Enable VAE slicing for memory optimization.
|
||||
@@ -76,5 +96,98 @@ Key arguments:
|
||||
|
||||
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
|
||||
|
||||
## Wan2.1 VACE Conditional Tasks
|
||||
|
||||
The shared script selects the VACE conditioning structure from the media inputs.
|
||||
No explicit mode parameter is required: the script constructs the source video,
|
||||
mask, or reference images consumed by the VACE pipeline.
|
||||
Download the same Hugging Face assets used by the original VACE example:
|
||||
|
||||
```bash
|
||||
wget -O astronaut.jpg https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg
|
||||
wget -O vace_first_frame.png https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_first_frame.png
|
||||
wget -O vace_last_frame.png https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_last_frame.png
|
||||
```
|
||||
|
||||
### Image-to-Video (I2V)
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--image astronaut.jpg \
|
||||
--prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \
|
||||
--seed 42 --height 480 --width 832 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_i2v_output.mp4
|
||||
```
|
||||
|
||||
### Video-to-Last-Frame (V2LF)
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--last-image astronaut.jpg \
|
||||
--prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \
|
||||
--seed 42 --height 480 --width 832 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_v2lf_output.mp4
|
||||
```
|
||||
|
||||
### First-Last-Frame-to-Video (FLF2V)
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--image vace_first_frame.png \
|
||||
--last-image vace_last_frame.png \
|
||||
--prompt "CG animation style, a small blue bird takes off from a branch and lands on another branch" \
|
||||
--seed 42 --height 512 --width 512 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_flf2v_output.mp4
|
||||
```
|
||||
|
||||
### Inpainting
|
||||
|
||||
Create a mask matching the original VACE example: a 160-pixel-wide white
|
||||
vertical stripe marks the region to regenerate.
|
||||
|
||||
```bash
|
||||
python - <<'PY'
|
||||
from PIL import Image
|
||||
|
||||
mask = Image.new("L", (832, 480), 0)
|
||||
mask.paste(255, (336, 0, 496, 480))
|
||||
mask.save("vace_center_mask.png")
|
||||
PY
|
||||
```
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--image astronaut.jpg \
|
||||
--mask-image vace_center_mask.png \
|
||||
--prompt "Shrek, the ogre, walks out of a building in a happy mood" \
|
||||
--seed 42 --height 480 --width 832 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_inpaint_output.mp4
|
||||
```
|
||||
|
||||
### Reference-to-Video (R2V)
|
||||
|
||||
Repeat `--reference-image` to provide more than one reference image.
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--reference-image astronaut.jpg \
|
||||
--prompt "Camera slowly zooms out from the character walking in a garden" \
|
||||
--seed 42 --height 480 --width 832 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_r2v_output.mp4
|
||||
```
|
||||
|
||||
The VACE T2V command is documented in the shared
|
||||
[`text_to_video.py`](../text_to_video/text_to_video.md#wan21-vace-t2v) example.
|
||||
|
||||
For Wan2.2 LightX2V-converted local Diffusers directories and related LoRA
|
||||
assets, see the [LoRA guide](../../../docs/user_guide/diffusion/lora.md#wan22-lightx2v-offline-assembly).
|
||||
@@ -2,13 +2,14 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, or HunyuanVideo-1.5.
|
||||
Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, HunyuanVideo-1.5, or Wan2.1 VACE.
|
||||
|
||||
Supports:
|
||||
- Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
|
||||
- Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B)
|
||||
- LTX2 image-to-video pipeline
|
||||
- HunyuanVideo-1.5 I2V: SigLIP + VAE dual image conditioning
|
||||
- Wan2.1 VACE: first/last-frame, inpainting, and reference conditioning
|
||||
|
||||
Usage:
|
||||
# Wan I2V-A14B (MoE)
|
||||
@@ -30,9 +31,16 @@ Usage:
|
||||
python image_to_video.py --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \
|
||||
--image input.jpg --prompt "A cat playing with yarn" \
|
||||
--flow-shift 5.0 --guidance-scale 6.0
|
||||
|
||||
# Cosmos3 I2V (image conditioning)
|
||||
python image_to_video.py --model nvidia/Cosmos3-Nano \
|
||||
--image input.jpg --prompt "The scene comes to life with smooth, natural motion." \
|
||||
--num-frames 189 --num-inference-steps 35 --guidance-scale 6.0 --fps 24 \
|
||||
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}'
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import functools
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
@@ -43,39 +51,59 @@ import PIL.Image
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.model_extras import (
|
||||
build_image_to_video_prompt,
|
||||
get_extra_body_params,
|
||||
get_model_class_name,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def parse_profiler_config(value: str) -> dict[str, Any]:
|
||||
def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]:
|
||||
"""Parse a CLI value as a JSON object, attributing errors to ``flag_name``."""
|
||||
try:
|
||||
config = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e
|
||||
if not isinstance(config, dict):
|
||||
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object")
|
||||
return config
|
||||
|
||||
|
||||
parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate a video from an image (Wan2.2, LTX2, HunyuanVideo-1.5).")
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate a video from one or more images (Wan2.2, LTX2, HunyuanVideo-1.5, or Wan2.1 VACE)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
help="Diffusers I2V model ID or local path (Wan2.2 or HunyuanVideo-1.5).",
|
||||
help="Diffusers I2V model ID or local path (Wan2.2, HunyuanVideo-1.5, or Wan2.1 VACE).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model-class-name",
|
||||
default=None,
|
||||
help="Override model class name (e.g., LTX2ImageToVideoPipeline).",
|
||||
)
|
||||
parser.add_argument("--image", required=True, help="Path to input image.")
|
||||
parser.add_argument("--image", help="Path to the first-frame or source image.")
|
||||
parser.add_argument("--last-image", help="Path to a last-frame condition (used by models such as VACE).")
|
||||
parser.add_argument("--mask-image", help="Path to an inpainting mask (used by models such as VACE).")
|
||||
parser.add_argument(
|
||||
"--reference-image",
|
||||
action="append",
|
||||
default=None,
|
||||
help="Path to a reference image. Repeat to provide multiple references.",
|
||||
)
|
||||
parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.")
|
||||
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
|
||||
parser.add_argument(
|
||||
"--guidance-scale-high", type=float, default=None, help="Optional separate CFG for high-noise (MoE only)."
|
||||
)
|
||||
@@ -83,8 +111,10 @@ def parse_args() -> argparse.Namespace:
|
||||
"--height", type=int, default=None, help="Video height (auto-calculated from image if not set)."
|
||||
)
|
||||
parser.add_argument("--width", type=int, default=None, help="Video width (auto-calculated from image if not set).")
|
||||
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps.")
|
||||
parser.add_argument("--num-frames", type=int, default=None, help="Number of frames. Default: model-specific.")
|
||||
parser.add_argument(
|
||||
"--num-inference-steps", type=int, default=None, help="Sampling steps. Default: model-specific."
|
||||
)
|
||||
parser.add_argument("--boundary-ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.")
|
||||
parser.add_argument(
|
||||
"--frame-rate",
|
||||
@@ -93,7 +123,10 @@ def parse_args() -> argparse.Namespace:
|
||||
help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
|
||||
"--flow-shift",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Scheduler flow_shift. Default: model-specific (Wan 5.0, Cosmos3 10.0).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample-solver",
|
||||
@@ -244,6 +277,19 @@ def parse_args() -> argparse.Namespace:
|
||||
default=None,
|
||||
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
|
||||
)
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=functools.partial(parse_json_object, flag_name="--extra-body"),
|
||||
default=None,
|
||||
help=(
|
||||
"Model-specific generation params as a JSON object. Keys are filtered "
|
||||
"against the model's declared extra_body_params (see vllm_omni/model_extras), "
|
||||
"so unknown keys for the chosen model are silently dropped. "
|
||||
'Cosmos3 V2V example: \'{"condition_frame_indexes_vision": [0, 1], '
|
||||
'"condition_video_keep": "first", "flow_shift": 10.0, '
|
||||
'"max_sequence_length": 4096, "guardrails": false}\'.'
|
||||
),
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -267,31 +313,63 @@ def main():
|
||||
model_name = str(args.model).lower() if args.model is not None else ""
|
||||
model_class_name = args.model_class_name
|
||||
is_ltx2 = "ltx2" in model_name or (model_class_name and "ltx2" in model_class_name.lower())
|
||||
is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower())
|
||||
if model_class_name is None and is_ltx2:
|
||||
model_class_name = "LTX2ImageToVideoPipeline"
|
||||
|
||||
# Load input image
|
||||
image = PIL.Image.open(args.image).convert("RGB")
|
||||
image = PIL.Image.open(args.image).convert("RGB") if args.image else None
|
||||
last_image = PIL.Image.open(args.last_image).convert("RGB") if args.last_image else None
|
||||
mask_image = PIL.Image.open(args.mask_image).convert("L") if args.mask_image else None
|
||||
reference_images = (
|
||||
[PIL.Image.open(path).convert("RGB") for path in args.reference_image] if args.reference_image else None
|
||||
)
|
||||
dimension_image = image or last_image or (reference_images[0] if reference_images else None)
|
||||
if dimension_image is None:
|
||||
raise ValueError("Provide --image, --last-image, or at least one --reference-image.")
|
||||
|
||||
fps = args.fps if args.fps is not None else (24 if is_ltx2 else 16)
|
||||
# Per-model generation defaults, applied only when the matching flag is omitted.
|
||||
# Cosmos3 would otherwise silently inherit the Wan2.2 defaults (wrong size/steps/shift).
|
||||
if is_cosmos:
|
||||
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
|
||||
24,
|
||||
6.0,
|
||||
189,
|
||||
35,
|
||||
10.0,
|
||||
1280 * 720,
|
||||
16,
|
||||
)
|
||||
elif is_ltx2:
|
||||
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 24, 4.0, 121, 40, 5.0, 512 * 768, 32
|
||||
else: # Wan2.2 / HunyuanVideo-1.5
|
||||
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 16, 5.0, 81, 50, 5.0, 480 * 832, 16
|
||||
|
||||
fps = args.fps if args.fps is not None else d_fps
|
||||
frame_rate = args.frame_rate if args.frame_rate is not None else float(fps)
|
||||
guidance_scale = args.guidance_scale if args.guidance_scale is not None else (4.0 if is_ltx2 else 5.0)
|
||||
num_frames = args.num_frames if args.num_frames is not None else (121 if is_ltx2 else 81)
|
||||
num_inference_steps = args.num_inference_steps if args.num_inference_steps is not None else (40 if is_ltx2 else 50)
|
||||
guidance_scale = args.guidance_scale if args.guidance_scale is not None else d_guidance
|
||||
num_frames = args.num_frames if args.num_frames is not None else d_num_frames
|
||||
num_inference_steps = args.num_inference_steps if args.num_inference_steps is not None else d_steps
|
||||
flow_shift = args.flow_shift if args.flow_shift is not None else d_flow_shift
|
||||
|
||||
# Calculate dimensions if not provided
|
||||
# Calculate dimensions if not provided (model-aware max area).
|
||||
height = args.height
|
||||
width = args.width
|
||||
if height is None or width is None:
|
||||
# Default to 480P area for Wan2.2 I2V, 512x768 area for LTX2
|
||||
max_area = 512 * 768 if is_ltx2 else 480 * 832
|
||||
mod_value = 32 if is_ltx2 else 16
|
||||
calc_height, calc_width = calculate_dimensions(image, max_area=max_area, mod_value=mod_value)
|
||||
calc_height, calc_width = calculate_dimensions(dimension_image, max_area=d_max_area, mod_value=d_mod)
|
||||
height = height or calc_height
|
||||
width = width or calc_width
|
||||
|
||||
# Resize image to target dimensions
|
||||
image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
media_inputs: dict[str, Any] = {}
|
||||
if image is not None:
|
||||
media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
if last_image is not None:
|
||||
media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
if mask_image is not None:
|
||||
media_inputs["mask"] = mask_image.resize((width, height), PIL.Image.Resampling.NEAREST)
|
||||
if reference_images is not None:
|
||||
media_inputs["reference_images"] = [
|
||||
reference.resize((width, height), PIL.Image.Resampling.LANCZOS) for reference in reference_images
|
||||
]
|
||||
|
||||
# Configure cache based on backend type
|
||||
cache_config = None
|
||||
@@ -342,7 +420,7 @@ def main():
|
||||
vae_use_slicing=args.vae_use_slicing,
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
boundary_ratio=args.boundary_ratio,
|
||||
flow_shift=args.flow_shift,
|
||||
flow_shift=flow_shift,
|
||||
diffusion_kv_cache_dtype=args.diffusion_kv_cache_dtype,
|
||||
diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps,
|
||||
diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers,
|
||||
@@ -357,7 +435,13 @@ def main():
|
||||
)
|
||||
if args.quantization is not None:
|
||||
omni_kwargs["quantization"] = args.quantization
|
||||
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
|
||||
# gate is an engine-level config (offline analog of the server's --no-guardrails).
|
||||
if args.extra_body and "guardrails" in args.extra_body:
|
||||
omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
|
||||
omni = Omni(**omni_kwargs)
|
||||
model_class_name = get_model_class_name(omni) or model_class_name
|
||||
declared_extra_body_params = get_extra_body_params(model_class_name)
|
||||
|
||||
if profiler_enabled:
|
||||
print("[Profiler] Starting profiling...")
|
||||
@@ -367,8 +451,8 @@ def main():
|
||||
print(f"\n{'=' * 60}")
|
||||
print("Generation Configuration:")
|
||||
print(f" Model: {args.model}")
|
||||
print(f" Inference steps: {args.num_inference_steps}")
|
||||
print(f" Frames: {args.num_frames}")
|
||||
print(f" Inference steps: {num_inference_steps}")
|
||||
print(f" Frames: {num_frames}")
|
||||
print(f" Solver: {args.sample_solver}")
|
||||
print(f" diffusion_kv_cache_dtype(config): {args.diffusion_kv_cache_dtype}")
|
||||
print(f" diffusion_kv_cache_skip_steps(config): {args.diffusion_kv_cache_skip_steps}")
|
||||
@@ -381,30 +465,43 @@ def main():
|
||||
print(f" Video size: {args.width}x{args.height}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
prompt_dict = build_image_to_video_prompt(
|
||||
model_class_name=model_class_name,
|
||||
prompt=args.prompt,
|
||||
negative_prompt=args.negative_prompt,
|
||||
media_inputs=media_inputs,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
)
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=height,
|
||||
width=width,
|
||||
generator=generator,
|
||||
guidance_scale=guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_high,
|
||||
boundary_ratio=args.boundary_ratio,
|
||||
num_inference_steps=num_inference_steps,
|
||||
num_frames=num_frames,
|
||||
frame_rate=frame_rate,
|
||||
extra_args={
|
||||
"sample_solver": args.sample_solver,
|
||||
"flow_shift": flow_shift,
|
||||
},
|
||||
)
|
||||
|
||||
# Route model-specific knobs through extra_body, filtered against the model's
|
||||
# declared extra_body_params. Models without a declaration only forward explicit
|
||||
# --extra-body JSON (e.g. Cosmos3 V2V's condition_frame_indexes_vision).
|
||||
extra_body = dict(args.extra_body or {})
|
||||
if declared_extra_body_params:
|
||||
apply_declared_extra_args(sampling_params, declared_extra_body_params, extra_body)
|
||||
elif extra_body:
|
||||
sampling_params.extra_args.update({k: v for k, v in extra_body.items() if v is not None})
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
# omni.generate() returns Generator[OmniRequestOutput, None, None]
|
||||
frames = omni.generate(
|
||||
{
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
"multi_modal_data": {"image": image},
|
||||
},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=height,
|
||||
width=width,
|
||||
generator=generator,
|
||||
guidance_scale=guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_high,
|
||||
boundary_ratio=args.boundary_ratio,
|
||||
num_inference_steps=num_inference_steps,
|
||||
num_frames=num_frames,
|
||||
frame_rate=frame_rate,
|
||||
extra_args={
|
||||
"sample_solver": args.sample_solver,
|
||||
"flow_shift": args.flow_shift,
|
||||
},
|
||||
),
|
||||
)
|
||||
frames = omni.generate(prompt_dict, sampling_params)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
import functools
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
@@ -38,16 +39,20 @@ def is_nextstep_model(model_name: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def parse_profiler_config(value: str) -> dict[str, Any]:
|
||||
def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]:
|
||||
"""Parse a CLI value as a JSON object, attributing errors to ``flag_name``."""
|
||||
try:
|
||||
config = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e
|
||||
if not isinstance(config, dict):
|
||||
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object")
|
||||
return config
|
||||
|
||||
|
||||
parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
|
||||
parser.add_argument(
|
||||
@@ -269,7 +274,7 @@ def parse_args() -> argparse.Namespace:
|
||||
)
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=parse_profiler_config,
|
||||
type=functools.partial(parse_json_object, flag_name="--extra-body"),
|
||||
default=None,
|
||||
help=(
|
||||
"Model-specific generation params as a JSON object, e.g. "
|
||||
@@ -425,6 +430,10 @@ def main():
|
||||
if use_nextstep:
|
||||
# NextStep-1.1 requires explicit pipeline class
|
||||
omni_kwargs["model_class_name"] = "NextStep11Pipeline"
|
||||
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
|
||||
# gate is an engine-level config (offline analog of the server's --no-guardrails).
|
||||
if args.extra_body and "guardrails" in args.extra_body:
|
||||
omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
|
||||
omni = Omni(**omni_kwargs)
|
||||
model_class_name = get_model_class_name(omni)
|
||||
declared_extra_body_params = get_extra_body_params(model_class_name)
|
||||
|
||||
@@ -23,6 +23,7 @@ list of supported architectures across all modalities, see
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
|
||||
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
|
||||
|
||||
## Common Quick Start
|
||||
@@ -436,6 +437,50 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [
|
||||
|
||||
---
|
||||
|
||||
## IndexTTS-2
|
||||
|
||||
2-stage TTS pipeline (GPT AR talker + S2Mel CFM DiT + BigVGAN vocoder) at 22.05 kHz. Every request requires reference audio for zero-shot voice cloning. Supports emotion conditioning via audio, text, or 8-dim vector.
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model IndexTeam/IndexTTS-2 \
|
||||
--text "你好,这是一个语音合成测试。" \
|
||||
--ref-audio /path/to/reference.wav
|
||||
```
|
||||
|
||||
### Emotion control
|
||||
```bash
|
||||
# Emotion from reference audio
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model IndexTeam/IndexTTS-2 \
|
||||
--text "今天天气真好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-audio /path/to/happy.wav
|
||||
|
||||
# Emotion from 8-dim vector (happy angry sad afraid disgusted melancholy surprised calm)
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model IndexTeam/IndexTTS-2 \
|
||||
--text "今天天气真好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-vector 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
|
||||
|
||||
# Emotion from text description
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model IndexTeam/IndexTTS-2 \
|
||||
--text "今天天气真好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-text "happy and excited"
|
||||
```
|
||||
|
||||
### Notes
|
||||
- `--ref-audio` is **required** — IndexTTS-2 does not support text-only synthesis.
|
||||
- Stage 0 (AR Talker): GPT-2 generates mel codes from text + reference audio.
|
||||
- Stage 1 (S2Mel + BigVGAN): CFM DiT converts mel codes to waveform at 22.05 kHz.
|
||||
- Deploy config: `vllm_omni/deploy/indextts2.yaml`. Stage 1 runs with `enforce_eager: true` (DiT has dynamic shapes).
|
||||
|
||||
---
|
||||
|
||||
## Voxtral TTS
|
||||
|
||||
Voxtral-4B-TTS (Mistral). Has its own argparse surface; uses voice presets and the `mistral_common` `SpeechRequest` protocol.
|
||||
@@ -476,3 +521,100 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
|
||||
- `--num-prompts N` replicates the prompt for performance measurement.
|
||||
- `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
|
||||
- Run `--help` for the full argument surface.
|
||||
|
||||
---
|
||||
|
||||
## SoulX-Singer
|
||||
|
||||
Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Script: `soulxsinger/end2end.py`. Deploy: `vllm_omni/deploy/soulxsinger_svs.yaml` or `soulxsinger_svc.yaml`.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Download DiT and preprocess weights, then set up separate SVS / SVC view directories. Copy `soulxsinger/utils/phoneme/phone_set.json` from upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) into the model weights dir as `phoneme/phone_set.json` — HuggingFace does not ship it.
|
||||
|
||||
```bash
|
||||
# 1. DiT weights
|
||||
export BASE=path/to/SoulX-Singer
|
||||
export PREPROCESS=path/to/SoulX-Singer-Preprocess
|
||||
export SVC_DIR=path/to/SoulX-Singer-svc
|
||||
|
||||
huggingface-cli download Soul-AILab/SoulX-Singer --local-dir "$BASE"
|
||||
|
||||
# 2. Preprocess weights (required)
|
||||
huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess --local-dir "$PREPROCESS"
|
||||
export SOULX_PREPROCESS_WEIGHTS_DIR="$PREPROCESS"
|
||||
|
||||
# 3. SVS / SVC view directories
|
||||
mkdir -p "$SVC_DIR"
|
||||
cp $BASE/{config.yaml,README.md,assets} $SVC_DIR
|
||||
mv $BASE/model-svc.pt $SVC_DIR/model-svc.pt
|
||||
|
||||
cat > "$BASE/config.json" <<'EOF'
|
||||
{
|
||||
"model_type": "soulxsinger",
|
||||
"architectures": ["SoulXSingerPipeline"],
|
||||
"max_num_seqs": 1
|
||||
}
|
||||
EOF
|
||||
|
||||
cat > "$SVC_DIR/config.json" <<'EOF'
|
||||
{
|
||||
"model_type": "soulxsinger",
|
||||
"architectures": ["SoulXSingerSVCPipeline"],
|
||||
"max_num_seqs": 1
|
||||
}
|
||||
EOF
|
||||
```
|
||||
|
||||
`config.yaml` hyper-parameters live under `$BASE`; each view's `config.json` `architectures` field is the single source of truth for SVS vs SVC. Point `--model` at the matching directory (`$BASE` for SVS, `$SVC_DIR` for SVC). Deploy YAML is chosen automatically from `config.json`; optional `--svs` / `--svc` only assert the mode matches.
|
||||
|
||||
**Online preprocess** is the default: pass `--prompt-audio` and `--target-audio`, and the worker runs vocal separation, F0, and (for SVS) lyrics/MIDI before DiT. Install only what your run needs:
|
||||
|
||||
```bash
|
||||
pip install "BS-RoFormer" # vocal sep + F0 on GPU — SVS and SVC
|
||||
```
|
||||
|
||||
Mandarin SVS also needs FunASR and Chinese G2P; `ffmpeg` must be on `PATH`:
|
||||
|
||||
```bash
|
||||
# install optional dependencies:
|
||||
pip install -e ".[soulx-svs]"
|
||||
```
|
||||
|
||||
English SVS adds NeMo ASR and NLTK data; pass `--language English`:
|
||||
|
||||
```bash
|
||||
pip install "nemo_toolkit[asr]==2.6.1" lhotse==1.32.2
|
||||
python -c "import nltk; nltk.download('cmudict'); nltk.download('averaged_perceptron_tagger_eng')"
|
||||
```
|
||||
|
||||
**Precomputed metadata** is the alternative: pass both `--prompt-metadata-path` and `--target-metadata-path` and skip online ASR/ROSVOT — none of the packages above are required. JSON can be produced by integrated preprocess on a prior run, or by upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) `preprocess/` scripts if you prefer to run that outside vLLM-Omni.
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
# SVS — default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3
|
||||
python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \
|
||||
--model "$BASE" \
|
||||
--preprocess-weights-dir "$PREPROCESS" \
|
||||
--control score \
|
||||
--num-inference-steps 32 \
|
||||
-o output.wav
|
||||
|
||||
python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \
|
||||
--model "$SVC_DIR" \
|
||||
--preprocess-weights-dir "$PREPROCESS" \
|
||||
--svc \
|
||||
--num-inference-steps 32 \
|
||||
-o output_svc.wav
|
||||
```
|
||||
|
||||
`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. Long SVS targets are handled in one request. See `end2end.py --help` for `--pitch-shift`, `--vocal-sep`, `--auto-shift`, and language/control options.
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 24 kHz mono WAV; batch only.
|
||||
- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on.
|
||||
- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only.
|
||||
|
||||
---
|
||||
@@ -0,0 +1,187 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline inference example for IndexTTS2 via vLLM-Omni.
|
||||
|
||||
Two-stage pipeline: GPT AR (Stage 0) → S2Mel + BigVGAN (Stage 1).
|
||||
Output is 22050 Hz mono WAV.
|
||||
|
||||
Usage:
|
||||
python end2end.py \
|
||||
--model /path/to/IndexTeam/IndexTTS-2 \
|
||||
--text "你好,这是一个语音合成测试。" \
|
||||
--ref-audio /path/to/ref.wav
|
||||
|
||||
# With emotion audio:
|
||||
python end2end.py \
|
||||
--model /path/to/IndexTeam/IndexTTS-2 \
|
||||
--text "今天天气真好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-audio /path/to/happy.wav
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from vllm import SamplingParams
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
|
||||
from vllm_omni import Omni # noqa: E402
|
||||
from vllm_omni.model_executor.models.indextts2.prompt_utils import ( # noqa: E402
|
||||
build_indextts2_prefill_prompt_ids,
|
||||
)
|
||||
|
||||
|
||||
def build_request(
|
||||
model: str,
|
||||
text: str,
|
||||
ref_audio_path: str | None = None,
|
||||
emo_audio_path: str | None = None,
|
||||
emo_text: str | None = None,
|
||||
emo_vector: list[float] | None = None,
|
||||
emo_alpha: float | None = None,
|
||||
use_emo_text: bool = False,
|
||||
use_random: bool = False,
|
||||
) -> dict:
|
||||
additional: dict = {"text": [text]}
|
||||
if ref_audio_path:
|
||||
additional["voice"] = [str(ref_audio_path)]
|
||||
if emo_audio_path:
|
||||
additional["emo_audio"] = [str(emo_audio_path)]
|
||||
if emo_text:
|
||||
additional["emo_text"] = [emo_text]
|
||||
if emo_vector is not None:
|
||||
additional["emo_vector"] = [emo_vector]
|
||||
if emo_alpha is not None:
|
||||
additional["emo_alpha"] = [emo_alpha]
|
||||
if use_emo_text:
|
||||
additional["use_emo_text"] = [True]
|
||||
if use_random:
|
||||
additional["use_random"] = [True]
|
||||
return {
|
||||
"prompt_token_ids": build_indextts2_prefill_prompt_ids(model, text),
|
||||
"additional_information": additional,
|
||||
}
|
||||
|
||||
|
||||
def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = 22050) -> None:
|
||||
audio_np = waveform.float().numpy()
|
||||
sf.write(path, audio_np, sample_rate)
|
||||
print(f" Saved {path} ({audio_np.shape}, {sample_rate} Hz)")
|
||||
|
||||
|
||||
def extract_audio(mm: dict) -> tuple[torch.Tensor | None, int]:
|
||||
audio = mm.get("audio")
|
||||
if audio is None:
|
||||
audio = mm.get("model_outputs")
|
||||
if isinstance(audio, list):
|
||||
chunks = [chunk.reshape(-1) for chunk in audio if isinstance(chunk, torch.Tensor) and chunk.numel() > 0]
|
||||
audio = torch.cat(chunks, dim=0) if chunks else None
|
||||
|
||||
sr_val = mm.get("sr")
|
||||
if isinstance(sr_val, list):
|
||||
sr_val = sr_val[-1] if sr_val else None
|
||||
if hasattr(sr_val, "item"):
|
||||
sample_rate = int(sr_val.item())
|
||||
else:
|
||||
sample_rate = int(sr_val) if sr_val is not None else 22050
|
||||
|
||||
return audio if isinstance(audio, torch.Tensor) else None, sample_rate
|
||||
|
||||
|
||||
def main(args) -> None:
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
deploy_config=args.deploy_config,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
# Stage 0: GPT AR (text → mel codes)
|
||||
gpt_sampling = SamplingParams(
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
top_k=30,
|
||||
max_tokens=1500,
|
||||
repetition_penalty=10.0,
|
||||
stop_token_ids=[8193],
|
||||
seed=args.seed if args.seed is not None else 42,
|
||||
detokenize=False,
|
||||
)
|
||||
# Stage 1: S2Mel + BigVGAN (non-AR, params mostly ignored)
|
||||
s2mel_sampling = SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=65536,
|
||||
detokenize=True,
|
||||
)
|
||||
sampling_params = [gpt_sampling, s2mel_sampling]
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"Synthesizing: {args.text!r}")
|
||||
if args.ref_audio:
|
||||
print(f" ref_audio: {args.ref_audio}")
|
||||
inputs = build_request(
|
||||
model=args.model,
|
||||
text=args.text,
|
||||
ref_audio_path=args.ref_audio,
|
||||
emo_audio_path=args.emo_audio,
|
||||
emo_text=args.emo_text,
|
||||
emo_vector=args.emo_vector,
|
||||
emo_alpha=args.emo_alpha,
|
||||
use_emo_text=args.use_emo_text,
|
||||
use_random=args.use_random,
|
||||
)
|
||||
|
||||
for i, omni_out in enumerate(omni.generate(inputs, sampling_params_list=sampling_params)):
|
||||
mm = omni_out.multimodal_output
|
||||
if not mm:
|
||||
print(f" [req {i}] No multimodal output.")
|
||||
continue
|
||||
audio, sr = extract_audio(mm)
|
||||
if audio is None:
|
||||
print(f" [req {i}] No waveform in multimodal_output.")
|
||||
continue
|
||||
out_path = str(output_dir / f"output_{i}.wav")
|
||||
save_audio(audio.cpu(), out_path, sr)
|
||||
|
||||
print("Done.")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="IndexTTS2 offline inference")
|
||||
parser.add_argument("--model", required=True, help="HF model path for IndexTTS2.")
|
||||
parser.add_argument("--text", default="你好,这是IndexTTS2语音合成测试。")
|
||||
parser.add_argument("--ref-audio", required=True, help="Reference audio for voice cloning.")
|
||||
parser.add_argument("--emo-audio", default=None, help="Emotion reference audio.")
|
||||
parser.add_argument("--emo-text", default=None, help="Emotion description text.")
|
||||
parser.add_argument(
|
||||
"--emo-vector",
|
||||
type=float,
|
||||
nargs=8,
|
||||
default=None,
|
||||
help="8-dim emotion vector: happy angry sad afraid disgusted melancholic surprised calm.",
|
||||
)
|
||||
parser.add_argument("--emo-alpha", type=float, default=None, help="Emotion weight in [0, 1].")
|
||||
parser.add_argument("--use-emo-text", action="store_true", help="Infer emotion vector from emo-text or text.")
|
||||
parser.add_argument("--use-random", action="store_true", help="Use random emotion prototypes.")
|
||||
parser.add_argument("--seed", type=int, default=None)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default=os.path.join(
|
||||
os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")),
|
||||
"indextts2_output",
|
||||
),
|
||||
)
|
||||
parser.add_argument("--deploy-config", default=None)
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=600)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(parse_args())
|
||||
@@ -0,0 +1,162 @@
|
||||
"""Repeatable SoulX-Singer offline benchmark (preprocess + DiT, RTF + stage timings).
|
||||
|
||||
Usage:
|
||||
python benchmark.py --model /path/to/SoulX-Singer --svs \\
|
||||
--prompt-audio ../SoulX-Singer/example/audio/zh_prompt.mp3 \\
|
||||
--target-audio ../SoulX-Singer/example/audio/music.mp3 \\
|
||||
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
|
||||
-o benchmark.wav
|
||||
|
||||
python benchmark.py ... --enable-diffusion-pipeline-profiler
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import statistics
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
from end2end import (
|
||||
SVC_DEPLOY_CONFIG,
|
||||
SVS_DEPLOY_CONFIG,
|
||||
add_inference_args,
|
||||
build_sampling,
|
||||
extract_audio,
|
||||
resolve_preprocess_weights_dir,
|
||||
)
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
|
||||
def _audio_duration_sec(outputs) -> float:
|
||||
audio_np, sr = extract_audio(outputs)
|
||||
return audio_np.size / sr
|
||||
|
||||
|
||||
def _stage_durations_from_output(outputs) -> dict[str, float]:
|
||||
omni_out = outputs[0]
|
||||
durations = getattr(omni_out, "stage_durations", None) or {}
|
||||
if durations:
|
||||
return dict(durations)
|
||||
ro = getattr(omni_out, "request_output", None)
|
||||
if ro is not None:
|
||||
return dict(getattr(ro, "stage_durations", None) or {})
|
||||
return {}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="SoulX-Singer offline benchmark")
|
||||
add_inference_args(parser)
|
||||
parser.add_argument("--warmup", type=int, default=1, help="Warmup runs (excluded from stats)")
|
||||
parser.add_argument("--runs", type=int, default=3, help="Measured runs")
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable low-overhead stage timing in DiffusionOutput.stage_durations",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enforce-eager",
|
||||
action="store_true",
|
||||
help="Disable torch.compile on diff_estimator (eager baseline for A/B)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["bfloat16", "bf16", "float16", "fp16", "half"],
|
||||
help="DiT trunk dtype (default from deploy YAML). Use float16 to match upstream --fp16.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
"-o",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Save WAV from the last measured run (after warmup)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save-each-run",
|
||||
action="store_true",
|
||||
help="With --output, also write measured runs as <stem>_run1.wav, _run2.wav, ...",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir)
|
||||
deploy_config = args.deploy_config or str(SVS_DEPLOY_CONFIG if args.svs else SVC_DEPLOY_CONFIG)
|
||||
mode = "SVS" if args.svs else "SVC"
|
||||
|
||||
omni_kwargs: dict = {
|
||||
"model": args.model,
|
||||
"deploy_config": deploy_config,
|
||||
"async_chunk": False, # SoulX-Singer currently supports only batch mode (pseudo-streaming was stashed)
|
||||
}
|
||||
if args.enable_diffusion_pipeline_profiler:
|
||||
omni_kwargs["enable_diffusion_pipeline_profiler"] = True
|
||||
if args.enforce_eager:
|
||||
omni_kwargs["enforce_eager"] = True
|
||||
if args.dtype is not None:
|
||||
omni_kwargs["dtype"] = args.dtype
|
||||
|
||||
compile_mode = "eager" if args.enforce_eager else "torch.compile"
|
||||
dtype_label = args.dtype or "from deploy YAML"
|
||||
print(f"Loading SoulX-Singer {mode} from {args.model} [{compile_mode}, dtype={dtype_label}]")
|
||||
|
||||
omni = Omni(**omni_kwargs)
|
||||
kind = "svs" if args.svs else "svc"
|
||||
sampling = build_sampling(args, preprocess_weights_dir=preprocess_dir, kind=kind)
|
||||
prompt = {"prompt_token_ids": [0]}
|
||||
|
||||
latencies_ms: list[float] = []
|
||||
rtfs: list[float] = []
|
||||
last_stages: dict[str, float] = {}
|
||||
last_measured_outputs = None
|
||||
|
||||
total_runs = args.warmup + args.runs
|
||||
for run_idx in range(total_runs):
|
||||
is_warmup = run_idx < args.warmup
|
||||
label = "warmup" if is_warmup else f"run {run_idx - args.warmup + 1}/{args.runs}"
|
||||
t0 = time.perf_counter()
|
||||
outputs = list(omni.generate([prompt], sampling))
|
||||
elapsed_ms = (time.perf_counter() - t0) * 1000.0
|
||||
audio_sec = _audio_duration_sec(outputs)
|
||||
rtf = (elapsed_ms / 1000.0) / audio_sec if audio_sec > 0 else float("inf")
|
||||
print(f"[{label}] client={elapsed_ms:.1f} ms, audio={audio_sec:.2f}s, RTF={rtf:.3f}")
|
||||
if not is_warmup:
|
||||
latencies_ms.append(elapsed_ms)
|
||||
rtfs.append(rtf)
|
||||
last_stages = _stage_durations_from_output(outputs)
|
||||
last_measured_outputs = outputs
|
||||
if args.output and args.save_each_run:
|
||||
measured_idx = run_idx - args.warmup + 1
|
||||
out_path = Path(args.output)
|
||||
run_path = out_path.with_name(f"{out_path.stem}_run{measured_idx}{out_path.suffix}")
|
||||
audio_np, sr = extract_audio(outputs)
|
||||
sf.write(str(run_path), audio_np, sr)
|
||||
print(f" saved {run_path} ({sr} Hz, {audio_np.size / sr:.2f}s)")
|
||||
|
||||
omni.close()
|
||||
|
||||
if args.output and last_measured_outputs is not None:
|
||||
audio_np, sr = extract_audio(last_measured_outputs)
|
||||
out_path = Path(args.output)
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
sf.write(str(out_path), audio_np, sr)
|
||||
print(f"\nSaved last measured run → {out_path} ({sr} Hz, {audio_np.size / sr:.2f}s)")
|
||||
|
||||
if latencies_ms:
|
||||
print("\n=== Summary (measured runs) ===")
|
||||
print(f"client_ms: mean={statistics.mean(latencies_ms):.1f}, stdev={statistics.pstdev(latencies_ms):.1f}")
|
||||
print(f"RTF: mean={statistics.mean(rtfs):.3f}, stdev={statistics.pstdev(rtfs):.3f}")
|
||||
if last_stages:
|
||||
print("\nStage durations (last measured run):")
|
||||
for name, value in sorted(last_stages.items()):
|
||||
if name.endswith("_ms"):
|
||||
print(f" {name}: {value:.1f} ms")
|
||||
else:
|
||||
print(f" {name}: {value * 1000.0:.1f} ms")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,285 @@
|
||||
"""Offline SoulX-Singer SVS / SVC: single-stage DiT (preprocess inline).
|
||||
|
||||
Pass raw prompt/target audio; vLLM-Omni runs integrated preprocess then inference.
|
||||
|
||||
Usage:
|
||||
# SVS (default)
|
||||
python end2end.py --model /path/to/SoulX-Singer \\
|
||||
--prompt-audio path/to/prompt/audio.mp3 \\
|
||||
--target-audio path/to/target/audio.mp3 \\
|
||||
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
|
||||
-o output.wav
|
||||
|
||||
# SVC (config.json must declare SoulXSingerSVCPipeline; --model points at SVC view)
|
||||
python end2end.py --model /path/to/SoulX-Singer-svc \\
|
||||
--prompt-audio path/to/prompt/audio.mp3 \\
|
||||
--target-audio path/to/target/audio.mp3 \\
|
||||
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
|
||||
-o output_svc.wav
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
from vllm_omni.diffusion.models.soulx_singer.utils import resolve_soulx_kind, validate_soulx_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
SVS_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svs.yaml"
|
||||
SVC_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svc.yaml"
|
||||
DEFAULT_PROMPT_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "zh_prompt.mp3"
|
||||
DEFAULT_TARGET_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "music.mp3"
|
||||
_SAMPLE_RATE = 24000
|
||||
|
||||
|
||||
def _require_paths(paths: dict[str, str | None]) -> None:
|
||||
missing = [name for name, p in paths.items() if not p or not os.path.isfile(p)]
|
||||
if missing:
|
||||
raise FileNotFoundError("Missing file(s): " + ", ".join(f"{k}={paths[k]!r}" for k in missing))
|
||||
|
||||
|
||||
def resolve_preprocess_weights_dir(explicit: str | None) -> str:
|
||||
if explicit:
|
||||
path = Path(explicit).expanduser().resolve()
|
||||
if not (path / "rmvpe" / "rmvpe.pt").is_file():
|
||||
raise FileNotFoundError(f"--preprocess-weights-dir must contain rmvpe/rmvpe.pt: {path}")
|
||||
return str(path)
|
||||
env = os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR")
|
||||
if env:
|
||||
path = Path(env).expanduser().resolve()
|
||||
if (path / "rmvpe" / "rmvpe.pt").is_file():
|
||||
return str(path)
|
||||
raise FileNotFoundError(
|
||||
"SoulX preprocess weights not found. Pass --preprocess-weights-dir or set "
|
||||
"SOULX_PREPROCESS_WEIGHTS_DIR to a directory containing rmvpe/rmvpe.pt "
|
||||
"(huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess)."
|
||||
)
|
||||
|
||||
|
||||
def build_sampling(
|
||||
args: argparse.Namespace,
|
||||
*,
|
||||
preprocess_weights_dir: str,
|
||||
kind: str,
|
||||
) -> OmniDiffusionSamplingParams:
|
||||
extra: dict = {
|
||||
"prompt_audio": os.path.abspath(args.prompt_audio),
|
||||
"target_audio": os.path.abspath(args.target_audio),
|
||||
"preprocess_weights_dir": preprocess_weights_dir,
|
||||
"vocal_sep": args.vocal_sep,
|
||||
"auto_shift": args.auto_shift,
|
||||
"pitch_shift": args.pitch_shift,
|
||||
}
|
||||
if kind == "svs":
|
||||
extra["language"] = args.language
|
||||
extra["control"] = args.control
|
||||
if args.prompt_metadata_path:
|
||||
extra["prompt_metadata_path"] = os.path.abspath(args.prompt_metadata_path)
|
||||
if args.target_metadata_path:
|
||||
extra["target_metadata_path"] = os.path.abspath(args.target_metadata_path)
|
||||
if args.prompt_metadata_path or args.target_metadata_path:
|
||||
extra["audio_path"] = os.path.abspath(args.prompt_wav_path or args.prompt_audio)
|
||||
|
||||
extra = validate_soulx_extra_args(kind, extra)
|
||||
|
||||
return OmniDiffusionSamplingParams(
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
seed=args.seed,
|
||||
extra_args=extra,
|
||||
)
|
||||
|
||||
|
||||
def extract_audio(outputs) -> tuple[np.ndarray, int]:
|
||||
out = outputs[0]
|
||||
mm = getattr(out, "multimodal_output", None)
|
||||
if mm is None:
|
||||
ro = getattr(out, "request_output", None)
|
||||
if ro is not None:
|
||||
mm = getattr(ro, "multimodal_output", None)
|
||||
if mm is None and ro.outputs:
|
||||
mm = getattr(ro.outputs[0], "multimodal_output", None)
|
||||
if not mm or "audio" not in mm:
|
||||
raise RuntimeError("No audio in multimodal_output")
|
||||
|
||||
audio_val = mm["audio"]
|
||||
if isinstance(audio_val, list):
|
||||
chunks = [
|
||||
np.asarray(chunk.detach().cpu().float().numpy() if hasattr(chunk, "detach") else chunk).reshape(-1)
|
||||
for chunk in audio_val
|
||||
if chunk is not None
|
||||
]
|
||||
audio_np = np.concatenate(chunks, axis=0) if chunks else np.array([], dtype=np.float32)
|
||||
elif hasattr(audio_val, "cpu"):
|
||||
audio_np = audio_val.detach().cpu().numpy().squeeze()
|
||||
else:
|
||||
audio_np = np.asarray(audio_val, dtype=np.float32).squeeze()
|
||||
|
||||
sr_val = mm.get("sr", mm.get("audio_sample_rate", _SAMPLE_RATE))
|
||||
if isinstance(sr_val, list) and sr_val:
|
||||
sr_val = sr_val[-1]
|
||||
if hasattr(sr_val, "item"):
|
||||
sr_val = sr_val.item()
|
||||
return audio_np.astype(np.float32, copy=False), int(sr_val)
|
||||
|
||||
|
||||
def add_inference_args(parser: argparse.ArgumentParser) -> None:
|
||||
mode = parser.add_mutually_exclusive_group()
|
||||
mode.add_argument(
|
||||
"--svc",
|
||||
action="store_true",
|
||||
help="Assert SVC mode (must match config.json architectures)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
required=True,
|
||||
help="SoulX-Singer weight directory ($BASE for SVS, $SVC_DIR for SVC)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override deploy YAML (default: soulxsinger_svs.yaml or soulxsinger_svc.yaml)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-audio",
|
||||
type=str,
|
||||
default=str(DEFAULT_PROMPT_AUDIO),
|
||||
help="Prompt / reference audio (mp3 or wav)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-audio",
|
||||
type=str,
|
||||
default=str(DEFAULT_TARGET_AUDIO),
|
||||
help="Target content audio (mp3 or wav)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-metadata-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="SVS precomputed prompt metadata.json (skips online ASR/ROSVOT for prompt)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-metadata-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="SVS precomputed target metadata.json (skips online ASR/ROSVOT for target)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-wav-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="SVS prompt vocal wav for metadata processor when using --prompt-metadata-path (default: --prompt-audio)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preprocess-weights-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="SoulX-Singer-Preprocess weights (or set SOULX_PREPROCESS_WEIGHTS_DIR)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--language",
|
||||
type=str,
|
||||
default="Mandarin",
|
||||
help="SVS lyrics language (Mandarin, Cantonese, English)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--control",
|
||||
type=str,
|
||||
default="melody",
|
||||
choices=["score", "melody"],
|
||||
help="SVS control mode",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vocal-sep",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=False,
|
||||
help="Run vocal separation during preprocess",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auto-shift",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Auto pitch shift to align target pitch with prompt",
|
||||
)
|
||||
parser.add_argument("--pitch-shift", type=int, default=0)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=32)
|
||||
parser.add_argument("--guidance-scale", type=float, default=3.0)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=42,
|
||||
help="Optional CFM RNG seed. Omit for non-deterministic sampling.",
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="SoulX-Singer offline SVS / SVC (preprocess + DiT)")
|
||||
add_inference_args(parser)
|
||||
parser.add_argument("--output", "-o", type=str, default="output.wav")
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
kind = resolve_soulx_kind(args.model)
|
||||
|
||||
if not (args.svc ^ (kind == "svs")):
|
||||
raise ValueError(f"mode mismatch: {args.model}/config.json (declares {kind!r}), --svc={args.svc}")
|
||||
|
||||
required = {
|
||||
"prompt_audio": args.prompt_audio,
|
||||
"target_audio": args.target_audio,
|
||||
}
|
||||
if kind == "svs" and bool(args.prompt_metadata_path) ^ bool(args.target_metadata_path):
|
||||
raise ValueError(
|
||||
"SVS precomputed metadata requires both --prompt-metadata-path and "
|
||||
"--target-metadata-path (or neither, for online preprocess)."
|
||||
)
|
||||
if args.prompt_metadata_path:
|
||||
required["prompt_metadata_path"] = args.prompt_metadata_path
|
||||
required["target_metadata_path"] = args.target_metadata_path
|
||||
if args.prompt_wav_path:
|
||||
required["prompt_wav_path"] = args.prompt_wav_path
|
||||
_require_paths(required)
|
||||
preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir)
|
||||
deploy_config = args.deploy_config or str(SVC_DEPLOY_CONFIG if kind == "svc" else SVS_DEPLOY_CONFIG)
|
||||
mode = kind.upper()
|
||||
|
||||
print(f"Loading SoulX-Singer {mode} from {args.model}")
|
||||
print(f" deploy: {deploy_config}")
|
||||
print(f" prompt: {args.prompt_audio}")
|
||||
print(f" target: {args.target_audio}")
|
||||
|
||||
omni = Omni(model=args.model, deploy_config=deploy_config, async_chunk=False) # SoulX-Singer is batch-only
|
||||
sampling = build_sampling(
|
||||
args,
|
||||
preprocess_weights_dir=preprocess_dir,
|
||||
kind=kind,
|
||||
)
|
||||
prompt = {"prompt_token_ids": [0], "multi_modal_data": {}}
|
||||
|
||||
print("Running generate (inline preprocess + DiT)...")
|
||||
outputs = list(omni.generate([prompt], sampling))
|
||||
if not outputs:
|
||||
raise RuntimeError("No output from omni.generate")
|
||||
|
||||
err = getattr(outputs[0], "error", None)
|
||||
if err is not None:
|
||||
raise RuntimeError(f"Pipeline error: {err}")
|
||||
|
||||
audio_np, sr = extract_audio(outputs)
|
||||
sf.write(args.output, audio_np, sr)
|
||||
print(f"Saved {args.output} ({sr} Hz, {len(audio_np) / sr:.2f}s)")
|
||||
omni.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -6,9 +6,11 @@ A unified script for text-to-video generation. Supports multiple models with mod
|
||||
|
||||
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|
||||
|---|---|---|---|---|---|
|
||||
| `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) |
|
||||
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
|
||||
| `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) |
|
||||
| `BestWishYsh/Helios-Base` / `Helios-Mid` / `Helios-Distilled` | 384x640 | 99 | 50 | 5.0 / 5.0 / 1.0 | — |
|
||||
|
||||
## Local CLI Usage
|
||||
@@ -30,6 +32,28 @@ python text_to_video.py \
|
||||
--output t2v_out.mp4
|
||||
```
|
||||
|
||||
### Wan2.1 VACE (T2V)
|
||||
|
||||
VACE text-to-video uses this shared entrypoint. Conditional VACE tasks use
|
||||
the shared [`image_to_video.py`](../image_to_video/README.md#wan21-vace-conditional-tasks)
|
||||
entrypoint, which constructs the pipeline-native conditioning data from the
|
||||
provided media inputs. No explicit mode parameter is required.
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--prompt "A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves." \
|
||||
--seed 0 \
|
||||
--height 480 \
|
||||
--width 832 \
|
||||
--num-frames 81 \
|
||||
--num-inference-steps 30 \
|
||||
--guidance-scale 5.0 \
|
||||
--flow-shift 5.0 \
|
||||
--vae-use-tiling \
|
||||
--output vace_t2v_output.mp4
|
||||
```
|
||||
|
||||
LTX2 example:
|
||||
|
||||
```bash
|
||||
@@ -101,6 +125,20 @@ python text_to_video.py \
|
||||
--output quick_test.mp4
|
||||
```
|
||||
|
||||
### Cosmos3
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model nvidia/Cosmos3-Nano \
|
||||
--prompt "A robot arm is cleaning a plate in the kitchen." \
|
||||
--negative-prompt "blurry, distorted, low quality, jittery, deformed" \
|
||||
--height 720 --width 1280 --num-frames 189 --fps 24 \
|
||||
--num-inference-steps 35 --guidance-scale 6.0 \
|
||||
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false,
|
||||
"use_resolution_template": false, "use_duration_template": false}' \
|
||||
--output cosmos3_t2v.mp4
|
||||
```
|
||||
|
||||
### Helios (T2V)
|
||||
|
||||
Helios ships three variants. Model-specific knobs (declared in
|
||||
@@ -167,6 +205,7 @@ python text_to_video.py \
|
||||
- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio).
|
||||
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
|
||||
- `--flow-shift`: scheduler flow_shift parameter.
|
||||
- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).
|
||||
|
||||
### Wan2.2-specific
|
||||
|
||||
|
||||
@@ -11,12 +11,23 @@ import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
_MODEL_PRESETS = {
|
||||
"vace": {
|
||||
"height": 480,
|
||||
"width": 832,
|
||||
"num_frames": 81,
|
||||
"num_inference_steps": 30,
|
||||
"guidance_scale": 5.0,
|
||||
"fps": 16,
|
||||
"output": "vace_t2v_output.mp4",
|
||||
},
|
||||
"wan": {
|
||||
"height": 720,
|
||||
"width": 1280,
|
||||
@@ -35,6 +46,16 @@ _MODEL_PRESETS = {
|
||||
"fps": 24,
|
||||
"output": "hunyuan_video_15_output.mp4",
|
||||
},
|
||||
"cosmos": {
|
||||
"height": 720,
|
||||
"width": 1280,
|
||||
"num_frames": 189,
|
||||
"num_inference_steps": 35,
|
||||
"guidance_scale": 6.0,
|
||||
"fps": 24,
|
||||
"flow_shift": 10.0,
|
||||
"output": "cosmos3_t2v_output.mp4",
|
||||
},
|
||||
"helios": {
|
||||
"height": 384,
|
||||
"width": 640,
|
||||
@@ -49,6 +70,10 @@ _MODEL_PRESETS = {
|
||||
|
||||
def _detect_preset(model: str) -> dict:
|
||||
model_lower = model.lower()
|
||||
if "vace" in model_lower:
|
||||
return _MODEL_PRESETS["vace"]
|
||||
if "cosmos" in model_lower:
|
||||
return _MODEL_PRESETS["cosmos"]
|
||||
if "hunyuan" in model_lower:
|
||||
return _MODEL_PRESETS["hunyuan"]
|
||||
if "helios" in model_lower:
|
||||
@@ -106,8 +131,12 @@ def parse_args() -> argparse.Namespace:
|
||||
type=parse_extra_body,
|
||||
default=None,
|
||||
help="JSON dict of model-specific extra_body params (declared in vllm_omni/model_extras/), "
|
||||
"merged into sampling extra_args. Example (Helios-Distilled): "
|
||||
'\'{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}\'.',
|
||||
"merged into sampling extra_args. Unknown keys for the chosen model are dropped. "
|
||||
'Cosmos3 example: \'{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false, '
|
||||
'"use_resolution_template": false, "use_duration_template": false}\' '
|
||||
'(add "generate_sound": true and optional "sound_duration" for synchronized audio). '
|
||||
'Helios-Distilled example: \'{"is_enable_stage2": true, '
|
||||
'"pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}\'.',
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
|
||||
@@ -357,7 +386,14 @@ def main():
|
||||
omni_kwargs["cache_config"] = cache_config
|
||||
omni_kwargs["enable_cache_dit_summary"] = args.enable_cache_dit_summary
|
||||
|
||||
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
|
||||
# gate is an engine-level config (offline analog of the server's --no-guardrails).
|
||||
if args.extra_body and "guardrails" in args.extra_body:
|
||||
omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
|
||||
|
||||
omni = Omni(**omni_kwargs)
|
||||
model_class_name = get_model_class_name(omni) or model_class_name
|
||||
declared_extra_body_params = get_extra_body_params(model_class_name)
|
||||
|
||||
if profiler_enabled:
|
||||
print("[Profiler] Starting profiling...")
|
||||
@@ -392,14 +428,21 @@ def main():
|
||||
)
|
||||
if args.guidance_scale_high is not None:
|
||||
sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
|
||||
if args.extra_body:
|
||||
# Model-specific knobs (declared in vllm_omni/model_extras/) routed via extra_args.
|
||||
sampling_kwargs["extra_args"] = dict(args.extra_body)
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(**sampling_kwargs)
|
||||
# Route model-specific knobs through extra_body, filtered against the model's
|
||||
# declared extra_body_params. Models without a declaration only forward explicit
|
||||
# --extra-body JSON (preserving the generic flags' legacy behavior).
|
||||
extra_body = dict(args.extra_body or {})
|
||||
if declared_extra_body_params:
|
||||
apply_declared_extra_args(sampling_params, declared_extra_body_params, extra_body)
|
||||
elif extra_body:
|
||||
sampling_params.extra_args.update({k: v for k, v in extra_body.items() if v is not None})
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
frames = omni.generate(
|
||||
prompt_dict,
|
||||
OmniDiffusionSamplingParams(**sampling_kwargs),
|
||||
sampling_params,
|
||||
)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
@@ -1,148 +0,0 @@
|
||||
# VACE Video Generation
|
||||
|
||||
Generate videos from text prompts, images, or video conditions using vLLM-Omni's VACE diffusion pipeline.
|
||||
|
||||
- `vace_video_generation.py`: command-line script for multi-mode video generation with advanced options.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Quick Start](#quick-start)
|
||||
- [Key Arguments](#key-arguments)
|
||||
- [More CLI Examples](#more-cli-examples)
|
||||
|
||||
## Overview
|
||||
|
||||
[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single unified model, including text-to-video, image-to-video, first-last-frame interpolation, inpainting, and reference image-guided generation.
|
||||
|
||||
### Supported Models
|
||||
|
||||
| Model | Architecture | Peak VRAM (GiB) * | Model Weights (GiB) | HuggingFace |
|
||||
|-------|-------------|-------------------|---------------------|-------------|
|
||||
| Wan2.1-VACE (1.3B) | Wan2.1 | TBD | ~10 | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
|
||||
| Wan2.1-VACE (14B) | Wan2.1 | TBD | ~38 | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
|
||||
|
||||
!!! info
|
||||
*Peak VRAM: based on basic single-card usage, 480x832 resolution, 81 frames, 30 inference steps, without any acceleration/optimization features.
|
||||
|
||||
Default model: `Wan-AI/Wan2.1-VACE-14B-diffusers`
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Python API
|
||||
|
||||
Text-to-video generation:
|
||||
|
||||
```python
|
||||
import torch
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="Wan-AI/Wan2.1-VACE-1.3B-diffusers")
|
||||
prompt = "A sleek robot stands in a vast warehouse filled with boxes"
|
||||
outputs = omni.generate(
|
||||
prompt,
|
||||
OmniDiffusionSamplingParams(
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=81,
|
||||
num_inference_steps=30,
|
||||
guidance_scale=5.0,
|
||||
),
|
||||
)
|
||||
video = outputs[0].images
|
||||
|
||||
from diffusers.utils import export_to_video
|
||||
export_to_video(list(video[0]), "t2v_output.mp4", fps=16)
|
||||
omni.close()
|
||||
```
|
||||
|
||||
### Local CLI Usage
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode t2v \
|
||||
--prompt "A sleek robot stands in a vast warehouse filled with boxes" \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--output t2v_output.mp4
|
||||
```
|
||||
|
||||
## Key Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| -------- | ---- | ------- | ----------- |
|
||||
| `--mode` | str | `"t2v"` | VACE task mode: `t2v`, `i2v`, `v2lf`, `flf2v`, `inpaint`, `r2v` |
|
||||
| `--model` | str | `"Wan-AI/Wan2.1-VACE-14B-diffusers"` | Model ID or local path |
|
||||
| `--prompt` | str | `"A cat walking in a garden"` | Text description of desired video |
|
||||
| `--negative-prompt` | str | `""` | Negative prompt for classifier-free guidance |
|
||||
| `--image` | str | `None` | Input image path (for I2V, R2V, FLF2V, inpaint modes) |
|
||||
| `--last-image` | str | `None` | Last frame image path (for FLF2V mode) |
|
||||
| `--height` | int | `480` | Output video height in pixels (should be a multiple of 16) |
|
||||
| `--width` | int | `832` | Output video width in pixels (should be a multiple of 16) |
|
||||
| `--num-frames` | int | `81` | Number of video frames to generate |
|
||||
| `--num-inference-steps` | int | `30` | Number of denoising steps (more steps = higher quality, slower) |
|
||||
| `--guidance-scale` | float | `5.0` | Classifier-free guidance scale |
|
||||
| `--flow-shift` | float | `5.0` | Scheduler flow shift parameter |
|
||||
| `--seed` | int | `42` | Random seed for deterministic sampling |
|
||||
| `--fps` | int | `16` | Frames per second for the saved MP4 |
|
||||
| `--output` | str | `"vace_output.mp4"` | Path to save the generated video |
|
||||
| `--vae-use-tiling` | flag | on | Enable VAE tiling for memory optimization |
|
||||
| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
|
||||
| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree for multi-GPU inference |
|
||||
| `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference |
|
||||
| `--tensor-parallel-size` | int | — | Tensor parallel size |
|
||||
| `--enforce-eager` | flag | off | Disable torch.compile |
|
||||
|
||||
> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options (`--ulysses-degree`, `--cfg-parallel-size`).
|
||||
|
||||
## More CLI Examples
|
||||
|
||||
### Image-to-Video (I2V)
|
||||
|
||||
First frame is kept, remaining frames are generated:
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode i2v \
|
||||
--image astronaut.jpg \
|
||||
--prompt "An astronaut emerging from a cracked egg on the moon" \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--output i2v_output.mp4
|
||||
```
|
||||
|
||||
### First-Last-Frame Interpolation (FLF2V)
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode flf2v \
|
||||
--image first_frame.jpg --last-image last_frame.jpg \
|
||||
--prompt "A bird takes off from a branch and lands on another" \
|
||||
--height 512 --width 512 --num-frames 81 \
|
||||
--output flf2v_output.mp4
|
||||
```
|
||||
|
||||
### Inpainting
|
||||
|
||||
Center vertical stripe is masked and regenerated:
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode inpaint \
|
||||
--image scene.jpg \
|
||||
--prompt "Shrek walks out of a building" \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--output inpaint_output.mp4
|
||||
```
|
||||
|
||||
### Reference Image-guided (R2V)
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode r2v \
|
||||
--image reference.jpg \
|
||||
--prompt "Camera slowly zooms out from the character" \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--output r2v_output.mp4
|
||||
```
|
||||
@@ -1,209 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""VACE video generation example.
|
||||
|
||||
VACE (Video All-in-one Creation Engine) supports multiple video tasks:
|
||||
- T2V: Text-to-Video
|
||||
- I2V: Image-to-Video (first frame conditioning)
|
||||
- V2LF: Video-to-Last-Frame
|
||||
- FLF2V: First-Last-Frame interpolation
|
||||
- Inpainting: Masked region generation
|
||||
- R2V: Reference image-guided generation
|
||||
|
||||
Usage examples:
|
||||
# T2V (text-to-video)
|
||||
python vace_video_generation.py --mode t2v --prompt "A robot in a warehouse"
|
||||
|
||||
# I2V (image-to-video, first frame kept)
|
||||
python vace_video_generation.py --mode i2v --image input.jpg --prompt "..."
|
||||
|
||||
# FLF2V (first-last frame interpolation)
|
||||
python vace_video_generation.py --mode flf2v --image first.jpg --last-image last.jpg
|
||||
|
||||
# R2V (reference image guided)
|
||||
python vace_video_generation.py --mode r2v --image ref.jpg --prompt "..."
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="VACE video generation.")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="Wan-AI/Wan2.1-VACE-14B-diffusers",
|
||||
help="VACE model ID or local path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
default="t2v",
|
||||
choices=["t2v", "i2v", "v2lf", "flf2v", "inpaint", "r2v"],
|
||||
help="Generation mode.",
|
||||
)
|
||||
parser.add_argument("--prompt", default="A cat walking in a garden", help="Text prompt.")
|
||||
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
|
||||
parser.add_argument("--image", type=str, default=None, help="Input image path (for I2V, R2V, FLF2V, inpaint).")
|
||||
parser.add_argument("--last-image", type=str, default=None, help="Last frame image path (for FLF2V).")
|
||||
parser.add_argument("--video-dir", type=str, default=None, help="Directory of video frames (for inpaint).")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG guidance scale.")
|
||||
parser.add_argument("--height", type=int, default=480, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=832, help="Video width.")
|
||||
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=30, help="Sampling steps.")
|
||||
parser.add_argument("--flow-shift", type=float, default=5.0, help="Scheduler flow_shift.")
|
||||
parser.add_argument("--output", type=str, default="vace_output.mp4", help="Output video path.")
|
||||
parser.add_argument("--fps", type=int, default=16, help="Output video FPS.")
|
||||
parser.add_argument("--vae-use-tiling", action="store_true", default=True, help="Enable VAE tiling.")
|
||||
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
|
||||
parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
|
||||
parser.add_argument("--ring-degree", type=int, default=1, help="Ring attention degree.")
|
||||
parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def build_prompts(args):
|
||||
"""Build prompt dict with multi_modal_data based on mode."""
|
||||
h, w, nf = args.height, args.width, args.num_frames
|
||||
|
||||
gray = PIL.Image.new("RGB", (w, h), (128, 128, 128))
|
||||
mask_black = PIL.Image.new("L", (w, h), 0)
|
||||
mask_white = PIL.Image.new("L", (w, h), 255)
|
||||
|
||||
prompt_data = {
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
}
|
||||
|
||||
if args.mode == "t2v":
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "r2v":
|
||||
assert args.image, "--image required for R2V mode"
|
||||
ref_img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
prompt_data["multi_modal_data"] = {"reference_images": [ref_img]}
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "i2v":
|
||||
assert args.image, "--image required for I2V mode"
|
||||
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
prompt_data["multi_modal_data"] = {
|
||||
"video": [img] + [gray] * (nf - 1),
|
||||
"mask": [mask_black] + [mask_white] * (nf - 1),
|
||||
}
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "v2lf":
|
||||
assert args.image, "--image required for V2LF mode"
|
||||
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
prompt_data["multi_modal_data"] = {
|
||||
"video": [gray] * (nf - 1) + [img],
|
||||
"mask": [mask_white] * (nf - 1) + [mask_black],
|
||||
}
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "flf2v":
|
||||
assert args.image and args.last_image, "--image and --last-image required for FLF2V"
|
||||
first = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
last = PIL.Image.open(args.last_image).convert("RGB").resize((w, h))
|
||||
prompt_data["multi_modal_data"] = {
|
||||
"video": [first] + [gray] * (nf - 2) + [last],
|
||||
"mask": [mask_black] + [mask_white] * (nf - 2) + [mask_black],
|
||||
}
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "inpaint":
|
||||
assert args.image, "--image required for inpaint mode"
|
||||
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
d = 80
|
||||
frames, masks = [], []
|
||||
for _ in range(nf):
|
||||
base = np.array(img).copy()
|
||||
mask = PIL.Image.new("L", (w, h), 0)
|
||||
stripe = PIL.Image.new("L", (2 * d, h), 255)
|
||||
mask.paste(stripe, (w // 2 - d, 0))
|
||||
base[np.array(mask) > 128] = 128
|
||||
frames.append(PIL.Image.fromarray(base))
|
||||
masks.append(mask)
|
||||
prompt_data["multi_modal_data"] = {"video": frames, "mask": masks}
|
||||
return prompt_data
|
||||
|
||||
raise ValueError(f"Unknown mode: {args.mode}")
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
)
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
flow_shift=args.flow_shift,
|
||||
enforce_eager=args.enforce_eager,
|
||||
parallel_config=parallel_config,
|
||||
)
|
||||
|
||||
prompt_data = build_prompts(args)
|
||||
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f"VACE {args.mode.upper()} Generation")
|
||||
print(f" Model: {args.model}")
|
||||
print(f" Size: {args.width}x{args.height}, {args.num_frames} frames, {args.num_inference_steps} steps")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
start = time.perf_counter()
|
||||
outputs = omni.generate(
|
||||
prompt_data,
|
||||
OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
generator=generator,
|
||||
),
|
||||
)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
video = outputs[0].images
|
||||
if isinstance(video, list):
|
||||
video = video[0]
|
||||
if isinstance(video, torch.Tensor):
|
||||
video = video.cpu().numpy()
|
||||
if video.ndim == 5:
|
||||
video = video[0]
|
||||
print(f"Output shape: {video.shape}, Time: {elapsed:.1f}s")
|
||||
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
if np.issubdtype(video.dtype, np.integer):
|
||||
video = video.astype(np.float32) / 255.0
|
||||
export_to_video(list(video), str(output_path), fps=args.fps)
|
||||
print(f"Saved to {output_path}")
|
||||
|
||||
omni.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -88,7 +88,7 @@ curl -s http://127.0.0.1:8091/v1/chat/completions \
|
||||
|
||||
### Optional knobs
|
||||
|
||||
Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_args.image_gen`:
|
||||
Pass image-gen overrides as flat keys on the diffusion-stage `sampling_params_list[1].extra_args`:
|
||||
|
||||
| Key | Default | Description |
|
||||
| --- | --- | --- |
|
||||
@@ -97,25 +97,48 @@ Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_
|
||||
| `cfg` | 2.0 | Classifier-free guidance scale. |
|
||||
| `seed` | 42 | Per-request RNG seed (deterministic when ByT5 is also seed-stable). |
|
||||
| `byte5_text` | (auto) | Override the glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "...". ` format. |
|
||||
| `negative_prompt` | empty | Real CFG negative conditioning (must be set on **stage-0 thinker** `sampling_params` so `expand_cfg_prompts` spawns the companion). |
|
||||
| `negative_prompt` | empty | Real CFG negative conditioning (set on **stage-0 thinker** `extra_args` so `expand_cfg_prompts` spawns the companion). |
|
||||
|
||||
Example with all the knobs:
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Jonathan1909/Ming-flash-omni-2.0",
|
||||
"modalities": ["image"],
|
||||
"sampling_params_list": [
|
||||
{"temperature":0.4,"top_p":0.9,"top_k":1,"max_tokens":1,"seed":42,
|
||||
"extra_args":{"image_gen":{"negative_prompt":"ugly, blurry, distorted"}}},
|
||||
{"seed":42,"extra_args":{"image_gen":{
|
||||
"steps":6,"cfg":1.5,"height":512,"width":512,"seed":123,
|
||||
"byte5_text":["理解与生成统一"]}}}
|
||||
],
|
||||
"messages": [{"role": "user", "content": "Draw a poster."}]
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen_knobs.png
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Jonathan1909/Ming-flash-omni-2.0",
|
||||
"modalities": ["image"],
|
||||
"sampling_params_list": [
|
||||
{
|
||||
"temperature": 0.4,
|
||||
"top_p": 0.9,
|
||||
"top_k": 1,
|
||||
"max_tokens": 1,
|
||||
"seed": 42,
|
||||
"extra_args": {
|
||||
"negative_prompt": "ugly, blurry, distorted"
|
||||
}
|
||||
},
|
||||
{
|
||||
"seed": 42,
|
||||
"extra_args": {
|
||||
"steps": 6,
|
||||
"cfg": 1.5,
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"seed": 123,
|
||||
"byte5_text": ["理解与生成统一"]
|
||||
}
|
||||
}
|
||||
],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Draw a poster."
|
||||
}
|
||||
]
|
||||
}' \
|
||||
| jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
|
||||
| base64 -d > ming_imagegen_knobs.png
|
||||
```
|
||||
|
||||
### img2img (reference image)
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
# Streaming Video Generation
|
||||
|
||||
This example uses the custom WebSocket endpoint `WS /v1/realtime/video` to receive a video byte stream as chunks are produced.
|
||||
It covers text-only video generation. Image/reference input is intentionally not included for now.
|
||||
|
||||
## Start The Server
|
||||
|
||||
Start a diffusion video model with streaming output enabled:
|
||||
|
||||
```bash
|
||||
vllm serve BestWishYsh/Helios-Distilled \
|
||||
--omni \
|
||||
--diffusion-streaming-output \
|
||||
--port 8000
|
||||
```
|
||||
|
||||
The `--diffusion-streaming-output` CLI flag is forwarded as `streaming_output=True` in the default diffusion stage `engine_args`, then loaded by `OmniDiffusionConfig.from_kwargs()`.
|
||||
|
||||
## WebSocket Protocol
|
||||
|
||||
| Direction | Message | Format | Description |
|
||||
| --- | --- | --- | --- |
|
||||
| Client to server | `session.start` | JSON text: `{"type":"session.start","model":"...","prompt":"...","format":"m4s"}` | Starts generation. `format` is optional and accepts `m4s` (default). Sampling fields such as `width`, `height`, `fps`, `num_frames`, and `extra_params` may be included. |
|
||||
| Server to client | `video.start` | JSON text: `{"type":"video.start","request_id":"...","format":"m4s","config":{...}}` | Confirms the session and mirrors the accepted `format`. |
|
||||
| Server to client | Video chunk | Binary WebSocket frame | Fragmented MP4 (`m4s`) video bytes. |
|
||||
| Client to server | `session.stop` | JSON text: `{"type":"session.stop"}` | Requests cancellation of the active session. |
|
||||
| Client to server | `session.ping` | JSON text: `{"type":"session.ping"}` | Optional keepalive; refreshes the server stall clock. |
|
||||
| Server to client | `session.done` | JSON text: `{"type":"session.done","request_id":"...","chunks":3,"stopped":false}` | Ends a completed or stopped session. |
|
||||
| Server to client | `session.pong` | JSON text: `{"type":"session.pong"}` | Reply to `session.ping`. |
|
||||
| Server to client | `error` | JSON text: `{"type":"error","message":"..."}` | Reports invalid input, unsupported formats, generation failures, control-message errors, or stall timeout. |
|
||||
|
||||
During generation the client normally sends only `session.start` and then receives binary chunks; silence on the client socket is expected. The server closes the session with a stall error only when there is no engine progress and no `session.ping` for about 60 seconds.
|
||||
|
||||
## Install Client Dependency
|
||||
|
||||
```bash
|
||||
pip install websockets
|
||||
# For the Gradio demo:
|
||||
pip install 'vllm-omni[demo]' websockets
|
||||
```
|
||||
|
||||
## Run The Client
|
||||
|
||||
```bash
|
||||
python streaming_video_client.py \
|
||||
--host 127.0.0.1 \
|
||||
--port 8000 \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--prompt "A serene lakeside sunrise with mist over the water." \
|
||||
--width 640 \
|
||||
--height 384 \
|
||||
--fps 16 \
|
||||
--num-frames 99 \
|
||||
--guidance-scale 1.0 \
|
||||
--seed 42 \
|
||||
--output helios_stream.mp4
|
||||
```
|
||||
|
||||
The client sends one `session.start` message, prints each received binary video chunk with its byte size and elapsed time, and saves the received bytes to `--output` after `session.done`.
|
||||
The client remuxes the gathered stream to a regular progressive MP4 file so that local playback knows the video duration.
|
||||
|
||||
## Run The Gradio Demo
|
||||
|
||||
```bash
|
||||
python gradio_demo.py \
|
||||
--host 127.0.0.1 \
|
||||
--port 7860
|
||||
```
|
||||
|
||||
The Gradio demo requests fMP4 (`m4s`) chunks and appends them directly in the browser with a Media Source Extensions player.
|
||||
|
||||
## Model Choice
|
||||
|
||||
### Helios
|
||||
|
||||
The example uses `BestWishYsh/Helios-Distilled` model by default.
|
||||
|
||||
To ensure streaming-level generation speed, `pyramid_num_inference_steps_list` is suggested to be as low as `[1, 1, 1]`. Both example clients uses the following Helios-Distilled preset by default:
|
||||
|
||||
```json
|
||||
{
|
||||
"is_enable_stage2": true,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [1, 1, 1],
|
||||
"is_amplify_first_chunk": true
|
||||
}
|
||||
```
|
||||
|
||||
Disable it in the CLI example with `--no-helios-distilled-preset`, or override/extend it with `--extra-params`:
|
||||
|
||||
```bash
|
||||
python streaming_video_client.py \
|
||||
--extra-params '{"pyramid_num_inference_steps_list":[2, 2, 2]}'
|
||||
```
|
||||
@@ -0,0 +1,301 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Gradio demo for `/v1/realtime/video`.
|
||||
|
||||
Usage:
|
||||
python gradio_demo.py --host 127.0.0.1 --port 7860
|
||||
|
||||
Start a vLLM-Omni server with streaming output enabled first, for example:
|
||||
vllm serve BestWishYsh/Helios-Distilled --omni --diffusion-streaming-output --port 8000
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
import gradio as gr
|
||||
except ImportError:
|
||||
raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
|
||||
|
||||
DEFAULT_PROMPT = "A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement."
|
||||
VIDEO_STREAM_VIEW_HTML = Path(__file__).with_name("video-stream-view.html")
|
||||
VIDEO_STREAM_VIEW_JS = Path(__file__).with_name("video-stream-view.js")
|
||||
|
||||
|
||||
def _maybe_set(payload: dict[str, Any], key: str, value: Any) -> None:
|
||||
if value is None:
|
||||
return
|
||||
if isinstance(value, str) and not value.strip():
|
||||
return
|
||||
payload[key] = value
|
||||
|
||||
|
||||
def _optional_int(value: int | float | None) -> int | None:
|
||||
return None if value is None else int(value)
|
||||
|
||||
|
||||
def _optional_float(value: int | float | None) -> float | None:
|
||||
return None if value is None else float(value)
|
||||
|
||||
|
||||
def _build_session_start(
|
||||
*,
|
||||
model: str,
|
||||
prompt: str,
|
||||
negative_prompt: str,
|
||||
width: int | float | None,
|
||||
height: int | float | None,
|
||||
fps: int | float | None,
|
||||
num_frames: int | float | None,
|
||||
num_inference_steps: int | float | None,
|
||||
guidance_scale: int | float | None,
|
||||
guidance_scale_2: int | float | None,
|
||||
boundary_ratio: int | float | None,
|
||||
flow_shift: int | float | None,
|
||||
true_cfg_scale: int | float | None,
|
||||
seed: int | float | None,
|
||||
) -> dict[str, Any]:
|
||||
payload: dict[str, Any] = {
|
||||
"type": "session.start",
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"format": "m4s",
|
||||
}
|
||||
|
||||
_maybe_set(payload, "negative_prompt", negative_prompt)
|
||||
_maybe_set(payload, "width", _optional_int(width))
|
||||
_maybe_set(payload, "height", _optional_int(height))
|
||||
_maybe_set(payload, "fps", _optional_int(fps))
|
||||
_maybe_set(payload, "num_frames", _optional_int(num_frames))
|
||||
_maybe_set(payload, "num_inference_steps", _optional_int(num_inference_steps))
|
||||
_maybe_set(payload, "guidance_scale", _optional_float(guidance_scale))
|
||||
_maybe_set(payload, "guidance_scale_2", _optional_float(guidance_scale_2))
|
||||
_maybe_set(payload, "boundary_ratio", _optional_float(boundary_ratio))
|
||||
_maybe_set(payload, "flow_shift", _optional_float(flow_shift))
|
||||
_maybe_set(payload, "true_cfg_scale", _optional_float(true_cfg_scale))
|
||||
_maybe_set(payload, "seed", _optional_int(seed))
|
||||
|
||||
if "Helios" in model:
|
||||
payload["extra_params"] = {
|
||||
"is_enable_stage2": True,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [1, 1, 1],
|
||||
"is_amplify_first_chunk": True,
|
||||
}
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
def _set_optional_field_enabled(enabled: bool) -> Any:
|
||||
return gr.update(interactive=enabled)
|
||||
|
||||
|
||||
def build_browser_stream_config(
|
||||
host: str,
|
||||
port: int | float,
|
||||
model: str,
|
||||
prompt: str,
|
||||
negative_prompt: str,
|
||||
width: int | float | None,
|
||||
height: int | float | None,
|
||||
fps: int | float,
|
||||
num_frames: int | float | None,
|
||||
num_inference_steps: int | float | None,
|
||||
include_guidance_scale: bool,
|
||||
guidance_scale: int | float | None,
|
||||
include_guidance_scale_2: bool,
|
||||
guidance_scale_2: int | float | None,
|
||||
include_seed: bool,
|
||||
seed: int | float | None,
|
||||
include_boundary_ratio: bool,
|
||||
boundary_ratio: int | float | None,
|
||||
include_flow_shift: bool,
|
||||
flow_shift: int | float | None,
|
||||
include_true_cfg_scale: bool,
|
||||
true_cfg_scale: int | float | None,
|
||||
) -> tuple[str, Any]:
|
||||
"""Build the browser-side WebSocket request config for the HTML/MSE player."""
|
||||
if not model.strip():
|
||||
raise gr.Error("Model is required.")
|
||||
if not prompt.strip():
|
||||
raise gr.Error("Prompt is required.")
|
||||
|
||||
payload = _build_session_start(
|
||||
model=model.strip(),
|
||||
prompt=prompt.strip(),
|
||||
negative_prompt=negative_prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
fps=fps,
|
||||
num_frames=num_frames,
|
||||
num_inference_steps=num_inference_steps,
|
||||
guidance_scale=guidance_scale if include_guidance_scale else None,
|
||||
guidance_scale_2=guidance_scale_2 if include_guidance_scale_2 else None,
|
||||
boundary_ratio=boundary_ratio if include_boundary_ratio else None,
|
||||
flow_shift=flow_shift if include_flow_shift else None,
|
||||
true_cfg_scale=true_cfg_scale if include_true_cfg_scale else None,
|
||||
seed=seed if include_seed else None,
|
||||
)
|
||||
config = {
|
||||
"url": f"ws://{host}:{int(port)}/v1/realtime/video",
|
||||
"payload": payload,
|
||||
}
|
||||
return json.dumps(config, ensure_ascii=False), gr.update(value="Streaming...", interactive=False)
|
||||
|
||||
|
||||
def create_demo() -> gr.Blocks:
|
||||
with gr.Blocks(title="Streaming Video Generation") as demo:
|
||||
gr.Markdown("# Streaming Video Generation")
|
||||
gr.Markdown(
|
||||
"Connects to `WS /v1/realtime/video` in the browser and appends fMP4 chunks "
|
||||
"directly to a Media Source Extensions video player."
|
||||
)
|
||||
|
||||
with gr.Row():
|
||||
with gr.Column(scale=1):
|
||||
host = gr.Textbox(label="Server Host", value="127.0.0.1")
|
||||
port = gr.Number(label="Server Port", value=8000, precision=0)
|
||||
model = gr.Textbox(label="Model", value="BestWishYsh/Helios-Distilled")
|
||||
prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=3)
|
||||
negative_prompt = gr.Textbox(label="Negative Prompt", value="", lines=2)
|
||||
|
||||
with gr.Row():
|
||||
width = gr.Number(label="Width", value=640, precision=0)
|
||||
height = gr.Number(label="Height", value=384, precision=0)
|
||||
|
||||
with gr.Row():
|
||||
fps = gr.Number(label="FPS", value=16, precision=0)
|
||||
num_frames = gr.Number(label="Num Frames", value=132, precision=0)
|
||||
num_inference_steps = gr.Number(label="Num Inference Steps", value=50, precision=0)
|
||||
|
||||
with gr.Row():
|
||||
with gr.Column():
|
||||
guidance_scale = gr.Number(
|
||||
label="Guidance Scale",
|
||||
value=1.0,
|
||||
interactive=True,
|
||||
scale=4,
|
||||
)
|
||||
include_guidance_scale = gr.Checkbox(label="Send", value=True, scale=1, min_width=70)
|
||||
with gr.Column():
|
||||
guidance_scale_2 = gr.Number(
|
||||
label="Guidance Scale 2",
|
||||
value=None,
|
||||
interactive=False,
|
||||
scale=4,
|
||||
)
|
||||
include_guidance_scale_2 = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
|
||||
with gr.Column():
|
||||
true_cfg_scale = gr.Number(
|
||||
label="True CFG Scale",
|
||||
value=None,
|
||||
interactive=False,
|
||||
scale=4,
|
||||
)
|
||||
include_true_cfg_scale = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
|
||||
# I want it to wrap after 3, but it auto wrap after 2. I'll rather fully delegate to auto wrapping.
|
||||
with gr.Column():
|
||||
boundary_ratio = gr.Number(
|
||||
label="Boundary Ratio",
|
||||
value=None,
|
||||
interactive=False,
|
||||
scale=4,
|
||||
)
|
||||
include_boundary_ratio = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
|
||||
with gr.Column():
|
||||
flow_shift = gr.Number(label="Flow Shift", value=None, interactive=False, scale=4)
|
||||
include_flow_shift = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
|
||||
with gr.Column():
|
||||
seed = gr.Number(label="Seed", value=42, precision=0, interactive=True, scale=4)
|
||||
include_seed = gr.Checkbox(label="Send", value=True, scale=1, min_width=70)
|
||||
|
||||
with gr.Row():
|
||||
start_button = gr.Button("Start", variant="primary", elem_id="streaming-video-start")
|
||||
hidden_stream_config = gr.Textbox(value="", visible=False)
|
||||
|
||||
with gr.Column(scale=1):
|
||||
gr.HTML(
|
||||
VIDEO_STREAM_VIEW_HTML.read_text(encoding="utf-8"),
|
||||
label="Streaming Preview",
|
||||
js_on_load=VIDEO_STREAM_VIEW_JS.read_text(encoding="utf-8"),
|
||||
)
|
||||
|
||||
inputs = [
|
||||
host,
|
||||
port,
|
||||
model,
|
||||
prompt,
|
||||
negative_prompt,
|
||||
width,
|
||||
height,
|
||||
fps,
|
||||
num_frames,
|
||||
num_inference_steps,
|
||||
include_guidance_scale,
|
||||
guidance_scale,
|
||||
include_guidance_scale_2,
|
||||
guidance_scale_2,
|
||||
include_seed,
|
||||
seed,
|
||||
include_boundary_ratio,
|
||||
boundary_ratio,
|
||||
include_flow_shift,
|
||||
flow_shift,
|
||||
include_true_cfg_scale,
|
||||
true_cfg_scale,
|
||||
]
|
||||
for toggle, field in [
|
||||
(include_guidance_scale, guidance_scale),
|
||||
(include_guidance_scale_2, guidance_scale_2),
|
||||
(include_seed, seed),
|
||||
(include_boundary_ratio, boundary_ratio),
|
||||
(include_flow_shift, flow_shift),
|
||||
(include_true_cfg_scale, true_cfg_scale),
|
||||
]:
|
||||
toggle.change(
|
||||
fn=_set_optional_field_enabled,
|
||||
inputs=toggle,
|
||||
outputs=field,
|
||||
)
|
||||
start_button.click(
|
||||
fn=build_browser_stream_config,
|
||||
inputs=inputs,
|
||||
outputs=[hidden_stream_config, start_button],
|
||||
).then(
|
||||
fn=lambda config: config,
|
||||
inputs=hidden_stream_config,
|
||||
outputs=hidden_stream_config,
|
||||
js="""
|
||||
(config) => {
|
||||
if (window.vllmStreamingVideoStart) {
|
||||
window.vllmStreamingVideoStart(config);
|
||||
}
|
||||
return config;
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
return demo
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Streaming video generation Gradio demo")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="Host/IP for gradio launch")
|
||||
parser.add_argument("--port", type=int, default=7860, help="Port for gradio launch")
|
||||
parser.add_argument("--share", action="store_true", help="Create a public Gradio link")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
demo = create_demo()
|
||||
demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,279 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Text-only WebSocket client for `/v1/realtime/video`.
|
||||
|
||||
The client prints a line for every binary video chunk it receives and writes
|
||||
the received bytes to disk when the session finishes or is interrupted.
|
||||
Image/reference input is intentionally not supported in this example yet.
|
||||
|
||||
Requirements:
|
||||
pip install websockets
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import av
|
||||
|
||||
try:
|
||||
from websockets.asyncio.client import connect # pyright: ignore[reportMissingImports]
|
||||
except ImportError:
|
||||
try:
|
||||
from websockets import connect # pyright: ignore[reportMissingImports]
|
||||
except ImportError:
|
||||
print("Please install websockets: pip install websockets")
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
HELIOS_DISTILLED_EXTRA_PARAMS = {
|
||||
"is_enable_stage2": True,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [1, 1, 1],
|
||||
"is_amplify_first_chunk": True,
|
||||
}
|
||||
|
||||
|
||||
def _json_object(value: str, *, argument_name: str) -> dict[str, Any]:
|
||||
try:
|
||||
parsed = json.loads(value)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise argparse.ArgumentTypeError(f"{argument_name} must be valid JSON: {exc}") from exc
|
||||
if not isinstance(parsed, dict):
|
||||
raise argparse.ArgumentTypeError(f"{argument_name} must be a JSON object")
|
||||
return parsed
|
||||
|
||||
|
||||
def _maybe_set(payload: dict[str, Any], key: str, value: Any) -> None:
|
||||
if value is not None:
|
||||
payload[key] = value
|
||||
|
||||
|
||||
def build_session_start(args: argparse.Namespace) -> dict[str, Any]:
|
||||
payload: dict[str, Any] = {
|
||||
"type": "session.start",
|
||||
"model": args.model,
|
||||
"prompt": args.prompt,
|
||||
}
|
||||
|
||||
_maybe_set(payload, "negative_prompt", args.negative_prompt)
|
||||
_maybe_set(payload, "width", args.width)
|
||||
_maybe_set(payload, "height", args.height)
|
||||
_maybe_set(payload, "size", args.size)
|
||||
_maybe_set(payload, "seconds", args.seconds)
|
||||
_maybe_set(payload, "fps", args.fps)
|
||||
_maybe_set(payload, "num_frames", args.num_frames)
|
||||
_maybe_set(payload, "num_inference_steps", args.num_inference_steps)
|
||||
_maybe_set(payload, "guidance_scale", args.guidance_scale)
|
||||
_maybe_set(payload, "guidance_scale_2", args.guidance_scale_2)
|
||||
_maybe_set(payload, "boundary_ratio", args.boundary_ratio)
|
||||
_maybe_set(payload, "flow_shift", args.flow_shift)
|
||||
_maybe_set(payload, "true_cfg_scale", args.true_cfg_scale)
|
||||
_maybe_set(payload, "seed", args.seed)
|
||||
|
||||
extra_params: dict[str, Any] = {}
|
||||
if "Helios-Distilled" in args.model and args.helios_distilled_preset:
|
||||
extra_params.update(HELIOS_DISTILLED_EXTRA_PARAMS)
|
||||
if args.extra_params:
|
||||
extra_params.update(args.extra_params)
|
||||
if extra_params:
|
||||
payload["extra_params"] = extra_params
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
def _count_decoded_frames(video_bytes: bytes, *, stream_format: str) -> int | None:
|
||||
"""Return the number of decodable video frames in a streamed byte sequence, if known."""
|
||||
if not video_bytes:
|
||||
return 0
|
||||
|
||||
av_format = "mp4" if stream_format == "m4s" else None
|
||||
try:
|
||||
with av.open(io.BytesIO(video_bytes), format=av_format, mode="r") as container:
|
||||
stream = container.streams.video[0]
|
||||
return sum(1 for _ in container.decode(stream))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def save_video(
|
||||
chunks: list[bytes],
|
||||
output_path: Path,
|
||||
*,
|
||||
fps: float,
|
||||
stream_format: str,
|
||||
) -> None:
|
||||
if not chunks:
|
||||
print("No video chunks received; nothing to save.")
|
||||
return
|
||||
|
||||
from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes
|
||||
|
||||
streamed_bytes = b"".join(chunks)
|
||||
playback_bytes = finalize_streaming_video_bytes(
|
||||
streamed_bytes,
|
||||
input_format="m4s",
|
||||
fps=fps,
|
||||
)
|
||||
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp_path = output_path.with_suffix(output_path.suffix + ".tmp")
|
||||
with tmp_path.open("wb") as f:
|
||||
f.write(playback_bytes)
|
||||
os.replace(tmp_path, output_path)
|
||||
print(
|
||||
f"Saved {len(chunks)} streamed chunk(s), {len(streamed_bytes)} streamed bytes "
|
||||
f"-> {len(playback_bytes)} playback bytes at {output_path}"
|
||||
)
|
||||
|
||||
|
||||
async def stream_video(args: argparse.Namespace) -> None:
|
||||
payload = build_session_start(args)
|
||||
output_path = Path(args.output)
|
||||
url = f"ws://{args.host}:{args.port}/v1/realtime/video"
|
||||
chunks: list[bytes] = []
|
||||
total_bytes = 0
|
||||
total_frames = 0
|
||||
accumulated_stream = bytearray()
|
||||
stream_format = "m4s"
|
||||
done = False
|
||||
ws = None
|
||||
|
||||
try:
|
||||
async with connect(url, max_size=None) as websocket:
|
||||
ws = websocket
|
||||
await websocket.send(json.dumps(payload))
|
||||
print(f"Connected: {url}")
|
||||
print(f"Sent session.start: {json.dumps(payload, ensure_ascii=False)}")
|
||||
started_at = time.perf_counter()
|
||||
last_chunk_at = started_at
|
||||
|
||||
while True:
|
||||
message = await websocket.recv()
|
||||
now = time.perf_counter()
|
||||
|
||||
if isinstance(message, bytes):
|
||||
chunks.append(message)
|
||||
chunk_bytes = len(message)
|
||||
chunk_elapsed = now - last_chunk_at
|
||||
last_chunk_at = now
|
||||
|
||||
total_bytes += chunk_bytes
|
||||
accumulated_stream.extend(message)
|
||||
decoded_total = _count_decoded_frames(bytes(accumulated_stream), stream_format=stream_format)
|
||||
if decoded_total is not None:
|
||||
chunk_frames = max(decoded_total - total_frames, 0)
|
||||
total_frames = decoded_total
|
||||
else:
|
||||
chunk_frames = 0
|
||||
|
||||
total_elapsed = now - started_at
|
||||
print(
|
||||
f"[chunk {len(chunks):03d}] "
|
||||
f"bytes={chunk_bytes} "
|
||||
f"frames={chunk_frames} "
|
||||
f"elapsed={chunk_elapsed:.2f}s "
|
||||
f"total_bytes={total_bytes} "
|
||||
f"total_frames={total_frames} "
|
||||
f"total_elapsed={total_elapsed:.2f}s"
|
||||
)
|
||||
continue
|
||||
|
||||
msg = json.loads(message)
|
||||
msg_type = msg.get("type")
|
||||
if msg_type == "video.start":
|
||||
stream_format = msg.get("format") or stream_format
|
||||
print(f"Video session started: request_id={msg.get('request_id')} format={msg.get('format')}")
|
||||
elif msg_type == "session.done":
|
||||
print(f"Session complete: {json.dumps(msg, ensure_ascii=False)}")
|
||||
done = True
|
||||
break
|
||||
elif msg_type == "error":
|
||||
print(f"ERROR: {msg.get('message', msg)}")
|
||||
else:
|
||||
print(f"Received control message: {msg}")
|
||||
except (KeyboardInterrupt, asyncio.CancelledError):
|
||||
print("\nInterrupted. Asking server to stop the active session...")
|
||||
if ws is not None:
|
||||
try:
|
||||
await ws.send(json.dumps({"type": "session.stop"}))
|
||||
await ws.close(code=1000, reason="client interrupted")
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
finally:
|
||||
print(
|
||||
"Saving video... (May take a while. Remuxing concatenated chunks into one progressive MP4 file with total duration metadata in file header)"
|
||||
)
|
||||
save_video(
|
||||
chunks,
|
||||
output_path,
|
||||
fps=float(payload.get("fps") or args.fps),
|
||||
stream_format=stream_format,
|
||||
)
|
||||
if not done:
|
||||
print("Saved partial output because the stream did not finish normally.")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Streaming video generation WebSocket client")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="API server host")
|
||||
parser.add_argument("--port", type=int, default=8000, help="API server port")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="BestWishYsh/Helios-Distilled",
|
||||
help="Model name served by the API server",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
default="A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement.",
|
||||
help="Text prompt for video generation",
|
||||
)
|
||||
parser.add_argument("--output", default="streaming_video_output.mp4", help="Path for the received video bytes")
|
||||
parser.add_argument("--width", type=int, default=640, help="Video width")
|
||||
parser.add_argument("--height", type=int, default=384, help="Video height")
|
||||
parser.add_argument("--size", default=None, help="Video size as WIDTHxHEIGHT; alternative to width/height")
|
||||
parser.add_argument("--seconds", default=None, help="Clip duration in seconds")
|
||||
parser.add_argument("--fps", type=int, default=16, help="Frames per second")
|
||||
parser.add_argument("--num-frames", type=int, default=99, help="Number of generated frames")
|
||||
|
||||
parser.add_argument("--negative-prompt", default=None, help="Negative prompt")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=None, help="Number of diffusion steps")
|
||||
parser.add_argument("--guidance-scale", type=float, default=1.0, help="Classifier-free guidance scale")
|
||||
parser.add_argument("--guidance-scale-2", type=float, default=None, help="High-noise CFG scale for video models")
|
||||
parser.add_argument("--boundary-ratio", type=float, default=None, help="Boundary split ratio for video models")
|
||||
parser.add_argument("--flow-shift", type=float, default=None, help="Scheduler flow_shift for video models")
|
||||
parser.add_argument("--true-cfg-scale", type=float, default=None, help="True CFG scale for models that support it")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed")
|
||||
|
||||
parser.add_argument(
|
||||
"--helios-distilled-preset",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Send the Helios-Distilled extra_params used by the bundled example.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--extra-params",
|
||||
type=lambda value: _json_object(value, argument_name="--extra-params"),
|
||||
default=None,
|
||||
help="JSON object merged into extra_params; overrides preset keys on conflict.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
try:
|
||||
asyncio.run(stream_video(args))
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,11 @@
|
||||
<!-- This file and the corresponding JS file are needed to display a streamable video view in the Gradio demo.
|
||||
Gradio's built-in video player does not support modern M4S streaming format. -->
|
||||
<div id="vllm-streaming-video-view" style="display:flex; flex-direction:column; gap:10px;">
|
||||
<video id="vllm-streaming-video-player" autoplay muted controls playsinline
|
||||
style="width:100%; min-height:320px; background:#111; border-radius:8px;"></video>
|
||||
<div style="display:flex; align-items:center; gap:10px;">
|
||||
<span id="vllm-streaming-video-status" style="font-weight:600;">Ready</span>
|
||||
</div>
|
||||
<pre id="vllm-streaming-video-log" style="height:220px; overflow:auto; margin:0; padding:10px; border-radius:8px;
|
||||
background:#0f172a; color:#dbeafe; font-size:12px; white-space:pre-wrap;">Ready.</pre>
|
||||
</div>
|
||||
@@ -0,0 +1,198 @@
|
||||
/**
|
||||
* This file and the corresponding HTML file are needed to display a streamable video view in the Gradio demo.
|
||||
* Gradio's built-in video player does not support modern M4S streaming format.
|
||||
*/
|
||||
(function () {
|
||||
const playerId = "vllm-streaming-video-player";
|
||||
const statusId = "vllm-streaming-video-status";
|
||||
const logId = "vllm-streaming-video-log";
|
||||
let ws = null;
|
||||
let mediaSource = null;
|
||||
let sourceBuffer = null;
|
||||
let queue = [];
|
||||
let chunkCount = 0;
|
||||
let totalBytes = 0;
|
||||
let done = false;
|
||||
let objectUrl = null;
|
||||
|
||||
function el(id) {
|
||||
return document.getElementById(id);
|
||||
}
|
||||
|
||||
function setStatus(text) {
|
||||
const node = el(statusId);
|
||||
if (node) node.textContent = text;
|
||||
}
|
||||
|
||||
function log(line) {
|
||||
const node = el(logId);
|
||||
if (!node) return;
|
||||
node.textContent += "\n" + line;
|
||||
node.scrollTop = node.scrollHeight;
|
||||
}
|
||||
|
||||
function setStartButton(disabled, label) {
|
||||
const btn = document.querySelector("#streaming-video-start button") || document.getElementById("streaming-video-start");
|
||||
if (!btn) return;
|
||||
btn.disabled = disabled;
|
||||
btn.textContent = label;
|
||||
}
|
||||
|
||||
function chooseMime() {
|
||||
const candidates = [
|
||||
'video/mp4; codecs="avc1.42E01E"',
|
||||
'video/mp4; codecs="avc1.4D401F"',
|
||||
'video/mp4; codecs="avc1.64001F"',
|
||||
"video/mp4"
|
||||
];
|
||||
for (const mime of candidates) {
|
||||
if (window.MediaSource && MediaSource.isTypeSupported(mime)) return mime;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
function pump() {
|
||||
if (!sourceBuffer || sourceBuffer.updating || queue.length === 0) return;
|
||||
try {
|
||||
sourceBuffer.appendBuffer(queue.shift());
|
||||
} catch (err) {
|
||||
done = true;
|
||||
setStatus("SourceBuffer error");
|
||||
setStartButton(false, "Restart");
|
||||
log("SourceBuffer append error: " + err.message);
|
||||
finishStream();
|
||||
}
|
||||
}
|
||||
|
||||
function finishStream() {
|
||||
if (!mediaSource || mediaSource.readyState !== "open") return;
|
||||
if (sourceBuffer && sourceBuffer.updating) return;
|
||||
if (queue.length > 0) return;
|
||||
try {
|
||||
mediaSource.endOfStream();
|
||||
} catch (_) { }
|
||||
}
|
||||
|
||||
function stopCurrent() {
|
||||
if (ws && ws.readyState === WebSocket.OPEN) {
|
||||
try { ws.send(JSON.stringify({ type: "session.stop" })); } catch (_) { }
|
||||
try { ws.close(); } catch (_) { }
|
||||
}
|
||||
ws = null;
|
||||
sourceBuffer = null;
|
||||
mediaSource = null;
|
||||
queue = [];
|
||||
if (objectUrl) URL.revokeObjectURL(objectUrl);
|
||||
objectUrl = null;
|
||||
}
|
||||
|
||||
window.vllmStreamingVideoStart = function (configJson) {
|
||||
stopCurrent();
|
||||
chunkCount = 0;
|
||||
totalBytes = 0;
|
||||
done = false;
|
||||
const logNode = el(logId);
|
||||
if (logNode) logNode.textContent = "Starting...";
|
||||
setStatus("Starting...");
|
||||
|
||||
let config;
|
||||
try {
|
||||
config = JSON.parse(configJson);
|
||||
} catch (err) {
|
||||
setStatus("Invalid request");
|
||||
setStartButton(false, "Restart");
|
||||
log("Invalid request config: " + err.message);
|
||||
return configJson;
|
||||
}
|
||||
|
||||
const mime = chooseMime();
|
||||
if (!mime) {
|
||||
setStatus("MSE unsupported");
|
||||
setStartButton(false, "Restart");
|
||||
log("This browser does not support MP4 Media Source Extensions playback.");
|
||||
return configJson;
|
||||
}
|
||||
|
||||
const video = el(playerId);
|
||||
if (!video) return configJson;
|
||||
mediaSource = new MediaSource();
|
||||
objectUrl = URL.createObjectURL(mediaSource);
|
||||
video.src = objectUrl;
|
||||
|
||||
mediaSource.addEventListener("sourceopen", () => {
|
||||
try {
|
||||
sourceBuffer = mediaSource.addSourceBuffer(mime);
|
||||
sourceBuffer.mode = "segments";
|
||||
sourceBuffer.addEventListener("updateend", () => {
|
||||
pump();
|
||||
if (done) finishStream();
|
||||
});
|
||||
} catch (err) {
|
||||
setStatus("SourceBuffer error");
|
||||
setStartButton(false, "Restart");
|
||||
log("SourceBuffer error: " + err.message);
|
||||
return;
|
||||
}
|
||||
|
||||
ws = new WebSocket(config.url);
|
||||
ws.binaryType = "arraybuffer";
|
||||
ws.onopen = () => {
|
||||
setStatus("Streaming...");
|
||||
setStartButton(true, "Streaming...");
|
||||
log("Connected: " + config.url);
|
||||
ws.send(JSON.stringify(config.payload));
|
||||
log("Sent session.start: " + JSON.stringify(config.payload));
|
||||
};
|
||||
ws.onmessage = (event) => {
|
||||
if (typeof event.data === "string") {
|
||||
const msg = JSON.parse(event.data);
|
||||
if (msg.type === "video.start") {
|
||||
log("Video session started: request_id=" + (msg.request_id || "") + " format=" + (msg.format || ""));
|
||||
} else if (msg.type === "session.done") {
|
||||
done = true;
|
||||
setStatus("Done");
|
||||
setStartButton(false, "Restart");
|
||||
log("Session complete: " + JSON.stringify(msg));
|
||||
finishStream();
|
||||
} else if (msg.type === "error") {
|
||||
done = true;
|
||||
setStatus("Error");
|
||||
setStartButton(false, "Restart");
|
||||
log("ERROR: " + (msg.message || JSON.stringify(msg)));
|
||||
finishStream();
|
||||
} else {
|
||||
log("Control message: " + JSON.stringify(msg));
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const data = event.data;
|
||||
queue.push(data);
|
||||
chunkCount += 1;
|
||||
totalBytes += data.byteLength;
|
||||
log(`[chunk ${String(chunkCount).padStart(3, "0")}] bytes=${data.byteLength} total_bytes=${totalBytes}`);
|
||||
pump();
|
||||
};
|
||||
ws.onerror = () => {
|
||||
done = true;
|
||||
setStatus("Error");
|
||||
setStartButton(false, "Restart");
|
||||
log("ERROR: WebSocket error");
|
||||
finishStream();
|
||||
};
|
||||
ws.onclose = () => {
|
||||
if (!done) {
|
||||
done = true;
|
||||
setStatus("Closed");
|
||||
setStartButton(false, "Restart");
|
||||
log("WebSocket closed before session.done");
|
||||
finishStream();
|
||||
}
|
||||
};
|
||||
}, { once: true });
|
||||
|
||||
return configJson;
|
||||
};
|
||||
|
||||
console.log("video-stream-view.html loaded");
|
||||
})();
|
||||
@@ -16,6 +16,7 @@ For the full list of supported architectures across all modalities, see
|
||||
|---|---|---|---|---|---|
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
|
||||
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | compat only, non-chunk | uploaded audio voice only; no presets | — |
|
||||
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | ✓ (`ref_audio` / `speaker_embedding`) | ✓ (PCM stream) | IP labels + structured `instructions` | — |
|
||||
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | — (caption-controlled) | — | caption fields (`instructions`) | — |
|
||||
| MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ |
|
||||
@@ -23,6 +24,7 @@ For the full list of supported architectures across all modalities, see
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
|
||||
| SoulX-Singer | `Soul-AILab/SoulX-Singer` | ✓ (prompt audio) | — (batch only) | — (prompt + target audio) | — (chat client) |
|
||||
|
||||
CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](../../offline_inference/text_to_speech/README.md#cosyvoice3) instead.
|
||||
|
||||
@@ -34,7 +36,7 @@ Launch the server (defaults shown — adjust `--port`, `--gpu-memory-utilization
|
||||
vllm serve <hf-repo-or-local-path> --omni --port 8091
|
||||
```
|
||||
|
||||
Send a TTS request via curl:
|
||||
Send a TTS request via curl. These generic snippets assume a model with a preset/default voice; voice-cloning-only models such as IndexTTS-2 require `ref_audio` or an uploaded audio `voice` (see model-specific sections below).
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
@@ -90,7 +92,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
|
||||
```
|
||||
|
||||
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 24 kHz for the others).
|
||||
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 22.05 kHz for IndexTTS-2, and 24 kHz for many others).
|
||||
|
||||
For full request-shape documentation (all parameters, response formats, error codes), see the [Speech API reference](../../../docs/serving/speech_api.md).
|
||||
|
||||
@@ -135,6 +137,40 @@ bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
|
||||
|
||||
---
|
||||
|
||||
## IndexTTS-2
|
||||
|
||||
2-stage TTS (GPT AR + S2Mel CFM DiT + BigVGAN) at 22.05 kHz. Requests use `ref_audio` for voice cloning, or an uploaded audio `voice` from `/v1/audio/voices`. Supports emotion conditioning via `emo_audio`, `emo_text`, or `emo_vector` passed in `extra_params`.
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092
|
||||
# or, to pass the bundled deploy config explicitly:
|
||||
bash examples/online_serving/text_to_speech/indextts2/run_server.sh
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
```bash
|
||||
# Voice cloning (ref_audio required)
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--text "你好,世界!" \
|
||||
--ref-audio /path/to/reference.wav
|
||||
|
||||
# With emotion audio
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--text "今天心情很好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-audio /path/to/happy.wav
|
||||
```
|
||||
|
||||
### Notes
|
||||
- Output: 22.05 kHz mono WAV.
|
||||
- Provide `ref_audio` on the documented raw request path, or pass `voice` only when it names an uploaded audio voice; IndexTTS-2 does not provide a built-in text-only preset voice.
|
||||
- Emotion params (`emo_audio`, `emo_text`, `emo_vector`, `emo_alpha`, `use_emo_text`, `use_random`) are passed via the `extra_params` field. Official precedence is `use_emo_text` > `emo_vector` > `emo_audio` > same emotion as the speaker reference.
|
||||
- `stream=true` is accepted as an OpenAI-compatible response path, but IndexTTS-2 is not async-chunk streaming; audio is produced after S2Mel receives the full mel-code sequence.
|
||||
- Deploy config: `vllm_omni/deploy/indextts2.yaml` (auto-loaded).
|
||||
|
||||
---
|
||||
|
||||
## Fish Speech S2 Pro
|
||||
|
||||
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
|
||||
@@ -626,3 +662,65 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
|
||||
- Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
|
||||
- Voice cloning is gated upstream and may require a recent `mistral_common`.
|
||||
- A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now.
|
||||
|
||||
---
|
||||
|
||||
## SoulX-Singer
|
||||
|
||||
Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Single-stage DiT with inline preprocess. Uses the `/v1/chat/completions` endpoint with multimodal input (`prompt_audio` + `target_audio`).
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Download DiT and preprocess weights, then set up separate SVS / SVC view directories and install dependencies as described in the [offline README](../../offline_inference/text_to_speech/README.md#soulx-singer). `config.json` `architectures` field is the single source of truth for SVS vs SVC — point `MODEL` at the matching directory.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
# SVS (default)
|
||||
export MODEL=/path/to/SoulX-Singer
|
||||
export PREPROCESS=/path/to/SoulX-Singer-Preprocess
|
||||
bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh
|
||||
|
||||
# SVC
|
||||
export MODE=svc
|
||||
export MODEL=/path/to/SoulX-Singer-svc
|
||||
bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh
|
||||
```
|
||||
|
||||
Or equivalently, set `SOULX_PREPROCESS_WEIGHTS_DIR` and launch directly:
|
||||
```bash
|
||||
export SOULX_PREPROCESS_WEIGHTS_DIR=$PREPROCESS
|
||||
vllm serve $MODEL --omni \
|
||||
--deploy-config vllm_omni/deploy/soulxsinger_${MODE}.yaml \
|
||||
--port 8192 --trust-remote-code --enforce-eager
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
|
||||
Audio paths must be reachable from the server host (local filesystem or data URL). The client sends prompt vocal via `input_audio` and target accompaniment via `extra_args['target_audio']`.
|
||||
|
||||
```bash
|
||||
# Default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3
|
||||
python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \
|
||||
--prompt-audio /path/on/server/zh_prompt.mp3 \
|
||||
--target-audio /path/on/server/music.mp3 \
|
||||
--preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \
|
||||
-o output.wav
|
||||
```
|
||||
|
||||
Use precomputed metadata to skip online preprocess with following command:
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \
|
||||
--prompt-metadata-path /path/on/server/zh_prompt.json \
|
||||
--target-metadata-path /path/on/server/music.json \
|
||||
--audio-path /path/on/server/zh_prompt.mp3 \
|
||||
-o output.wav
|
||||
```
|
||||
|
||||
`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. See `openai_chat_client.py --help` for `--vocal-sep`, `--language`, `--num-inference-steps`, `--guidance-scale`, and `--seed`.
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 24 kHz mono WAV; batch only.
|
||||
- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on.
|
||||
- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only.
|
||||
@@ -0,0 +1,25 @@
|
||||
#!/bin/bash
|
||||
# Launch vLLM-Omni server for IndexTTS2
|
||||
#
|
||||
# Usage from repository root:
|
||||
# examples/online_serving/text_to_speech/indextts2/run_server.sh
|
||||
# CUDA_VISIBLE_DEVICES=0 PORT=8092 MODEL=/path/to/IndexTeam/IndexTTS-2 examples/online_serving/text_to_speech/indextts2/run_server.sh
|
||||
|
||||
set -e
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
ROOT_DIR="$(cd -- "$SCRIPT_DIR/../../../.." && pwd)"
|
||||
|
||||
MODEL="${MODEL:-IndexTeam/IndexTTS-2}"
|
||||
PORT="${PORT:-8092}"
|
||||
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$ROOT_DIR/vllm_omni/deploy/indextts2.yaml}"
|
||||
|
||||
echo "Starting IndexTTS2 server with model: $MODEL"
|
||||
|
||||
FLASHINFER_DISABLE_VERSION_CHECK=1 \
|
||||
vllm serve "$MODEL" \
|
||||
--host 0.0.0.0 \
|
||||
--port "$PORT" \
|
||||
--omni \
|
||||
--trust-remote-code \
|
||||
--deploy-config "$DEPLOY_CONFIG"
|
||||
@@ -0,0 +1,126 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""OpenAI-compatible client for IndexTTS2 TTS via /v1/audio/speech endpoint.
|
||||
|
||||
Examples:
|
||||
# With reference audio for voice cloning
|
||||
python speech_client.py --text "你好,世界!" \
|
||||
--ref-audio /path/to/reference.wav
|
||||
|
||||
# With emotion audio
|
||||
python speech_client.py --text "今天心情很好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-audio /path/to/happy.wav
|
||||
|
||||
Server setup:
|
||||
vllm serve IndexTeam/IndexTTS-2 --omni --host 0.0.0.0 --port 8092
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import os
|
||||
|
||||
import httpx
|
||||
|
||||
DEFAULT_API_BASE = "http://localhost:8092"
|
||||
DEFAULT_API_KEY = "sk-empty"
|
||||
|
||||
|
||||
def encode_audio_to_base64(audio_path: str) -> str:
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
ext = audio_path.lower().rsplit(".", 1)[-1]
|
||||
mime = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac"}.get(ext, "audio/wav")
|
||||
with open(audio_path, "rb") as f:
|
||||
b64 = base64.b64encode(f.read()).decode("utf-8")
|
||||
return f"data:{mime};base64,{b64}"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="IndexTTS2 OpenAI speech client")
|
||||
parser.add_argument("--text", type=str, required=True)
|
||||
parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio for voice cloning")
|
||||
parser.add_argument("--emo-audio", type=str, default=None, help="Emotion reference audio")
|
||||
parser.add_argument("--emo-text", type=str, default=None, help="Emotion description text")
|
||||
parser.add_argument(
|
||||
"--emo-vector",
|
||||
type=float,
|
||||
nargs=8,
|
||||
default=None,
|
||||
help="8-dim emotion vector: happy angry sad afraid disgusted melancholic surprised calm",
|
||||
)
|
||||
parser.add_argument("--emo-alpha", type=float, default=None, help="Emotion weight in [0, 1]")
|
||||
parser.add_argument("--use-emo-text", action="store_true", help="Infer emotion vector from emo-text or text")
|
||||
parser.add_argument("--use-random", action="store_true", help="Use random emotion prototypes")
|
||||
parser.add_argument("--model", type=str, default="IndexTeam/IndexTTS-2")
|
||||
parser.add_argument("--voice", type=str, default=None, help="Uploaded voice name to use instead of --ref-audio")
|
||||
parser.add_argument("--output", type=str, default="output.wav")
|
||||
parser.add_argument("--api-base", type=str, default=DEFAULT_API_BASE)
|
||||
parser.add_argument("--api-key", type=str, default=DEFAULT_API_KEY)
|
||||
parser.add_argument("--response-format", type=str, default="wav")
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.ref_audio and not args.voice:
|
||||
parser.error("IndexTTS2 requires --ref-audio or --voice for voice cloning")
|
||||
|
||||
payload: dict = {
|
||||
"model": args.model,
|
||||
"input": args.text,
|
||||
"response_format": args.response_format,
|
||||
}
|
||||
if args.voice:
|
||||
payload["voice"] = args.voice
|
||||
|
||||
if args.ref_audio:
|
||||
ref = args.ref_audio
|
||||
if ref.startswith(("http://", "https://", "data:")):
|
||||
payload["ref_audio"] = ref
|
||||
else:
|
||||
payload["ref_audio"] = encode_audio_to_base64(ref)
|
||||
|
||||
extra_params = {}
|
||||
if args.emo_audio:
|
||||
emo = args.emo_audio
|
||||
if emo.startswith(("http://", "https://", "data:")):
|
||||
extra_params["emo_audio"] = emo
|
||||
else:
|
||||
extra_params["emo_audio"] = encode_audio_to_base64(emo)
|
||||
if args.emo_text:
|
||||
extra_params["emo_text"] = args.emo_text
|
||||
if args.emo_vector is not None:
|
||||
extra_params["emo_vector"] = args.emo_vector
|
||||
if args.emo_alpha is not None:
|
||||
extra_params["emo_alpha"] = args.emo_alpha
|
||||
if args.use_emo_text:
|
||||
extra_params["use_emo_text"] = True
|
||||
if args.use_random:
|
||||
extra_params["use_random"] = True
|
||||
if extra_params:
|
||||
payload["extra_params"] = extra_params
|
||||
|
||||
url = f"{args.api_base}/v1/audio/speech"
|
||||
print(f"POST {url}")
|
||||
print(f" text: {args.text}")
|
||||
if args.ref_audio:
|
||||
print(f" ref_audio: {args.ref_audio[:80]}...")
|
||||
|
||||
with httpx.Client(timeout=300) as client:
|
||||
resp = client.post(
|
||||
url,
|
||||
json=payload,
|
||||
headers={"Authorization": f"Bearer {args.api_key}"},
|
||||
)
|
||||
|
||||
if resp.status_code != 200:
|
||||
print(f"Error {resp.status_code}: {resp.text[:500]}")
|
||||
return
|
||||
|
||||
with open(args.output, "wb") as f:
|
||||
f.write(resp.content)
|
||||
print(f"Saved: {args.output} ({len(resp.content):,} bytes)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,186 @@
|
||||
#!/usr/bin/env python3
|
||||
"""SoulX-Singer OpenAI-compatible chat client (SVS / SVC).
|
||||
|
||||
Sends prompt audio via ``input_audio`` and target accompaniment via
|
||||
``extra_args['target_audio']`` (server-local path). For integrated preprocess,
|
||||
also pass ``preprocess_weights_dir`` in ``extra_args``.
|
||||
|
||||
Usage:
|
||||
python openai_chat_client.py \\
|
||||
--prompt-audio /path/on/server/zh_prompt.mp3 \\
|
||||
--target-audio /path/on/server/music.mp3 \\
|
||||
--preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \\
|
||||
-o output.wav
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import io
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
import soundfile
|
||||
import torch
|
||||
|
||||
|
||||
def _audio_to_data_url(path: Path) -> str:
|
||||
with path.open("rb") as handle:
|
||||
data = base64.b64encode(handle.read()).decode("ascii")
|
||||
return f"data:audio/mpeg;base64,{data}"
|
||||
|
||||
|
||||
def _save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None:
|
||||
audio = audio.to(torch.float32)
|
||||
peak = audio.abs().max().clamp(min=1e-8)
|
||||
audio = audio / peak
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16")
|
||||
|
||||
|
||||
def _decode_audio_from_response(body: dict) -> tuple[torch.Tensor, int]:
|
||||
for choice in body.get("choices", []):
|
||||
audio_obj = choice.get("message", {}).get("audio")
|
||||
if isinstance(audio_obj, dict) and audio_obj.get("data"):
|
||||
data, sr = soundfile.read(
|
||||
io.BytesIO(base64.b64decode(audio_obj["data"])),
|
||||
dtype="float32",
|
||||
always_2d=True,
|
||||
)
|
||||
return torch.from_numpy(data).transpose(0, 1), sr
|
||||
brief = {k: v for k, v in body.items() if k != "choices"}
|
||||
raise RuntimeError(f"no audio in response message.audio: {brief}")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
repo_root = Path(__file__).resolve().parents[4]
|
||||
default_assets = repo_root / "tests" / "assets" / "soulxsinger"
|
||||
|
||||
parser = argparse.ArgumentParser(description="SoulX-Singer online chat client")
|
||||
parser.add_argument("--port", type=int, default=8192)
|
||||
parser.add_argument("--model", default="Soul-AILab/SoulX-Singer")
|
||||
parser.add_argument(
|
||||
"--prompt-audio",
|
||||
default=str(default_assets / "zh_prompt.mp3"),
|
||||
help="Prompt vocal audio (path on server if using extra_args, or local for input_audio)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-audio",
|
||||
default=str(default_assets / "music.mp3"),
|
||||
help="Target accompaniment path on the server (extra_args['target_audio'])",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-metadata-path",
|
||||
default=None,
|
||||
help="SVS precomputed prompt metadata.json",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-metadata-path",
|
||||
default=None,
|
||||
help="SVS precomputed target metadata.json",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-path",
|
||||
default=None,
|
||||
help="SVS prompt vocal wav for precomputed metadata",
|
||||
)
|
||||
parser.add_argument("--preprocess-weights-dir", default=None)
|
||||
parser.add_argument("--output", "-o", default="soulxsinger_out.wav")
|
||||
parser.add_argument("--svc", action="store_true", help="Use SVC mode knobs")
|
||||
parser.add_argument("--language", default="Mandarin")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=32)
|
||||
parser.add_argument("--guidance-scale", type=float, default=3.0)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=42,
|
||||
help="Optional CFM seed. Omit for non-deterministic sampling.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auto-shift",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Auto pitch shift (default: on, original upstream infer.sh)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--control",
|
||||
default="melody",
|
||||
choices=["melody", "score"],
|
||||
help="SVS control mode",
|
||||
)
|
||||
parser.add_argument("--vocal-sep", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
meta_paths = (args.prompt_metadata_path, args.target_metadata_path, args.audio_path)
|
||||
if any(meta_paths) and not all(meta_paths):
|
||||
print(
|
||||
"ERROR: precomputed metadata requires --prompt-metadata-path, "
|
||||
"--target-metadata-path, and --audio-path together.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 2
|
||||
|
||||
extra_args: dict = {
|
||||
"vocal_sep": args.vocal_sep,
|
||||
"auto_shift": args.auto_shift,
|
||||
"pitch_shift": 0,
|
||||
}
|
||||
if all(meta_paths):
|
||||
extra_args.update(
|
||||
{
|
||||
"prompt_metadata_path": str(Path(args.prompt_metadata_path).expanduser().resolve()),
|
||||
"target_metadata_path": str(Path(args.target_metadata_path).expanduser().resolve()),
|
||||
"audio_path": str(Path(args.audio_path).expanduser().resolve()),
|
||||
}
|
||||
)
|
||||
content = [{"type": "text", "text": "soulx-singer"}]
|
||||
else:
|
||||
prompt_path = Path(args.prompt_audio).expanduser().resolve()
|
||||
if not prompt_path.is_file():
|
||||
print(f"ERROR: prompt audio not found: {prompt_path}", file=sys.stderr)
|
||||
return 2
|
||||
extra_args["prompt_audio"] = str(prompt_path)
|
||||
extra_args["target_audio"] = str(Path(args.target_audio).expanduser().resolve())
|
||||
if args.preprocess_weights_dir:
|
||||
extra_args["preprocess_weights_dir"] = str(Path(args.preprocess_weights_dir).expanduser().resolve())
|
||||
content = [
|
||||
{"type": "text", "text": "soulx-singer"},
|
||||
{
|
||||
"type": "input_audio",
|
||||
"input_audio": {"data": _audio_to_data_url(prompt_path), "format": "mp3"},
|
||||
},
|
||||
]
|
||||
if not args.svc:
|
||||
extra_args["language"] = args.language
|
||||
extra_args["control"] = args.control
|
||||
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"modalities": ["audio"],
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"guidance_scale": args.guidance_scale,
|
||||
"extra_args": extra_args,
|
||||
}
|
||||
if args.seed is not None:
|
||||
payload["seed"] = args.seed
|
||||
|
||||
print(f"POST http://localhost:{args.port}/v1/chat/completions")
|
||||
response = requests.post(
|
||||
f"http://localhost:{args.port}/v1/chat/completions",
|
||||
headers={"Content-Type": "application/json"},
|
||||
json=payload,
|
||||
timeout=1800,
|
||||
)
|
||||
response.raise_for_status()
|
||||
audio, sample_rate = _decode_audio_from_response(response.json())
|
||||
_save_wav(audio, Path(args.output), sample_rate)
|
||||
duration = audio.shape[-1] / sample_rate
|
||||
print(f"saved {args.output} sr={sample_rate}Hz duration={duration:.2f}s")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,39 @@
|
||||
#!/bin/bash
|
||||
# Launch vLLM-Omni server for SoulX-Singer (single-stage DiT, preprocess inline).
|
||||
#
|
||||
# Usage:
|
||||
# MODEL=/path/to/SoulX-Singer PREPROCESS=/path/to/Preprocess \
|
||||
#
|
||||
# Audio paths in client extra_args must be readable on the server host.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)"
|
||||
|
||||
MODEL="${MODEL:-Soul-AILab/SoulX-Singer}"
|
||||
MODE="${MODE:-svs}"
|
||||
PORT="${PORT:-8192}"
|
||||
GPUS="${GPUS:-0}"
|
||||
|
||||
if [[ "$MODE" == "svc" ]]; then
|
||||
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svc.yaml}"
|
||||
else
|
||||
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svs.yaml}"
|
||||
fi
|
||||
|
||||
echo "Starting SoulX-Singer server"
|
||||
echo " MODEL=$MODEL"
|
||||
echo " MODE=$MODE"
|
||||
echo " PORT=$PORT"
|
||||
echo " DEPLOY_CONFIG=$DEPLOY_CONFIG"
|
||||
echo " CUDA_VISIBLE_DEVICES=$GPUS"
|
||||
|
||||
CUDA_VISIBLE_DEVICES="$GPUS" \
|
||||
vllm serve "$MODEL" \
|
||||
--omni \
|
||||
--deploy-config "$DEPLOY_CONFIG" \
|
||||
--host 0.0.0.0 \
|
||||
--port "$PORT" \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
@@ -40,6 +40,7 @@ dev = [
|
||||
"pytest-asyncio>=0.21.0",
|
||||
"pytest-cov>=4.0.0",
|
||||
"pytest-mock>=3.10.0",
|
||||
"websockets>=15.0.1",
|
||||
"datasets>=2.14.0",
|
||||
"mypy==1.11.1",
|
||||
"pre-commit==4.0.1",
|
||||
@@ -118,6 +119,12 @@ docs = [
|
||||
"pydantic",
|
||||
]
|
||||
|
||||
soulx-svs = [
|
||||
"funasr>=1.0.0",
|
||||
"g2pM>=0.1.2.5",
|
||||
"g2p-en>=2.1.0",
|
||||
"ToJyutping>=3.2.0",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://github.com/vllm-project/vllm-omni"
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
# IndexTTS-2 for voice-cloned TTS on 1x GPU
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: IndexTeam
|
||||
- Model: `IndexTeam/IndexTTS-2`
|
||||
- Task: Text-to-speech with voice cloning and optional emotion control
|
||||
- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe to serve IndexTTS-2 as a two-stage TTS system for
|
||||
voice-cloned speech synthesis. Each request provides synthesis text plus either
|
||||
a reference audio clip or an uploaded audio voice; the model returns 22.05 kHz
|
||||
mono speech. IndexTTS-2 also supports optional emotion conditioning through
|
||||
model-specific `extra_params`.
|
||||
|
||||
IndexTTS-2 is **not an async-chunk streaming model** in the bundled deployment
|
||||
configuration: the S2Mel flow-matching stage consumes the full mel-code sequence
|
||||
from the AR talker before it produces audio. The OpenAI-compatible endpoint can
|
||||
still be called with `stream=true` as an HTTP compatibility path, but that should
|
||||
not be interpreted as incremental model streaming.
|
||||
|
||||
## References
|
||||
|
||||
- Upstream model:
|
||||
[IndexTeam/IndexTTS-2 on Hugging Face](https://huggingface.co/IndexTeam/IndexTTS-2)
|
||||
- Online serving example:
|
||||
[`examples/online_serving/text_to_speech/README.md#indextts-2`](../../examples/online_serving/text_to_speech/README.md#indextts-2)
|
||||
- Offline inference example:
|
||||
[`examples/offline_inference/text_to_speech/README.md#indextts-2`](../../examples/offline_inference/text_to_speech/README.md#indextts-2)
|
||||
- OpenAI-compatible client:
|
||||
[`examples/online_serving/text_to_speech/indextts2/speech_client.py`](../../examples/online_serving/text_to_speech/indextts2/speech_client.py)
|
||||
- Deploy config:
|
||||
[`vllm_omni/deploy/indextts2.yaml`](../../vllm_omni/deploy/indextts2.yaml)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
### GPU
|
||||
|
||||
### 1x L4 24GB or larger CUDA GPU
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.10+
|
||||
- Driver / runtime: NVIDIA CUDA environment
|
||||
- vLLM version: Match the repository requirements for your checkout
|
||||
- vLLM-Omni version or commit: Use the commit you are deploying from
|
||||
|
||||
#### Command
|
||||
|
||||
Start the server from the repository root:
|
||||
|
||||
```bash
|
||||
vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092
|
||||
```
|
||||
|
||||
The bundled deployment config is selected automatically for the `indextts2`
|
||||
model type. To make the config choice explicit, pass it directly:
|
||||
|
||||
```bash
|
||||
vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092 \
|
||||
--deploy-config vllm_omni/deploy/indextts2.yaml
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
Run the bundled client. Local audio paths passed to this client are converted to
|
||||
base64 data URLs before they are sent to the HTTP API:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--api-base http://localhost:8092 \
|
||||
--text "你好,这是 IndexTTS-2 语音合成测试。" \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--output indextts2.wav
|
||||
```
|
||||
|
||||
For a raw HTTP request, pass `ref_audio` as an HTTP(S) URL, a base64 data URL,
|
||||
or a `file://` URI allowed by the server's `--allowed-local-media-path`. Do not
|
||||
pass an arbitrary local filesystem path directly in JSON unless your client
|
||||
converts it first:
|
||||
|
||||
```bash
|
||||
curl http://localhost:8092/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "IndexTeam/IndexTTS-2",
|
||||
"input": "你好,这是 IndexTTS-2 语音合成测试。",
|
||||
"response_format": "wav",
|
||||
"ref_audio": "data:audio/wav;base64,<BASE64_ENCODED_AUDIO>"
|
||||
}' \
|
||||
--output indextts2.wav
|
||||
```
|
||||
|
||||
Optional emotion conditioning is passed through `extra_params`. Use one emotion
|
||||
source mode at a time:
|
||||
|
||||
```bash
|
||||
curl http://localhost:8092/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "IndexTeam/IndexTTS-2",
|
||||
"input": "今天心情很好!",
|
||||
"response_format": "wav",
|
||||
"ref_audio": "data:audio/wav;base64,<BASE64_ENCODED_REFERENCE_AUDIO>",
|
||||
"extra_params": {
|
||||
"emo_vector": [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
|
||||
"emo_alpha": 0.8
|
||||
}
|
||||
}' \
|
||||
--output indextts2_happy.wav
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- Hardware scope: the default config is a single-GPU CUDA deployment. Stage 0
|
||||
and Stage 1 share GPU 0 by default.
|
||||
- Audio output: 22.05 kHz mono WAV.
|
||||
- Voice cloning: `ref_audio` is required for the documented raw request path.
|
||||
Alternatively, pass `voice` only when it names an uploaded audio voice from
|
||||
`/v1/audio/voices`; IndexTTS-2 does not provide a built-in text-only preset
|
||||
voice in this serving recipe.
|
||||
- Reference audio format: use an HTTP(S) URL, a base64 data URL, or a `file://`
|
||||
URI allowed by the server's `--allowed-local-media-path` for raw API calls.
|
||||
The bundled Python client accepts local paths and converts them to base64 data
|
||||
URLs.
|
||||
- Emotion control: `extra_params.emo_audio`, `extra_params.emo_vector`, and
|
||||
`extra_params.use_emo_text` are alternative conditioning modes. Official
|
||||
precedence is `use_emo_text` > `emo_vector` > `emo_audio` > same emotion as
|
||||
the speaker reference. The `emo_vector` order is `[happy, angry, sad, afraid,
|
||||
disgusted, melancholic, surprised, calm]`.
|
||||
- Streaming limitation: `async_chunk` is disabled in
|
||||
`vllm_omni/deploy/indextts2.yaml`; S2Mel requires the full mel-code sequence
|
||||
before audio generation.
|
||||
- First request can be slower because reference encoders and auxiliary models
|
||||
are lazily loaded and cached.
|
||||
@@ -5,7 +5,7 @@
|
||||
## Summary
|
||||
|
||||
- Vendor: JD (Joy Future Academy)
|
||||
- Model: [`jdopensource/JoyAI-VL-Interaction-Preview`](https://huggingface.co/jdopensource/JoyAI-VL-Interaction-Preview) (8B, vanilla Qwen3-VL weights)
|
||||
- Model: [`jdopensource/JoyAI-VL-Interaction-Preview`](https://huggingface.co/jdopensource/JoyAI-VL-Interaction-Preview) (8B, Qwen3-VL architecture with weights retrained for streaming interaction)
|
||||
- Task: Per-tick proactive interaction over a live video stream — the model decides
|
||||
on its own each second to speak, stay silent, or delegate a hard question
|
||||
- Mode: Online serving — an OpenAI-compatible interaction orchestrator in front of
|
||||
@@ -25,7 +25,8 @@ framework internals and how to add another full-duplex model, see
|
||||
- OS: Linux
|
||||
- Python: 3.10+
|
||||
- Hardware: 1x GPU. The default `T_s=100` frame window + `--max-model-len 131072` wants
|
||||
≈48GB+; for ~24GB, lower `--chunk-frames`, `--max-model-len`, and the image limit together
|
||||
≈48GB+; for ~24GB, lower `--chunk-frames`, `--max-model-len`, and the image limit together,
|
||||
and/or load the weights in fp8 (see "Smaller GPU: fp8 weights" below — ~9.9 GiB vs 16.8 GiB)
|
||||
- vLLM / vLLM-Omni: versions from your current checkout
|
||||
|
||||
## Start server
|
||||
@@ -33,7 +34,7 @@ framework internals and how to add another full-duplex model, see
|
||||
From repository root:
|
||||
|
||||
```bash
|
||||
# 1. Serve the model (plain `vllm serve`, NOT --omni; it is vanilla Qwen3-VL).
|
||||
# 1. Serve the model (plain `vllm serve`, NOT --omni; it uses the Qwen3-VL architecture).
|
||||
# The image limit must cover the short-term frame window (chunk_frames, default 100 = T_s);
|
||||
# prefix caching keeps the accumulating window cheap. Lower both for smaller GPUs.
|
||||
vllm serve jdopensource/JoyAI-VL-Interaction-Preview \
|
||||
@@ -46,6 +47,26 @@ python -m vllm_omni.experimental.fullduplex.joyvl.serving.server --port 8070 \
|
||||
--main-backend-url http://127.0.0.1:8061/v1 --main-model JoyAI-VL-Interaction-Preview
|
||||
```
|
||||
|
||||
### Smaller GPU: fp8 weights
|
||||
|
||||
Add `--quantization fp8` to the step-1 `vllm serve` to load the weights in fp8 online
|
||||
from the bf16 checkpoint (no separate quantized checkpoint needed; uses vLLM's generic
|
||||
fp8 path, since the model is a standard Qwen3-VL VLM):
|
||||
|
||||
```bash
|
||||
vllm serve jdopensource/JoyAI-VL-Interaction-Preview \
|
||||
--served-model-name JoyAI-VL-Interaction-Preview --port 8061 \
|
||||
--quantization fp8 --max-model-len 131072 --enable-prefix-caching \
|
||||
--limit-mm-per-prompt '{"image":256,"video":1}'
|
||||
```
|
||||
|
||||
Measured here: model weights **9.9 GiB** vs **16.8 GiB** bf16 (**−41%**), which lets the 8B
|
||||
fit a much smaller card. This is a **memory** option, not a speed one: at one request per
|
||||
tick (the streaming regime) fp8 was **not** faster (≈19% slower in our test) because
|
||||
single-stream decode is memory-latency-bound and pays fp8's dynamic-scale overhead; its
|
||||
throughput gain needs batching. Decisions stay coherent, but fp8 can shift the
|
||||
speak/silence boundary, so for timing-critical alerting validate output against bf16 first.
|
||||
|
||||
Optional one-shot launch (model + orchestrator + JD webui + ASR/TTS/background,
|
||||
all env-configurable). The JD webui frontend is external — set `WEBUI_DIR` to point at it:
|
||||
|
||||
@@ -176,14 +197,22 @@ the audio-track caveat.
|
||||
|
||||
## Notes
|
||||
|
||||
- `--omni` is **not** used: the model is standard Qwen3-VL, so stock `vllm serve`
|
||||
runs the forward pass; this recipe only adds the interaction/serving layer.
|
||||
- `--omni` is **not** used: the model keeps the Qwen3-VL architecture (only the
|
||||
weights are retrained), so stock `vllm serve` runs the forward pass; this recipe
|
||||
only adds the interaction/serving layer.
|
||||
- On a host without `nvcc` / `ninja`, `vllm serve` of the 8B can crash engine-core in the
|
||||
FlashInfer sampler JIT (`FileNotFoundError: 'ninja'`) during `profile_run`. Set
|
||||
`VLLM_USE_FLASHINFER_SAMPLER=0` (or install `ninja`) to work around it.
|
||||
- `force_silence_before_query` is on by default — the model stays silent until an
|
||||
instruction arrives; give a standing task (e.g. "translate the on-screen text")
|
||||
to arm proactive output.
|
||||
- **Frame resolution is the main latency knob.** The system/memory prefix is prefix-cached,
|
||||
so per-tick *new* compute is dominated by the new frame's vision tokens. Measured on the
|
||||
8B: ~256×192 frames hit Qwen3-VL's min-pixel floor (~72 vision tokens, ~17 ms/tick) vs
|
||||
~302 tokens / ~38 ms at 640×480 — about 2× cheaper. Downsample frames to ~256×192 for the
|
||||
tightest latency / highest concurrency; one GPU then sustains ~150–180 concurrent 1 fps
|
||||
streams with p95 < 200 ms. Per-tick latency is already far inside the 1 fps budget, so
|
||||
serving is rarely the bottleneck — resolution is the lever if you need more headroom.
|
||||
- Speech is external and pluggable: point `ASR_URL` / `TTS_URL` at the bridges in
|
||||
`examples/online_serving/joyvl_interaction/bridges/` or any compatible service.
|
||||
- The decision prompts, sampling, and 3-tier summary memory (`T_s=100`, mid→long every
|
||||
|
||||
@@ -43,6 +43,7 @@ recipes/
|
||||
| [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
|
||||
| [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
|
||||
| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
|
||||
| [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
|
||||
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
|
||||
| [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
|
||||
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
|
||||
@@ -57,10 +58,12 @@ recipes/
|
||||
| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
|
||||
| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB |
|
||||
| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) |
|
||||
| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB |
|
||||
| [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
|
||||
| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
|
||||
| [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
|
||||
| [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
|
||||
| [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) |
|
||||
| [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
|
||||
| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
|
||||
| [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
|
||||
|
||||
@@ -0,0 +1,170 @@
|
||||
# Stable-Audio-Open Text-To-Audio Generation on 1x GPU
|
||||
|
||||
> Text-to-audio recipe for Stable Audio Open with offline inference and
|
||||
> OpenAI-compatible online serving on 1x RTX 4090 24GB.
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Stability AI
|
||||
- Model: `stabilityai/stable-audio-open-1.0`
|
||||
- Task: Text-to-audio generation
|
||||
- Mode: Offline inference and online serving
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when you want to run Stable Audio Open on a single RTX 4090
|
||||
24GB GPU for music or sound-effect generation. The recipe covers a 10-second
|
||||
offline validation sample with TeaCache and online serving through the
|
||||
`/v1/audio/generate` endpoint.
|
||||
|
||||
## References
|
||||
|
||||
- Model: <https://huggingface.co/stabilityai/stable-audio-open-1.0>
|
||||
- Offline example:
|
||||
[`examples/offline_inference/text_to_audio`](../../examples/offline_inference/text_to_audio)
|
||||
- Online example:
|
||||
[`examples/online_serving/stable_audio`](../../examples/online_serving/stable_audio)
|
||||
- Audio generation API:
|
||||
[`docs/serving/audio_generate_api.md`](../../docs/serving/audio_generate_api.md)
|
||||
- Related issue:
|
||||
[#2645](https://github.com/vllm-project/vllm-omni/issues/2645)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
### GPU
|
||||
|
||||
#### 1x RTX 4090 24GB
|
||||
|
||||
##### Environment
|
||||
|
||||
- OS: Ubuntu 22.04.5
|
||||
- Python: 3.12
|
||||
- GPU: NVIDIA GeForce RTX 4090, 24564 MiB VRAM
|
||||
- Driver / runtime: NVIDIA driver 595.80, CUDA-capable runtime matching the
|
||||
repository build
|
||||
- vLLM version: 0.22.0
|
||||
- vLLM-Omni version: source checkout
|
||||
- PyTorch: 2.11.0+cu130
|
||||
- Model path used in the commands below: `/path/to/stable-audio-open-1.0`
|
||||
|
||||
Stable Audio Open is a gated Hugging Face model. Accept the model license on
|
||||
the Hugging Face model card before downloading the checkpoint.
|
||||
|
||||
```bash
|
||||
hf auth login
|
||||
|
||||
hf download stabilityai/stable-audio-open-1.0 \
|
||||
--local-dir /path/to/stable-audio-open-1.0
|
||||
```
|
||||
|
||||
##### Commands
|
||||
|
||||
Run a 10-second offline validation sample from the repository root:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_audio/text_to_audio.py \
|
||||
--model /path/to/stable-audio-open-1.0 \
|
||||
--prompt "A gentle piano melody with soft room ambience" \
|
||||
--negative-prompt "Low quality, distorted, noisy" \
|
||||
--seed 42 \
|
||||
--guidance-scale 7.0 \
|
||||
--audio-length 10.0 \
|
||||
--num-inference-steps 50 \
|
||||
--cache-backend tea_cache \
|
||||
--output examples/offline_inference/text_to_audio/stable_audio_10s.wav
|
||||
```
|
||||
|
||||
Start the online serving endpoint:
|
||||
|
||||
```bash
|
||||
vllm-omni serve /path/to/stable-audio-open-1.0 \
|
||||
--host 0.0.0.0 \
|
||||
--port 8091 \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--omni
|
||||
```
|
||||
|
||||
Generate a 10-second WAV file from the repository root in another terminal:
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/health
|
||||
|
||||
curl -X POST http://localhost:8091/v1/audio/generate \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "A gentle piano melody with soft room ambience",
|
||||
"audio_length": 10.0,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 7.0,
|
||||
"negative_prompt": "Low quality, distorted, noisy",
|
||||
"seed": 42,
|
||||
"response_format": "wav"
|
||||
}' \
|
||||
--output examples/online_serving/stable_audio/piano_10s.wav
|
||||
```
|
||||
|
||||
##### Verification
|
||||
|
||||
Check that:
|
||||
|
||||
- The offline command writes a valid WAV file.
|
||||
- The server responds on `http://localhost:8091/health`.
|
||||
- The online request writes a valid WAV file.
|
||||
- The generated audio sample rate is 44.1 kHz.
|
||||
- The generated duration is approximately 10 seconds.
|
||||
- Peak sampled GPU memory is within the 24GB RTX 4090 budget. In the validated
|
||||
run, offline and online generation each peaked at about 12.6 GiB.
|
||||
|
||||
Validate the offline outputs:
|
||||
|
||||
```bash
|
||||
ls -lh examples/offline_inference/text_to_audio/stable_audio_10s.wav
|
||||
|
||||
python - <<'PY'
|
||||
import soundfile as sf
|
||||
|
||||
path = "examples/offline_inference/text_to_audio/stable_audio_10s.wav"
|
||||
audio, sample_rate = sf.read(path)
|
||||
print("sample_rate:", sample_rate)
|
||||
print("shape:", audio.shape)
|
||||
print("duration:", len(audio) / sample_rate)
|
||||
PY
|
||||
```
|
||||
|
||||
Validate the online output:
|
||||
|
||||
```bash
|
||||
ls -lh examples/online_serving/stable_audio/piano_10s.wav
|
||||
|
||||
python - <<'PY'
|
||||
import soundfile as sf
|
||||
|
||||
path = "examples/online_serving/stable_audio/piano_10s.wav"
|
||||
audio, sample_rate = sf.read(path)
|
||||
print("sample_rate:", sample_rate)
|
||||
print("shape:", audio.shape)
|
||||
print("duration:", len(audio) / sample_rate)
|
||||
PY
|
||||
```
|
||||
|
||||
##### Notes
|
||||
|
||||
- `stable-audio-open-1.0` can generate up to about 47 seconds of 44.1 kHz
|
||||
stereo audio. This recipe validates 10-second WAV outputs.
|
||||
- `--cache-backend tea_cache` is supported and was used for the 10-second
|
||||
offline validation command.
|
||||
- The model is gated on Hugging Face and requires license acceptance before
|
||||
download.
|
||||
- If online serving fails while importing `torchaudio`, make sure the
|
||||
`torchaudio` wheel matches the installed PyTorch and CUDA build. The
|
||||
validated environment used `torch==2.11.0+cu130` and
|
||||
`torchaudio==2.11.0+cu130`.
|
||||
- The `NIXL is not available`, `GLOO_SOCKET_IFNAME`, and `torchsde` boundary
|
||||
warnings observed during validation did not prevent successful generation.
|
||||
- This recipe was validated on one RTX 4090 24GB GPU only. Other GPU counts,
|
||||
ROCm, XPU, and NPU setups are not covered here.
|
||||
- Long generations, higher inference-step counts, and non-WAV response formats
|
||||
were not benchmarked in this recipe.
|
||||
@@ -0,0 +1,147 @@
|
||||
# Wan2.1 VACE Unified Video Generation
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Wan-AI
|
||||
- Models: `Wan2.1-VACE-1.3B-diffusers`, `Wan2.1-VACE-14B-diffusers`
|
||||
- Tasks: T2V, I2V, V2LF, FLF2V, inpainting, and R2V
|
||||
- Mode: offline inference through the shared task examples
|
||||
|
||||
The shared examples construct VACE's pipeline-native conditioning data from the
|
||||
provided media inputs. They do not require a mode parameter or a model-specific
|
||||
example script.
|
||||
|
||||
## References
|
||||
|
||||
- [VACE project](https://github.com/ali-vilab/VACE)
|
||||
- [Hugging Face 14B checkpoint](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers)
|
||||
- [ModelScope 14B checkpoint](https://modelscope.cn/models/Wan-AI/Wan2.1-VACE-14B-diffusers)
|
||||
- [Shared text-to-video example](../../examples/offline_inference/text_to_video/text_to_video.md)
|
||||
- [Shared image-to-video example](../../examples/offline_inference/image_to_video/README.md)
|
||||
- [Migration RFC #4539](https://github.com/vllm-project/vllm-omni/issues/4539)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
The 1.3B checkpoint was validated on an RTX 5090; the 81-frame T2V command below
|
||||
used approximately 20 GiB of peak reserved GPU memory with VAE tiling. The 14B
|
||||
checkpoint requires approximately 70 GiB of disk space and was validated on a
|
||||
single 48 GiB L40S with VAE tiling and layerwise offload. Start with the smoke
|
||||
configuration before attempting an 81-frame run.
|
||||
|
||||
## Setup
|
||||
|
||||
The public examples use the Hugging Face model ID by default:
|
||||
|
||||
```bash
|
||||
MODEL="${MODEL:-Wan-AI/Wan2.1-VACE-14B-diffusers}"
|
||||
ASSET_DIR="${ASSET_DIR:-$HOME/datasets/vace}"
|
||||
OUT="${OUT:-$PWD/vace-results/14b}"
|
||||
mkdir -p "$ASSET_DIR" "$OUT"
|
||||
```
|
||||
|
||||
For environments where Hugging Face downloads are unavailable, download the
|
||||
same checkpoint from ModelScope and point `MODEL` at the local directory:
|
||||
|
||||
```bash
|
||||
MODEL="$HOME/models/Wan2.1-VACE-14B-diffusers"
|
||||
modelscope download --model Wan-AI/Wan2.1-VACE-14B-diffusers --local_dir "$MODEL"
|
||||
```
|
||||
|
||||
Download the public conditioning images and create a mask matching the output
|
||||
size. White pixels are regenerated and black pixels are preserved.
|
||||
|
||||
```bash
|
||||
wget -O "$ASSET_DIR/astronaut.jpg" \
|
||||
https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg
|
||||
wget -O "$ASSET_DIR/vace_first_frame.png" \
|
||||
https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_first_frame.png
|
||||
wget -O "$ASSET_DIR/vace_last_frame.png" \
|
||||
https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_last_frame.png
|
||||
|
||||
python - "$ASSET_DIR/vace_center_mask.png" <<'PY'
|
||||
import sys
|
||||
from PIL import Image
|
||||
|
||||
mask = Image.new("L", (832, 480), 0)
|
||||
mask.paste(255, (336, 0, 496, 480))
|
||||
mask.save(sys.argv[1])
|
||||
PY
|
||||
|
||||
ASTRONAUT="$ASSET_DIR/astronaut.jpg"
|
||||
FIRST="$ASSET_DIR/vace_first_frame.png"
|
||||
LAST="$ASSET_DIR/vace_last_frame.png"
|
||||
MASK="$ASSET_DIR/vace_center_mask.png"
|
||||
```
|
||||
|
||||
## Commands
|
||||
|
||||
The commands below reproduce the original VACE examples with `81 frames / 30
|
||||
steps`. Most tasks use `480x832`; FLF2V uses `512x512`. For a smoke test,
|
||||
replace the frame and step counts with `5 frames / 2 steps`.
|
||||
|
||||
```bash
|
||||
# T2V
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model "$MODEL" \
|
||||
--prompt "A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves." \
|
||||
--seed 0 \
|
||||
--height 480 --width 832 --num-frames 81 --num-inference-steps 30 \
|
||||
--guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \
|
||||
--output "$OUT/t2v.mp4"
|
||||
|
||||
# I2V
|
||||
python examples/offline_inference/image_to_video/image_to_video.py \
|
||||
--model "$MODEL" --image "$ASTRONAUT" \
|
||||
--prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \
|
||||
--seed 42 \
|
||||
--height 480 --width 832 --num-frames 81 --num-inference-steps 30 \
|
||||
--guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \
|
||||
--output "$OUT/i2v.mp4"
|
||||
|
||||
# V2LF
|
||||
python examples/offline_inference/image_to_video/image_to_video.py \
|
||||
--model "$MODEL" --last-image "$ASTRONAUT" \
|
||||
--prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \
|
||||
--seed 42 \
|
||||
--height 480 --width 832 --num-frames 81 --num-inference-steps 30 \
|
||||
--guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \
|
||||
--output "$OUT/v2lf.mp4"
|
||||
|
||||
# FLF2V
|
||||
python examples/offline_inference/image_to_video/image_to_video.py \
|
||||
--model "$MODEL" --image "$FIRST" --last-image "$LAST" \
|
||||
--prompt "CG animation style, a small blue bird takes off from a branch and lands on another branch" \
|
||||
--seed 42 --height 512 --width 512 --num-frames 81 --num-inference-steps 30 \
|
||||
--guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \
|
||||
--output "$OUT/flf2v.mp4"
|
||||
|
||||
# Inpaint
|
||||
python examples/offline_inference/image_to_video/image_to_video.py \
|
||||
--model "$MODEL" --image "$ASTRONAUT" --mask-image "$MASK" \
|
||||
--prompt "Shrek, the ogre, walks out of a building in a happy mood" \
|
||||
--seed 42 \
|
||||
--height 480 --width 832 --num-frames 81 --num-inference-steps 30 \
|
||||
--guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \
|
||||
--output "$OUT/inpaint.mp4"
|
||||
|
||||
# R2V; repeat --reference-image to add references
|
||||
python examples/offline_inference/image_to_video/image_to_video.py \
|
||||
--model "$MODEL" --reference-image "$ASTRONAUT" \
|
||||
--prompt "Camera slowly zooms out from the character walking in a garden" \
|
||||
--seed 42 \
|
||||
--height 480 --width 832 --num-frames 81 --num-inference-steps 30 \
|
||||
--guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \
|
||||
--output "$OUT/r2v.mp4"
|
||||
```
|
||||
|
||||
## Verification
|
||||
|
||||
```bash
|
||||
for video in "$OUT"/*.mp4; do
|
||||
ffprobe -v error -select_streams v:0 \
|
||||
-show_entries stream=codec_name,width,height,nb_frames,r_frame_rate,duration \
|
||||
-of default=noprint_wrappers=1 "$video"
|
||||
done
|
||||
```
|
||||
|
||||
For the 1.3B checkpoint, `--enable-layerwise-offload` is generally unnecessary.
|
||||
@@ -367,60 +367,58 @@ vllm serve nvidia/Cosmos3-Nano-Policy-DROID \
|
||||
|
||||
#### Command
|
||||
|
||||
```python
|
||||
# cosmos3_offline.py — run with: python cosmos3_offline.py
|
||||
import torch
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
Cosmos3 runs through the standard task examples; pass model-specific knobs via
|
||||
`--extra-body`. Guardrails are on by default — pass `"guardrails": false` for a
|
||||
quick local run (install `cosmos-guardrail` + accept the gated repo to enable them).
|
||||
|
||||
```bash
|
||||
# Text-to-image -> examples/offline_inference/text_to_image
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model nvidia/Cosmos3-Nano \
|
||||
--prompt "A photorealistic red sports car at golden hour, cinematic lighting." \
|
||||
--negative-prompt "blurry, distorted, low quality" \
|
||||
--height 1024 --width 1024 --num-inference-steps 50 --guidance-scale 7.0 \
|
||||
--extra-body '{"flow_shift": 3.0, "guardrails": false}' \
|
||||
--output cosmos3_t2i.png
|
||||
|
||||
def main():
|
||||
omni = Omni(
|
||||
model="nvidia/Cosmos3-Nano",
|
||||
model_class_name="Cosmos3OmniDiffusersPipeline",
|
||||
trust_remote_code=True,
|
||||
enforce_eager=True,
|
||||
# Guardrails are disabled here for a quick local run; install
|
||||
# cosmos-guardrail + gated-repo access and drop this to enable them.
|
||||
model_config={"guardrails": False},
|
||||
)
|
||||
gen = torch.Generator(device="cpu").manual_seed(42)
|
||||
# Text-to-video -> examples/offline_inference/text_to_video
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model nvidia/Cosmos3-Nano \
|
||||
--prompt "A robot arm is cleaning a plate in the kitchen." \
|
||||
--negative-prompt "blurry, distorted, low quality, jittery, deformed" \
|
||||
--height 720 --width 1280 --num-frames 189 --fps 24 \
|
||||
--num-inference-steps 35 --guidance-scale 6.0 \
|
||||
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false,
|
||||
"use_resolution_template": false, "use_duration_template": false}' \
|
||||
--output cosmos3_t2v.mp4
|
||||
|
||||
# Text-to-image (modalities=["image"]). For T2V use modalities=["video"]
|
||||
# plus num_frames/fps; for I2V add multi_modal_data={"image": <PIL.Image>}.
|
||||
outputs = omni.generate(
|
||||
{
|
||||
"prompt": "A photorealistic red sports car at golden hour, cinematic lighting.",
|
||||
"negative_prompt": "blurry, distorted, low quality",
|
||||
"modalities": ["image"],
|
||||
},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=1024, width=1024, generator=gen,
|
||||
guidance_scale=7.0, num_inference_steps=50, num_outputs_per_prompt=1,
|
||||
),
|
||||
)
|
||||
outputs[0].request_output.images[0].save("cosmos3_t2i.png")
|
||||
omni.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
# Image-to-video -> examples/offline_inference/image_to_video
|
||||
# (Cosmos3 bundles example frames under assets/; any RGB image works too.)
|
||||
python examples/offline_inference/image_to_video/image_to_video.py \
|
||||
--model nvidia/Cosmos3-Nano \
|
||||
--image /path/to/Cosmos3-Nano/assets/example_i2v_input.jpg \
|
||||
--prompt "The scene comes to life with smooth, natural motion." \
|
||||
--height 720 --width 1280 --num-frames 189 --fps 24 \
|
||||
--num-inference-steps 35 --guidance-scale 6.0 \
|
||||
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}' \
|
||||
--output cosmos3_i2v.mp4
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
```bash
|
||||
python cosmos3_offline.py
|
||||
python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image', im.size, im.mode)"
|
||||
ffprobe -v error -show_entries stream=codec_type,nb_frames,width,height cosmos3_t2v.mp4
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- Same `Cosmos3OmniDiffusersPipeline` as online; mode is chosen by
|
||||
`prompt["modalities"]` (`["image"]` → T2I, `["video"]` → T2V) plus
|
||||
`num_frames`/`fps`, `multi_modal_data={"image": ...}` for I2V, and
|
||||
`multi_modal_data={"video": [<PIL frames>]}` or a video tensor/array for V2V.
|
||||
For video, frames are returned in `outputs[0].request_output.images` as an
|
||||
`(B, F, H, W, 3)` array.
|
||||
- The offline entry must be guarded by `if __name__ == "__main__":` — the engine
|
||||
spawns workers with the `spawn` start method.
|
||||
- A single `Cosmos3OmniDiffusersPipeline` serves every mode; the standard examples
|
||||
select it automatically from `model_index.json`. T2I is chosen by the
|
||||
`text_to_image` prompt builder (which marks `modalities=["image"]`); `text_to_video`
|
||||
defaults to T2V; `image_to_video` adds `multi_modal_data={"image": ...}` (I2V).
|
||||
V2V is served online (`/v1/videos/sync`).
|
||||
- Model-specific knobs (`flow_shift`, `max_sequence_length`, `condition_*`,
|
||||
`generate_sound`/`sound_duration`, `guardrails`, `action_*`, ...) are declared
|
||||
once in `vllm_omni/model_extras/cosmos3.py` and forwarded through `--extra-body`;
|
||||
unknown keys for the model are dropped.
|
||||
Binary file not shown.
@@ -0,0 +1,16 @@
|
||||
[
|
||||
{
|
||||
"index": "vocal_240_51240",
|
||||
"language": "Mandarin",
|
||||
"time": [
|
||||
240,
|
||||
51240
|
||||
],
|
||||
"duration": "0.21 0.18 0.34 0.38 0.26 0.22 0.50 0.20 0.33 0.13 0.60 0.22 0.38 0.52 0.30 0.12 1.82 0.98 0.20 0.26 0.38 0.36 0.58 0.54 0.26 0.50 1.46 3.03 0.24 0.24 0.29 0.25 0.20 0.14 0.80 0.22 0.54 0.30 0.24 0.16 0.58 0.28 0.38 1.73 0.79 0.24 0.30 0.34 0.30 0.30 0.34 0.21 0.21 0.36 0.34 0.23 1.85 1.90 0.23 0.39 0.68 0.50 0.31 0.43 0.76 0.38 2.00 1.87 0.68 0.72 0.56 0.62 0.80 0.40 0.42 1.68 1.79 0.70 0.66 0.54 0.24 0.48 0.68 0.40 2.34 0.14",
|
||||
"text": "<SP> 只 是 因 为 为 在 人 群 群 中 多 看 了 你 一 眼 <SP> 再 也 没 能 忘 掉 你 容 颜 <SP> 梦 想 着 着 偶 偶 然 然 有 <SP> 一 一 天 再 相 见 <SP> 从 此 我 开 始 始 孤 孤 单 思 念 念 <SP> 想 想 你 时 你 你 在 天 边 <SP> 想 你 时 你 在 眼 前 前 <SP> 想 你 时 你 你 在 脑 海 <SP>",
|
||||
"phoneme": "<SP> zh_zhi3 zh_shi4 zh_yin1 zh_wei4 zh_wei2 zh_zai4 zh_ren2 zh_qun2 zh_qun2 zh_zhong1 zh_duo1 zh_kan4 zh_le5 zh_ni3 zh_yi1 zh_yan3 <SP> zh_zai4 zh_ye3 zh_mei2 zh_neng2 zh_wang4 zh_diao4 zh_ni3 zh_rong2 zh_yan2 <SP> zh_meng4 zh_xiang3 zh_zhe5 zh_zhe5 zh_ou3 zh_ou3 zh_ran2 zh_ran2 zh_you3 <SP> zh_yi1 zh_yi1 zh_tian1 zh_zai4 zh_xiang1 zh_jian4 <SP> zh_cong2 zh_ci3 zh_wo3 zh_kai1 zh_shi3 zh_shi3 zh_gu1 zh_gu1 zh_dan1 zh_si1 zh_nian4 zh_nian4 <SP> zh_xiang3 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_tian1 zh_bian1 <SP> zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_zai4 zh_yan3 zh_qian2 zh_qian2 <SP> zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_nao3 zh_hai3 <SP>",
|
||||
"note_pitch": "0 64 64 64 66 68 66 66 66 64 64 64 64 66 64 60 61 0 63 63 63 64 66 63 61 59 56 0 68 66 66 68 68 66 66 64 64 0 64 66 61 61 61 64 0 62 63 63 64 64 66 65 66 61 58 59 56 0 69 71 66 68 68 71 66 64 61 0 66 61 68 66 64 63 61 59 0 71 66 68 68 71 66 64 61 0",
|
||||
"note_type": "1 2 2 2 2 3 2 2 2 3 2 2 2 2 2 2 2 1 2 2 2 2 2 2 2 2 2 1 2 2 2 3 2 3 2 3 2 1 2 3 2 2 2 2 1 2 2 2 2 2 3 2 3 2 2 2 3 1 2 3 2 2 2 3 2 2 2 1 2 2 2 2 2 2 2 3 1 2 2 2 2 3 2 2 2 1",
|
||||
"f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 318.9 332.2 331.2 323.9 312.3 0.0 0.0 0.0 0.0 0.0 0.0 340.9 342.3 340.2 337.4 333.5 329.3 327.8 328.3 329.1 330.5 331.9 331.5 329.2 327.6 327.4 327.7 329.6 330.5 329.5 328.1 328.5 330.1 330.4 329.5 328.9 331.0 333.6 332.8 331.7 330.2 330.6 330.9 329.5 327.6 332.4 343.5 363.9 368.7 369.3 368.2 366.7 371.3 382.0 392.9 402.3 409.8 413.8 416.3 416.4 415.4 414.1 413.4 415.3 417.4 418.3 413.6 416.0 0.0 0.0 0.0 0.0 368.2 362.5 361.3 362.0 362.9 364.3 366.4 367.5 368.4 368.9 368.1 367.6 369.7 371.1 371.1 369.5 368.0 367.3 366.6 367.1 369.5 372.0 371.3 368.2 368.1 369.6 369.9 373.4 376.5 360.2 285.6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 384.0 376.7 372.0 371.7 373.9 376.3 375.5 361.2 333.0 317.8 319.0 329.1 333.5 333.2 333.2 333.3 328.0 314.5 0.0 0.0 0.0 320.6 326.5 332.9 335.7 334.1 329.9 327.4 326.3 326.6 329.1 330.5 330.0 328.0 326.7 328.6 331.1 332.4 332.1 332.8 333.2 333.1 331.8 328.6 318.9 293.2 323.6 327.5 330.0 332.8 331.1 319.4 272.2 0.0 0.0 0.0 0.0 278.5 294.5 310.0 318.5 322.7 327.4 334.1 337.2 336.3 328.5 321.0 324.2 341.5 362.2 374.7 375.8 370.5 366.0 365.2 368.7 371.7 374.4 373.0 368.7 370.4 374.8 375.1 372.5 368.7 363.0 357.1 359.0 366.8 377.1 379.5 371.5 359.9 351.0 358.4 371.7 377.9 375.8 367.8 359.7 363.6 375.3 379.9 377.5 372.9 359.8 345.2 337.1 334.5 333.4 332.8 329.6 319.2 292.8 261.7 253.4 262.8 273.0 278.1 279.3 278.6 277.6 277.9 278.3 278.0 277.3 276.7 275.6 275.3 276.1 277.8 278.2 277.7 277.6 277.2 276.4 275.5 273.9 271.8 270.7 274.7 282.7 285.5 281.4 273.1 264.4 262.2 268.9 278.4 284.3 283.6 277.3 268.5 262.2 263.2 270.7 278.8 285.0 287.3 282.0 272.1 265.2 266.0 272.1 278.2 283.8 285.1 281.4 272.6 267.2 269.1 276.1 282.7 289.3 290.5 285.6 273.9 267.5 267.4 271.8 277.7 282.6 283.1 277.8 269.8 265.6 268.6 273.8 281.1 286.2 287.2 282.3 270.4 264.5 263.3 263.0 268.5 268.3 268.7 270.9 277.2 277.4 284.3 285.7 282.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 311.5 306.4 306.0 306.8 305.8 305.0 307.2 310.8 312.4 313.5 314.4 314.6 313.7 312.2 311.1 311.1 312.8 314.6 312.7 306.5 299.7 302.7 308.7 315.2 314.2 310.9 309.4 309.8 310.6 310.7 310.8 311.4 312.2 313.0 312.5 312.2 312.2 311.9 312.0 312.2 311.4 306.1 303.3 309.3 315.0 325.3 324.4 323.9 325.0 326.4 327.7 327.7 328.1 328.6 327.8 328.3 328.3 327.6 328.6 336.1 347.2 359.4 371.0 374.9 371.3 366.5 369.5 375.5 376.2 370.2 365.7 367.3 372.9 376.9 375.3 368.4 360.1 358.0 365.4 380.8 383.5 380.4 374.4 367.6 365.9 372.0 376.3 377.9 375.7 372.6 368.6 361.2 353.7 0.0 309.4 302.3 299.2 300.6 305.7 308.5 308.6 309.0 310.9 310.9 309.5 308.6 309.0 311.3 313.0 314.4 314.5 313.0 312.1 311.2 308.0 299.5 295.8 295.0 285.7 278.0 280.1 281.2 279.9 278.3 278.3 279.5 279.7 272.2 259.1 0.0 0.0 0.0 240.7 236.0 233.8 233.5 234.6 235.7 235.6 236.2 239.8 245.4 248.8 250.1 250.4 249.9 249.1 247.2 241.7 231.7 216.0 197.6 192.1 197.2 204.3 206.5 205.4 201.7 200.8 203.5 208.2 210.9 210.3 207.1 203.5 202.3 202.1 202.8 205.1 207.3 208.1 207.3 205.4 202.1 199.3 200.5 206.1 212.9 212.9 206.5 197.1 191.1 191.2 196.8 203.6 207.8 208.7 206.4 201.7 197.3 197.4 200.7 206.5 209.9 209.9 206.9 203.3 201.2 203.4 207.7 210.5 208.9 207.6 206.5 204.2 202.0 203.6 205.7 210.8 213.1 214.3 210.6 204.1 199.7 202.3 211.9 217.7 215.1 215.0 215.3 213.3 0.0 0.0 216.3 0.0 0.0 195.2 209.0 205.3 201.3 196.8 195.7 195.8 195.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 431.9 425.8 421.0 417.7 419.3 422.7 423.4 421.1 417.9 411.9 408.6 411.5 414.9 419.6 426.7 424.2 0.0 0.0 0.0 280.8 279.4 369.7 371.8 370.6 369.8 370.5 370.5 371.4 374.6 374.0 361.0 332.7 328.0 0.0 359.0 365.0 371.1 373.8 371.9 369.7 369.6 371.0 378.3 390.4 403.4 414.1 417.9 417.5 417.0 416.6 415.7 414.9 414.5 413.8 413.5 413.6 413.8 415.2 416.9 417.3 415.4 415.0 414.1 411.4 409.2 403.5 392.7 375.6 367.5 365.8 368.8 370.7 370.8 369.0 367.0 366.0 366.2 367.8 369.2 368.8 367.2 366.7 367.3 367.3 368.8 368.5 366.8 365.3 364.0 363.0 365.0 367.3 367.7 365.5 364.0 364.9 367.4 370.3 371.3 369.9 366.5 364.1 368.3 385.8 406.7 409.4 399.7 376.0 357.2 359.9 372.4 377.5 366.9 345.1 320.6 315.8 321.8 329.0 329.9 326.4 324.5 324.8 324.9 325.3 325.1 324.7 325.9 326.8 326.6 326.6 326.4 325.4 323.2 322.5 323.9 326.2 328.6 329.9 329.5 329.0 328.2 327.7 327.9 329.6 331.2 331.4 334.2 340.1 335.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.Line truncated
|
||||
}
|
||||
]
|
||||
Binary file not shown.
@@ -0,0 +1,16 @@
|
||||
[
|
||||
{
|
||||
"index": "vocal_320_10687",
|
||||
"language": "Mandarin",
|
||||
"time": [
|
||||
320,
|
||||
10687
|
||||
],
|
||||
"duration": "0.23 0.34 0.26 0.70 0.52 0.46 0.36 0.44 0.14 0.24 0.64 0.47 0.51 1.10 0.28 0.38 0.32 0.32 0.38 0.32 0.31 0.19 1.45",
|
||||
"text": "<SP> 除 了 想 你 你 <SP> 除 了 了 爱 你 你 <SP> 我 什 么 什 么 都 愿 愿 意",
|
||||
"phoneme": "<SP> zh_chu2 zh_le5 zh_xiang3 zh_ni3 zh_ni3 <SP> zh_chu2 zh_le5 zh_le5 zh_ai4 zh_ni3 zh_ni3 <SP> zh_wo3 zh_shen2 zh_me5 zh_shen2 zh_me5 zh_dou1 zh_yuan4 zh_yuan4 zh_yi4",
|
||||
"note_pitch": "0 62 65 67 67 69 0 67 69 67 65 67 69 67 67 66 64 64 60 60 65 67 0",
|
||||
"note_type": "1 2 2 2 2 3 1 2 2 3 2 2 3 1 2 2 2 2 2 2 2 3 2",
|
||||
"f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 294.1 288.2 290.6 294.6 295.7 292.8 291.5 294.4 295.4 294.7 293.5 292.2 294.4 295.8 293.2 297.7 320.1 338.1 348.3 348.7 344.4 342.6 346.2 354.8 356.9 353.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 401.1 403.8 405.0 402.1 398.4 395.6 393.3 392.0 391.4 390.3 390.2 390.3 390.3 391.3 392.8 393.6 391.7 390.5 391.4 391.6 391.5 393.5 393.8 390.8 387.4 387.8 389.3 390.8 392.2 391.6 390.2 389.8 389.1 388.4 390.0 395.5 397.2 396.7 395.5 395.1 394.6 394.9 395.6 395.2 394.6 395.4 395.9 394.0 391.7 390.7 391.7 392.6 391.6 395.7 405.8 441.7 462.5 463.8 450.2 430.9 414.5 415.2 426.7 439.8 454.4 462.9 447.8 422.5 400.6 403.2 423.5 451.3 482.3 492.8 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 435.9 414.3 406.2 402.8 398.2 396.7 396.7 395.2 398.3 416.9 441.6 446.9 441.5 437.3 435.3 432.9 431.6 432.3 434.8 436.7 433.8 422.8 406.7 390.9 382.0 381.7 384.7 382.7 368.1 357.4 355.6 355.1 352.4 348.1 346.4 348.7 351.6 355.0 354.4 351.7 349.9 349.2 348.1 346.0 345.4 344.2 344.4 345.5 346.6 349.0 349.7 349.1 349.5 349.6 349.7 349.4 349.5 352.2 354.5 355.7 355.6 356.9 359.4 361.5 363.8 360.4 354.3 357.2 363.9 372.4 382.6 399.1 402.9 400.6 395.5 390.5 388.9 390.2 391.1 391.9 391.5 390.4 390.2 391.3 391.6 391.0 388.6 386.2 389.1 403.8 430.2 441.8 449.5 448.1 443.2 438.3 432.9 430.7 434.6 442.0 447.6 446.0 440.3 434.7 431.1 435.7 442.3 445.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 423.8 402.1 398.5 397.6 393.4 390.5 391.4 402.5 427.6 442.5 442.5 435.1 430.1 430.8 439.9 447.4 442.2 426.1 412.3 399.8 391.8 389.5 388.5 387.8 386.4 384.7 384.5 387.9 391.2 391.8 392.9 393.8 392.0 392.0 395.4 398.1 398.2 396.3 393.1 391.1 388.9 386.6 383.1 381.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 357.8 367.6 370.9 365.4 359.3 355.7 358.1 372.7 396.8 404.0 398.4 392.6 389.2 388.8 383.6 362.3 341.1 325.1 326.3 327.9 331.3 333.3 326.0 319.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 356.4 353.5 349.0 346.5 338.4 328.4 323.7 326.2 334.2 338.4 331.6 309.8 280.6 256.4 252.8 256.1 258.3 256.9 257.6 261.1 259.7 258.7 258.6 260.2 262.2 262.4 262.9 264.2 263.4 260.2 256.4 253.2 238.7 223.3 231.6 257.7 258.1 258.4 258.8 258.0 256.8 255.8 254.0 255.1 258.1 261.9 263.7 262.6 256.7 253.1 250.2 246.7 258.7 294.4 327.1 342.6 346.0 344.1 341.2 342.3 345.2 350.1 364.7 382.5 396.1 396.1 389.2 381.8 381.1 387.2 397.0 399.3 390.7 374.3 360.1 350.6 346.6 347.4 350.7 354.4 354.3 351.7 349.8 348.4 346.9 347.0 348.4 349.5 351.0 352.3 353.6 353.3 350.5 348.3 345.5 344.3 344.4 347.0 350.6 352.0 351.0 350.8 350.1 347.6 345.7 347.0 350.3 351.7 350.7 348.5 346.9 347.7 349.0 349.1 348.5 346.8 346.3 348.4 349.0 349.2 351.1 349.6 348.3 350.5 351.1 348.0 347.6 349.1 351.3 356.0 361.3 360.6 354.0 341.0 316.2 302.9 302.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0"
|
||||
}
|
||||
]
|
||||
Binary file not shown.
@@ -1,146 +1,68 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the central pipeline registry (2.5/N)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
"""Tests for out of tree registration to OMNI_PIPELINES."""
|
||||
|
||||
import pytest
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES
|
||||
from vllm_omni.config.stage_config import (
|
||||
_PIPELINE_REGISTRY,
|
||||
PipelineConfig,
|
||||
StageExecutionType,
|
||||
StagePipelineConfig,
|
||||
register_pipeline,
|
||||
)
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline
|
||||
from vllm_omni.config.stage_config import PipelineConfig, pipeline_cfg_resolver
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class TestCentralRegistryDeclarations:
|
||||
"""Every in-tree pipeline must be declared exactly once in the central registry."""
|
||||
@pytest.fixture
|
||||
def custom_resolver():
|
||||
"""Build a reusable custom resolver for PipelineConfigs."""
|
||||
|
||||
def test_omni_entries_visible_in_registry(self):
|
||||
for key in _OMNI_PIPELINES:
|
||||
assert key in _PIPELINE_REGISTRY
|
||||
class CustomConfigType(PretrainedConfig):
|
||||
pass
|
||||
|
||||
def test_expected_omni_pipelines_present(self):
|
||||
# Guard against accidental removal during future refactors.
|
||||
assert "qwen2_5_omni" in _OMNI_PIPELINES
|
||||
assert "qwen2_5_omni_thinker_only" in _OMNI_PIPELINES
|
||||
assert "qwen3_omni_moe" in _OMNI_PIPELINES
|
||||
assert "aura_omni" in _OMNI_PIPELINES
|
||||
assert "qwen3_tts" in _OMNI_PIPELINES
|
||||
@pipeline_cfg_resolver(config_type=CustomConfigType)
|
||||
def custom_resolver(
|
||||
hf_config: CustomConfigType,
|
||||
) -> PipelineConfig:
|
||||
return PipelineConfig(model_type="resolved_type")
|
||||
|
||||
return custom_resolver
|
||||
|
||||
|
||||
class TestLazyLoading:
|
||||
"""Pipelines are imported only on first access."""
|
||||
|
||||
def test_contains_without_import(self):
|
||||
# ``in`` hits the lazy map, not the loaded cache.
|
||||
assert "qwen3_omni_moe" in _PIPELINE_REGISTRY
|
||||
|
||||
def test_getitem_loads_correct_pipeline(self):
|
||||
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
|
||||
assert pipeline.model_type == "qwen3_omni_moe"
|
||||
assert pipeline.model_arch == "Qwen3OmniMoeForConditionalGeneration"
|
||||
|
||||
def test_getitem_loads_aura_omni_pipeline(self):
|
||||
pipeline = _PIPELINE_REGISTRY["aura_omni"]
|
||||
assert pipeline.model_type == "aura_omni"
|
||||
assert [stage.model_stage for stage in pipeline.stages] == [
|
||||
"asr",
|
||||
"aura",
|
||||
"qwen3_tts",
|
||||
"code2wav",
|
||||
]
|
||||
assert pipeline.stages[0].model_arch == "Qwen3ASRForConditionalGeneration"
|
||||
assert pipeline.stages[1].model_arch == "AuraQwen3VLForConditionalGeneration"
|
||||
assert pipeline.stages[2].model_arch == "Qwen3TTSTalkerForConditionalGeneration"
|
||||
assert pipeline.stages[3].model_arch == "Qwen3TTSCode2Wav"
|
||||
|
||||
def test_unknown_model_type_returns_none_via_get(self):
|
||||
assert _PIPELINE_REGISTRY.get("not_a_real_pipeline") is None
|
||||
|
||||
def test_unknown_model_type_raises_keyerror_via_getitem(self):
|
||||
with pytest.raises(KeyError):
|
||||
_PIPELINE_REGISTRY["not_a_real_pipeline"]
|
||||
|
||||
def test_iteration_yields_registered_pipelines(self):
|
||||
keys = set(_PIPELINE_REGISTRY)
|
||||
assert "qwen2_5_omni" in keys
|
||||
assert "qwen3_omni_moe" in keys
|
||||
|
||||
def test_iteration_resilience(self, monkeypatch, caplog):
|
||||
# Synthetic failure to verify iteration robustness against load errors.
|
||||
target_logger = logging.getLogger("vllm_omni.config.stage_config")
|
||||
target_logger.addHandler(caplog.handler)
|
||||
prev_level = target_logger.level
|
||||
target_logger.setLevel(logging.ERROR)
|
||||
try:
|
||||
monkeypatch.setattr(
|
||||
_PIPELINE_REGISTRY,
|
||||
"_lazy_map",
|
||||
{
|
||||
"healthy": ("vllm_omni.model_executor.models.qwen2_5_omni.pipeline", "QWEN2_5_OMNI_PIPELINE"),
|
||||
"broken": ("vllm_omni.non_existent_module", "SomeConfig"),
|
||||
"attr_fail": ("vllm_omni.config.stage_config", "NON_EXISTENT_VAR"),
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(_PIPELINE_REGISTRY, "_loaded", {})
|
||||
|
||||
# Assert that list(_PIPELINE_REGISTRY.values()) and dict(_PIPELINE_REGISTRY.items())
|
||||
# both yield only healthy entries.
|
||||
healthy_instance = _PIPELINE_REGISTRY["healthy"]
|
||||
assert list(_PIPELINE_REGISTRY.values()) == [healthy_instance]
|
||||
assert dict(_PIPELINE_REGISTRY.items()) == {"healthy": healthy_instance}
|
||||
assert "Failed to import pipeline module" in caplog.text
|
||||
assert "Pipeline variable 'NON_EXISTENT_VAR' not found" in caplog.text
|
||||
|
||||
finally:
|
||||
target_logger.removeHandler(caplog.handler)
|
||||
target_logger.setLevel(prev_level)
|
||||
def test_register_pipeline_config(clean_pipeline_registry):
|
||||
"""Ensure that we can register a custom pipeline config to OMNI_PIPELINES."""
|
||||
new_model_type = "new_model_type"
|
||||
pipe_cfg = PipelineConfig(model_type=new_model_type)
|
||||
assert new_model_type not in OMNI_PIPELINES
|
||||
register_pipeline(pipe_cfg)
|
||||
assert new_model_type in OMNI_PIPELINES
|
||||
assert OMNI_PIPELINES[new_model_type] is pipe_cfg
|
||||
|
||||
|
||||
class TestDynamicRegistration:
|
||||
"""``register_pipeline()`` still works for plugins and tests."""
|
||||
def test_register_pipeline_config_with_model_type(clean_pipeline_registry):
|
||||
"""Ensure that we can register a custom pipeline config with an explicit model_type to OMNI_PIPELINES."""
|
||||
new_model_type = "new_model_type"
|
||||
unused_model_type = "foo"
|
||||
pipe_cfg = PipelineConfig(model_type=unused_model_type)
|
||||
assert new_model_type not in OMNI_PIPELINES
|
||||
assert unused_model_type not in OMNI_PIPELINES
|
||||
|
||||
def test_register_adds_to_registry(self):
|
||||
custom = PipelineConfig(
|
||||
model_type="_test_dynamic_registration",
|
||||
model_arch="DynamicTestModel",
|
||||
stages=(
|
||||
StagePipelineConfig(
|
||||
stage_id=0,
|
||||
model_stage="test",
|
||||
execution_type=StageExecutionType.LLM_AR,
|
||||
input_sources=(),
|
||||
final_output=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
register_pipeline(custom)
|
||||
try:
|
||||
assert "_test_dynamic_registration" in _PIPELINE_REGISTRY
|
||||
assert _PIPELINE_REGISTRY["_test_dynamic_registration"] is custom
|
||||
finally:
|
||||
# Don't leak the test registration into other tests.
|
||||
if "_test_dynamic_registration" in _PIPELINE_REGISTRY:
|
||||
del _PIPELINE_REGISTRY["_test_dynamic_registration"]
|
||||
# Registering with an explicitly provided model_type uses
|
||||
# the passed value instead of the pipeline_cfg.model_type
|
||||
register_pipeline(pipe_cfg, new_model_type)
|
||||
assert new_model_type in OMNI_PIPELINES
|
||||
assert unused_model_type not in OMNI_PIPELINES
|
||||
assert OMNI_PIPELINES[new_model_type] is pipe_cfg
|
||||
|
||||
def test_dynamic_registration_overrides_lazy_entry(self):
|
||||
# Build a substitute for qwen3_omni_moe that we can distinguish.
|
||||
original = _PIPELINE_REGISTRY["qwen3_omni_moe"]
|
||||
override = PipelineConfig(
|
||||
model_type="qwen3_omni_moe",
|
||||
model_arch="OverriddenArch",
|
||||
stages=original.stages,
|
||||
)
|
||||
register_pipeline(override)
|
||||
try:
|
||||
assert _PIPELINE_REGISTRY["qwen3_omni_moe"].model_arch == "OverriddenArch"
|
||||
finally:
|
||||
# Remove the dynamic override so later tests see the original.
|
||||
if "qwen3_omni_moe" in _PIPELINE_REGISTRY._loaded:
|
||||
del _PIPELINE_REGISTRY["qwen3_omni_moe"]
|
||||
|
||||
def test_register_resolver(custom_resolver, clean_pipeline_registry):
|
||||
"""Ensure that we can register a custom resolver to OMNI_PIPELINES."""
|
||||
new_model_type = "new_model_type"
|
||||
assert new_model_type not in OMNI_PIPELINES
|
||||
register_pipeline(custom_resolver, new_model_type)
|
||||
assert new_model_type in OMNI_PIPELINES
|
||||
assert OMNI_PIPELINES[new_model_type] is custom_resolver
|
||||
|
||||
|
||||
def test_register_resolver_requires_model_type(custom_resolver, clean_pipeline_registry):
|
||||
"""Ensure that registering a custom resolver to OMNI_PIPELINES requires an explicit model_type."""
|
||||
with pytest.raises(ValueError):
|
||||
register_pipeline(custom_resolver)
|
||||
@@ -17,6 +17,7 @@ from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules
|
||||
bootstrap_vllm_layer_custom_op_modules()
|
||||
|
||||
pytest_plugins = (
|
||||
"tests.helpers.fixtures.config",
|
||||
"tests.helpers.fixtures.env",
|
||||
"tests.helpers.fixtures.log",
|
||||
"tests.helpers.fixtures.run_args",
|
||||
|
||||
@@ -115,6 +115,7 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase):
|
||||
("MiMoAudioModel", "code2wav"),
|
||||
("Qwen3TTSCode2Wav", "code2wav"),
|
||||
("CosyVoice3Model", "cosyvoice3_code2wav"),
|
||||
("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),
|
||||
("DyninOmniForConditionalGeneration", "token2image"),
|
||||
("DyninOmniForConditionalGeneration", "token2audio"),
|
||||
}
|
||||
|
||||
@@ -621,7 +621,6 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
|
||||
{"consent": "c1", "name": "v1"},
|
||||
("audio_sample", "speaker_embedding", "provided"),
|
||||
id="neither_audio_nor_embedding",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
pytest.param(
|
||||
"wav_ok",
|
||||
@@ -673,7 +672,6 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
|
||||
{"consent": "consent_ok", "name": "evil/name"},
|
||||
("voice name", "invalid voice name"),
|
||||
id="invalid_name_path_sep",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
pytest.param(
|
||||
"wav_ok",
|
||||
@@ -698,21 +696,18 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
|
||||
{"consent": "consent_ok", "name": "v_bad_mime_{uuid}"},
|
||||
("mime", "unsupported", "audio/mp4"),
|
||||
id="audio_unsupported_mime",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
pytest.param(
|
||||
"wav_short",
|
||||
{"consent": "consent_ok", "name": "v_short_{uuid}"},
|
||||
("too short", "reference audio"),
|
||||
id="audio_too_short",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
pytest.param(
|
||||
"wav_bad_body",
|
||||
{"consent": "consent_ok", "name": "v_bad_wav_{uuid}"},
|
||||
("decode", "audio", "Format not recognised"),
|
||||
id="audio_decode_error",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
@@ -723,7 +718,6 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
|
||||
},
|
||||
("speaker_embedding", "valid JSON"),
|
||||
id="speaker_embedding_invalid_json",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
@@ -734,7 +728,6 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
|
||||
},
|
||||
("speaker_embedding", "non-empty"),
|
||||
id="speaker_embedding_empty",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
@@ -745,7 +738,6 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
|
||||
},
|
||||
("speaker_embedding", "finite"),
|
||||
id="speaker_embedding_nan",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
@@ -756,7 +748,6 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
|
||||
},
|
||||
("expected", "dimensions", "1024"),
|
||||
id="speaker_embedding_wrong_dim",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
pytest.param(
|
||||
"wav_ok",
|
||||
@@ -767,7 +758,6 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
|
||||
},
|
||||
("mutually exclusive", "audio_sample", "speaker_embedding"),
|
||||
id="audio_and_embedding_mutually_exclusive",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
@@ -19,9 +19,19 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
class StubScheduler:
|
||||
def __init__(self, timesteps: list[int] | None = None, *, flow_shift: float = 1.0) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
timesteps: list[int] | None = None,
|
||||
*,
|
||||
flow_shift: float = 1.0,
|
||||
use_karras_sigmas: bool = True,
|
||||
) -> None:
|
||||
self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64)
|
||||
self.config = SimpleNamespace(num_train_timesteps=1000, flow_shift=flow_shift)
|
||||
self.config = SimpleNamespace(
|
||||
num_train_timesteps=1000,
|
||||
flow_shift=flow_shift,
|
||||
use_karras_sigmas=use_karras_sigmas,
|
||||
)
|
||||
self.set_timesteps_calls: list[tuple[int, torch.device]] = []
|
||||
self.step_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
|
||||
|
||||
@@ -194,6 +204,8 @@ def make_cosmos3_pipeline():
|
||||
pipeline._base_scheduler_config = pipeline.scheduler.config
|
||||
pipeline._engine_init_flow_shift = 1.0
|
||||
pipeline._current_flow_shift = 1.0
|
||||
pipeline._base_scheduler_use_karras_sigmas = True
|
||||
pipeline._current_scheduler_use_karras_sigmas = True
|
||||
pipeline._guidance_scale = None
|
||||
pipeline._num_timesteps = None
|
||||
pipeline._cosmos3_branch_caches = None
|
||||
@@ -292,8 +304,8 @@ def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch):
|
||||
class _StubDiffusersScheduler:
|
||||
from_config_calls: list[dict[str, Any]] = []
|
||||
|
||||
def __init__(self, *, flow_shift: float = 1.0) -> None:
|
||||
self.config = SimpleNamespace(flow_shift=flow_shift)
|
||||
def __init__(self, *, flow_shift: float = 1.0, use_karras_sigmas: bool = True) -> None:
|
||||
self.config = SimpleNamespace(flow_shift=flow_shift, use_karras_sigmas=use_karras_sigmas)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, *args, **kwargs):
|
||||
@@ -302,7 +314,13 @@ def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch):
|
||||
@classmethod
|
||||
def from_config(cls, config, **kwargs):
|
||||
cls.from_config_calls.append({"config": config, "kwargs": dict(kwargs)})
|
||||
return cls(flow_shift=kwargs.get("flow_shift", getattr(config, "flow_shift", 1.0)))
|
||||
return cls(
|
||||
flow_shift=kwargs.get("flow_shift", getattr(config, "flow_shift", 1.0)),
|
||||
use_karras_sigmas=kwargs.get(
|
||||
"use_karras_sigmas",
|
||||
getattr(config, "use_karras_sigmas", True),
|
||||
),
|
||||
)
|
||||
|
||||
class _StubVideoProcessor:
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
@@ -349,10 +367,13 @@ def test_pipeline_init_skips_tokenizer_when_sound_disabled(stub_real_pipeline_in
|
||||
|
||||
pipeline = Cosmos3OmniDiffusersPipeline(od_config=_make_od_config(sound_gen=False))
|
||||
|
||||
assert stub_real_pipeline_init.from_config_calls == []
|
||||
assert pipeline._sound_tokenizer is None
|
||||
assert pipeline.transformer.sound_gen is False
|
||||
assert not hasattr(pipeline.transformer, "audio_proj_in")
|
||||
assert not hasattr(pipeline.transformer, "audio_proj_out")
|
||||
assert pipeline._base_scheduler_use_karras_sigmas is True
|
||||
assert pipeline._current_scheduler_use_karras_sigmas is True
|
||||
|
||||
|
||||
def test_pipeline_init_rebuilds_scheduler_with_cosmos3_defaults(stub_real_pipeline_init) -> None:
|
||||
@@ -366,8 +387,52 @@ def test_pipeline_init_rebuilds_scheduler_with_cosmos3_defaults(stub_real_pipeli
|
||||
assert stub_real_pipeline_init.from_config_calls
|
||||
call = stub_real_pipeline_init.from_config_calls[-1]
|
||||
assert getattr(call["config"], "flow_shift") == 1.0
|
||||
assert call["kwargs"] == {"flow_shift": 2.5, "use_karras_sigmas": False}
|
||||
assert call["kwargs"] == {"flow_shift": 2.5}
|
||||
assert pipeline._engine_init_flow_shift == 2.5
|
||||
assert pipeline._base_scheduler_use_karras_sigmas is True
|
||||
|
||||
|
||||
def test_set_flow_shift_restores_checkpoint_scheduler_mode(
|
||||
make_cosmos3_pipeline,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3
|
||||
|
||||
class RecordingScheduler:
|
||||
from_config_calls: list[dict[str, Any]] = []
|
||||
|
||||
def __init__(self, *, flow_shift: float, use_karras_sigmas: bool) -> None:
|
||||
self.config = SimpleNamespace(flow_shift=flow_shift, use_karras_sigmas=use_karras_sigmas)
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config, **kwargs):
|
||||
cls.from_config_calls.append({"config": config, "kwargs": dict(kwargs)})
|
||||
return cls(
|
||||
flow_shift=kwargs.get("flow_shift", getattr(config, "flow_shift", 1.0)),
|
||||
use_karras_sigmas=kwargs.get(
|
||||
"use_karras_sigmas",
|
||||
getattr(config, "use_karras_sigmas", True),
|
||||
),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(pipeline_cosmos3, "UniPCMultistepScheduler", RecordingScheduler)
|
||||
pipeline = make_cosmos3_pipeline()
|
||||
|
||||
pipeline._set_flow_shift(10.0, use_karras_sigmas=False)
|
||||
assert RecordingScheduler.from_config_calls[-1]["kwargs"] == {
|
||||
"flow_shift": 10.0,
|
||||
"use_karras_sigmas": False,
|
||||
}
|
||||
assert pipeline._current_flow_shift == 10.0
|
||||
assert pipeline._current_scheduler_use_karras_sigmas is False
|
||||
|
||||
pipeline._set_flow_shift(10.0)
|
||||
assert RecordingScheduler.from_config_calls[-1]["kwargs"] == {"flow_shift": 10.0}
|
||||
assert pipeline._current_flow_shift == 10.0
|
||||
assert pipeline._current_scheduler_use_karras_sigmas is True
|
||||
|
||||
pipeline._set_flow_shift(10.0)
|
||||
assert len(RecordingScheduler.from_config_calls) == 2
|
||||
|
||||
|
||||
def test_pipeline_init_passes_tokenizer_attrs_into_transformer(
|
||||
@@ -1022,7 +1087,12 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint
|
||||
pipeline._encode_video_tensor = fake_encode
|
||||
pipeline._prepare_transfer_latents = fake_prepare
|
||||
pipeline.diffuse_transfer = fake_diffuse_transfer
|
||||
pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target)
|
||||
|
||||
def fake_set_flow_shift(target, *, use_karras_sigmas=None):
|
||||
captured.setdefault("flow_shifts", []).append(target)
|
||||
captured.setdefault("scheduler_use_karras_sigmas", []).append(use_karras_sigmas)
|
||||
|
||||
pipeline._set_flow_shift = fake_set_flow_shift
|
||||
pipeline._decode_latents = lambda latents: torch.zeros(1, 3, 5, 16, 16, device="meta")
|
||||
|
||||
control = torch.zeros(3, 5, 16, 16, dtype=torch.uint8)
|
||||
@@ -1055,6 +1125,8 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint
|
||||
assert captured["format_frame_rate"] == expected_fps
|
||||
assert captured["shared_kwargs"]["fps"] == expected_fps
|
||||
assert captured["flow_shifts"] == [10.0]
|
||||
# Transfer uses the V2V flow-sigma schedule (karras off) so flow_shift applies.
|
||||
assert captured["scheduler_use_karras_sigmas"] == [False]
|
||||
assert output.custom_output["fps"] == expected_fps
|
||||
assert output.output["video"].device.type == "meta"
|
||||
|
||||
@@ -1068,7 +1140,7 @@ def test_forward_transfer_runs_multichunk_overlap_path(
|
||||
|
||||
pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16)
|
||||
pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (_ids(2), _mask(), _ids(1), _mask())
|
||||
pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target)
|
||||
pipeline._set_flow_shift = lambda target, **_kwargs: captured.setdefault("flow_shifts", []).append(target)
|
||||
|
||||
original_prepare = pipeline._prepare_transfer_latents
|
||||
|
||||
@@ -1213,7 +1285,12 @@ class TestForwardRouting:
|
||||
|
||||
pipeline._format_and_tokenize_prompts = fake_format
|
||||
pipeline._prepare_latents = fake_prepare
|
||||
pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target)
|
||||
|
||||
def fake_set_flow_shift(target, *, use_karras_sigmas=None):
|
||||
captured.setdefault("flow_shifts", []).append(target)
|
||||
captured.setdefault("scheduler_use_karras_sigmas", []).append(use_karras_sigmas)
|
||||
|
||||
pipeline._set_flow_shift = fake_set_flow_shift
|
||||
pipeline.diffuse = fake_diffuse
|
||||
pipeline._decode_latents = lambda latents: latents
|
||||
return captured
|
||||
@@ -1227,7 +1304,6 @@ class TestForwardRouting:
|
||||
{
|
||||
"key": "image",
|
||||
"is_t2i": True,
|
||||
"use_system_prompt": True,
|
||||
"flow": [3.0],
|
||||
"steps": [50, 50],
|
||||
"frames": 1,
|
||||
@@ -1239,7 +1315,6 @@ class TestForwardRouting:
|
||||
{
|
||||
"key": "video",
|
||||
"is_t2i": False,
|
||||
"use_system_prompt": True,
|
||||
"flow": [1.0],
|
||||
"steps": [35],
|
||||
"frames": 189,
|
||||
@@ -1261,9 +1336,9 @@ class TestForwardRouting:
|
||||
|
||||
assert expected["key"] in output.output
|
||||
assert captured["format"]["is_t2i"] is expected["is_t2i"]
|
||||
assert captured["format"]["use_system_prompt"] is expected["use_system_prompt"]
|
||||
assert captured["format"]["num_frames"] == expected["frames"]
|
||||
assert captured["flow_shifts"] == expected["flow"]
|
||||
assert captured["scheduler_use_karras_sigmas"] == [None]
|
||||
assert [call[0] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"]
|
||||
|
||||
def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> None:
|
||||
@@ -1290,7 +1365,6 @@ class TestForwardRouting:
|
||||
)
|
||||
)
|
||||
assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask
|
||||
assert captured["format"]["use_system_prompt"] is True
|
||||
|
||||
video_tensor = torch.zeros(1, 3, 5, 16, 16)
|
||||
v2v_condition = torch.full((1, 2, 2, 1, 1), 4.0)
|
||||
@@ -1316,8 +1390,8 @@ class TestForwardRouting:
|
||||
)
|
||||
)
|
||||
assert captured["flow_shifts"][-1] == 10.0
|
||||
assert captured["scheduler_use_karras_sigmas"][-1] is False
|
||||
assert captured["format"]["negative_prompt"] == ""
|
||||
assert captured["format"]["use_system_prompt"] is True
|
||||
assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is v2v_mask
|
||||
assert captured["diffuse_calls"][-1]["condition_latents"] is v2v_condition
|
||||
|
||||
@@ -1333,7 +1407,6 @@ class TestForwardRouting:
|
||||
)
|
||||
)
|
||||
assert captured["diffuse_calls"][-1]["sound_latents"] is sound_latents
|
||||
assert captured["format"]["use_system_prompt"] is True
|
||||
assert output.output["audio_sample_rate"] == 10
|
||||
|
||||
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4)
|
||||
@@ -1359,7 +1432,6 @@ class TestForwardRouting:
|
||||
)
|
||||
)
|
||||
assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
|
||||
assert captured["format"]["use_system_prompt"] is False
|
||||
assert output.custom_output["action"].shape == (1, 2, 2)
|
||||
assert "action_only_output" not in output.custom_output
|
||||
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for HunyuanImage3 AR image RGB conversion."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3 import (
|
||||
HunyuanImage3Processor,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _FixedResolutionGroup:
|
||||
def get_target_size(self, width: int, height: int):
|
||||
return width, height
|
||||
|
||||
def get_base_size_and_ratio_index(self, width: int, height: int):
|
||||
return 1024, 0
|
||||
|
||||
|
||||
def test_ar_processor_uses_pil_rgb_conversion_for_rgba_images():
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_vit_processor(image: Image.Image):
|
||||
captured["mode"] = image.mode
|
||||
captured["pixel"] = image.getpixel((0, 0))
|
||||
return {
|
||||
"pixel_values": torch.zeros((1, 1, 3), dtype=torch.float32),
|
||||
"pixel_attention_mask": torch.ones((1, 1), dtype=torch.bool),
|
||||
"spatial_shapes": torch.tensor([[1, 1]], dtype=torch.long),
|
||||
}
|
||||
|
||||
def fake_vae_processor(_image: Image.Image):
|
||||
return torch.zeros((1, 3, 1, 1), dtype=torch.float32)
|
||||
|
||||
processor = object.__new__(HunyuanImage3Processor)
|
||||
processor.hf_config = SimpleNamespace(
|
||||
vit={"num_channels": 3},
|
||||
vae_downsample_factor=(1, 1),
|
||||
patch_size=1,
|
||||
)
|
||||
processor.reso_group = _FixedResolutionGroup()
|
||||
processor.vision_encoder_processor = fake_vit_processor
|
||||
processor.vae_processor = fake_vae_processor
|
||||
|
||||
processor.process_image(Image.new("RGBA", (1, 1), color=(255, 0, 0, 0)))
|
||||
|
||||
assert captured["mode"] == "RGB"
|
||||
assert captured["pixel"] == (255, 0, 0)
|
||||
@@ -20,6 +20,7 @@ RATIO_START = 200
|
||||
RATIO_END = 210
|
||||
RATIO_OTHER_START = 220
|
||||
RATIO_OTHER_END = 223
|
||||
TIMESTEP = 224
|
||||
|
||||
|
||||
class FakeSamplerModel:
|
||||
@@ -60,6 +61,67 @@ class FakeSamplerModel:
|
||||
_apply_ratio_restriction = _Real._apply_ratio_restriction
|
||||
|
||||
|
||||
class FakeTokenEmbedding:
|
||||
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return torch.stack(
|
||||
[
|
||||
input_ids.to(dtype=torch.float32),
|
||||
input_ids.to(dtype=torch.float32) + 0.5,
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
|
||||
|
||||
class FakeEmbeddingModel:
|
||||
def __init__(self):
|
||||
self.model = FakeTokenEmbedding()
|
||||
self._timestep_token_id = TIMESTEP
|
||||
|
||||
from vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3 import (
|
||||
HunyuanImage3ForConditionalGeneration as _Real,
|
||||
)
|
||||
|
||||
embed_input_ids = _Real.embed_input_ids
|
||||
|
||||
def _timestep_encode(self, timestep: torch.Tensor) -> torch.Tensor:
|
||||
return torch.full(
|
||||
(1, 2),
|
||||
-7.0,
|
||||
device=timestep.device,
|
||||
dtype=timestep.dtype,
|
||||
)
|
||||
|
||||
|
||||
class TestTimestepEmbedding:
|
||||
def test_timestep_embedding_avoids_scalar_item_sync(self, monkeypatch):
|
||||
model = FakeEmbeddingModel()
|
||||
input_ids = torch.tensor([1, TIMESTEP, 2, TIMESTEP], dtype=torch.long)
|
||||
|
||||
def fail_item(_self):
|
||||
raise AssertionError("embed_input_ids should not synchronize timestep count")
|
||||
|
||||
monkeypatch.setattr(torch.Tensor, "item", fail_item)
|
||||
|
||||
actual = model.embed_input_ids(input_ids)
|
||||
expected = model.model.embed_input_ids(input_ids)
|
||||
expected[input_ids == TIMESTEP] = -7.0
|
||||
|
||||
torch.testing.assert_close(actual, expected)
|
||||
|
||||
def test_timestep_embedding_preserves_no_timestep_inputs(self, monkeypatch):
|
||||
model = FakeEmbeddingModel()
|
||||
input_ids = torch.tensor([1, 2, 3], dtype=torch.long)
|
||||
|
||||
def fail_item(_self):
|
||||
raise AssertionError("embed_input_ids should not synchronize timestep count")
|
||||
|
||||
monkeypatch.setattr(torch.Tensor, "item", fail_item)
|
||||
|
||||
actual = model.embed_input_ids(input_ids)
|
||||
|
||||
torch.testing.assert_close(actual, model.model.embed_input_ids(input_ids))
|
||||
|
||||
|
||||
class TestGetForcedToken:
|
||||
"""Tests for the stateless _get_forced_token method."""
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler
|
||||
@@ -33,6 +34,7 @@ def _make_engine() -> DiffusionEngine:
|
||||
engine._rpc_lock = threading.RLock()
|
||||
engine._cv = threading.Condition(engine._rpc_lock)
|
||||
engine._out_queue = {}
|
||||
engine._out_queue_streaming = {}
|
||||
engine._closed = False
|
||||
engine._shutdown_complete = False
|
||||
engine.abort_queue = queue.Queue()
|
||||
@@ -57,6 +59,23 @@ def test_close_completes_pending_async_waiters() -> None:
|
||||
event_loop.close()
|
||||
|
||||
|
||||
def test_close_completes_pending_streaming_waiters() -> None:
|
||||
engine = _make_engine()
|
||||
event_loop = asyncio.new_event_loop()
|
||||
try:
|
||||
engine.main_loop = event_loop
|
||||
queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
|
||||
engine._out_queue_streaming["pending-stream"] = queue
|
||||
|
||||
engine.close()
|
||||
|
||||
output = queue.get_nowait()
|
||||
assert output.error == "DiffusionEngine is closed."
|
||||
assert output.finished is True
|
||||
finally:
|
||||
event_loop.close()
|
||||
|
||||
|
||||
def test_handle_finished_requests_ignores_already_drained_waiter() -> None:
|
||||
class RacingOutQueue(dict):
|
||||
def __contains__(self, key):
|
||||
|
||||
@@ -9,6 +9,7 @@ import torch
|
||||
|
||||
import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
|
||||
|
||||
pytestmark = [pytest.mark.diffusion]
|
||||
@@ -31,6 +32,46 @@ class _DummyPipeline:
|
||||
return self._output
|
||||
|
||||
|
||||
class _ChunkStepPipeline:
|
||||
device = torch.device("cpu")
|
||||
supports_step_execution = True
|
||||
|
||||
def __init__(self, outputs):
|
||||
self._outputs = outputs
|
||||
self.prepare_calls = 0
|
||||
self.decode_calls = 0
|
||||
|
||||
def prepare_encode(self, state):
|
||||
self.prepare_calls += 1
|
||||
state.prompt_embeds = torch.zeros(1, 1, 1)
|
||||
state.latents = torch.zeros(1, 1)
|
||||
state.timesteps = torch.tensor([1.0, 0.0])
|
||||
state.step_index = 0
|
||||
state.step_in_chunk = 0
|
||||
state.chunk_num_steps = 2
|
||||
state.total_chunks = len(self._outputs)
|
||||
return state
|
||||
|
||||
def denoise_step(self, input_batch, states):
|
||||
del states
|
||||
return torch.ones_like(input_batch.latents)
|
||||
|
||||
def step_scheduler(self, state, noise_pred):
|
||||
state.latents = noise_pred
|
||||
state.step_index += 1
|
||||
state.step_in_chunk += 1
|
||||
|
||||
def post_decode(self, state):
|
||||
output = self._outputs[self.decode_calls]
|
||||
self.decode_calls += 1
|
||||
state.chunk_index += 1
|
||||
state.step_index = 0
|
||||
state.step_in_chunk = 0
|
||||
if not state.request_denoise_completed:
|
||||
state.latents = torch.zeros(1, 1)
|
||||
return output
|
||||
|
||||
|
||||
def _make_request(skip_cache_refresh: bool = True):
|
||||
sampling_params = SimpleNamespace(
|
||||
generator=None,
|
||||
@@ -52,11 +93,13 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
|
||||
runner.pipeline = _DummyPipeline(output=SimpleNamespace(output="ok"))
|
||||
runner.cache_backend = cache_backend
|
||||
runner.offload_backend = None
|
||||
runner.state_cache = {}
|
||||
runner.prompt_embed_cache = None
|
||||
runner.od_config = SimpleNamespace(
|
||||
cache_backend=cache_backend_name,
|
||||
enable_cache_dit_summary=enable_cache_dit_summary,
|
||||
parallel_config=SimpleNamespace(use_hsdp=False),
|
||||
streaming_output=False,
|
||||
)
|
||||
runner.kv_transfer_manager = SimpleNamespace(
|
||||
receive_kv_cache=lambda req, target_device=None: None,
|
||||
@@ -66,6 +109,49 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
|
||||
return runner
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):
|
||||
"""Step streaming returns empty step results until a chunk decode boundary."""
|
||||
chunks = [
|
||||
DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2),
|
||||
DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2),
|
||||
]
|
||||
runner = _make_runner(cache_backend=None, cache_backend_name=None)
|
||||
runner.pipeline = _ChunkStepPipeline(chunks)
|
||||
runner.od_config.streaming_output = True
|
||||
runner.od_config.step_execution = True
|
||||
req = _make_request(skip_cache_refresh=True)
|
||||
req.request_id = "req"
|
||||
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
|
||||
scheduler_output = SimpleNamespace(
|
||||
finished_req_ids=set(),
|
||||
scheduled_new_reqs=[SimpleNamespace(request_id="req", req=req)],
|
||||
scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
|
||||
)
|
||||
|
||||
first = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
|
||||
assert first.get_request_output("req").result is None
|
||||
|
||||
scheduler_output = SimpleNamespace(
|
||||
finished_req_ids=set(),
|
||||
scheduled_new_reqs=[],
|
||||
scheduled_cached_reqs=SimpleNamespace(request_ids=["req"]),
|
||||
)
|
||||
second = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
|
||||
assert second.get_request_output("req").result == chunks[0]
|
||||
assert second.get_request_output("req").finished is False
|
||||
|
||||
DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
|
||||
fourth = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
|
||||
assert fourth.get_request_output("req").result == chunks[1]
|
||||
assert fourth.get_request_output("req").finished is True
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypatch):
|
||||
@@ -76,6 +162,9 @@ def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypa
|
||||
cache_summary_calls = []
|
||||
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
|
||||
monkeypatch.setattr(
|
||||
model_runner_module,
|
||||
"cache_summary",
|
||||
@@ -101,6 +190,9 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp
|
||||
cache_summary_calls = []
|
||||
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
|
||||
monkeypatch.setattr(
|
||||
model_runner_module,
|
||||
"cache_summary",
|
||||
@@ -157,6 +249,7 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
|
||||
cache_config={},
|
||||
model_class_name="NextStep11Pipeline",
|
||||
enforce_eager=True,
|
||||
streaming_output=False,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
import queue
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
@@ -238,6 +239,35 @@ class TestRequestScheduler:
|
||||
assert finished == {req_id}
|
||||
assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
|
||||
|
||||
def test_streaming_output_keeps_request_running_until_final_chunk(self) -> None:
|
||||
scheduler = StepScheduler()
|
||||
scheduler.initialize(SimpleNamespace())
|
||||
req_id = scheduler.add_request(_make_request("stream"))
|
||||
|
||||
sched_output = scheduler.schedule()
|
||||
chunk = RunnerOutput(
|
||||
request_id=req_id,
|
||||
step_index=1,
|
||||
finished=False,
|
||||
result=DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2),
|
||||
)
|
||||
finished = scheduler.update_from_output(sched_output, chunk)
|
||||
|
||||
assert finished == set()
|
||||
assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.RUNNING
|
||||
assert scheduler.has_requests() is True
|
||||
|
||||
final_chunk = RunnerOutput(
|
||||
request_id=req_id,
|
||||
step_index=2,
|
||||
finished=True,
|
||||
result=DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2),
|
||||
)
|
||||
finished = scheduler.update_from_output(sched_output, final_chunk)
|
||||
|
||||
assert finished == {req_id}
|
||||
assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
|
||||
|
||||
def test_fifo_single_request_scheduling(self) -> None:
|
||||
req_id_a = self.scheduler.add_request(_make_request("a"))
|
||||
req_id_b = self.scheduler.add_request(_make_request("b"))
|
||||
@@ -369,6 +399,7 @@ class TestRequestScheduler:
|
||||
class TestDiffusionEngine:
|
||||
def test_add_req_and_wait_for_response_single_path(self, mocker: MockerFixture) -> None:
|
||||
engine = DiffusionEngine.__new__(DiffusionEngine)
|
||||
engine.od_config = SimpleNamespace(streaming_output=False)
|
||||
engine.scheduler = RequestScheduler()
|
||||
engine.scheduler.initialize(SimpleNamespace())
|
||||
engine._rpc_lock = threading.RLock()
|
||||
@@ -391,6 +422,7 @@ class TestDiffusionEngine:
|
||||
scheduler = _StubScheduler(request, runner_output)
|
||||
|
||||
engine = DiffusionEngine.__new__(DiffusionEngine)
|
||||
engine.od_config = SimpleNamespace(streaming_output=False)
|
||||
engine.scheduler = scheduler
|
||||
engine._rpc_lock = threading.RLock()
|
||||
engine._cv = threading.Condition(engine._rpc_lock)
|
||||
@@ -410,7 +442,7 @@ class TestDiffusionEngine:
|
||||
) -> None:
|
||||
request = _make_request("init")
|
||||
scheduler = _StubScheduler(request, DiffusionOutput(output=None))
|
||||
od_config = SimpleNamespace(model_class_name="mock_model")
|
||||
od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False)
|
||||
fake_executor_cls = mocker.Mock(return_value=mocker.Mock())
|
||||
|
||||
monkeypatch.setattr(
|
||||
@@ -472,7 +504,7 @@ class TestDiffusionEngine:
|
||||
|
||||
def test_finalize_finished_request_returns_aborted_output(self) -> None:
|
||||
engine = DiffusionEngine.__new__(DiffusionEngine)
|
||||
engine.scheduler = RequestScheduler()
|
||||
engine.scheduler = StepScheduler()
|
||||
engine.scheduler.initialize(SimpleNamespace())
|
||||
|
||||
req_id = engine.scheduler.add_request(_make_request("req-finalize"))
|
||||
@@ -483,12 +515,77 @@ class TestDiffusionEngine:
|
||||
assert output.aborted is True
|
||||
assert output.abort_message == "Request req-finalize aborted."
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_runner_output_notifies_each_chunk(self) -> None:
|
||||
engine = DiffusionEngine.__new__(DiffusionEngine)
|
||||
engine.scheduler = StepScheduler()
|
||||
engine.scheduler.initialize(SimpleNamespace())
|
||||
engine._rpc_lock = threading.RLock()
|
||||
engine._cv = threading.Condition(engine._rpc_lock)
|
||||
engine._out_queue_streaming = {}
|
||||
engine.main_loop = asyncio.get_running_loop()
|
||||
|
||||
req_id = engine.scheduler.add_request(_make_request("stream-engine"))
|
||||
queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
|
||||
engine._out_queue_streaming[req_id] = queue
|
||||
sched_output = engine.scheduler.schedule()
|
||||
|
||||
chunk = RunnerOutput(
|
||||
request_id=req_id,
|
||||
step_index=1,
|
||||
finished=False,
|
||||
result=DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2),
|
||||
)
|
||||
finished_req_ids = engine.scheduler.update_from_output(sched_output, chunk)
|
||||
engine._handle_step_streaming_runner_output(finished_req_ids, sched_output.scheduled_request_ids, chunk)
|
||||
|
||||
notified_chunk = await asyncio.wait_for(queue.get(), timeout=1)
|
||||
assert notified_chunk.output == "chunk-0"
|
||||
assert notified_chunk.finished is False
|
||||
assert engine.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.RUNNING
|
||||
|
||||
final_chunk = RunnerOutput(
|
||||
request_id=req_id,
|
||||
step_index=2,
|
||||
finished=True,
|
||||
result=DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2),
|
||||
)
|
||||
finished_req_ids = engine.scheduler.update_from_output(sched_output, final_chunk)
|
||||
engine._handle_step_streaming_runner_output(finished_req_ids, sched_output.scheduled_request_ids, final_chunk)
|
||||
|
||||
notified_final = await asyncio.wait_for(queue.get(), timeout=1)
|
||||
assert notified_final.output == "chunk-1"
|
||||
assert notified_final.finished is True
|
||||
assert engine.scheduler.get_request_state(req_id) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finished_streaming_request_without_runner_output_notifies_waiter(self) -> None:
|
||||
engine = DiffusionEngine.__new__(DiffusionEngine)
|
||||
engine.scheduler = RequestScheduler()
|
||||
engine.scheduler.initialize(SimpleNamespace())
|
||||
engine._rpc_lock = threading.RLock()
|
||||
engine._cv = threading.Condition(engine._rpc_lock)
|
||||
engine._out_queue_streaming = {}
|
||||
engine.main_loop = asyncio.get_running_loop()
|
||||
|
||||
req_id = engine.scheduler.add_request(_make_request("stream-abort"))
|
||||
queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
|
||||
engine._out_queue_streaming[req_id] = queue
|
||||
engine.scheduler.finish_requests(req_id, DiffusionRequestStatus.FINISHED_ABORTED)
|
||||
|
||||
engine._handle_empty_streaming_requests({req_id})
|
||||
|
||||
output = await asyncio.wait_for(queue.get(), timeout=1)
|
||||
assert output.aborted is True
|
||||
assert output.finished is True
|
||||
assert engine.scheduler.get_request_state(req_id) is None
|
||||
|
||||
def test_initializes_step_scheduler_when_step_execution_enabled(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
mocker: MockerFixture,
|
||||
) -> None:
|
||||
od_config = SimpleNamespace(model_class_name="mock_model")
|
||||
od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False)
|
||||
od_config.step_execution = True
|
||||
fake_executor = mocker.Mock()
|
||||
fake_executor_cls = mocker.Mock(return_value=fake_executor)
|
||||
|
||||
@@ -249,6 +249,7 @@ def _make_runner():
|
||||
runner.od_config = SimpleNamespace(
|
||||
cache_backend=None,
|
||||
parallel_config=SimpleNamespace(use_hsdp=False),
|
||||
streaming_output=False,
|
||||
)
|
||||
runner.device = torch.device("cpu")
|
||||
runner.pipeline = _StepPipeline()
|
||||
@@ -267,6 +268,7 @@ def _make_distributed_runner(mode: str, device: torch.device):
|
||||
runner.od_config = SimpleNamespace(
|
||||
cache_backend=None,
|
||||
parallel_config=SimpleNamespace(use_hsdp=False),
|
||||
streaming_output=False,
|
||||
)
|
||||
runner.device = device
|
||||
runner.pipeline = _DistributedStepPipeline(mode=mode, device=device)
|
||||
@@ -327,7 +329,7 @@ def _make_cached_scheduler_output(request_id="req-1", step_id=1, finished_req_id
|
||||
|
||||
def _make_engine(scheduler, execute_fn=None) -> DiffusionEngine:
|
||||
engine = object.__new__(DiffusionEngine)
|
||||
engine.od_config = SimpleNamespace(model_class_name="QwenImagePipeline")
|
||||
engine.od_config = SimpleNamespace(model_class_name="QwenImagePipeline", streaming_output=False)
|
||||
engine.pre_process_func = None
|
||||
engine.post_process_func = None
|
||||
engine.scheduler = scheduler
|
||||
@@ -563,6 +565,7 @@ class TestRunner:
|
||||
step_execution=True,
|
||||
model_class_name="RequestOnlyPipeline",
|
||||
parallel_config=SimpleNamespace(use_hsdp=False),
|
||||
streaming_output=False,
|
||||
)
|
||||
runner.device = torch.device("cpu")
|
||||
runner.pipeline = None
|
||||
@@ -710,6 +713,7 @@ class TestExecutor:
|
||||
|
||||
def test_execute_step_passes_through_runner_output(self, mocker: MockerFixture):
|
||||
executor = object.__new__(MultiprocDiffusionExecutor)
|
||||
executor.od_config = SimpleNamespace(streaming_output=False)
|
||||
executor._ensure_open = lambda: None
|
||||
expected = RunnerOutput(request_id="req-step", step_index=1, finished=False, result=None)
|
||||
executor.collective_rpc = mocker.Mock(return_value=expected)
|
||||
|
||||
@@ -0,0 +1,525 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Integration tests for diffusion pipeline streaming output mode.
|
||||
|
||||
Covers:
|
||||
- Mock pipeline -> StageDiffusionClient (ZMQ subprocess path)
|
||||
- Mock pipeline -> InlineStageDiffusionClient -> Orchestrator -> AsyncOmni
|
||||
- Mock pipeline -> InlineStageDiffusionClient -> Orchestrator -> AsyncOmni -> `/v1/realtime/video`
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import queue
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
import zmq
|
||||
from fastapi import FastAPI, WebSocket
|
||||
from pytest_mock import MockerFixture
|
||||
from starlette.testclient import TestClient
|
||||
from vllm.utils.network_utils import get_open_zmq_ipc_path
|
||||
|
||||
import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
|
||||
from tests.engine.test_orchestrator import OrchestratorFixture, _build_harness, _wait_for
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
|
||||
from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
|
||||
from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
|
||||
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
|
||||
from vllm_omni.distributed.omni_connectors.utils.serialization import (
|
||||
OmniMsgpackDecoder,
|
||||
OmniMsgpackEncoder,
|
||||
)
|
||||
from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
|
||||
from vllm_omni.engine.messages import ShutdownRequestMessage, StageSubmissionMessage
|
||||
from vllm_omni.engine.stage_init_utils import StageMetadata
|
||||
from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni
|
||||
from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _StepStreamingPipeline:
|
||||
supports_step_execution = True
|
||||
|
||||
def __init__(self, outputs: list[DiffusionOutput]) -> None:
|
||||
self.outputs = outputs
|
||||
self.requests: list[Any] = []
|
||||
|
||||
def step_outputs(self, request):
|
||||
self.requests.append(request)
|
||||
return list(self.outputs)
|
||||
|
||||
|
||||
class _FailingStreamingPipeline:
|
||||
supports_step_execution = True
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.requests: list[Any] = []
|
||||
|
||||
def step_outputs(self, request):
|
||||
self.requests.append(request)
|
||||
raise RuntimeError("stream failed midway")
|
||||
|
||||
|
||||
class _PipelineBackedEngine:
|
||||
"""DiffusionEngine stand-in that exposes step-streaming chunks from a mock pipeline."""
|
||||
|
||||
def __init__(self, pipeline: _StepStreamingPipeline | _FailingStreamingPipeline) -> None:
|
||||
self.pipeline = pipeline
|
||||
self.executor = SimpleNamespace(
|
||||
register_failure_callback=MagicMock(),
|
||||
check_health=MagicMock(),
|
||||
)
|
||||
|
||||
async def step_streaming(self, request):
|
||||
try:
|
||||
outputs = self.pipeline.step_outputs(request)
|
||||
except Exception as exc:
|
||||
outputs = [
|
||||
DiffusionOutput(custom_output={"chunk": 0}, finished=False),
|
||||
DiffusionOutput(error=str(exc), finished=True),
|
||||
]
|
||||
for output in outputs:
|
||||
if output.error is not None:
|
||||
yield [OmniRequestOutput.from_error(request_id=request.request_id, error_message=output.error)]
|
||||
continue
|
||||
custom_output = output.custom_output or {}
|
||||
yield [
|
||||
OmniRequestOutput.from_diffusion(
|
||||
request_id=request.request_id,
|
||||
images=list(custom_output.get("images") or []),
|
||||
custom_output=custom_output,
|
||||
finished=output.finished,
|
||||
)
|
||||
]
|
||||
|
||||
def abort(self, request_id: str) -> None:
|
||||
del request_id
|
||||
|
||||
|
||||
class TestPipelineStreamingOutputToStageDiffusionClient:
|
||||
"""Streaming pipeline output over ZMQ into ``StageDiffusionClient``."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_output_reaches_stage_client_from_mock_pipeline(self) -> None:
|
||||
"""Mock pipeline chunks over ZMQ reach StageDiffusionClient with correct finished flags."""
|
||||
pipeline = _StepStreamingPipeline(
|
||||
[
|
||||
DiffusionOutput(custom_output={"chunk": 0}, finished=False),
|
||||
DiffusionOutput(custom_output={"chunk": 1}, finished=True),
|
||||
]
|
||||
)
|
||||
|
||||
outputs = await self._run_streaming_client_proc(pipeline, request_id="req-stream")
|
||||
|
||||
assert [output.request_id for output in outputs] == ["req-stream", "req-stream"]
|
||||
assert [output.finished for output in outputs] == [False, True]
|
||||
assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
|
||||
assert pipeline.requests[0].request_id == "req-stream"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_midway_error_reaches_stage_client_after_prior_chunk(self) -> None:
|
||||
"""A pipeline error after the first chunk still delivers that chunk then a terminal error output."""
|
||||
pipeline = _FailingStreamingPipeline()
|
||||
|
||||
outputs = await self._run_streaming_client_proc(pipeline, request_id="req-error")
|
||||
|
||||
assert len(outputs) == 2
|
||||
assert outputs[0].request_id == "req-error"
|
||||
assert outputs[0].finished is False
|
||||
assert outputs[0].custom_output["chunk"] == 0
|
||||
assert outputs[1].request_id == "req-error"
|
||||
assert outputs[1].finished is True
|
||||
assert outputs[1].error is not None
|
||||
assert "stream failed midway" in outputs[1].error
|
||||
|
||||
@classmethod
|
||||
async def _run_streaming_client_proc(
|
||||
cls,
|
||||
pipeline: _StepStreamingPipeline | _FailingStreamingPipeline,
|
||||
*,
|
||||
request_id: str,
|
||||
) -> list[OmniRequestOutput]:
|
||||
request_address = get_open_zmq_ipc_path()
|
||||
response_address = get_open_zmq_ipc_path()
|
||||
proc = object.__new__(StageDiffusionProc)
|
||||
proc._od_config = SimpleNamespace(streaming_output=True)
|
||||
proc._engine = _PipelineBackedEngine(pipeline)
|
||||
proc._closed = False
|
||||
client = cls._make_client(request_address, response_address)
|
||||
proc_task = asyncio.create_task(proc.run_loop(request_address, response_address))
|
||||
try:
|
||||
await asyncio.sleep(0.05)
|
||||
await client.add_request_async(
|
||||
request_id,
|
||||
"prompt",
|
||||
OmniDiffusionSamplingParams(),
|
||||
)
|
||||
deadline = asyncio.get_running_loop().time() + 2.0
|
||||
outputs: list[OmniRequestOutput] = []
|
||||
while (not outputs or not outputs[-1].finished) and asyncio.get_running_loop().time() < deadline:
|
||||
output = client.get_diffusion_output_nowait()
|
||||
if output is not None:
|
||||
outputs.append(output)
|
||||
else:
|
||||
await asyncio.sleep(0.01)
|
||||
return outputs
|
||||
finally:
|
||||
try:
|
||||
client._request_socket.send(client._encoder.encode({"type": "shutdown"}))
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
await asyncio.wait_for(proc_task, timeout=2)
|
||||
finally:
|
||||
cls._close_client(client)
|
||||
|
||||
@staticmethod
|
||||
def _make_client(request_address: str, response_address: str) -> StageDiffusionClient:
|
||||
client = object.__new__(StageDiffusionClient)
|
||||
client.stage_id = 0
|
||||
client.replica_id = 0
|
||||
client.final_output = True
|
||||
client.final_output_type = "image"
|
||||
client.default_sampling_params = {}
|
||||
client.requires_multimodal_data = False
|
||||
client.custom_process_input_func = None
|
||||
client.engine_input_source = []
|
||||
client._proc = None
|
||||
client._proc_manager = None
|
||||
client._owns_process = False
|
||||
client._zmq_ctx = zmq.Context()
|
||||
client._request_socket = client._zmq_ctx.socket(zmq.PUSH)
|
||||
client._request_socket.bind(request_address)
|
||||
client._response_socket = client._zmq_ctx.socket(zmq.PULL)
|
||||
client._response_socket.bind(response_address)
|
||||
client._encoder = OmniMsgpackEncoder()
|
||||
client._decoder = OmniMsgpackDecoder()
|
||||
client._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
|
||||
client._rpc_results = {}
|
||||
client._pending_rpcs = set()
|
||||
client._tasks = {}
|
||||
client._shutting_down = False
|
||||
client._engine_dead = False
|
||||
return client
|
||||
|
||||
@staticmethod
|
||||
def _close_client(client: StageDiffusionClient) -> None:
|
||||
client._request_socket.close(linger=0)
|
||||
client._response_socket.close(linger=0)
|
||||
client._zmq_ctx.term()
|
||||
|
||||
|
||||
class TestPipelineStreamingOutputToEntrypoint:
|
||||
"""Streaming pipeline output through inline stage, orchestrator, and entrypoints."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_output_reaches_async_omni_from_mock_pipeline(self) -> None:
|
||||
"""Mock pipeline streaming chunks reach AsyncOmni.generate() via inline stage and orchestrator."""
|
||||
pipeline = _StepStreamingPipeline(
|
||||
[
|
||||
DiffusionOutput(custom_output={"chunk": 0}, finished=False),
|
||||
DiffusionOutput(custom_output={"chunk": 1}, finished=True),
|
||||
]
|
||||
)
|
||||
inline_client = self._make_inline_pipeline_client(pipeline)
|
||||
fixture = _build_harness([inline_client])
|
||||
omni = self._make_async_omni(self._OrchestratorBridgeEngine(fixture))
|
||||
|
||||
try:
|
||||
outputs: list[OmniRequestOutput] = []
|
||||
async for output in omni.generate(
|
||||
prompt={"prompt": "a cat"},
|
||||
request_id="req-omni",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
output_modalities=["image"],
|
||||
):
|
||||
outputs.append(output)
|
||||
|
||||
await _wait_for(lambda: len(pipeline.requests) == 1)
|
||||
assert pipeline.requests[0].request_id.startswith("req-omni-")
|
||||
assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
|
||||
assert [getattr(output.request_output, "finished", output.finished) for output in outputs] == [False, True]
|
||||
finally:
|
||||
await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_output_reaches_api_websocket_from_mock_pipeline(self, mocker: MockerFixture) -> None:
|
||||
"""Mock pipeline streaming chunks reach the video WebSocket handler through AsyncOmni.generate()."""
|
||||
frames = [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(2)]
|
||||
pipeline = _StepStreamingPipeline(
|
||||
[
|
||||
DiffusionOutput(custom_output={"chunk": 0, "images": [frames]}, finished=False),
|
||||
DiffusionOutput(custom_output={"chunk": 1, "images": [frames]}, finished=True),
|
||||
]
|
||||
)
|
||||
inline_client = self._make_inline_pipeline_client(pipeline)
|
||||
fixture = _build_harness([inline_client])
|
||||
bridge_engine = self._OrchestratorBridgeEngine(fixture)
|
||||
omni = self._make_async_omni(bridge_engine)
|
||||
|
||||
encoded = [b"pipeline-fmp4-0", b"pipeline-fmp4-1"]
|
||||
|
||||
class FakeStreamingVideoEncoder:
|
||||
encode_calls: list[int] = []
|
||||
|
||||
def encode(self, video):
|
||||
del video
|
||||
idx = len(self.encode_calls)
|
||||
self.encode_calls.append(idx)
|
||||
return encoded[idx]
|
||||
|
||||
def close(self):
|
||||
return b""
|
||||
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
|
||||
lambda *, output_format, fps, video_codec_options=None: FakeStreamingVideoEncoder(),
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
|
||||
return_value="diffusion",
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
|
||||
return_value=[OmniDiffusionSamplingParams()],
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list",
|
||||
return_value=[OmniDiffusionSamplingParams()],
|
||||
)
|
||||
|
||||
handler = OmniStreamingVideoOutputHandler(
|
||||
engine_client=omni,
|
||||
model_name="test-model",
|
||||
stage_configs=bridge_engine.stage_configs,
|
||||
stall_timeout=5.0,
|
||||
start_timeout=5.0,
|
||||
)
|
||||
app = FastAPI()
|
||||
|
||||
@app.websocket("/v1/realtime/video")
|
||||
async def ws_endpoint(websocket: WebSocket):
|
||||
await handler.handle_session(websocket)
|
||||
|
||||
try:
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect("/v1/realtime/video") as ws:
|
||||
ws.send_json({"type": "session.start", "prompt": "integration test"})
|
||||
assert ws.receive_json()["type"] == "video.start"
|
||||
assert ws.receive_bytes() == b"pipeline-fmp4-0"
|
||||
assert ws.receive_bytes() == b"pipeline-fmp4-1"
|
||||
done = ws.receive_json()
|
||||
assert done["type"] == "session.done"
|
||||
assert done["chunks"] == 2
|
||||
|
||||
await _wait_for(lambda: len(pipeline.requests) == 1)
|
||||
assert pipeline.requests[0].request_id
|
||||
finally:
|
||||
await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client)
|
||||
|
||||
class _OrchestratorBridgeEngine:
|
||||
"""Minimal AsyncOmni engine facade backed by a live Orchestrator harness."""
|
||||
|
||||
def __init__(self, fixture: OrchestratorFixture) -> None:
|
||||
self._fixture = fixture
|
||||
self.stage_metadata = [
|
||||
StageRuntimeInfo(
|
||||
stage_type="diffusion",
|
||||
final_output=True,
|
||||
final_output_type="image",
|
||||
)
|
||||
]
|
||||
self.stage_configs: list[Any] = [SimpleNamespace(stage_type="diffusion")]
|
||||
self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
|
||||
self.num_stages = 1
|
||||
self.supported_tasks = ("generate",)
|
||||
self._alive = True
|
||||
|
||||
async def add_request_async(
|
||||
self,
|
||||
*,
|
||||
request_id: str,
|
||||
prompt: Any,
|
||||
sampling_params_list: list[Any],
|
||||
final_stage_id: int,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
self._fixture.request_sync_q.put_nowait(
|
||||
StageSubmissionMessage(
|
||||
type="add_request",
|
||||
request_id=request_id,
|
||||
prompt=prompt,
|
||||
original_prompt=prompt,
|
||||
output_prompt_text=None,
|
||||
sampling_params_list=sampling_params_list,
|
||||
final_stage_id=final_stage_id,
|
||||
final_output_stage_ids=kwargs.get("final_output_stage_ids"),
|
||||
preprocess_ms=0.0,
|
||||
request_timestamp=kwargs.get("arrival_time", time.time()),
|
||||
enqueue_ts=time.perf_counter(),
|
||||
)
|
||||
)
|
||||
|
||||
async def try_get_output_async(self) -> Any | None:
|
||||
try:
|
||||
return self._fixture.output_sync_q.get_nowait()
|
||||
except queue.Empty:
|
||||
return None
|
||||
|
||||
def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo:
|
||||
return self.stage_metadata[stage_id]
|
||||
|
||||
def is_alive(self) -> bool:
|
||||
return self._fixture.thread.is_alive()
|
||||
|
||||
async def abort_async(self, request_ids: list[str]) -> None:
|
||||
del request_ids
|
||||
|
||||
@classmethod
|
||||
def _make_inline_pipeline_client(cls, pipeline: _StepStreamingPipeline) -> InlineStageDiffusionClient:
|
||||
"""Build an inline diffusion stage client whose engine runs the given mock pipeline."""
|
||||
metadata = StageMetadata(
|
||||
stage_id=0,
|
||||
stage_type="diffusion",
|
||||
engine_output_type="image",
|
||||
is_comprehension=False,
|
||||
requires_multimodal_data=False,
|
||||
engine_input_source=[],
|
||||
final_output=True,
|
||||
final_output_type="image",
|
||||
default_sampling_params=OmniDiffusionSamplingParams(),
|
||||
custom_process_input_func=None,
|
||||
model_stage=None,
|
||||
runtime_cfg=None,
|
||||
)
|
||||
pipeline_engine = _PipelineBackedEngine(pipeline)
|
||||
with patch.object(InlineStageDiffusionClient, "_enrich_config"):
|
||||
with patch(
|
||||
"vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine.make_engine",
|
||||
return_value=pipeline_engine,
|
||||
):
|
||||
od_config = MagicMock(spec=OmniDiffusionConfig)
|
||||
od_config.streaming_output = True
|
||||
return InlineStageDiffusionClient(
|
||||
model="test_model",
|
||||
od_config=od_config,
|
||||
metadata=metadata,
|
||||
batch_size=1,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _make_async_omni(engine: _OrchestratorBridgeEngine) -> AsyncOmni:
|
||||
omni = object.__new__(AsyncOmni)
|
||||
omni.engine = engine
|
||||
omni.log_stats = False
|
||||
omni._pause_cond = asyncio.Condition()
|
||||
omni._paused = False
|
||||
omni.request_states = {}
|
||||
omni.final_output_task = None
|
||||
omni.event_resolver = AsyncEventResolver()
|
||||
omni._enable_ar_profiler = False
|
||||
omni._is_sleeping = False
|
||||
omni.prom_metrics = MagicMock()
|
||||
omni.mod_metrics = MagicMock()
|
||||
omni.resolve_sampling_params_list = lambda params, allow_delta_coercion: params
|
||||
omni._compute_final_stage_id = lambda output_modalities: 0
|
||||
omni._compute_final_output_stage_ids = lambda output_modalities: [0]
|
||||
omni.default_sampling_params_list = engine.default_sampling_params_list
|
||||
omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None)
|
||||
return omni
|
||||
|
||||
@staticmethod
|
||||
async def _shutdown_pipeline_omni_harness(
|
||||
omni: AsyncOmni,
|
||||
fixture: OrchestratorFixture,
|
||||
inline_client: InlineStageDiffusionClient,
|
||||
) -> None:
|
||||
if omni.final_output_task is not None:
|
||||
omni.final_output_task.cancel()
|
||||
await asyncio.gather(omni.final_output_task, return_exceptions=True)
|
||||
inline_client.shutdown()
|
||||
fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
|
||||
await asyncio.to_thread(fixture.thread.join, 5)
|
||||
|
||||
|
||||
def _make_vllm_config():
|
||||
@contextlib.contextmanager
|
||||
def set_priority(*args, **kwargs):
|
||||
del args, kwargs
|
||||
yield
|
||||
|
||||
return SimpleNamespace(
|
||||
kernel_config=SimpleNamespace(ir_op_priority=SimpleNamespace(set_priority=set_priority)),
|
||||
compilation_config=SimpleNamespace(ir_enable_torch_wrap=True),
|
||||
)
|
||||
|
||||
|
||||
class TestSupportedPipelines:
|
||||
"""Streaming-output protocol checks for supported pipelines."""
|
||||
|
||||
def test_helios_supports_step_execution_for_streaming_output(self) -> None:
|
||||
from vllm_omni.diffusion.models.helios.pipeline_helios import HeliosPipeline
|
||||
from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
|
||||
|
||||
# Avoid loading model weights; protocol membership depends on the class contract.
|
||||
pipeline = object.__new__(HeliosPipeline)
|
||||
|
||||
assert pipeline.supports_step_execution is True
|
||||
assert supports_step_execution(pipeline) is True
|
||||
assert isinstance(pipeline, SupportsStepExecution) is True
|
||||
|
||||
def test_load_model_rejects_streaming_output_without_step_execution(self, monkeypatch) -> None:
|
||||
class _NoStepPipeline:
|
||||
def forward(self): ...
|
||||
|
||||
class _FakeDiffusersPipelineLoader:
|
||||
def __init__(self, *args, **kwargs): ...
|
||||
|
||||
def load_model(self, **kwargs):
|
||||
return _NoStepPipeline()
|
||||
|
||||
class _FakeDeviceMemoryProfiler:
|
||||
consumed_memory = 0
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
runner = object.__new__(DiffusionModelRunner)
|
||||
runner.vllm_config = _make_vllm_config()
|
||||
runner.od_config = SimpleNamespace(
|
||||
enable_cpu_offload=False,
|
||||
enable_layerwise_offload=False,
|
||||
enforce_eager=True,
|
||||
cache_backend=None,
|
||||
cache_config=None,
|
||||
streaming_output=True,
|
||||
step_execution=False,
|
||||
model_class_name="NoStepPipeline",
|
||||
parallel_config=SimpleNamespace(use_hsdp=False),
|
||||
)
|
||||
runner.device = torch.device("cpu")
|
||||
runner.pipeline = None
|
||||
runner.cache_backend = None
|
||||
runner.offload_backend = None
|
||||
runner.kv_transfer_manager = SimpleNamespace()
|
||||
|
||||
monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeDiffusersPipelineLoader)
|
||||
monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeDeviceMemoryProfiler)
|
||||
monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None)
|
||||
|
||||
with pytest.raises(ValueError, match="NoStepPipeline"):
|
||||
DiffusionModelRunner.load_model(runner)
|
||||
@@ -46,6 +46,7 @@ def client(mock_engine):
|
||||
)
|
||||
with patch.object(InlineStageDiffusionClient, "_enrich_config"):
|
||||
od_config = MagicMock(spec=OmniDiffusionConfig)
|
||||
od_config.streaming_output = False
|
||||
c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1)
|
||||
yield c
|
||||
c.shutdown()
|
||||
@@ -72,6 +73,35 @@ async def test_inline_dispatch_request_success(client, mock_engine):
|
||||
mock_engine.step.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inline_dispatch_request_streaming_success(client, mock_engine):
|
||||
"""Inline StageDiffusionClient correctly receives streaming chunk output from Diffusion Engine"""
|
||||
chunks = [
|
||||
OmniRequestOutput.from_diffusion(request_id="req-stream", images=[], finished=False),
|
||||
OmniRequestOutput.from_diffusion(request_id="req-stream", images=[], finished=True),
|
||||
]
|
||||
|
||||
async def _step_streaming(_request):
|
||||
for chunk in chunks:
|
||||
yield [chunk]
|
||||
|
||||
client.od_config.streaming_output = True
|
||||
mock_engine.step_streaming = _step_streaming
|
||||
|
||||
await client.add_request_async("req-stream", "A test prompt", OmniDiffusionSamplingParams())
|
||||
|
||||
outputs = []
|
||||
for _ in range(20):
|
||||
output = client.get_diffusion_output_nowait()
|
||||
if output is not None:
|
||||
outputs.append(output)
|
||||
if output.finished:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert outputs == chunks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inline_dispatch_request_error(client, mock_engine):
|
||||
# Setup mock engine step to raise an exception
|
||||
|
||||
@@ -21,6 +21,7 @@ from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE
|
||||
from vllm_omni.diffusion.sched import RequestScheduler
|
||||
from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
|
||||
from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc
|
||||
from vllm_omni.diffusion.worker.utils import RunnerOutput
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
@@ -49,7 +50,7 @@ def _make_executor(num_gpus: int = 1):
|
||||
|
||||
Returns ``(executor, request_queue, result_queue)``.
|
||||
"""
|
||||
od_cfg = SimpleNamespace(num_gpus=num_gpus)
|
||||
od_cfg = SimpleNamespace(num_gpus=num_gpus, streaming_output=False)
|
||||
monkeypatch = pytest.MonkeyPatch()
|
||||
monkeypatch.setattr(MultiprocDiffusionExecutor, "_init_executor", lambda self: None)
|
||||
executor = MultiprocDiffusionExecutor(od_cfg)
|
||||
@@ -75,6 +76,7 @@ def _make_engine(num_gpus: int = 1):
|
||||
"""Create a lightweight ``DiffusionEngine`` wired to mocked executor."""
|
||||
executor, req_q, res_q = _make_executor(num_gpus)
|
||||
engine = DiffusionEngine.__new__(DiffusionEngine)
|
||||
engine.od_config = SimpleNamespace(streaming_output=False)
|
||||
sched = RequestScheduler()
|
||||
sched.initialize(SimpleNamespace())
|
||||
engine.scheduler = sched
|
||||
@@ -581,6 +583,58 @@ class TestMultiprocExecutorRaisesEngineDeadError:
|
||||
)
|
||||
|
||||
|
||||
class TestMultiprocExecutorStepStreamingOutput:
|
||||
"""Streaming output uses step execution and one worker reply per step."""
|
||||
|
||||
def test_execute_step_allows_streaming_output_mode(self):
|
||||
executor, req_q, res_q = _make_executor()
|
||||
executor.od_config = SimpleNamespace(streaming_output=True) # pyright: ignore[reportAttributeAccessIssue]
|
||||
runner_outputs = [
|
||||
RunnerOutput(
|
||||
request_id="sched-stream",
|
||||
step_index=1,
|
||||
finished=False,
|
||||
result=DiffusionOutput(output={"chunk": 0}, finished=False, chunk_index=0, total_chunks=2),
|
||||
),
|
||||
RunnerOutput(
|
||||
request_id="sched-stream",
|
||||
step_index=2,
|
||||
finished=True,
|
||||
result=DiffusionOutput(output={"chunk": 1}, finished=True, chunk_index=1, total_chunks=2),
|
||||
),
|
||||
]
|
||||
scheduler_output = SimpleNamespace(
|
||||
scheduled_request_ids=["sched-stream"],
|
||||
)
|
||||
|
||||
def _worker():
|
||||
for runner_output in runner_outputs:
|
||||
req_q.get(timeout=10)
|
||||
res_q.put(runner_output)
|
||||
|
||||
thread = threading.Thread(target=_worker, daemon=True)
|
||||
thread.start()
|
||||
|
||||
first: RunnerOutput = MultiprocDiffusionExecutor.execute_step( # pyright: ignore[reportAssignmentType]
|
||||
executor,
|
||||
scheduler_output, # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
second: RunnerOutput = MultiprocDiffusionExecutor.execute_step( # pyright: ignore[reportAssignmentType]
|
||||
executor,
|
||||
scheduler_output, # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
|
||||
assert first is runner_outputs[0]
|
||||
assert first.result is not None
|
||||
assert first.result.output == {"chunk": 0}
|
||||
assert first.finished is False
|
||||
assert second is runner_outputs[1]
|
||||
assert second.result is not None
|
||||
assert second.result.output == {"chunk": 1}
|
||||
assert second.finished is True
|
||||
thread.join(timeout=2)
|
||||
|
||||
|
||||
class TestDiffusionEngineDeadErrorPassthrough:
|
||||
"""``DiffusionEngine.add_req_and_wait_for_response`` re-raises
|
||||
``EngineDeadError`` from executor and wraps other errors."""
|
||||
|
||||
@@ -11,6 +11,7 @@ import pytest
|
||||
|
||||
from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
@@ -54,6 +55,7 @@ def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
|
||||
|
||||
proc = object.__new__(StageDiffusionProc)
|
||||
proc._engine = SimpleNamespace(step=step)
|
||||
proc._od_config = SimpleNamespace(streaming_output=False)
|
||||
proc._executor = ThreadPoolExecutor(max_workers=1)
|
||||
|
||||
try:
|
||||
@@ -102,6 +104,40 @@ class MockDiffusionEngine:
|
||||
return [MockOmniRequestOutput(request_id=request.request_id)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proc_streaming_request_yields_each_engine_chunk():
|
||||
"""Ensure that the streaming output chunks from DiffusionEngine reaches StageDiffusionProc"""
|
||||
captured = {}
|
||||
chunks = [
|
||||
OmniRequestOutput.from_diffusion(request_id="", images=[], finished=False),
|
||||
OmniRequestOutput.from_diffusion(request_id="", images=[], finished=True),
|
||||
]
|
||||
|
||||
class _StreamingEngine:
|
||||
async def step_streaming(self, request):
|
||||
captured["request"] = request
|
||||
for chunk in chunks:
|
||||
yield [chunk]
|
||||
|
||||
stage_proc = object.__new__(StageDiffusionProc)
|
||||
stage_proc._engine = _StreamingEngine()
|
||||
|
||||
outputs = [
|
||||
output
|
||||
async for output in stage_proc._process_streaming_request(
|
||||
request_id="req-stream",
|
||||
prompt="prompt",
|
||||
sampling_params_dict=asdict(OmniDiffusionSamplingParams()),
|
||||
kv_sender_info={0: {"host": "127.0.0.1"}},
|
||||
)
|
||||
]
|
||||
|
||||
assert outputs == chunks
|
||||
assert [output.request_id for output in outputs] == ["req-stream", "req-stream"]
|
||||
assert [output.finished for output in outputs] == [False, True]
|
||||
assert captured["request"].kv_sender_info == {0: {"host": "127.0.0.1"}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proc_process_request_with_batching_async_output():
|
||||
stage_proc = object.__new__(StageDiffusionProc)
|
||||
|
||||
@@ -305,7 +305,7 @@ def assert_video_similarity_metrics(
|
||||
offline_path: Path,
|
||||
ssim_threshold: float,
|
||||
psnr_threshold: float,
|
||||
) -> None:
|
||||
) -> tuple[float, float]:
|
||||
ssim_output = run_ffmpeg_similarity("ssim", online_path, offline_path)
|
||||
psnr_output = run_ffmpeg_similarity("psnr", online_path, offline_path)
|
||||
ssim_score = parse_ssim_score(ssim_output)
|
||||
@@ -339,6 +339,7 @@ def assert_video_similarity_metrics(
|
||||
f"PSNR below threshold: got {psnr_score:.6f}, expected >= {psnr_threshold:.6f}. "
|
||||
f"online={online_path} offline={offline_path}"
|
||||
)
|
||||
return ssim_score, psnr_score
|
||||
|
||||
|
||||
def assert_similarity(
|
||||
|
||||
@@ -33,7 +33,7 @@ Prerequisites
|
||||
|
||||
* **Seed-TTS** optional extras for WER/SIM/UTMOS::
|
||||
|
||||
pip install 'vllm-omni[seed-tts-eval]'
|
||||
pip install 'vllm-omni[dev]'
|
||||
|
||||
Examples
|
||||
--------
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
"""
|
||||
Online serving smoke test for streaming video output similarity.
|
||||
"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.e2e.accuracy.helpers import assert_video_similarity_metrics, probe_video
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL = "BestWishYsh/Helios-Distilled"
|
||||
PROMPT = "A serene lakeside sunrise with mist over the water."
|
||||
WIDTH = 640
|
||||
HEIGHT = 384
|
||||
FPS = 16
|
||||
NUM_FRAMES = 65
|
||||
NUM_INFERENCE_STEPS = 5
|
||||
SEED = 42
|
||||
MIN_PSNR_DB = 28.0
|
||||
MIN_SSIM = 0.95
|
||||
|
||||
SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||
|
||||
HELIOS_DISTILLED_EXTRA_PARAMS = {
|
||||
"is_enable_stage2": True,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [1, 1, 1],
|
||||
"is_amplify_first_chunk": True,
|
||||
}
|
||||
|
||||
|
||||
def _server_args(*, streaming_output: bool) -> list[str]:
|
||||
args = [
|
||||
"--stage-init-timeout",
|
||||
"600",
|
||||
"--init-timeout",
|
||||
"900",
|
||||
"--log-stats",
|
||||
]
|
||||
if streaming_output:
|
||||
args.append("--diffusion-streaming-output")
|
||||
return args
|
||||
|
||||
|
||||
def _request_config(model: str, *, serialize_extra_params: bool) -> dict:
|
||||
extra_params = (
|
||||
json.dumps(HELIOS_DISTILLED_EXTRA_PARAMS) if serialize_extra_params else HELIOS_DISTILLED_EXTRA_PARAMS
|
||||
)
|
||||
return {
|
||||
"model": model,
|
||||
"form_data": {
|
||||
"prompt": PROMPT,
|
||||
"width": WIDTH,
|
||||
"height": HEIGHT,
|
||||
"fps": FPS,
|
||||
"num_frames": NUM_FRAMES,
|
||||
"num_inference_steps": NUM_INFERENCE_STEPS,
|
||||
"guidance_scale": 1.0,
|
||||
"seed": SEED,
|
||||
"extra_params": extra_params,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _client_for(server: OmniServer, *, run_level: str) -> OpenAIClientHandler:
|
||||
return OpenAIClientHandler(
|
||||
host=server.host,
|
||||
port=server.port,
|
||||
api_key="EMPTY",
|
||||
run_level=run_level,
|
||||
log_stats=server.log_stats,
|
||||
)
|
||||
|
||||
|
||||
def _write_video_artifact(response: DiffusionResponse, output_path: Path) -> Path:
|
||||
assert response.videos, "No video artifact returned"
|
||||
output_path.write_bytes(response.videos[0])
|
||||
return output_path
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
pytest.param(
|
||||
MODEL,
|
||||
id="helios_distilled",
|
||||
marks=SINGLE_CARD_MARKS,
|
||||
)
|
||||
],
|
||||
)
|
||||
def test_helios_streaming_output_matches_non_streaming(
|
||||
model: str,
|
||||
model_prefix: str,
|
||||
run_level: str,
|
||||
tmp_path: Path,
|
||||
record_property,
|
||||
):
|
||||
model = model_prefix + model
|
||||
|
||||
with OmniServer(model, _server_args(streaming_output=True), use_omni=True) as server:
|
||||
client = _client_for(server, run_level=run_level)
|
||||
streaming_response = client.send_streaming_video_diffusion_request(
|
||||
_request_config(server.model, serialize_extra_params=False)
|
||||
)[0]
|
||||
|
||||
streaming_path = _write_video_artifact(streaming_response, tmp_path / "helios_streaming.mp4")
|
||||
record_property("helios_streaming_video", str(streaming_path))
|
||||
record_property("helios_streaming_latency_s", streaming_response.e2e_latency)
|
||||
|
||||
with OmniServer(model, _server_args(streaming_output=False), use_omni=True) as server:
|
||||
client = _client_for(server, run_level=run_level)
|
||||
non_streaming_response = client.send_video_diffusion_request(
|
||||
_request_config(server.model, serialize_extra_params=True)
|
||||
)[0]
|
||||
|
||||
non_streaming_path = _write_video_artifact(non_streaming_response, tmp_path / "helios_non_streaming.mp4")
|
||||
record_property("helios_non_streaming_video", str(non_streaming_path))
|
||||
record_property("helios_non_streaming_latency_s", non_streaming_response.e2e_latency)
|
||||
|
||||
streaming_metadata = probe_video(streaming_path)
|
||||
non_streaming_metadata = probe_video(non_streaming_path)
|
||||
assert streaming_metadata == non_streaming_metadata, (
|
||||
f"Video metadata mismatch for helios_streaming_output:\n"
|
||||
f"streaming={streaming_metadata}\n"
|
||||
f"non_streaming={non_streaming_metadata}\n"
|
||||
f"streaming_path={streaming_path}\n"
|
||||
f"non_streaming_path={non_streaming_path}"
|
||||
)
|
||||
|
||||
ssim, psnr = assert_video_similarity_metrics(
|
||||
label="helios_streaming_output",
|
||||
online_path=streaming_path,
|
||||
offline_path=non_streaming_path,
|
||||
ssim_threshold=MIN_SSIM,
|
||||
psnr_threshold=MIN_PSNR_DB,
|
||||
)
|
||||
record_property("helios_streaming_vs_non_streaming_psnr_db", psnr)
|
||||
record_property("helios_streaming_vs_non_streaming_ssim_percent", ssim * 100.0)
|
||||
@@ -0,0 +1,96 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""E2E offline inference tests for IndexTTS2 two-stage pipeline.
|
||||
|
||||
Stage 0 (GPT AR) → Stage 1 (S2Mel + BigVGAN). Output is 22050 Hz mono WAV.
|
||||
Covers: basic TTS.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.model_executor.models.indextts2.prompt_utils import (
|
||||
build_indextts2_prefill_prompt_ids,
|
||||
)
|
||||
|
||||
MODEL_NAME = "IndexTeam/IndexTTS-2"
|
||||
STAGE_CONFIG = get_deploy_config_path("indextts2.yaml")
|
||||
|
||||
_OMNI_RUNNER_PARAM = (
|
||||
MODEL_NAME,
|
||||
STAGE_CONFIG,
|
||||
)
|
||||
|
||||
ASSETS_DIR = Path(__file__).resolve().parents[2] / "assets" / "indextts2"
|
||||
REF_AUDIO = str(ASSETS_DIR / "ref_audio.wav")
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.advanced_model,
|
||||
pytest.mark.tts,
|
||||
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
|
||||
]
|
||||
|
||||
SAMPLE_RATE = 22050
|
||||
|
||||
|
||||
def _audio_from_mm(mm: dict) -> torch.Tensor | None:
|
||||
audio = mm.get("audio")
|
||||
if audio is None:
|
||||
audio = mm.get("model_outputs")
|
||||
if isinstance(audio, list):
|
||||
chunks = [chunk.reshape(-1) for chunk in audio if isinstance(chunk, torch.Tensor) and chunk.numel() > 0]
|
||||
return torch.cat(chunks, dim=0) if chunks else None
|
||||
return audio if isinstance(audio, torch.Tensor) else None
|
||||
|
||||
|
||||
def _sample_rate_from_mm(mm: dict) -> int:
|
||||
sr = mm.get("sr")
|
||||
if isinstance(sr, list):
|
||||
sr = sr[-1] if sr else None
|
||||
if hasattr(sr, "item"):
|
||||
return int(sr.item())
|
||||
return int(sr) if sr is not None else SAMPLE_RATE
|
||||
|
||||
|
||||
def _build_request(text: str, ref_audio: str, **extra) -> dict:
|
||||
info = {
|
||||
"text": [text],
|
||||
"voice": [ref_audio],
|
||||
}
|
||||
info.update(extra)
|
||||
return {
|
||||
"prompt_token_ids": build_indextts2_prefill_prompt_ids(MODEL_NAME, text),
|
||||
"additional_information": info,
|
||||
}
|
||||
|
||||
|
||||
def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
|
||||
for omni_out in omni.generate(request):
|
||||
mm = omni_out.multimodal_output
|
||||
assert mm is not None, "Expected multimodal_output"
|
||||
audio = _audio_from_mm(mm)
|
||||
assert audio is not None, "Expected audio output"
|
||||
return audio.cpu(), _sample_rate_from_mm(mm)
|
||||
raise AssertionError("No outputs received")
|
||||
|
||||
|
||||
def _assert_valid_audio(audio: torch.Tensor, sr: int) -> None:
|
||||
assert sr == SAMPLE_RATE
|
||||
assert audio.numel() > 0
|
||||
assert torch.any(audio != 0).item()
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"})
|
||||
def test_basic_english(omni_runner: OmniRunner) -> None:
|
||||
"""English TTS: text + ref_audio, no emotion conditioning."""
|
||||
req = _build_request("Hello, this is a voice synthesis test.", REF_AUDIO)
|
||||
audio, sr = _collect_audio(omni_runner.omni, req)
|
||||
_assert_valid_audio(audio, sr)
|
||||
@@ -14,7 +14,7 @@ from tests.helpers.stage_config import get_deploy_config_path, modify_stage_conf
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
|
||||
|
||||
thinker_only_models = ["Qwen/Qwen3-Omni-30B-A3B-Captioner"]
|
||||
|
||||
# Single CI deploy YAML; rocm/xpu deltas are picked automatically via the
|
||||
# platforms: section. Only CUDA needs an extra enforce_eager tweak.
|
||||
@@ -40,6 +40,10 @@ else:
|
||||
|
||||
# Create parameter combinations for model and stage config
|
||||
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
|
||||
# we can use the same config for a model that only has thinker (i.e., does not
|
||||
# enable audio output) because the resolver should figure out that it doesn't
|
||||
# need the full pipeline based on the HF config.
|
||||
thinker_test_params = [(model, stage_config) for model in thinker_only_models for stage_config in stage_configs]
|
||||
|
||||
|
||||
def get_question(prompt_type="video"):
|
||||
@@ -61,3 +65,15 @@ def test_video_to_audio(omni_runner, omni_runner_handler) -> None:
|
||||
|
||||
# Test single completion
|
||||
omni_runner_handler.send_omni_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_runner", thinker_test_params, indirect=True)
|
||||
def test_thinker_only_model_request(omni_runner, omni_runner_handler) -> None:
|
||||
"""Test that we can load and run a request through a model that only has the thinker stage."""
|
||||
request_config = {"prompts": "what color is the sky?", "modalities": ["text"]}
|
||||
|
||||
# Test single completion
|
||||
omni_runner_handler.send_omni_request(request_config)
|
||||
@@ -0,0 +1,241 @@
|
||||
"""E2E offline inference tests for SoulX-Singer (single-stage, preprocess inline)."""
|
||||
|
||||
import functools
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.media import get_asset_path
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
PROMPT_AUDIO = get_asset_path("soulxsinger/zh_prompt.mp3")
|
||||
TARGET_AUDIO = get_asset_path("soulxsinger/music.mp3")
|
||||
PROMPT_METADATA = get_asset_path("soulxsinger/zh_prompt.json")
|
||||
TARGET_METADATA = get_asset_path("soulxsinger/music.json")
|
||||
SAMPLE_RATE = 24_000
|
||||
|
||||
# phone_set.json (phoneme vocab) is not shipped on HuggingFace; SVS loads it
|
||||
# from the model dir. Stage it from this pinned upstream commit so the SVS DiT
|
||||
# can run in CI; the SVS tests skip gracefully if it cannot be fetched.
|
||||
_PHONE_SET_URL = (
|
||||
"https://raw.githubusercontent.com/Soul-AILab/SoulX-Singer/"
|
||||
"81aeb3ae772c70093c3de74dc23c92d983801ae4/soulxsinger/utils/phoneme/phone_set.json"
|
||||
)
|
||||
|
||||
if not PROMPT_AUDIO.is_file() or not TARGET_AUDIO.is_file():
|
||||
pytest.skip(
|
||||
f"Missing SoulX-Singer audio assets: {PROMPT_AUDIO.name}, {TARGET_AUDIO.name}",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.tts]
|
||||
|
||||
_CASES = (
|
||||
pytest.param(
|
||||
"SoulXSingerPipeline",
|
||||
"soulxsinger_svs.yaml",
|
||||
{
|
||||
"language": "Mandarin",
|
||||
"vocal_sep": False,
|
||||
"control": "score",
|
||||
"auto_shift": False,
|
||||
"pitch_shift": 0,
|
||||
},
|
||||
("g2pM", "g2p_en"),
|
||||
id="svs",
|
||||
),
|
||||
pytest.param(
|
||||
"SoulXSingerSVCPipeline",
|
||||
"soulxsinger_svc.yaml",
|
||||
{"vocal_sep": False, "auto_shift": False, "pitch_shift": 0},
|
||||
(),
|
||||
id="svc",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _resolve_weights() -> tuple[Path, Path, Path]:
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
base = Path(snapshot_download("Soul-AILab/SoulX-Singer", allow_patterns=["*"]))
|
||||
|
||||
# phone_set.json is not on HF; best-effort stage it from pinned upstream so
|
||||
# SVS can load. SVS tests skip (not fail) if both this and a manual copy are absent.
|
||||
phone_set = base / "phoneme" / "phone_set.json"
|
||||
if not phone_set.is_file() and not (base / "phone_set.json").is_file():
|
||||
try:
|
||||
from urllib.request import urlretrieve
|
||||
|
||||
phone_set.parent.mkdir(exist_ok=True)
|
||||
urlretrieve(_PHONE_SET_URL, phone_set)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# make a temporary svc directory
|
||||
svc_dir = base.parent / "SoulX-Singer-SVC"
|
||||
svc_dir.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copyfile(base / "config.yaml", svc_dir / "config.yaml")
|
||||
shutil.move(base / "model-svc.pt", svc_dir / "model-svc.pt")
|
||||
|
||||
if raw := os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR"):
|
||||
pre = Path(raw).expanduser().resolve()
|
||||
if (pre / "rmvpe" / "rmvpe.pt").is_file():
|
||||
return base, svc_dir, pre
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
pre = Path(snapshot_download("Soul-AILab/SoulX-Singer-Preprocess", allow_patterns=["*"]))
|
||||
return base, svc_dir, pre
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def soulx_weights() -> tuple[Path, Path, Path]:
|
||||
try:
|
||||
return _resolve_weights()
|
||||
except Exception as exc:
|
||||
pytest.skip(f"Set SOULXSINGER_MODEL_DIR / SOULX_PREPROCESS_WEIGHTS_DIR. ({exc})")
|
||||
|
||||
|
||||
def _flatten_audio(audio_val) -> np.ndarray:
|
||||
import torch
|
||||
|
||||
if isinstance(audio_val, list):
|
||||
chunks = [c.detach().cpu().float().numpy().reshape(-1) for c in audio_val if c is not None]
|
||||
return np.concatenate(chunks) if chunks else np.array([], dtype=np.float32)
|
||||
if isinstance(audio_val, torch.Tensor):
|
||||
return audio_val.detach().cpu().float().numpy().reshape(-1)
|
||||
return np.asarray(audio_val, dtype=np.float32).reshape(-1)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("architecture,deploy_yaml,extra_args,py_deps", _CASES)
|
||||
def test_soulxsinger_multistage_from_audio(
|
||||
soulx_weights: tuple[Path, Path, Path],
|
||||
architecture: str,
|
||||
deploy_yaml: str,
|
||||
extra_args: dict,
|
||||
py_deps: tuple[str, ...],
|
||||
) -> None:
|
||||
for mod in py_deps:
|
||||
try:
|
||||
importlib.import_module(mod)
|
||||
except ImportError as exc:
|
||||
pytest.fail(f"SoulX SVS requires {mod}: {exc}")
|
||||
|
||||
base_dir, svc_dir, preprocess_dir = soulx_weights
|
||||
|
||||
# SVS mode requires phone_set.json in the model directory
|
||||
if architecture == "SoulXSingerPipeline":
|
||||
json.dump(
|
||||
{"model_type": "soulxsinger", "architectures": ["SoulXSingerPipeline"], "max_num_seqs": 1},
|
||||
open(base_dir / "config.json", "w"),
|
||||
indent=4,
|
||||
)
|
||||
model = str(base_dir)
|
||||
if not (base_dir / "phoneme" / "phone_set.json").is_file() and not (base_dir / "phone_set.json").is_file():
|
||||
pytest.skip(
|
||||
"SoulX-Singer SVS test requires phoneme/phone_set.json. "
|
||||
"Copy it from github.com/Soul-AILab/SoulX-Singer into the model dir. "
|
||||
"See `examples/offline_inference/text_to_speech/README.md` for details."
|
||||
)
|
||||
if architecture == "SoulXSingerSVCPipeline":
|
||||
json.dump(
|
||||
{"model_type": "soulxsinger", "architectures": ["SoulXSingerSVCPipeline"], "max_num_seqs": 1},
|
||||
open(svc_dir / "config.json", "w"),
|
||||
indent=4,
|
||||
)
|
||||
model = str(svc_dir)
|
||||
|
||||
with OmniRunner(
|
||||
model,
|
||||
stage_configs_path=get_deploy_config_path(deploy_yaml),
|
||||
async_chunk=False,
|
||||
) as runner:
|
||||
sampling = OmniDiffusionSamplingParams(
|
||||
num_inference_steps=4,
|
||||
guidance_scale=3.0,
|
||||
seed=42,
|
||||
extra_args={
|
||||
"prompt_audio": str(PROMPT_AUDIO),
|
||||
"target_audio": str(TARGET_AUDIO),
|
||||
"preprocess_weights_dir": str(preprocess_dir),
|
||||
**extra_args,
|
||||
},
|
||||
)
|
||||
prompt = {"prompt_token_ids": [0]}
|
||||
outputs = runner.generate([prompt], sampling)
|
||||
|
||||
assert outputs and outputs[0].error is None, outputs[0].error if outputs else "no output"
|
||||
mm = outputs[0].multimodal_output
|
||||
assert isinstance(mm, dict) and "audio" in mm
|
||||
audio = _flatten_audio(mm["audio"])
|
||||
assert 12_000 <= audio.size
|
||||
assert np.isfinite(audio).all() and float(np.max(np.abs(audio))) > 1e-4
|
||||
duration_s = audio.size / SAMPLE_RATE
|
||||
assert 50.0 <= duration_s <= 52.0, f"duration={duration_s:.1f}s"
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_soulxsinger_svs_precomputed(soulx_weights: tuple[Path, Path, Path]) -> None:
|
||||
"""SVS from precomputed note/lyric metadata.
|
||||
|
||||
Exercises the SVS DiT + vocoder serving path while skipping the heavy
|
||||
upstream preprocess (funasr ASR + ROSVOT note transcription), so it runs in
|
||||
the merge gate without those non-pip dependencies.
|
||||
"""
|
||||
base_dir, _svc_dir, preprocess_dir = soulx_weights
|
||||
|
||||
if not PROMPT_METADATA.is_file() or not TARGET_METADATA.is_file():
|
||||
pytest.skip(f"Missing SoulX-Singer metadata fixtures: {PROMPT_METADATA.name}, {TARGET_METADATA.name}")
|
||||
if not (base_dir / "phoneme" / "phone_set.json").is_file() and not (base_dir / "phone_set.json").is_file():
|
||||
pytest.skip(
|
||||
"SoulX-Singer SVS requires phoneme/phone_set.json (upstream fetch unavailable). "
|
||||
"See `examples/offline_inference/text_to_speech/README.md`."
|
||||
)
|
||||
|
||||
json.dump(
|
||||
{"model_type": "soulxsinger", "architectures": ["SoulXSingerPipeline"], "max_num_seqs": 1},
|
||||
open(base_dir / "config.json", "w"),
|
||||
indent=4,
|
||||
)
|
||||
|
||||
with OmniRunner(
|
||||
str(base_dir),
|
||||
stage_configs_path=get_deploy_config_path("soulxsinger_svs.yaml"),
|
||||
async_chunk=False,
|
||||
) as runner:
|
||||
sampling = OmniDiffusionSamplingParams(
|
||||
num_inference_steps=4,
|
||||
guidance_scale=3.0,
|
||||
seed=42,
|
||||
extra_args={
|
||||
"language": "Mandarin",
|
||||
"control": "score",
|
||||
"prompt_metadata_path": str(PROMPT_METADATA),
|
||||
"target_metadata_path": str(TARGET_METADATA),
|
||||
"audio_path": str(PROMPT_AUDIO),
|
||||
"preprocess_weights_dir": str(preprocess_dir),
|
||||
},
|
||||
)
|
||||
prompt = {"prompt_token_ids": [0]}
|
||||
outputs = runner.generate([prompt], sampling)
|
||||
|
||||
assert outputs and outputs[0].error is None, outputs[0].error if outputs else "no output"
|
||||
mm = outputs[0].multimodal_output
|
||||
assert isinstance(mm, dict) and "audio" in mm
|
||||
audio = _flatten_audio(mm["audio"])
|
||||
assert 12_000 <= audio.size
|
||||
assert np.isfinite(audio).all() and float(np.max(np.abs(audio))) > 1e-4
|
||||
duration_s = audio.size / SAMPLE_RATE
|
||||
assert 50.0 <= duration_s <= 52.0, f"duration={duration_s:.1f}s"
|
||||
@@ -0,0 +1,82 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""E2E online serving tests for IndexTTS2 via /v1/audio/speech endpoint.
|
||||
|
||||
Two-stage pipeline: GPT AR → S2Mel + BigVGAN. Output is 22050 Hz mono WAV.
|
||||
Covers: basic TTS and the speech endpoint's stream=True compatibility path.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.media import load_test_audio_data_url
|
||||
from tests.helpers.runtime import OmniServerParams
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
pytestmark = [pytest.mark.advanced_model, pytest.mark.tts]
|
||||
|
||||
MODEL = "IndexTeam/IndexTTS-2"
|
||||
|
||||
REF_AUDIO_URL = load_test_audio_data_url("indextts2/ref_audio.wav")
|
||||
|
||||
tts_server_params = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
stage_config_path=get_deploy_config_path("indextts2.yaml"),
|
||||
server_args=["--disable-log-stats"],
|
||||
),
|
||||
id="indextts2",
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
|
||||
def test_basic_english(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Basic English TTS: text + ref_audio, no emotion.
|
||||
Deploy Setting: default yaml
|
||||
Input Modal: text + reference audio
|
||||
Output Modal: audio (22050 Hz, WAV)
|
||||
Input Setting: stream=False
|
||||
"""
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"input": "Hello, this is a voice cloning demo.",
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"ref_audio": REF_AUDIO_URL,
|
||||
"min_audio_bytes": 1024,
|
||||
}
|
||||
openai_client.send_audio_speech_request(request_config)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
|
||||
def test_streaming(omni_server, openai_client) -> None:
|
||||
"""
|
||||
stream=True compatibility: IndexTTS2 itself is non-async-chunking, but the
|
||||
OpenAI speech endpoint should still accept the streaming request path and
|
||||
return a valid PCM response.
|
||||
|
||||
Deploy Setting: default yaml
|
||||
Input Modal: text + reference audio
|
||||
Output Modal: audio (22050 Hz, PCM)
|
||||
Input Setting: stream=True
|
||||
"""
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"input": "Hello, this is a voice cloning demo.",
|
||||
"stream": True,
|
||||
"response_format": "pcm",
|
||||
"ref_audio": REF_AUDIO_URL,
|
||||
"min_audio_bytes": 1024,
|
||||
}
|
||||
openai_client.send_audio_speech_request(request_config)
|
||||
@@ -0,0 +1,95 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Expansion tests for LTX-2 two-stage pipelines.
|
||||
|
||||
Coverage:
|
||||
- HSDP with LTX2TwoStagesPipeline (text-to-video)
|
||||
- HSDP with LTX2ImageToVideoTwoStagesPipeline (image-to-video)
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.media import generate_synthetic_image
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL = os.getenv("VLLM_TEST_LTX2_MODEL", "rootonchair/LTX-2-19b-distilled")
|
||||
PROMPT = "A serene lake at sunset with mountains in the background."
|
||||
PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
|
||||
HSDP_ARGS = ["--use-hsdp", "--hsdp-shard-size", "2"]
|
||||
|
||||
|
||||
def _cases():
|
||||
cases = []
|
||||
|
||||
# T2V: LTX2TwoStagesPipeline
|
||||
cases.append(
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
server_args=[
|
||||
*HSDP_ARGS,
|
||||
"--model-class-name",
|
||||
"LTX2TwoStagesPipeline",
|
||||
],
|
||||
),
|
||||
id="t2v_hsdp",
|
||||
marks=PARALLEL_MARKS,
|
||||
)
|
||||
)
|
||||
|
||||
# I2V: LTX2ImageToVideoTwoStagesPipeline
|
||||
cases.append(
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
server_args=[
|
||||
*HSDP_ARGS,
|
||||
"--model-class-name",
|
||||
"LTX2ImageToVideoTwoStagesPipeline",
|
||||
],
|
||||
),
|
||||
id="i2v_hsdp",
|
||||
marks=PARALLEL_MARKS,
|
||||
)
|
||||
)
|
||||
|
||||
return cases
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _cases(), indirect=True)
|
||||
def test_ltx2_two_stage_hsdp(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
):
|
||||
is_i2v = any("ImageToVideo" in arg for arg in omni_server.serve_args)
|
||||
|
||||
# The two-stage pipeline generates at the requested resolution then 2x
|
||||
# upsamples via LTX2LatentUpsamplerModel, so output dimensions differ
|
||||
# from the request. Omit height/width from form_data so the assertion
|
||||
# helper skips the dimension check (it only asserts when the key is
|
||||
# present). The pipeline falls back to its own defaults.
|
||||
form_data = {
|
||||
"prompt": PROMPT,
|
||||
"model": omni_server.model,
|
||||
"num_frames": 9,
|
||||
"fps": 8,
|
||||
"num_inference_steps": 2,
|
||||
"guidance_scale": 1.0,
|
||||
"seed": 42,
|
||||
}
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"form_data": form_data,
|
||||
}
|
||||
|
||||
if is_i2v:
|
||||
request_config["image_reference"] = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
|
||||
|
||||
openai_client.send_video_diffusion_request(request_config)
|
||||
@@ -9,13 +9,13 @@ from vllm.inputs import PromptType
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni.config.stage_config import _PIPELINE_REGISTRY
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
SEED = 42
|
||||
pipeline = _PIPELINE_REGISTRY["qwen2_5_omni_thinker_only"]
|
||||
pipeline = OMNI_PIPELINES["qwen2_5_omni_thinker_only"]
|
||||
|
||||
# Single-stage thinker-only deploy, materialized from tests.helpers.stage_config._CI_OVERLAYS.
|
||||
stage_config = get_deploy_config_path("ci/qwen2_5_omni_thinker_only.yaml")
|
||||
|
||||
@@ -523,6 +523,67 @@ async def test_run_single_stage_diffusion(orchestrator_factory) -> None:
|
||||
await _shutdown_orchestrator(orchestrator_fixture)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_single_stage_diffusion_streaming_forwards_intermediate_chunks(orchestrator_factory) -> None:
|
||||
"""Intermediate diffusion chunks (finished=False) reach the frontend before the final chunk."""
|
||||
stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
|
||||
orchestrator_fixture = orchestrator_factory([stage0])
|
||||
params = OmniDiffusionSamplingParams()
|
||||
|
||||
try:
|
||||
await _enqueue_add_request(
|
||||
orchestrator_fixture,
|
||||
request_id="req-stream",
|
||||
prompt={"prompt": "draw a cat"},
|
||||
original_prompt={"prompt": "draw a cat"},
|
||||
sampling_params_list=[params],
|
||||
final_stage_id=0,
|
||||
)
|
||||
|
||||
await _wait_for(lambda: len(stage0.add_request_calls) == 1)
|
||||
stage0.push_diffusion_output(
|
||||
OmniRequestOutput.from_diffusion(
|
||||
request_id="req-stream",
|
||||
images=[],
|
||||
final_output_type="image",
|
||||
custom_output={"chunk": 0},
|
||||
finished=False,
|
||||
)
|
||||
)
|
||||
stage0.push_diffusion_output(
|
||||
OmniRequestOutput.from_diffusion(
|
||||
request_id="req-stream",
|
||||
images=[],
|
||||
final_output_type="image",
|
||||
custom_output={"chunk": 1},
|
||||
finished=True,
|
||||
)
|
||||
)
|
||||
|
||||
output_msgs: list[OutputMessage] = []
|
||||
deadline = time.monotonic() + 2.0
|
||||
while not output_msgs or not output_msgs[-1].finished:
|
||||
if time.monotonic() >= deadline:
|
||||
raise AssertionError(
|
||||
f"Timed out waiting for finished orchestrator output, got {len(output_msgs)} message(s)"
|
||||
)
|
||||
try:
|
||||
msg = orchestrator_fixture.output_sync_q.get_nowait()
|
||||
except queue.Empty:
|
||||
await asyncio.sleep(0.01)
|
||||
continue
|
||||
if isinstance(msg, OutputMessage):
|
||||
output_msgs.append(msg)
|
||||
|
||||
assert [msg.request_id for msg in output_msgs] == ["req-stream", "req-stream"]
|
||||
assert [msg.finished for msg in output_msgs] == [False, True]
|
||||
assert [msg.engine_outputs.finished for msg in output_msgs] == [False, True]
|
||||
assert [msg.engine_outputs.custom_output["chunk"] for msg in output_msgs] == [0, 1]
|
||||
await _wait_for(lambda: "req-stream" not in orchestrator_fixture.orchestrator.request_states)
|
||||
finally:
|
||||
await _shutdown_orchestrator(orchestrator_fixture)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_llm_to_diffusion(orchestrator_factory) -> None:
|
||||
stage0 = FakeStageClient(stage_type="llm", final_output=False)
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.outputs import PoolingRequestOutput
|
||||
@@ -250,6 +251,25 @@ def test_no_detokenizer_make_request_output():
|
||||
assert AUDIO in result.outputs[0].multimodal_output
|
||||
|
||||
|
||||
def test_no_detokenizer_make_request_output_with_routed_experts():
|
||||
"""make_request_output accepts the routed_experts arg that the multimodal
|
||||
output channel (_process_mm_only_outputs) passes positionally, and attaches
|
||||
it to the completion output.
|
||||
|
||||
Regression: the call site passes 6 positional args
|
||||
(..., kv_transfer_params, routed_experts) but the override previously took
|
||||
only 5, raising TypeError on every generation-stage output.
|
||||
"""
|
||||
s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
|
||||
s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
|
||||
routed_experts = np.zeros((2, 3), dtype=np.int32)
|
||||
# Mirror the exact call shape of _process_mm_only_outputs.
|
||||
result = s.make_request_output([], None, FinishReason.STOP, None, None, routed_experts)
|
||||
assert result is not None
|
||||
assert not isinstance(result, PoolingRequestOutput)
|
||||
assert result.outputs[0].routed_experts is routed_experts
|
||||
|
||||
|
||||
def test_no_detokenizer_stream_interval_skipped():
|
||||
"""stream_interval > 1 does not crash when detokenizer is None."""
|
||||
kwargs = {**_NO_DETOK_STATE_KWARGS, "stream_interval": 5}
|
||||
|
||||
@@ -1301,6 +1301,58 @@ def test_generate_images_rejects_model_mismatch(test_client):
|
||||
assert "model mismatch" in response.json()["detail"].lower()
|
||||
|
||||
|
||||
def test_image_file_response_format_multiple(test_client):
|
||||
"""Test response_format=file with n>1 returns ZIP archive"""
|
||||
response = test_client.post(
|
||||
"/v1/images/generations",
|
||||
json={
|
||||
"prompt": "a dog",
|
||||
"n": 3,
|
||||
"response_format": "file",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "application/zip"
|
||||
assert "attachment" in response.headers.get("content-disposition", "")
|
||||
assert ".zip" in response.headers.get("content-disposition", "")
|
||||
|
||||
# Verify it's a valid ZIP with 3 PNG files
|
||||
import zipfile
|
||||
|
||||
zip_buffer = io.BytesIO(response.content)
|
||||
with zipfile.ZipFile(zip_buffer, "r") as zf:
|
||||
files = zf.namelist()
|
||||
assert len(files) == 3
|
||||
assert all(f.endswith(".png") for f in files)
|
||||
|
||||
# Verify each file is a valid PNG
|
||||
for filename in files:
|
||||
img_bytes = zf.read(filename)
|
||||
img = Image.open(io.BytesIO(img_bytes))
|
||||
assert img.format == "PNG"
|
||||
|
||||
|
||||
def test_image_file_response_format_single(test_client):
|
||||
"""Test response_format=file with n=1 returns a single image file."""
|
||||
response = test_client.post(
|
||||
"/v1/images/generations",
|
||||
json={
|
||||
"prompt": "a dog",
|
||||
"n": 1,
|
||||
"response_format": "file",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "image/png"
|
||||
assert "attachment" in response.headers.get("content-disposition", "")
|
||||
assert ".png" in response.headers.get("content-disposition", "")
|
||||
|
||||
img = Image.open(io.BytesIO(response.content))
|
||||
assert img.format == "PNG"
|
||||
|
||||
|
||||
def make_test_image_bytes(size=(64, 64)) -> bytes:
|
||||
img = Image.new(
|
||||
"RGB",
|
||||
|
||||
@@ -29,6 +29,7 @@ from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
|
||||
from vllm_omni.entrypoints.openai.protocol.audio import (
|
||||
BatchSpeechRequest,
|
||||
CreateAudio,
|
||||
OpenAICreateAudioGenerateRequest,
|
||||
OpenAICreateSpeechRequest,
|
||||
SpeechBatchItem,
|
||||
StreamingSpeechSessionConfig,
|
||||
@@ -2292,26 +2293,260 @@ def test_api_server_create_speech_wraps_error_response_status(mocker: MockerFixt
|
||||
)
|
||||
)
|
||||
|
||||
raw_request = _make_api_server_request(handler, path="/v1/audio/speech")
|
||||
request = OpenAICreateSpeechRequest(input="Hello")
|
||||
|
||||
response = asyncio.run(api_server_module.create_speech(request, raw_request))
|
||||
|
||||
_assert_openai_error_response(response, status_code=400, message="bad request")
|
||||
|
||||
|
||||
def _make_api_server_request(handler, *, method: str = "POST", path: str = "/v1/audio/voices") -> Request:
|
||||
app = FastAPI()
|
||||
app.state.openai_serving_speech = handler
|
||||
scope = {
|
||||
"type": "http",
|
||||
"app": app,
|
||||
"method": "POST",
|
||||
"path": "/v1/audio/speech",
|
||||
"method": method,
|
||||
"path": path,
|
||||
"headers": [],
|
||||
"query_string": b"",
|
||||
"client": ("127.0.0.1", 12345),
|
||||
"server": ("testserver", 80),
|
||||
"scheme": "http",
|
||||
}
|
||||
raw_request = Request(scope)
|
||||
return Request(scope)
|
||||
|
||||
|
||||
def _patch_api_server_base(mocker: MockerFixture):
|
||||
def _fake_create_error_response(message, err_type="BadRequestError", status_code=400, param=None):
|
||||
return ErrorResponse(
|
||||
error=ErrorInfo(
|
||||
message=message,
|
||||
type=err_type,
|
||||
param=param,
|
||||
code=getattr(status_code, "value", status_code),
|
||||
)
|
||||
)
|
||||
|
||||
fake_base = mocker.MagicMock()
|
||||
fake_base.create_error_response.side_effect = _fake_create_error_response
|
||||
mocker.patch.object(api_server_module, "base", return_value=fake_base)
|
||||
return fake_base
|
||||
|
||||
|
||||
def _assert_openai_error_response(
|
||||
response: JSONResponse,
|
||||
*,
|
||||
status_code: int,
|
||||
message: str,
|
||||
err_type: str = "BadRequestError",
|
||||
) -> None:
|
||||
assert isinstance(response, JSONResponse)
|
||||
assert response.status_code == status_code
|
||||
body = json.loads(response.body)
|
||||
assert body["error"]["code"] == status_code
|
||||
assert body["error"]["type"] == err_type
|
||||
assert message in body["error"]["message"]
|
||||
|
||||
|
||||
def test_api_server_list_voices_without_speech_handler_returns_404(mocker: MockerFixture):
|
||||
_patch_api_server_base(mocker)
|
||||
raw_request = _make_api_server_request(None, method="GET")
|
||||
|
||||
response = asyncio.run(api_server_module.list_voices(raw_request))
|
||||
|
||||
_assert_openai_error_response(
|
||||
response, status_code=404, message="does not support Speech API", err_type="NotFoundError"
|
||||
)
|
||||
|
||||
|
||||
def test_api_server_upload_voice_value_error_returns_400(mocker: MockerFixture):
|
||||
_patch_api_server_base(mocker)
|
||||
handler = mocker.MagicMock()
|
||||
handler.upload_voice = mocker.AsyncMock(side_effect=ValueError("Unsupported MIME type: audio/x-m4a"))
|
||||
raw_request = _make_api_server_request(handler)
|
||||
|
||||
response = asyncio.run(
|
||||
api_server_module.upload_voice(
|
||||
raw_request,
|
||||
audio_sample=mocker.MagicMock(),
|
||||
speaker_embedding=None,
|
||||
consent="cons_test",
|
||||
name="probe",
|
||||
)
|
||||
)
|
||||
|
||||
_assert_openai_error_response(response, status_code=400, message="Unsupported MIME type")
|
||||
|
||||
|
||||
def test_api_server_upload_voice_without_speech_handler_returns_404(mocker: MockerFixture):
|
||||
_patch_api_server_base(mocker)
|
||||
raw_request = _make_api_server_request(None)
|
||||
|
||||
response = asyncio.run(
|
||||
api_server_module.upload_voice(
|
||||
raw_request,
|
||||
consent="cons_test",
|
||||
name="probe",
|
||||
)
|
||||
)
|
||||
|
||||
_assert_openai_error_response(
|
||||
response, status_code=404, message="does not support Speech API", err_type="NotFoundError"
|
||||
)
|
||||
|
||||
|
||||
def test_api_server_upload_voice_without_input_returns_400(mocker: MockerFixture):
|
||||
_patch_api_server_base(mocker)
|
||||
raw_request = _make_api_server_request(mocker.MagicMock())
|
||||
|
||||
response = asyncio.run(
|
||||
api_server_module.upload_voice(
|
||||
raw_request,
|
||||
audio_sample=None,
|
||||
speaker_embedding=None,
|
||||
consent="cons_test",
|
||||
name="probe",
|
||||
)
|
||||
)
|
||||
|
||||
_assert_openai_error_response(response, status_code=400, message="must be provided")
|
||||
|
||||
|
||||
def test_api_server_upload_voice_with_audio_and_embedding_returns_400(mocker: MockerFixture):
|
||||
_patch_api_server_base(mocker)
|
||||
raw_request = _make_api_server_request(mocker.MagicMock())
|
||||
|
||||
response = asyncio.run(
|
||||
api_server_module.upload_voice(
|
||||
raw_request,
|
||||
audio_sample=mocker.MagicMock(),
|
||||
speaker_embedding="[0.1]",
|
||||
consent="cons_test",
|
||||
name="probe",
|
||||
)
|
||||
)
|
||||
|
||||
_assert_openai_error_response(response, status_code=400, message="mutually exclusive")
|
||||
|
||||
|
||||
def test_api_server_upload_voice_exception_returns_500(mocker: MockerFixture):
|
||||
_patch_api_server_base(mocker)
|
||||
handler = mocker.MagicMock()
|
||||
handler.upload_voice = mocker.AsyncMock(side_effect=RuntimeError("disk failed"))
|
||||
raw_request = _make_api_server_request(handler)
|
||||
|
||||
response = asyncio.run(
|
||||
api_server_module.upload_voice(
|
||||
raw_request,
|
||||
audio_sample=mocker.MagicMock(),
|
||||
speaker_embedding=None,
|
||||
consent="cons_test",
|
||||
name="probe",
|
||||
)
|
||||
)
|
||||
|
||||
_assert_openai_error_response(
|
||||
response,
|
||||
status_code=500,
|
||||
message="Failed to upload voice",
|
||||
err_type="InternalServerError",
|
||||
)
|
||||
|
||||
|
||||
def test_api_server_delete_voice_without_speech_handler_returns_404(mocker: MockerFixture):
|
||||
_patch_api_server_base(mocker)
|
||||
raw_request = _make_api_server_request(None, method="DELETE", path="/v1/audio/voices/probe")
|
||||
|
||||
response = asyncio.run(api_server_module.delete_voice("probe", raw_request))
|
||||
|
||||
_assert_openai_error_response(
|
||||
response, status_code=404, message="does not support Speech API", err_type="NotFoundError"
|
||||
)
|
||||
|
||||
|
||||
def test_api_server_delete_voice_value_error_returns_400(mocker: MockerFixture):
|
||||
_patch_api_server_base(mocker)
|
||||
handler = mocker.MagicMock()
|
||||
handler.delete_voice = mocker.AsyncMock(side_effect=ValueError("Invalid voice name"))
|
||||
raw_request = _make_api_server_request(handler, method="DELETE", path="/v1/audio/voices/probe")
|
||||
|
||||
response = asyncio.run(api_server_module.delete_voice("probe", raw_request))
|
||||
|
||||
_assert_openai_error_response(response, status_code=400, message="Invalid voice name")
|
||||
|
||||
|
||||
def test_api_server_delete_voice_not_found_returns_404(mocker: MockerFixture):
|
||||
_patch_api_server_base(mocker)
|
||||
handler = mocker.MagicMock()
|
||||
handler.delete_voice = mocker.AsyncMock(return_value=False)
|
||||
raw_request = _make_api_server_request(handler, method="DELETE", path="/v1/audio/voices/missing")
|
||||
|
||||
response = asyncio.run(api_server_module.delete_voice("missing", raw_request))
|
||||
|
||||
_assert_openai_error_response(
|
||||
response,
|
||||
status_code=404,
|
||||
message="Voice 'missing' not found",
|
||||
err_type="NotFoundError",
|
||||
)
|
||||
|
||||
|
||||
def test_api_server_delete_voice_exception_returns_500(mocker: MockerFixture):
|
||||
_patch_api_server_base(mocker)
|
||||
handler = mocker.MagicMock()
|
||||
handler.delete_voice = mocker.AsyncMock(side_effect=RuntimeError("disk failed"))
|
||||
raw_request = _make_api_server_request(handler, method="DELETE", path="/v1/audio/voices/probe")
|
||||
|
||||
response = asyncio.run(api_server_module.delete_voice("probe", raw_request))
|
||||
|
||||
_assert_openai_error_response(
|
||||
response,
|
||||
status_code=500,
|
||||
message="Failed to delete voice",
|
||||
err_type="InternalServerError",
|
||||
)
|
||||
|
||||
|
||||
def test_api_server_create_speech_without_handler_returns_404(mocker: MockerFixture):
|
||||
fake_base = _patch_api_server_base(mocker)
|
||||
raw_request = _make_api_server_request(None, path="/v1/audio/speech")
|
||||
raw_request.app.state.openai_serving_tokenization = fake_base
|
||||
request = OpenAICreateSpeechRequest(input="Hello")
|
||||
|
||||
response = asyncio.run(api_server_module.create_speech(request, raw_request))
|
||||
|
||||
assert isinstance(response, JSONResponse)
|
||||
assert response.status_code == 400
|
||||
_assert_openai_error_response(
|
||||
response, status_code=404, message="does not support Speech API", err_type="NotFoundError"
|
||||
)
|
||||
|
||||
|
||||
def test_api_server_create_speech_batch_without_handler_returns_404(mocker: MockerFixture):
|
||||
fake_base = _patch_api_server_base(mocker)
|
||||
raw_request = _make_api_server_request(None, path="/v1/audio/speech/batch")
|
||||
raw_request.app.state.openai_serving_tokenization = fake_base
|
||||
request = BatchSpeechRequest(items=[SpeechBatchItem(input="hi")])
|
||||
|
||||
response = asyncio.run(api_server_module.create_speech_batch(request, raw_request))
|
||||
|
||||
_assert_openai_error_response(
|
||||
response, status_code=404, message="does not support Speech API", err_type="NotFoundError"
|
||||
)
|
||||
|
||||
|
||||
def test_api_server_create_audio_generate_without_handler_returns_404(mocker: MockerFixture):
|
||||
fake_base = _patch_api_server_base(mocker)
|
||||
raw_request = _make_api_server_request(None, path="/v1/audio/generate")
|
||||
raw_request.app.state.openai_serving_audio_generate = None
|
||||
raw_request.app.state.openai_serving_tokenization = fake_base
|
||||
request = OpenAICreateAudioGenerateRequest(input="a bird singing")
|
||||
|
||||
response = asyncio.run(api_server_module.create_audio_generate(request, raw_request))
|
||||
|
||||
_assert_openai_error_response(
|
||||
response, status_code=404, message="does not support Audio Generate API", err_type="NotFoundError"
|
||||
)
|
||||
|
||||
|
||||
def test_api_server_create_speech_engine_error_response_includes_request_and_stage_id(mocker: MockerFixture):
|
||||
@@ -2325,26 +2560,13 @@ def test_api_server_create_speech_engine_error_response_includes_request_and_sta
|
||||
|
||||
terminate_mock = mocker.patch.object(api_server_module, "terminate_if_errored")
|
||||
|
||||
app = FastAPI()
|
||||
app.state.args = SimpleNamespace(log_error_stack=False)
|
||||
app.state.openai_serving_speech = handler
|
||||
app.state.engine_client = SimpleNamespace(
|
||||
raw_request = _make_api_server_request(handler, path="/v1/audio/speech")
|
||||
raw_request.app.state.args = SimpleNamespace(log_error_stack=False)
|
||||
raw_request.app.state.engine_client = SimpleNamespace(
|
||||
engine=SimpleNamespace(is_alive=lambda: False),
|
||||
errored=True,
|
||||
)
|
||||
app.state.server = SimpleNamespace()
|
||||
scope = {
|
||||
"type": "http",
|
||||
"app": app,
|
||||
"method": "POST",
|
||||
"path": "/v1/audio/speech",
|
||||
"headers": [],
|
||||
"query_string": b"",
|
||||
"client": ("127.0.0.1", 12345),
|
||||
"server": ("testserver", 80),
|
||||
"scheme": "http",
|
||||
}
|
||||
raw_request = Request(scope)
|
||||
raw_request.app.state.server = SimpleNamespace()
|
||||
raw_request.state.request_metadata = SimpleNamespace(request_id="speech-req-1")
|
||||
request = OpenAICreateSpeechRequest(input="Hello")
|
||||
|
||||
|
||||
@@ -0,0 +1,446 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for `/v1/realtime/video` WebSocket video output streaming."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncGenerator, Callable
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from fastapi import FastAPI, WebSocket
|
||||
from pytest_mock import MockerFixture
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _fake_video_frames(num_frames: int = 2) -> list[np.ndarray]:
|
||||
return [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(num_frames)]
|
||||
|
||||
|
||||
def _build_test_app(
|
||||
*,
|
||||
mocker: MockerFixture,
|
||||
streaming_chunks: list[tuple[bytes, bool]],
|
||||
mock_generate: Callable[..., AsyncGenerator[OmniRequestOutput, None]],
|
||||
final_chunk: bytes = b"",
|
||||
stall_timeout: float = 5.0,
|
||||
) -> tuple[FastAPI, OmniStreamingVideoOutputHandler, MagicMock]:
|
||||
engine_client = mocker.MagicMock()
|
||||
engine_client.abort = mocker.AsyncMock()
|
||||
engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
|
||||
engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
|
||||
engine_client.generate = mock_generate
|
||||
|
||||
class FakeStreamingVideoEncoder:
|
||||
encode_calls: list[int] = []
|
||||
|
||||
def encode(self, video):
|
||||
del video
|
||||
idx = len(self.encode_calls)
|
||||
self.encode_calls.append(idx)
|
||||
return streaming_chunks[idx][0]
|
||||
|
||||
def close(self):
|
||||
return final_chunk
|
||||
|
||||
def _make_encoder(*, output_format, fps, video_codec_options=None):
|
||||
del output_format, fps, video_codec_options
|
||||
return FakeStreamingVideoEncoder()
|
||||
|
||||
encoder_factory = mocker.MagicMock(side_effect=_make_encoder)
|
||||
engine_client.streaming_encoder_factory = encoder_factory
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
|
||||
encoder_factory,
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
|
||||
return_value="diffusion",
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
|
||||
return_value=[OmniDiffusionSamplingParams()],
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list",
|
||||
return_value=[OmniDiffusionSamplingParams()],
|
||||
)
|
||||
|
||||
handler = OmniStreamingVideoOutputHandler(
|
||||
engine_client=engine_client,
|
||||
model_name="test-model",
|
||||
stage_configs=engine_client.stage_configs,
|
||||
stall_timeout=stall_timeout,
|
||||
start_timeout=5.0,
|
||||
)
|
||||
app = FastAPI()
|
||||
|
||||
@app.websocket("/v1/realtime/video")
|
||||
async def ws_endpoint(websocket: WebSocket):
|
||||
await handler.handle_session(websocket)
|
||||
|
||||
return app, handler, engine_client
|
||||
|
||||
|
||||
class TestStreamingVideoOutputWebSocket:
|
||||
"""WebSocket protocol tests for streaming generated video output."""
|
||||
|
||||
def test_streaming_session_emits_video_start_binary_chunks_and_done(self, mocker: MockerFixture):
|
||||
"""A full session delivers video.start, binary chunks, then session.done."""
|
||||
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
for frames, finished in [
|
||||
(_fake_video_frames(2), False),
|
||||
(_fake_video_frames(3), True),
|
||||
]:
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[frames],
|
||||
final_output_type="image",
|
||||
finished=finished,
|
||||
)
|
||||
|
||||
app, _handler, engine_client = _build_test_app(
|
||||
mocker=mocker,
|
||||
streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)],
|
||||
mock_generate=mock_generate,
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect("/v1/realtime/video") as ws:
|
||||
ws.send_json({"type": "session.start", "prompt": "A cat walking in the rain"})
|
||||
|
||||
start = ws.receive_json()
|
||||
assert start["type"] == "video.start"
|
||||
assert start["format"] == "m4s"
|
||||
assert "format" not in start["config"]
|
||||
assert "request_id" in start
|
||||
|
||||
assert ws.receive_bytes() == b"mp4-chunk-0"
|
||||
assert ws.receive_bytes() == b"mp4-chunk-1"
|
||||
|
||||
done = ws.receive_json()
|
||||
assert done["type"] == "session.done"
|
||||
assert done["chunks"] == 2
|
||||
assert done["stopped"] is False
|
||||
|
||||
engine_client.abort.assert_not_called()
|
||||
engine_client.streaming_encoder_factory.assert_called_once()
|
||||
assert engine_client.streaming_encoder_factory.call_args.kwargs["output_format"] == "m4s"
|
||||
|
||||
def test_streaming_session_rejects_unknown_format(self, mocker: MockerFixture):
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
if False:
|
||||
yield None
|
||||
|
||||
app, _handler, engine_client = _build_test_app(
|
||||
mocker=mocker,
|
||||
streaming_chunks=[],
|
||||
mock_generate=mock_generate,
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect("/v1/realtime/video") as ws:
|
||||
ws.send_json({"type": "session.start", "prompt": "hello", "format": "webm"})
|
||||
err = ws.receive_json()
|
||||
assert err["type"] == "error"
|
||||
assert "format" in err["message"]
|
||||
|
||||
engine_client.abort.assert_not_called()
|
||||
|
||||
def test_streaming_session_emits_final_encoder_delta_before_done(self, mocker: MockerFixture):
|
||||
"""A non-empty encoder close delta is delivered as the final binary chunk."""
|
||||
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[_fake_video_frames(2)],
|
||||
final_output_type="image",
|
||||
finished=True,
|
||||
)
|
||||
|
||||
app, _handler, engine_client = _build_test_app(
|
||||
mocker=mocker,
|
||||
streaming_chunks=[(b"mp4-chunk-0", True)],
|
||||
mock_generate=mock_generate,
|
||||
final_chunk=b"mp4-trailer",
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect("/v1/realtime/video") as ws:
|
||||
ws.send_json({"type": "session.start", "prompt": "A cat walking in the rain"})
|
||||
assert ws.receive_json()["type"] == "video.start"
|
||||
assert ws.receive_bytes() == b"mp4-chunk-0"
|
||||
assert ws.receive_bytes() == b"mp4-trailer"
|
||||
|
||||
done = ws.receive_json()
|
||||
assert done["type"] == "session.done"
|
||||
assert done["chunks"] == 2
|
||||
assert done["stopped"] is False
|
||||
|
||||
engine_client.abort.assert_not_called()
|
||||
|
||||
def test_generation_error_output_returns_websocket_error(self, mocker: MockerFixture):
|
||||
"""An erroneous OmniRequestOutput from generate is sent to the client as a WebSocket error."""
|
||||
error_message = "diffusion engine exploded"
|
||||
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[_fake_video_frames(2)],
|
||||
final_output_type="image",
|
||||
finished=False,
|
||||
)
|
||||
yield OmniRequestOutput.from_error("req-ws", error_message)
|
||||
|
||||
app, _handler, engine_client = _build_test_app(
|
||||
mocker=mocker,
|
||||
streaming_chunks=[(b"mp4-chunk-0", False)],
|
||||
mock_generate=mock_generate,
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect("/v1/realtime/video") as ws:
|
||||
ws.send_json({"type": "session.start", "prompt": "hello"})
|
||||
assert ws.receive_json()["type"] == "video.start"
|
||||
assert ws.receive_bytes() == b"mp4-chunk-0"
|
||||
err = ws.receive_json()
|
||||
assert err["type"] == "error"
|
||||
assert error_message in err["message"]
|
||||
|
||||
engine_client.abort.assert_not_called()
|
||||
|
||||
|
||||
class _MockVideoOutputWebSocket:
|
||||
"""Feeds scripted client messages and records server sends."""
|
||||
|
||||
def __init__(self, messages: list[str]) -> None:
|
||||
self._messages: asyncio.Queue[str] = asyncio.Queue()
|
||||
for message in messages:
|
||||
self._messages.put_nowait(message)
|
||||
self.sent: list[dict | bytes] = []
|
||||
|
||||
def enqueue(self, message: str) -> None:
|
||||
self._messages.put_nowait(message)
|
||||
|
||||
async def accept(self) -> None:
|
||||
return None
|
||||
|
||||
async def receive_text(self) -> str:
|
||||
return await self._messages.get()
|
||||
|
||||
async def send_json(self, data: dict) -> None:
|
||||
self.sent.append(data)
|
||||
|
||||
async def send_bytes(self, data: bytes) -> None:
|
||||
self.sent.append(data)
|
||||
|
||||
@property
|
||||
def sent_json(self) -> list[dict]:
|
||||
return [m for m in self.sent if isinstance(m, dict)]
|
||||
|
||||
|
||||
def _build_handler_for_async_tests(
|
||||
mocker: MockerFixture,
|
||||
*,
|
||||
mock_generate: Callable[..., AsyncGenerator[OmniRequestOutput, None]],
|
||||
streaming_chunks: list[tuple[bytes, bool]] | None = None,
|
||||
stall_timeout: float = 5.0,
|
||||
) -> tuple[OmniStreamingVideoOutputHandler, MagicMock, _MockVideoOutputWebSocket]:
|
||||
engine_client = mocker.MagicMock()
|
||||
engine_client.abort = mocker.AsyncMock()
|
||||
engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
|
||||
engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
|
||||
engine_client.generate = mock_generate
|
||||
|
||||
chunks = streaming_chunks or [(b"mp4-chunk-0", True)]
|
||||
|
||||
class FakeStreamingVideoEncoder:
|
||||
def __init__(self) -> None:
|
||||
self.encode_calls = 0
|
||||
|
||||
def encode(self, video):
|
||||
del video
|
||||
idx = self.encode_calls
|
||||
self.encode_calls += 1
|
||||
return chunks[idx][0]
|
||||
|
||||
def close(self):
|
||||
return b""
|
||||
|
||||
encoder_factory = mocker.MagicMock(side_effect=lambda **_: FakeStreamingVideoEncoder())
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
|
||||
encoder_factory,
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
|
||||
return_value="diffusion",
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
|
||||
return_value=[OmniDiffusionSamplingParams()],
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list",
|
||||
return_value=[OmniDiffusionSamplingParams()],
|
||||
)
|
||||
|
||||
handler = OmniStreamingVideoOutputHandler(
|
||||
engine_client=engine_client,
|
||||
model_name="test-model",
|
||||
stage_configs=engine_client.stage_configs,
|
||||
stall_timeout=stall_timeout,
|
||||
start_timeout=5.0,
|
||||
)
|
||||
start_msg = json.dumps({"type": "session.start", "prompt": "async test"})
|
||||
ws = _MockVideoOutputWebSocket([start_msg])
|
||||
return handler, engine_client, ws
|
||||
|
||||
|
||||
class TestStreamingVideoOutputStallTimeout:
|
||||
"""Stall-timeout and session.ping behavior."""
|
||||
|
||||
def test_long_generation_without_client_messages_succeeds(self, mocker: MockerFixture):
|
||||
"""Engine progress between yields keeps the session alive without client pings."""
|
||||
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[_fake_video_frames(2)],
|
||||
final_output_type="image",
|
||||
finished=False,
|
||||
)
|
||||
await asyncio.sleep(0.3)
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[_fake_video_frames(2)],
|
||||
final_output_type="image",
|
||||
finished=True,
|
||||
)
|
||||
|
||||
app, _handler, engine_client = _build_test_app(
|
||||
mocker=mocker,
|
||||
streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)],
|
||||
mock_generate=mock_generate,
|
||||
stall_timeout=0.5,
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect("/v1/realtime/video") as ws:
|
||||
ws.send_json({"type": "session.start", "prompt": "slow generation"})
|
||||
assert ws.receive_json()["type"] == "video.start"
|
||||
assert ws.receive_bytes() == b"mp4-chunk-0"
|
||||
assert ws.receive_bytes() == b"mp4-chunk-1"
|
||||
done = ws.receive_json()
|
||||
assert done["type"] == "session.done"
|
||||
assert done["stopped"] is False
|
||||
|
||||
engine_client.abort.assert_not_called()
|
||||
|
||||
def test_stall_timeout_aborts_when_engine_silent(self, mocker: MockerFixture):
|
||||
"""No engine output for longer than stall_timeout aborts the request."""
|
||||
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
await asyncio.sleep(1.0)
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[_fake_video_frames(2)],
|
||||
final_output_type="image",
|
||||
finished=True,
|
||||
)
|
||||
|
||||
app, _handler, engine_client = _build_test_app(
|
||||
mocker=mocker,
|
||||
streaming_chunks=[(b"mp4-chunk-0", True)],
|
||||
mock_generate=mock_generate,
|
||||
stall_timeout=0.2,
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect("/v1/realtime/video") as ws:
|
||||
ws.send_json({"type": "session.start", "prompt": "stalled"})
|
||||
assert ws.receive_json()["type"] == "video.start"
|
||||
err = ws.receive_json()
|
||||
assert err["type"] == "error"
|
||||
assert "Stall timeout" in err["message"]
|
||||
|
||||
engine_client.abort.assert_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_ping_returns_pong(self, mocker: MockerFixture):
|
||||
"""session.ping during generation receives session.pong."""
|
||||
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
await asyncio.sleep(0.5)
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[_fake_video_frames(2)],
|
||||
final_output_type="image",
|
||||
finished=True,
|
||||
)
|
||||
|
||||
handler, _engine_client, ws = _build_handler_for_async_tests(
|
||||
mocker,
|
||||
mock_generate=mock_generate,
|
||||
stall_timeout=5.0,
|
||||
)
|
||||
|
||||
async def _send_ping_after_delay() -> None:
|
||||
await asyncio.sleep(0.05)
|
||||
ws.enqueue(json.dumps({"type": "session.ping"}))
|
||||
|
||||
ping_task = asyncio.create_task(_send_ping_after_delay())
|
||||
await handler.handle_session(ws) # type: ignore[arg-type]
|
||||
await ping_task
|
||||
|
||||
assert any(m.get("type") == "session.pong" for m in ws.sent_json)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ping_prevents_stall_during_silent_engine(self, mocker: MockerFixture):
|
||||
"""Client pings refresh the stall clock while the engine produces no output."""
|
||||
|
||||
release = asyncio.Event()
|
||||
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
await release.wait()
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[_fake_video_frames(2)],
|
||||
final_output_type="image",
|
||||
finished=True,
|
||||
)
|
||||
|
||||
handler, engine_client, ws = _build_handler_for_async_tests(
|
||||
mocker,
|
||||
mock_generate=mock_generate,
|
||||
stall_timeout=0.3,
|
||||
)
|
||||
|
||||
async def _pinger() -> None:
|
||||
while not release.is_set():
|
||||
ws.enqueue(json.dumps({"type": "session.ping"}))
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
pinger_task = asyncio.create_task(_pinger())
|
||||
session_task = asyncio.create_task(handler.handle_session(ws)) # type: ignore[arg-type]
|
||||
await asyncio.sleep(0.8)
|
||||
release.set()
|
||||
await session_task
|
||||
await pinger_task
|
||||
|
||||
errors = [m for m in ws.sent_json if m.get("type") == "error"]
|
||||
assert not any("Stall timeout" in m.get("message", "") for m in errors)
|
||||
assert any(m.get("type") == "session.pong" for m in ws.sent_json)
|
||||
assert any(m.get("type") == "session.done" for m in ws.sent_json)
|
||||
engine_client.abort.assert_not_called()
|
||||
@@ -48,6 +48,107 @@ def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch):
|
||||
assert mux_calls[0]["audio"] is None
|
||||
|
||||
|
||||
def test_fragmented_mp4_video_encoder_reuses_single_muxer(monkeypatch):
|
||||
muxers = []
|
||||
|
||||
class FakeFragmentedMP4Muxer:
|
||||
def __init__(self, *, width, height, fps, video_codec_options=None):
|
||||
self.calls = []
|
||||
muxers.append(
|
||||
{
|
||||
"width": width,
|
||||
"height": height,
|
||||
"fps": fps,
|
||||
"video_codec_options": video_codec_options,
|
||||
"instance": self,
|
||||
}
|
||||
)
|
||||
|
||||
def mux_video_frames(self, frames):
|
||||
self.calls.append(frames)
|
||||
return f"fragment-{len(self.calls)}".encode()
|
||||
|
||||
def close(self):
|
||||
return b"trailer"
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.utils.media_utils.FragmentedMP4Muxer",
|
||||
FakeFragmentedMP4Muxer,
|
||||
)
|
||||
|
||||
encoder = video_api_utils.FragmentedMP4VideoEncoder(
|
||||
fps=12,
|
||||
video_codec_options={"preset": "ultrafast"},
|
||||
)
|
||||
assert encoder.encode([np.zeros((4, 6, 3), dtype=np.float32)]) == b"fragment-1"
|
||||
assert encoder.encode([np.ones((4, 6, 3), dtype=np.float32)]) == b"fragment-2"
|
||||
assert encoder.close() == b"trailer"
|
||||
|
||||
assert len(muxers) == 1
|
||||
assert muxers[0]["width"] == 6
|
||||
assert muxers[0]["height"] == 4
|
||||
assert muxers[0]["fps"] == 12.0
|
||||
assert muxers[0]["video_codec_options"] == {"preset": "ultrafast"}
|
||||
assert len(muxers[0]["instance"].calls) == 2
|
||||
|
||||
|
||||
def test_create_streaming_video_encoder_selects_requested_format():
|
||||
assert isinstance(
|
||||
video_api_utils.create_streaming_video_encoder(output_format="m4s", fps=12),
|
||||
video_api_utils.FragmentedMP4VideoEncoder,
|
||||
)
|
||||
|
||||
|
||||
def test_finalize_streaming_mp4_bytes_produces_progressive_mp4():
|
||||
"""Fragment MP4 chunks are remuxed into decodable progressive MP4 bytes."""
|
||||
import numpy as np
|
||||
|
||||
from vllm_omni.diffusion.utils.media_utils import FragmentedMP4Muxer, finalize_streaming_video_bytes
|
||||
|
||||
def _read_mp4_video_info(mp4_bytes: bytes) -> tuple[int, float, float | None, int, int]:
|
||||
import io
|
||||
|
||||
import av
|
||||
|
||||
with av.open(io.BytesIO(mp4_bytes), mode="r", format="mp4") as container:
|
||||
stream = container.streams.video[0]
|
||||
frame_count = sum(1 for _ in container.decode(stream))
|
||||
assert stream.average_rate is not None
|
||||
fps = float(stream.average_rate)
|
||||
duration_sec = None
|
||||
if stream.duration is not None:
|
||||
assert stream.time_base is not None
|
||||
duration_sec = float(stream.duration * stream.time_base)
|
||||
return frame_count, fps, duration_sec, stream.width, stream.height
|
||||
|
||||
width = 32
|
||||
height = 32
|
||||
fps = 16.0
|
||||
input_frame_count = 2
|
||||
|
||||
muxer = FragmentedMP4Muxer(width=width, height=height, fps=fps)
|
||||
frames = np.zeros((input_frame_count, height, width, 3), dtype=np.uint8)
|
||||
streamed = muxer.mux_video_frames(frames) + muxer.close()
|
||||
|
||||
finalized = finalize_streaming_video_bytes(streamed, input_format="m4s", fps=fps)
|
||||
assert finalized
|
||||
assert finalized != streamed
|
||||
|
||||
streamed_info = _read_mp4_video_info(streamed)
|
||||
final_info = _read_mp4_video_info(finalized)
|
||||
expected_duration = input_frame_count / fps
|
||||
|
||||
assert streamed_info[0] == input_frame_count
|
||||
assert final_info[0] == input_frame_count
|
||||
assert streamed_info[1] == pytest.approx(fps)
|
||||
assert final_info[1] == pytest.approx(fps)
|
||||
assert streamed_info[3:] == (width, height)
|
||||
assert final_info[3:] == (width, height)
|
||||
|
||||
assert final_info[2] == pytest.approx(expected_duration, rel=0.05)
|
||||
assert streamed_info[2] == pytest.approx(expected_duration, rel=0.05)
|
||||
|
||||
|
||||
def test_rife_model_inference_runs_on_dummy_tensors():
|
||||
model = rife_interpolator.Model().eval()
|
||||
img0 = torch.rand(1, 3, 32, 32)
|
||||
|
||||
@@ -1334,7 +1334,7 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture)
|
||||
extra_params = {
|
||||
"is_enable_stage2": True,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [20, 20, 20],
|
||||
"pyramid_num_inference_steps_list": [1, 1, 1],
|
||||
"use_cfg_zero_star": True,
|
||||
}
|
||||
response = test_client.post(
|
||||
@@ -1352,7 +1352,7 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture)
|
||||
captured = engine.captured_sampling_params_list[0]
|
||||
assert captured.extra_args["is_enable_stage2"] is True
|
||||
assert captured.extra_args["pyramid_num_stages"] == 3
|
||||
assert captured.extra_args["pyramid_num_inference_steps_list"] == [20, 20, 20]
|
||||
assert captured.extra_args["pyramid_num_inference_steps_list"] == [1, 1, 1]
|
||||
assert captured.extra_args["use_cfg_zero_star"] is True
|
||||
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model]
|
||||
|
||||
@@ -198,6 +199,48 @@ def test_generate_accepts_request_after_repeated_cancellations():
|
||||
asyncio.run(run_test())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_generate_yields_streaming_diffusion_chunks_before_final():
|
||||
"""AsyncOmni.generate yields every intermediate diffusion chunk before the final one."""
|
||||
|
||||
async def streaming_process_results(request_id, metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts):
|
||||
del metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id=request_id,
|
||||
images=[],
|
||||
final_output_type="image",
|
||||
custom_output={"chunk": 0},
|
||||
finished=False,
|
||||
)
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id=request_id,
|
||||
images=[],
|
||||
final_output_type="image",
|
||||
custom_output={"chunk": 1},
|
||||
finished=True,
|
||||
)
|
||||
|
||||
async def run_test():
|
||||
omni = get_async_omni_instance()
|
||||
omni._process_orchestrator_results = streaming_process_results
|
||||
|
||||
outputs = []
|
||||
async for output in AsyncOmni.generate(
|
||||
omni,
|
||||
prompt={"prompt": "a cat"},
|
||||
request_id="req-stream",
|
||||
sampling_params_list=[SimpleNamespace()],
|
||||
output_modalities=["image"],
|
||||
):
|
||||
outputs.append(output)
|
||||
|
||||
assert len(outputs) == 2
|
||||
assert [output.finished for output in outputs] == [False, True]
|
||||
assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
|
||||
|
||||
asyncio.run(run_test())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize(
|
||||
"output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY, RequestOutputKind.CUMULATIVE]
|
||||
|
||||
@@ -24,7 +24,7 @@ def test_dreamzero_vla_resolves_to_dreamzero_config(monkeypatch):
|
||||
|
||||
def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.config.stage_config.StageConfigFactory._auto_detect_model_type",
|
||||
"vllm_omni.config.config_factory.StageConfigFactory._auto_detect_model_type",
|
||||
classmethod(lambda _cls, _model, trust_remote_code=True: ("vla", None)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
|
||||
@@ -1,11 +1,9 @@
|
||||
"""
|
||||
Offline inference tests: vace video generation.
|
||||
See examples/offline_inference/vace/vace_video_generation.md
|
||||
"""
|
||||
"""VACE coverage through the shared text-to-video and image-to-video examples."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet
|
||||
from tests.helpers.assertions import assert_video_valid
|
||||
@@ -14,22 +12,83 @@ from tests.helpers.mark import hardware_marks
|
||||
pytestmark = [
|
||||
pytest.mark.usefixtures("clean_gpu_memory_between_tests"),
|
||||
pytest.mark.full_model,
|
||||
pytest.mark.diffusion,
|
||||
pytest.mark.example,
|
||||
*hardware_marks(res={"cuda": "H100"}),
|
||||
]
|
||||
|
||||
VACE_SCRIPT = EXAMPLES / "offline_inference" / "vace" / "vace_video_generation.py"
|
||||
README_PATH = VACE_SCRIPT.with_name("vace_video_generation.md")
|
||||
TEXT_TO_VIDEO_README = EXAMPLES / "offline_inference" / "text_to_video" / "text_to_video.md"
|
||||
IMAGE_TO_VIDEO_README = EXAMPLES / "offline_inference" / "image_to_video" / "README.md"
|
||||
EXAMPLE_OUTPUT_SUBFOLDER = "example_offline_vace"
|
||||
|
||||
README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH)
|
||||
TEXT_TO_VIDEO_SNIPPET = next(
|
||||
snippet
|
||||
for snippet in ReadmeSnippet.extract_readme_snippets(TEXT_TO_VIDEO_README)
|
||||
if "vace_t2v_output.mp4" in snippet.code
|
||||
)
|
||||
CONDITIONAL_OUTPUTS = (
|
||||
"vace_i2v_output.mp4",
|
||||
"vace_v2lf_output.mp4",
|
||||
"vace_flf2v_output.mp4",
|
||||
"vace_inpaint_output.mp4",
|
||||
"vace_r2v_output.mp4",
|
||||
)
|
||||
IMAGE_TO_VIDEO_SNIPPETS = [
|
||||
snippet
|
||||
for snippet in ReadmeSnippet.extract_readme_snippets(IMAGE_TO_VIDEO_README)
|
||||
if snippet.h2_title == "Wan2.1 VACE Conditional Tasks"
|
||||
and any(output in snippet.code for output in CONDITIONAL_OUTPUTS)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id)
|
||||
def test_vace_video_generation(snippet: ReadmeSnippet, example_runner: ExampleRunner):
|
||||
should_skip, reason = snippet.skip
|
||||
if should_skip:
|
||||
pytest.skip(reason)
|
||||
def _smoke_snippet(snippet: ReadmeSnippet, asset_dir: Path) -> ReadmeSnippet:
|
||||
# Six fresh model launches share a 90-minute X2V job. Keep the documented
|
||||
# quality settings intact while reducing only the CI executions.
|
||||
replacements = {
|
||||
"astronaut.jpg": str(asset_dir / "astronaut.jpg"),
|
||||
"vace_first_frame.png": str(asset_dir / "vace_first_frame.png"),
|
||||
"vace_last_frame.png": str(asset_dir / "vace_last_frame.png"),
|
||||
"vace_center_mask.png": str(asset_dir / "vace_center_mask.png"),
|
||||
"--num-frames 81": "--num-frames 5",
|
||||
"--num-inference-steps 30": "--num-inference-steps 2",
|
||||
}
|
||||
code = snippet.code
|
||||
for source, target in replacements.items():
|
||||
code = code.replace(source, target)
|
||||
return snippet._replace(code=code)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def vace_asset_dir(tmp_path: Path) -> Path:
|
||||
asset_dir = tmp_path / "vace_assets"
|
||||
asset_dir.mkdir(parents=True, exist_ok=True)
|
||||
Image.new("RGB", (832, 480), (180, 180, 180)).save(asset_dir / "astronaut.jpg")
|
||||
Image.new("RGB", (512, 512), (80, 140, 220)).save(asset_dir / "vace_first_frame.png")
|
||||
Image.new("RGB", (512, 512), (220, 140, 80)).save(asset_dir / "vace_last_frame.png")
|
||||
mask = Image.new("L", (832, 480), 0)
|
||||
mask.paste(255, (336, 0, 496, 480))
|
||||
mask.save(asset_dir / "vace_center_mask.png")
|
||||
return asset_dir
|
||||
|
||||
|
||||
def test_vace_text_to_video_shared_example(example_runner: ExampleRunner) -> None:
|
||||
snippet = TEXT_TO_VIDEO_SNIPPET._replace(
|
||||
code=TEXT_TO_VIDEO_SNIPPET.code.replace("--num-frames 81", "--num-frames 5").replace(
|
||||
"--num-inference-steps 30", "--num-inference-steps 2"
|
||||
)
|
||||
)
|
||||
result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER))
|
||||
for asset in result.assets:
|
||||
assert_video_valid(asset)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("snippet", IMAGE_TO_VIDEO_SNIPPETS, ids=lambda snippet: snippet.test_id)
|
||||
def test_vace_conditional_shared_examples(
|
||||
snippet: ReadmeSnippet,
|
||||
example_runner: ExampleRunner,
|
||||
vace_asset_dir: Path,
|
||||
) -> None:
|
||||
snippet = _smoke_snippet(snippet, vace_asset_dir)
|
||||
result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER))
|
||||
for asset in result.assets:
|
||||
assert_video_valid(asset)
|
||||
@@ -17,8 +17,8 @@ class _FakeSummarizer:
|
||||
self.chunks += 1
|
||||
return f"summary-of-chunk-{chunk_index}"
|
||||
|
||||
async def compress_to_long_term(self, existing, mids):
|
||||
return (existing + " | " if existing else "") + f"compressed({len(mids)})"
|
||||
async def compress_to_long_term(self, mids):
|
||||
return f"compressed({len(mids)})"
|
||||
|
||||
|
||||
class _FakeDelegation:
|
||||
@@ -42,7 +42,7 @@ def test_query_freshness_and_archive_on_change():
|
||||
b.record_response("1 bottle")
|
||||
assert b.update_query("now describe") is True
|
||||
assert b.memory.qa_history[-1].query == "count bottles"
|
||||
assert b.memory.qa_history[-1].responses == [("0.0s", "1 bottle")]
|
||||
assert b.memory.qa_history[-1].responses == [("0.0 seconds", "1 bottle")]
|
||||
|
||||
|
||||
def test_query_in_frame_on_arrival_then_prefix_on_later_chunks():
|
||||
@@ -71,7 +71,7 @@ async def test_flush_archives_qa_and_summarizes():
|
||||
assert summ.chunks == 1
|
||||
assert b.chunk_index == 2
|
||||
assert b.should_flush() is False
|
||||
assert b.memory.qa_history[-1].responses == [("2.0s", "a person enters")]
|
||||
assert b.memory.qa_history[-1].responses == [("2.0 seconds", "a person enters")]
|
||||
assert len(b.memory.mid_term_summaries) == 1
|
||||
|
||||
|
||||
@@ -85,6 +85,20 @@ async def test_long_term_compression_after_n_chunks():
|
||||
assert b.memory.mid_term_summaries == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_long_term_window_caps_blocks():
|
||||
# Long-term memory must not grow unbounded over hours: oldest blocks are dropped
|
||||
# once the sliding window is full (reference window = 15; here 2 for the test).
|
||||
summ = _FakeSummarizer()
|
||||
b = InteractionBrain(
|
||||
summarizer=summ, chunk_frames=1, long_term_every_n_chunks=1, long_term_window=2, frame_seconds=1.0
|
||||
)
|
||||
for _ in range(4):
|
||||
await b.flush([("0.0s", "u")])
|
||||
assert len(b.memory.long_term_blocks) == 2 # capped, not 4
|
||||
assert b.memory.long_term_memory.count("compressed(") == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delegation_submit_and_fold():
|
||||
b = InteractionBrain(delegation=_FakeDelegation())
|
||||
|
||||
@@ -38,7 +38,7 @@ def test_commit_records_only_spoken():
|
||||
assert p.commit("</silence>").action.value == "silence"
|
||||
assert p.brain.response_records == []
|
||||
assert p.commit("</response> hello").action.value == "response"
|
||||
assert p.brain.response_records == [("0.0s", "hello")]
|
||||
assert p.brain.response_records == [("0.0 seconds", "hello")]
|
||||
|
||||
|
||||
def test_commit_dedup_default_is_exact_match_only():
|
||||
|
||||
@@ -115,6 +115,11 @@ def assert_video_diffusion_response(
|
||||
expected_height = _maybe_int(form_data.get("height"))
|
||||
expected_fps = _maybe_int(form_data.get("fps"))
|
||||
|
||||
# Skip num_frames assertion for Helios models because they round up frames
|
||||
model = request_config.get("model", "")
|
||||
if "Helios" in model:
|
||||
expected_frames = None
|
||||
|
||||
for vid_bytes in response.videos:
|
||||
assert_video_valid(
|
||||
vid_bytes,
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
"""Fixtures config creation & out of tree registry management."""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clean_pipeline_registry():
|
||||
"""Ensure the OMNI_PIPELINES are in a clean state for a test that mutates it."""
|
||||
snapshot = dict(OMNI_PIPELINES)
|
||||
yield
|
||||
OMNI_PIPELINES.clear()
|
||||
OMNI_PIPELINES.update(snapshot)
|
||||
@@ -19,7 +19,7 @@ from collections.abc import Generator
|
||||
from dataclasses import asdict, dataclass
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Any, NamedTuple
|
||||
from typing import Any, NamedTuple, cast
|
||||
from urllib.parse import quote
|
||||
|
||||
import psutil
|
||||
@@ -2168,6 +2168,102 @@ class OpenAIClientHandler:
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
def send_streaming_video_diffusion_request(
|
||||
self,
|
||||
request_config: dict[str, Any],
|
||||
request_num: int = 1,
|
||||
*,
|
||||
timeout_seconds: float = 600.0,
|
||||
) -> list[DiffusionResponse]:
|
||||
"""
|
||||
Send a native ``/v1/realtime/video`` WebSocket request and return one
|
||||
finalized MP4 artifact assembled from the streamed binary fragments.
|
||||
"""
|
||||
if request_num != 1:
|
||||
raise NotImplementedError("Concurrent streaming video diffusion requests are not currently implemented")
|
||||
|
||||
response = asyncio.run(
|
||||
self._send_streaming_video_diffusion_request_once(
|
||||
request_config,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
)
|
||||
assert_diffusion_response(response, request_config, run_level=self.run_level)
|
||||
if response.e2e_latency is not None:
|
||||
self._print_client_stat(f"[diffusion.stream] request#1 success in {response.e2e_latency:.3f}s")
|
||||
else:
|
||||
self._print_client_stat("[diffusion.stream] request#1 completed")
|
||||
return [response]
|
||||
|
||||
async def _send_streaming_video_diffusion_request_once(
|
||||
self,
|
||||
request_config: dict[str, Any],
|
||||
*,
|
||||
timeout_seconds: float,
|
||||
) -> DiffusionResponse:
|
||||
form_data = request_config.get("form_data")
|
||||
if not isinstance(form_data, dict):
|
||||
raise ValueError("Video request_config must contain 'form_data'")
|
||||
payload: dict[str, Any] = {
|
||||
"type": "session.start",
|
||||
**{key: value for key, value in form_data.items() if value is not None},
|
||||
}
|
||||
model = request_config.get("model")
|
||||
if model is not None:
|
||||
payload["model"] = model
|
||||
payload.setdefault("format", "m4s")
|
||||
|
||||
fps = float(payload.get("fps") or 16)
|
||||
stream_format = payload["format"]
|
||||
url = self._build_ws_url("/v1/realtime/video")
|
||||
|
||||
result = DiffusionResponse()
|
||||
chunks: list[bytes] = []
|
||||
start_time = time.perf_counter()
|
||||
deadline = start_time + timeout_seconds
|
||||
|
||||
import websockets
|
||||
|
||||
async with websockets.connect(url, max_size=None) as websocket:
|
||||
await websocket.send(json.dumps(payload))
|
||||
|
||||
while True:
|
||||
remaining = deadline - time.perf_counter()
|
||||
if remaining <= 0:
|
||||
raise TimeoutError(f"Streaming video request did not complete within {timeout_seconds}s")
|
||||
|
||||
message = await asyncio.wait_for(websocket.recv(), timeout=remaining)
|
||||
if isinstance(message, bytes):
|
||||
chunks.append(message)
|
||||
continue
|
||||
|
||||
msg = json.loads(message)
|
||||
msg_type = msg.get("type")
|
||||
if msg_type == "video.start":
|
||||
stream_format = msg.get("format") or stream_format
|
||||
continue
|
||||
if msg_type == "session.done":
|
||||
break
|
||||
if msg_type == "error":
|
||||
raise RuntimeError(str(msg.get("message", msg)))
|
||||
|
||||
from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes
|
||||
from vllm_omni.entrypoints.openai.video_api_utils import StreamingVideoFormat
|
||||
|
||||
streamed_bytes = b"".join(chunks)
|
||||
if not streamed_bytes:
|
||||
raise RuntimeError("Streaming video request completed without binary video chunks")
|
||||
result.videos = [
|
||||
finalize_streaming_video_bytes(
|
||||
streamed_bytes,
|
||||
input_format=cast(StreamingVideoFormat, stream_format),
|
||||
fps=fps,
|
||||
)
|
||||
]
|
||||
result.e2e_latency = time.perf_counter() - start_time
|
||||
result.success = True
|
||||
return result
|
||||
|
||||
def _wait_until_video_completed(
|
||||
self, video_id: str, poll_interval_seconds: int = 2, timeout_seconds: int = 300
|
||||
) -> None:
|
||||
|
||||
@@ -1,76 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for GLM-Image model __init__.py lazy import pattern.
|
||||
|
||||
The ``__init__.py`` uses ``__getattr__`` for lazy loading to avoid importing
|
||||
``transformers.models.glm_image`` at module init, which may not be available
|
||||
in all environments.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class TestLazyImport:
|
||||
"""Test the __getattr__ lazy import pattern in glm_image __init__.py."""
|
||||
|
||||
def test_getattr_exists_and_is_callable(self):
|
||||
"""Verify __getattr__ exists and is callable."""
|
||||
import vllm_omni.model_executor.models.glm_image as glm_image_pkg
|
||||
|
||||
assert hasattr(glm_image_pkg, "__getattr__")
|
||||
assert callable(glm_image_pkg.__getattr__)
|
||||
|
||||
def test_getattr_returns_class_for_known_attribute(self):
|
||||
"""Verify __getattr__ returns GlmImageForConditionalGeneration for known attribute."""
|
||||
import vllm_omni.model_executor.models.glm_image as glm_image_pkg
|
||||
|
||||
# Call __getattr__ directly to test the lazy import logic
|
||||
result = glm_image_pkg.__getattr__("GlmImageForConditionalGeneration")
|
||||
|
||||
# Verify we get a class (the actual GlmImageForConditionalGeneration)
|
||||
assert result is not None
|
||||
assert isinstance(result, type)
|
||||
assert result.__name__ == "GlmImageForConditionalGeneration"
|
||||
|
||||
def test_getattr_raises_for_unknown_attribute(self):
|
||||
"""Verify __getattr__ raises AttributeError for unknown attributes."""
|
||||
import vllm_omni.model_executor.models.glm_image as glm_image_pkg
|
||||
|
||||
# Test unknown attribute via __getattr__ directly
|
||||
with pytest.raises(AttributeError, match="has no attribute"):
|
||||
glm_image_pkg.__getattr__("UnknownClass")
|
||||
|
||||
def test___all___exports_correct_symbols(self):
|
||||
"""Verify __all__ contains the expected exported symbols."""
|
||||
import vllm_omni.model_executor.models.glm_image as glm_image_pkg
|
||||
|
||||
assert hasattr(glm_image_pkg, "__all__")
|
||||
assert "GlmImageForConditionalGeneration" in glm_image_pkg.__all__
|
||||
|
||||
|
||||
class TestLazyImportDoesNotImportTransformersAtInit:
|
||||
"""Verify that importing the package does not eagerly load transformers."""
|
||||
|
||||
def test_glm_image_module_has_getattr(self):
|
||||
"""Test that the module has __getattr__ for lazy loading."""
|
||||
import vllm_omni.model_executor.models.glm_image as glm_image_pkg
|
||||
|
||||
# The module should have __getattr__
|
||||
assert hasattr(glm_image_pkg, "__getattr__")
|
||||
assert callable(glm_image_pkg.__getattr__)
|
||||
|
||||
# And __all__
|
||||
assert hasattr(glm_image_pkg, "__all__")
|
||||
|
||||
def test_module_has_proper_structure(self):
|
||||
"""Test that the module has proper Python module structure."""
|
||||
import vllm_omni.model_executor.models.glm_image as glm_image_pkg
|
||||
|
||||
# Should be a module
|
||||
assert hasattr(glm_image_pkg, "__name__")
|
||||
assert glm_image_pkg.__name__ == "vllm_omni.model_executor.models.glm_image"
|
||||
|
||||
# Should have __file__ (even if it's a package)
|
||||
assert hasattr(glm_image_pkg, "__file__")
|
||||
@@ -19,9 +19,8 @@ from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
|
||||
from vllm_omni.config.stage_config import (
|
||||
_PIPELINE_REGISTRY,
|
||||
PipelineConfig,
|
||||
StageExecutionType,
|
||||
)
|
||||
@@ -34,13 +33,13 @@ _PIPELINE_KEY = "minicpmo_4_5"
|
||||
|
||||
class TestRegistryDeclaration:
|
||||
def test_declared_in_omni_pipelines(self) -> None:
|
||||
assert _PIPELINE_KEY in _OMNI_PIPELINES
|
||||
assert _PIPELINE_KEY in OMNI_PIPELINES
|
||||
|
||||
def test_visible_in_central_registry(self) -> None:
|
||||
assert _PIPELINE_KEY in _PIPELINE_REGISTRY
|
||||
assert _PIPELINE_KEY in OMNI_PIPELINES
|
||||
|
||||
def test_lazy_load_returns_pipeline_config(self) -> None:
|
||||
pipeline = _PIPELINE_REGISTRY[_PIPELINE_KEY]
|
||||
pipeline = OMNI_PIPELINES[_PIPELINE_KEY]
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
assert pipeline.model_type == _PIPELINE_KEY
|
||||
assert pipeline.model_arch == "MiniCPMO45OmniForConditionalGeneration"
|
||||
@@ -49,7 +48,7 @@ class TestRegistryDeclaration:
|
||||
class TestPipelineTopology:
|
||||
@pytest.fixture(scope="class")
|
||||
def pipeline(self) -> PipelineConfig:
|
||||
return _PIPELINE_REGISTRY[_PIPELINE_KEY]
|
||||
return OMNI_PIPELINES[_PIPELINE_KEY]
|
||||
|
||||
def test_two_stages(self, pipeline: PipelineConfig) -> None:
|
||||
assert len(pipeline.stages) == 2
|
||||
@@ -99,7 +98,7 @@ class TestArchAliases:
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def pipeline(self) -> PipelineConfig:
|
||||
return _PIPELINE_REGISTRY[_PIPELINE_KEY]
|
||||
return OMNI_PIPELINES[_PIPELINE_KEY]
|
||||
|
||||
def test_shared_minicpmo_alias_present(self, pipeline: PipelineConfig) -> None:
|
||||
# MiniCPM-o 4.5 ships ``architectures=["MiniCPMO"]`` in its HF config.
|
||||
@@ -123,7 +122,7 @@ class TestHfConfigPredicate:
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def predicate(self):
|
||||
pipeline = _PIPELINE_REGISTRY[_PIPELINE_KEY]
|
||||
pipeline = OMNI_PIPELINES[_PIPELINE_KEY]
|
||||
assert pipeline.hf_config_predicate is not None, (
|
||||
"MiniCPM-o 4.5 pipeline must declare hf_config_predicate to "
|
||||
"avoid misrouting MiniCPM-o 2.6 checkpoints into the 4.5 path."
|
||||
|
||||
@@ -375,6 +375,22 @@ def test_malformed_codec_length_warning_is_rate_limited():
|
||||
assert out2.multimodal_outputs["model_outputs"][0].numel() == 0
|
||||
|
||||
|
||||
def test_forward_emits_zero_samples_for_empty_codec_payload():
|
||||
"""Consumer side of #4463.
|
||||
|
||||
The empty-but-finished payload that ``talker2code2wav_full_payload`` now
|
||||
returns on a degenerate take carries a zero-length ``codes.audio``. Code2Wav
|
||||
must turn that into a 0-sample output (finishing the request) so the Stage-1
|
||||
wait gate releases, rather than stalling the pipeline to the connector
|
||||
timeout.
|
||||
"""
|
||||
model = _make_model()
|
||||
|
||||
out = model.forward(input_ids=torch.zeros(0, dtype=torch.long))
|
||||
|
||||
assert out.multimodal_outputs["model_outputs"][0].numel() == 0
|
||||
|
||||
|
||||
def test_forward_batches_equal_length_requests_in_one_decoder_call():
|
||||
model = _make_model()
|
||||
|
||||
|
||||
@@ -95,7 +95,7 @@ def test_aura2tts_builds_qwen3_tts_prompt_information():
|
||||
|
||||
[tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt])
|
||||
|
||||
assert len(tts_input["prompt_token_ids"]) >= 32
|
||||
assert len(tts_input["prompt_token_ids"]) > 0
|
||||
assert tts_input["additional_information"]["text"] == ["Hello."]
|
||||
assert PRECOMPUTED_TEXT_IDS_KEY not in tts_input["additional_information"]
|
||||
assert tts_input["additional_information"]["task_type"] == ["Base"]
|
||||
@@ -166,6 +166,7 @@ def test_aura2tts_supports_custom_voice_mode():
|
||||
assert tts_input["additional_information"]["task_type"] == ["CustomVoice"]
|
||||
assert tts_input["additional_information"]["speaker"] == ["Vivian"]
|
||||
assert "ref_audio" not in tts_input["additional_information"]
|
||||
assert len(tts_input["prompt_token_ids"]) == 14
|
||||
|
||||
|
||||
def test_aura2tts_passes_token_ids_to_qwen3_tts_when_enabled():
|
||||
|
||||
@@ -0,0 +1,336 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.stage_input_processors.indextts2 import (
|
||||
_strip_stop_token,
|
||||
talker2s2mel,
|
||||
talker2s2mel_full_payload,
|
||||
talker2s2mel_token_only,
|
||||
)
|
||||
|
||||
STOP_MEL_TOKEN = 8193
|
||||
LATENT_DIM = 16
|
||||
|
||||
|
||||
def _make_talker_output(
|
||||
*,
|
||||
mel_codes: torch.Tensor,
|
||||
latent: torch.Tensor,
|
||||
s_ref: torch.Tensor | None = None,
|
||||
ref_mel: torch.Tensor | None = None,
|
||||
style: torch.Tensor | None = None,
|
||||
finished: bool = True,
|
||||
):
|
||||
meta = {}
|
||||
if s_ref is not None:
|
||||
meta["S_ref"] = s_ref
|
||||
if ref_mel is not None:
|
||||
meta["ref_mel"] = ref_mel
|
||||
if style is not None:
|
||||
meta["style"] = style
|
||||
|
||||
mm = {
|
||||
"codes": {"mel": mel_codes},
|
||||
"hidden_states": {"latent": latent},
|
||||
"meta": meta,
|
||||
}
|
||||
return SimpleNamespace(
|
||||
finished=finished,
|
||||
outputs=[SimpleNamespace(multimodal_output=mm)],
|
||||
)
|
||||
|
||||
|
||||
# ── _strip_stop_token ──
|
||||
|
||||
|
||||
def test_strip_stop_token_basic():
|
||||
codes = torch.tensor([10, 20, 30, STOP_MEL_TOKEN, 99])
|
||||
latent = torch.randn(5, LATENT_DIM)
|
||||
|
||||
c, lat, lens = _strip_stop_token(codes, latent)
|
||||
|
||||
assert c.shape == (1, 3)
|
||||
assert lat.shape == (1, 3, LATENT_DIM)
|
||||
assert lens.tolist() == [3]
|
||||
assert (c[0] == torch.tensor([10, 20, 30])).all()
|
||||
|
||||
|
||||
def test_strip_stop_token_no_stop():
|
||||
codes = torch.tensor([10, 20, 30])
|
||||
latent = torch.randn(3, LATENT_DIM)
|
||||
|
||||
c, lat, lens = _strip_stop_token(codes, latent)
|
||||
|
||||
assert c.shape == (1, 3)
|
||||
assert lens.tolist() == [3]
|
||||
|
||||
|
||||
def test_strip_stop_token_at_start():
|
||||
codes = torch.tensor([STOP_MEL_TOKEN, 10, 20])
|
||||
latent = torch.randn(3, LATENT_DIM)
|
||||
|
||||
c, lat, lens = _strip_stop_token(codes, latent)
|
||||
|
||||
assert c.shape == (1, 0)
|
||||
assert lat.shape == (1, 0, LATENT_DIM)
|
||||
assert lens.tolist() == [0]
|
||||
|
||||
|
||||
def test_strip_stop_token_batch_with_padding():
|
||||
codes = torch.tensor(
|
||||
[
|
||||
[10, 20, STOP_MEL_TOKEN, 0],
|
||||
[30, 40, 50, STOP_MEL_TOKEN],
|
||||
]
|
||||
)
|
||||
latent = torch.randn(2, 4, LATENT_DIM)
|
||||
|
||||
c, lat, lens = _strip_stop_token(codes, latent)
|
||||
|
||||
assert c.shape == (2, 3)
|
||||
assert lat.shape == (2, 3, LATENT_DIM)
|
||||
assert lens.tolist() == [2, 3]
|
||||
assert c[0, 2].item() == 0 # padded with a valid semantic code
|
||||
|
||||
|
||||
def test_strip_stop_token_all_stop():
|
||||
codes = torch.tensor([STOP_MEL_TOKEN])
|
||||
latent = torch.randn(1, LATENT_DIM)
|
||||
|
||||
c, lat, lens = _strip_stop_token(codes, latent)
|
||||
|
||||
assert lens.tolist() == [0]
|
||||
assert c.shape[1] == 0
|
||||
assert lat.shape[1] == 0
|
||||
|
||||
|
||||
# ── talker2s2mel ──
|
||||
|
||||
|
||||
def test_talker2s2mel_basic():
|
||||
mel_codes = torch.tensor([10, 20, 30, STOP_MEL_TOKEN])
|
||||
latent = torch.randn(4, LATENT_DIM)
|
||||
s_ref = torch.randn(1, 128)
|
||||
ref_mel = torch.randn(1, 80, 200)
|
||||
style = torch.randn(1, 256)
|
||||
|
||||
outputs = talker2s2mel(
|
||||
source_outputs=[
|
||||
_make_talker_output(
|
||||
mel_codes=mel_codes,
|
||||
latent=latent,
|
||||
s_ref=s_ref,
|
||||
ref_mel=ref_mel,
|
||||
style=style,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert len(outputs) == 1
|
||||
info = outputs[0]["additional_information"]
|
||||
assert info["mel_codes"].shape == (1, 3)
|
||||
assert info["latent"].shape == (1, 3, LATENT_DIM)
|
||||
assert info["code_lens"].tolist() == [3]
|
||||
assert "S_ref" in info
|
||||
assert "ref_mel" in info
|
||||
assert "style" in info
|
||||
assert outputs[0]["prompt_token_ids"] == [0]
|
||||
|
||||
|
||||
def test_talker2s2mel_skips_unfinished():
|
||||
mel_codes = torch.tensor([10, 20])
|
||||
latent = torch.randn(2, LATENT_DIM)
|
||||
|
||||
outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, finished=False)])
|
||||
|
||||
assert len(outputs) == 0
|
||||
|
||||
|
||||
def test_talker2s2mel_skips_empty_mel():
|
||||
outputs = talker2s2mel(
|
||||
source_outputs=[
|
||||
_make_talker_output(
|
||||
mel_codes=torch.tensor([]),
|
||||
latent=torch.randn(0, LATENT_DIM),
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert len(outputs) == 0
|
||||
|
||||
|
||||
def test_talker2s2mel_missing_optional_meta():
|
||||
mel_codes = torch.tensor([10, 20])
|
||||
latent = torch.randn(2, LATENT_DIM)
|
||||
|
||||
outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent)])
|
||||
|
||||
assert len(outputs) == 1
|
||||
info = outputs[0]["additional_information"]
|
||||
assert "S_ref" not in info
|
||||
assert "ref_mel" not in info
|
||||
assert "style" not in info
|
||||
|
||||
|
||||
def test_talker2s2mel_batch():
|
||||
out0 = _make_talker_output(
|
||||
mel_codes=torch.tensor([1, 2, STOP_MEL_TOKEN]),
|
||||
latent=torch.randn(3, LATENT_DIM),
|
||||
s_ref=torch.randn(1, 64),
|
||||
)
|
||||
out1 = _make_talker_output(
|
||||
mel_codes=torch.tensor([5, 6, 7, 8]),
|
||||
latent=torch.randn(4, LATENT_DIM),
|
||||
ref_mel=torch.randn(1, 80, 100),
|
||||
)
|
||||
|
||||
outputs = talker2s2mel(source_outputs=[out0, out1])
|
||||
|
||||
assert len(outputs) == 2
|
||||
assert outputs[0]["additional_information"]["code_lens"].tolist() == [2]
|
||||
assert outputs[1]["additional_information"]["code_lens"].tolist() == [4]
|
||||
|
||||
|
||||
def test_talker2s2mel_tensors_on_cpu():
|
||||
device = torch.device("cpu")
|
||||
mel_codes = torch.tensor([10, 20], device=device)
|
||||
latent = torch.randn(2, LATENT_DIM, device=device)
|
||||
s_ref = torch.randn(1, 64, device=device)
|
||||
|
||||
outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, s_ref=s_ref)])
|
||||
|
||||
info = outputs[0]["additional_information"]
|
||||
assert info["latent"].device.type == "cpu"
|
||||
assert info["mel_codes"].device.type == "cpu"
|
||||
assert info["S_ref"].device.type == "cpu"
|
||||
|
||||
|
||||
def test_talker2s2mel_no_multimodal_dict():
|
||||
output = SimpleNamespace(
|
||||
finished=True,
|
||||
outputs=[SimpleNamespace(multimodal_output="not_a_dict")],
|
||||
)
|
||||
|
||||
results = talker2s2mel(source_outputs=[output])
|
||||
|
||||
assert len(results) == 0
|
||||
|
||||
|
||||
# ── full-payload connector path ──
|
||||
|
||||
|
||||
def _conditioning_tensors():
|
||||
return {
|
||||
"S_ref": torch.randn(1, 3, LATENT_DIM),
|
||||
"ref_mel": torch.randn(1, 80, 5),
|
||||
"style": torch.randn(1, 4),
|
||||
}
|
||||
|
||||
|
||||
def test_talker2s2mel_token_only_placeholder_finished_only():
|
||||
finished = _make_talker_output(mel_codes=torch.tensor([1]), latent=torch.randn(1, LATENT_DIM), finished=True)
|
||||
unfinished = _make_talker_output(mel_codes=torch.tensor([2]), latent=torch.randn(1, LATENT_DIM), finished=False)
|
||||
|
||||
outputs = talker2s2mel_token_only([unfinished, finished])
|
||||
|
||||
assert len(outputs) == 1
|
||||
assert outputs[0]["prompt_token_ids"] == [0]
|
||||
assert outputs[0]["additional_information"] is None
|
||||
assert outputs[0]["multi_modal_data"] is None
|
||||
|
||||
|
||||
def test_talker2s2mel_full_payload_flat_keys_builds_s2mel_contract():
|
||||
cond = _conditioning_tensors()
|
||||
latent = torch.randn(3, LATENT_DIM)
|
||||
payload = {
|
||||
"codes.mel": torch.tensor([[10], [20], [30]]),
|
||||
"hidden_states.latent": latent,
|
||||
"meta.S_ref": cond["S_ref"],
|
||||
"meta.ref_mel": cond["ref_mel"],
|
||||
"meta.style": cond["style"],
|
||||
}
|
||||
|
||||
result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-flat"))
|
||||
|
||||
assert result is not None
|
||||
assert result["mel_codes"].shape == (1, 3)
|
||||
assert result["mel_codes"].tolist() == [[10, 20, 30]]
|
||||
assert result["latent"].shape == (1, 3, LATENT_DIM)
|
||||
assert result["code_lens"].tolist() == [3]
|
||||
assert result["S_ref"].device.type == "cpu"
|
||||
assert result["ref_mel"].device.type == "cpu"
|
||||
assert result["style"].device.type == "cpu"
|
||||
assert result["S_ref"].data_ptr() != cond["S_ref"].data_ptr()
|
||||
|
||||
|
||||
def test_talker2s2mel_full_payload_nested_fallback_input():
|
||||
cond = _conditioning_tensors()
|
||||
payload = {
|
||||
"codes": {"mel": torch.tensor([[4], [5]])},
|
||||
"hidden_states": {"latent": torch.randn(2, LATENT_DIM)},
|
||||
"meta": cond,
|
||||
}
|
||||
|
||||
result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-nested"))
|
||||
|
||||
assert result is not None
|
||||
assert result["mel_codes"].tolist() == [[4, 5]]
|
||||
assert result["latent"].shape == (1, 2, LATENT_DIM)
|
||||
|
||||
|
||||
def test_talker2s2mel_full_payload_normalizes_one_by_t_mel_row():
|
||||
cond = _conditioning_tensors()
|
||||
payload = {
|
||||
"codes.mel": torch.tensor([[7, 8, 9]]),
|
||||
"hidden_states.latent": torch.randn(3, LATENT_DIM),
|
||||
"meta": cond,
|
||||
}
|
||||
|
||||
result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-row"))
|
||||
|
||||
assert result is not None
|
||||
assert result["mel_codes"].tolist() == [[7, 8, 9]]
|
||||
|
||||
|
||||
def test_talker2s2mel_full_payload_trims_stop_token_and_matching_latent():
|
||||
cond = _conditioning_tensors()
|
||||
latent = torch.arange(4 * LATENT_DIM, dtype=torch.float32).reshape(4, LATENT_DIM)
|
||||
payload = {
|
||||
"codes.mel": torch.tensor([[1], [2], [STOP_MEL_TOKEN], [99]]),
|
||||
"hidden_states.latent": latent,
|
||||
"meta": cond,
|
||||
}
|
||||
|
||||
result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-stop"))
|
||||
|
||||
assert result is not None
|
||||
assert result["mel_codes"].tolist() == [[1, 2]]
|
||||
assert result["latent"].shape == (1, 2, LATENT_DIM)
|
||||
assert result["latent"][0, 0].tolist() == latent[0].tolist()
|
||||
assert result["latent"][0, 1].tolist() == latent[1].tolist()
|
||||
assert result["code_lens"].tolist() == [2]
|
||||
|
||||
|
||||
def test_talker2s2mel_full_payload_crops_mel_latent_to_common_length():
|
||||
cond = _conditioning_tensors()
|
||||
payload = {
|
||||
"codes.mel": torch.tensor([[1], [2], [3], [4]]),
|
||||
"hidden_states.latent": torch.randn(2, LATENT_DIM),
|
||||
"meta": cond,
|
||||
}
|
||||
|
||||
result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-crop"))
|
||||
|
||||
assert result is not None
|
||||
assert result["mel_codes"].tolist() == [[1, 2]]
|
||||
assert result["latent"].shape == (1, 2, LATENT_DIM)
|
||||
assert result["code_lens"].tolist() == [2]
|
||||
|
||||
|
||||
def test_talker2s2mel_full_payload_missing_required_fields_returns_none():
|
||||
assert talker2s2mel_full_payload(None, {"hidden_states.latent": torch.randn(1, LATENT_DIM)}, None) is None
|
||||
assert talker2s2mel_full_payload(None, {"codes.mel": torch.tensor([[1]])}, None) is None
|
||||
@@ -486,3 +486,32 @@ def test_full_payload_omits_left_context_size_without_ref():
|
||||
assert payload["meta"]["finished"].item() is True
|
||||
assert "left_context_size" not in payload["meta"]
|
||||
assert len(payload["codes"]["audio"]) == _Q * 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pooling_output",
|
||||
[
|
||||
pytest.param(SimpleNamespace(codes="not-a-dict"), id="non_dict_output"),
|
||||
pytest.param({}, id="missing_codes_audio"),
|
||||
pytest.param({"codes.audio": torch.zeros((3, _Q), dtype=torch.long)}, id="all_codes_filtered"),
|
||||
],
|
||||
)
|
||||
def test_full_payload_emits_empty_finished_payload_on_degenerate_take(pooling_output):
|
||||
"""Regression for #4463.
|
||||
|
||||
A degenerate talker take used to return ``None`` from
|
||||
``talker2code2wav_full_payload``. The connector treats ``None`` as "drop the
|
||||
request", but Stage-1 was already scheduled to receive it, so its wait gate
|
||||
polls to ``connector_get_max_wait`` (~300s) and the orchestrator aborts — one
|
||||
stuck request stalls the whole two-stage pipeline. Each degenerate case
|
||||
(non-dict pooling_output, missing ``codes.audio``, all codec frames dropped by
|
||||
the filter) must instead return an empty-but-finished payload so the gate
|
||||
releases and the request finishes immediately with zero-length audio.
|
||||
"""
|
||||
request = SimpleNamespace(request_id="r", output_token_ids=[0, 1, 2])
|
||||
|
||||
payload = talker2code2wav_full_payload(transfer_manager=None, pooling_output=pooling_output, request=request)
|
||||
|
||||
assert payload is not None
|
||||
assert payload["meta"]["finished"].item() is True
|
||||
assert payload["codes"]["audio"].numel() == 0
|
||||
@@ -10,6 +10,7 @@ from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.model_extras import (
|
||||
build_image_to_image_prompt,
|
||||
build_image_to_video_prompt,
|
||||
build_text_to_image_prompt,
|
||||
get_extra_body_params,
|
||||
get_extra_output_params,
|
||||
@@ -57,6 +58,57 @@ def test_sensenova_extra_registry_declares_request_and_response_params() -> None
|
||||
assert should_init_extra_args_for_non_diffusion_stages("SenseNovaU1Pipeline") is False
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_cosmos3_extra_registry_declares_request_and_response_params() -> None:
|
||||
assert get_extra_body_params("Cosmos3OmniDiffusersPipeline") == frozenset(
|
||||
{
|
||||
"flow_shift",
|
||||
"max_sequence_length",
|
||||
"use_resolution_template",
|
||||
"use_duration_template",
|
||||
"use_system_prompt",
|
||||
"system_prompt",
|
||||
"negative_prompt",
|
||||
"guardrails",
|
||||
"condition_frame_indexes_vision",
|
||||
"condition_video_keep",
|
||||
"generate_sound",
|
||||
"sound_gen",
|
||||
"sound_duration",
|
||||
"audio_duration",
|
||||
"action_mode",
|
||||
"action",
|
||||
"domain_name",
|
||||
"domain_id",
|
||||
"raw_action_dim",
|
||||
"action_chunk_size",
|
||||
"action_space",
|
||||
"action_fps",
|
||||
"image_height",
|
||||
"image_width",
|
||||
"history_length",
|
||||
"conditioning_fps",
|
||||
"resolution",
|
||||
"image_size",
|
||||
"use_state",
|
||||
"observation",
|
||||
"robot_obs",
|
||||
"deterministic_seed",
|
||||
"session_id",
|
||||
}
|
||||
)
|
||||
assert get_extra_output_params("Cosmos3OmniDiffusersPipeline") == frozenset(
|
||||
{
|
||||
"action",
|
||||
"raw_action_dim",
|
||||
"domain_id",
|
||||
"action_mode",
|
||||
}
|
||||
)
|
||||
assert should_init_extra_args_for_non_diffusion_stages("Cosmos3OmniDiffusersPipeline") is False
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_magi_human_extra_registry_declares_request_and_response_params() -> None:
|
||||
@@ -74,6 +126,27 @@ def test_magi_human_extra_registry_declares_request_and_response_params() -> Non
|
||||
assert should_init_extra_args_for_non_diffusion_stages("MagiHumanPipeline") is False
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_cosmos3_text_to_image_prompt_builder_selects_image_modality() -> None:
|
||||
assert build_text_to_image_prompt(
|
||||
"Cosmos3OmniDiffusersPipeline",
|
||||
prompt="a red sports car at golden hour",
|
||||
negative_prompt="blurry, distorted",
|
||||
height=1024,
|
||||
width=1024,
|
||||
) == {
|
||||
"prompt": "a red sports car at golden hour",
|
||||
"modalities": ["image"],
|
||||
"negative_prompt": "blurry, distorted",
|
||||
}
|
||||
assert build_text_to_image_prompt(
|
||||
"Cosmos3OmniDiffusersPipeline",
|
||||
prompt="a red sports car",
|
||||
negative_prompt=None,
|
||||
) == {"prompt": "a red sports car", "modalities": ["image"]}
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_helios_extra_registry_declares_request_and_response_params() -> None:
|
||||
@@ -104,6 +177,14 @@ def test_helios_extra_registry_declares_request_and_response_params() -> None:
|
||||
assert should_init_extra_args_for_non_diffusion_stages(cls) is False
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_vace_extra_registry_has_no_pipeline_params() -> None:
|
||||
assert get_extra_body_params("WanVACEPipeline") == frozenset()
|
||||
assert get_extra_output_params("WanVACEPipeline") == frozenset()
|
||||
assert should_init_extra_args_for_non_diffusion_stages("WanVACEPipeline") is False
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_unknown_pipeline_has_empty_extra_registry() -> None:
|
||||
@@ -181,6 +262,58 @@ def test_unknown_pipeline_uses_default_image_to_image_prompt() -> None:
|
||||
}
|
||||
|
||||
|
||||
def _build_vace_prompt(media_inputs: dict[str, object], *, num_frames: int = 5) -> dict:
|
||||
return build_image_to_video_prompt(
|
||||
"WanVACEPipeline",
|
||||
prompt="a bird flying",
|
||||
negative_prompt=None,
|
||||
media_inputs=media_inputs,
|
||||
height=16,
|
||||
width=320,
|
||||
num_frames=num_frames,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize(
|
||||
"media_inputs",
|
||||
[
|
||||
{"image": Image.new("RGB", (320, 16), "red")},
|
||||
{"last_image": Image.new("RGB", (320, 16), "blue")},
|
||||
{
|
||||
"image": Image.new("RGB", (320, 16), "red"),
|
||||
"last_image": Image.new("RGB", (320, 16), "blue"),
|
||||
},
|
||||
{"image": Image.new("RGB", (320, 16), "red"), "mask": Image.new("L", (320, 16), 0)},
|
||||
{"reference_images": [Image.new("RGB", (64, 64), "red")]},
|
||||
],
|
||||
ids=["i2v", "v2lf", "flf2v", "inpaint", "r2v"],
|
||||
)
|
||||
def test_vace_image_to_video_prompt_builder(media_inputs: dict[str, object]) -> None:
|
||||
result = _build_vace_prompt(media_inputs)
|
||||
mmd = result["multi_modal_data"]
|
||||
if "reference_images" in mmd:
|
||||
assert mmd["reference_images"] is media_inputs["reference_images"]
|
||||
else:
|
||||
assert len(mmd["video"]) == len(mmd["mask"]) == 5
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize(
|
||||
("media_inputs", "message"),
|
||||
[
|
||||
({}, "requires a conditioning media input"),
|
||||
({"mask": Image.new("L", (320, 16))}, "mask input requires an image"),
|
||||
({"control_image": Image.new("RGB", (320, 16))}, "Unsupported VACE media input"),
|
||||
],
|
||||
)
|
||||
def test_vace_rejects_invalid_media_combinations(media_inputs: dict[str, object], message: str) -> None:
|
||||
with pytest.raises(ValueError, match=message):
|
||||
_build_vace_prompt(media_inputs)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_declared_extra_args_apply_to_existing_sampling_params() -> None:
|
||||
|
||||
+252
-73
@@ -12,14 +12,15 @@ from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from transformers import PretrainedConfig, Qwen3OmniMoeConfig
|
||||
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni.config.config_factory import StageConfigFactory
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline
|
||||
from vllm_omni.config.stage_config import (
|
||||
_PIPELINE_REGISTRY,
|
||||
DeployConfig,
|
||||
PipelineConfig,
|
||||
StageConfig,
|
||||
StageConfigFactory,
|
||||
StageDeployConfig,
|
||||
StageExecutionType,
|
||||
StagePipelineConfig,
|
||||
@@ -30,14 +31,16 @@ from vllm_omni.config.stage_config import (
|
||||
build_stage_runtime_overrides,
|
||||
load_deploy_config,
|
||||
merge_pipeline_deploy,
|
||||
register_pipeline,
|
||||
pipeline_cfg_resolver,
|
||||
strip_parent_engine_args,
|
||||
)
|
||||
from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True)
|
||||
Q3_OMNI_THINKER_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=False)
|
||||
|
||||
|
||||
class TestStageType:
|
||||
@@ -511,26 +514,39 @@ class TestStageResolutionHelpers:
|
||||
|
||||
|
||||
class TestPipelineDiscovery:
|
||||
"""Tests for the central pipeline registry (``pipeline_registry._OMNI_PIPELINES``)."""
|
||||
"""Tests for the central pipeline registry (``OMNI_PIPELINES``)."""
|
||||
|
||||
def test_registry_has_known_models(self):
|
||||
"""Built-in pipelines are lazy-loaded from the central declaration
|
||||
on first access; no eager import or discovery walk needed."""
|
||||
# ``in`` triggers the lazy-map lookup without forcing a load.
|
||||
assert "qwen2_5_omni" in _PIPELINE_REGISTRY
|
||||
assert "qwen3_omni_moe" in _PIPELINE_REGISTRY
|
||||
assert "qwen3_tts" in _PIPELINE_REGISTRY
|
||||
"""Check that specific models are in OMNI_PIPELINES."""
|
||||
assert "qwen2_5_omni" in OMNI_PIPELINES
|
||||
assert "qwen3_omni_moe" in OMNI_PIPELINES
|
||||
assert "qwen3_tts" in OMNI_PIPELINES
|
||||
|
||||
def test_registry_loads_pipeline_on_getitem(self):
|
||||
"""Looking up a registered model_type returns the matching PipelineConfig."""
|
||||
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
|
||||
def test_registry_resolver_qwen3_omni_all_stages(self):
|
||||
"""Test that providing the HF config for qwen3 omni with audio enabled uses all stages."""
|
||||
pipeline = StageConfigFactory.resolve_pipeline_config(
|
||||
"qwen3_omni_moe",
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
assert pipeline.model_type == "qwen3_omni_moe"
|
||||
assert len(pipeline.stages) == 3 # thinker + talker + code2wav
|
||||
|
||||
def test_registry_resolver_qwen3_omni_thinker_only(self):
|
||||
"""Test that providing the HF config for qwen3 omni without audio is thinker only."""
|
||||
pipeline = StageConfigFactory.resolve_pipeline_config(
|
||||
"qwen3_omni_moe",
|
||||
Q3_OMNI_THINKER_HF_CONFIG,
|
||||
)
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
|
||||
assert len(pipeline.stages) == 1 # thinker only
|
||||
|
||||
def test_registry_returns_none_for_unknown(self):
|
||||
"""Unknown model_types aren't found; ``get()`` returns None."""
|
||||
assert "definitely_not_a_real_model" not in _PIPELINE_REGISTRY
|
||||
assert _PIPELINE_REGISTRY.get("definitely_not_a_real_model") is None
|
||||
"""Unknown model_types aren't found and resolve to `None`."""
|
||||
assert "definitely_not_a_real_model" not in OMNI_PIPELINES
|
||||
assert OMNI_PIPELINES.get("definitely_not_a_real_model") is None
|
||||
assert StageConfigFactory.resolve_pipeline_config("definitely_not_a_real_model") is None
|
||||
|
||||
def test_pipeline_config_supports_hf_architectures(self):
|
||||
"""PipelineConfig accepts hf_architectures for HF-arch fallback
|
||||
@@ -579,6 +595,67 @@ class TestPipelineConfigNew:
|
||||
assert any("no stages" in e.lower() for e in p.validate())
|
||||
|
||||
|
||||
class TestPipelineRegistration:
|
||||
def test_pipeline_registration(self, clean_pipeline_registry):
|
||||
"""Ensure that we can register and create a custom pipeline config."""
|
||||
new_model_type = "new_model_type"
|
||||
pipe_cfg = PipelineConfig(model_type=new_model_type)
|
||||
|
||||
# Register the new PipelineConfig
|
||||
assert new_model_type not in OMNI_PIPELINES
|
||||
register_pipeline(pipe_cfg)
|
||||
assert new_model_type in OMNI_PIPELINES
|
||||
assert OMNI_PIPELINES[new_model_type] is pipe_cfg
|
||||
|
||||
class FakeConfig(PretrainedConfig):
|
||||
model_type = new_model_type
|
||||
|
||||
# Create the model
|
||||
with (
|
||||
patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()),
|
||||
patch.object(StageConfigFactory, "_create_from_registry") as mock_create,
|
||||
):
|
||||
StageConfigFactory.create_from_model("fake/model")
|
||||
mock_create.assert_called_once()
|
||||
# Ensure that the passed pipeline_config is the right type
|
||||
pipeline_cfg = mock_create.call_args[0][1]
|
||||
assert isinstance(pipeline_cfg, PipelineConfig)
|
||||
assert pipe_cfg.model_type == new_model_type
|
||||
|
||||
def test_resolver_registration(self, clean_pipeline_registry):
|
||||
"""Ensure that we can register and create a custom resolver."""
|
||||
new_model_type = "new_model_type"
|
||||
resolved_type = "resolved_type"
|
||||
|
||||
class FakeConfig(PretrainedConfig):
|
||||
model_type = new_model_type
|
||||
|
||||
@pipeline_cfg_resolver(config_type=FakeConfig)
|
||||
def custom_resolver(
|
||||
hf_config: FakeConfig,
|
||||
) -> PipelineConfig:
|
||||
return PipelineConfig(model_type=resolved_type)
|
||||
|
||||
# Register the new PipelineConfig
|
||||
assert new_model_type not in OMNI_PIPELINES
|
||||
|
||||
register_pipeline(custom_resolver, model_type=new_model_type)
|
||||
assert new_model_type in OMNI_PIPELINES
|
||||
assert OMNI_PIPELINES[new_model_type] is custom_resolver
|
||||
|
||||
# Create the model
|
||||
with (
|
||||
patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()),
|
||||
patch.object(StageConfigFactory, "_create_from_registry") as mock_create,
|
||||
):
|
||||
StageConfigFactory.create_from_model("fake/model")
|
||||
mock_create.assert_called_once()
|
||||
# Ensure that the passed pipeline_config is the right type
|
||||
pipeline_cfg = mock_create.call_args[0][1]
|
||||
assert isinstance(pipeline_cfg, PipelineConfig)
|
||||
assert pipeline_cfg.model_type == resolved_type
|
||||
|
||||
|
||||
class TestResolveScheduler:
|
||||
def test_all_execution_types_handled(self):
|
||||
for et in StageExecutionType:
|
||||
@@ -603,18 +680,6 @@ class TestResolveScheduler:
|
||||
assert _resolve_scheduler(StageExecutionType.DIFFUSION) is None
|
||||
|
||||
|
||||
class TestPipelineRegistry:
|
||||
def test_register_and_lookup(self):
|
||||
p = PipelineConfig(
|
||||
model_type="__test_only__",
|
||||
model_arch="A",
|
||||
stages=(StagePipelineConfig(stage_id=0, model_stage="a"),),
|
||||
)
|
||||
register_pipeline(p)
|
||||
assert _PIPELINE_REGISTRY["__test_only__"] is p
|
||||
del _PIPELINE_REGISTRY["__test_only__"]
|
||||
|
||||
|
||||
class TestDeployConfigLoading:
|
||||
def test_custom_voice_dir_is_pipeline_wide_engine_arg(self, tmp_path):
|
||||
deploy_path = tmp_path / "qwen3_tts_custom_voice.yaml"
|
||||
@@ -633,8 +698,9 @@ stages:
|
||||
)
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
|
||||
pipeline = StageConfigFactory.resolve_pipeline_config("qwen3_tts")
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
|
||||
assert deploy.custom_voice_dir == str(custom_voice_dir)
|
||||
assert {s.yaml_engine_args.get("custom_voice_dir") for s in stages} == {str(custom_voice_dir)}
|
||||
@@ -664,11 +730,16 @@ stages:
|
||||
assert "compilation_config" not in deploy.stages[0].engine_extras
|
||||
|
||||
def test_merge_pipeline_deploy(self):
|
||||
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
|
||||
deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
|
||||
if not deploy_path.exists():
|
||||
pytest.skip("Deploy config not found")
|
||||
|
||||
pipeline = StageConfigFactory.resolve_pipeline_config(
|
||||
"qwen3_omni_moe",
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
@@ -680,7 +751,12 @@ stages:
|
||||
assert s0.yaml_extras["default_sampling_params"]["detokenize"] is True
|
||||
|
||||
def test_merge_pipeline_deploy_preserves_num_replicas(self, tmp_path):
|
||||
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
|
||||
pipeline = StageConfigFactory.resolve_pipeline_config(
|
||||
"qwen3_omni_moe",
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
|
||||
base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
|
||||
if not base.exists():
|
||||
pytest.skip("Deploy config not found")
|
||||
@@ -894,14 +970,24 @@ stages:
|
||||
|
||||
class TestQwen3OmniPipeline:
|
||||
def test_registered(self):
|
||||
p = _PIPELINE_REGISTRY.get("qwen3_omni_moe")
|
||||
assert p is not None
|
||||
p = StageConfigFactory.resolve_pipeline_config(
|
||||
"qwen3_omni_moe",
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "Qwen3OmniMoeForConditionalGeneration"
|
||||
assert len(p.stages) == 3
|
||||
assert p.validate() == []
|
||||
|
||||
def test_thinker(self):
|
||||
s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(0)
|
||||
p = StageConfigFactory.resolve_pipeline_config(
|
||||
"qwen3_omni_moe",
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
assert isinstance(p, PipelineConfig)
|
||||
|
||||
s = p.get_stage(0)
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
assert s.model_stage == "thinker"
|
||||
assert s.execution_type == StageExecutionType.LLM_AR
|
||||
assert s.owns_tokenizer is True
|
||||
@@ -909,14 +995,28 @@ class TestQwen3OmniPipeline:
|
||||
assert s.sampling_constraints["detokenize"] is True
|
||||
|
||||
def test_talker(self):
|
||||
s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(1)
|
||||
p = StageConfigFactory.resolve_pipeline_config(
|
||||
"qwen3_omni_moe",
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
assert isinstance(p, PipelineConfig)
|
||||
|
||||
s = p.get_stage(1)
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
assert s.input_sources == (0,)
|
||||
assert s.sampling_constraints["stop_token_ids"] == [2150]
|
||||
assert s.custom_process_input_func is not None
|
||||
assert s.custom_process_next_stage_input_func is not None
|
||||
|
||||
def test_code2wav(self):
|
||||
s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(2)
|
||||
p = StageConfigFactory.resolve_pipeline_config(
|
||||
"qwen3_omni_moe",
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
assert isinstance(p, PipelineConfig)
|
||||
|
||||
s = p.get_stage(2)
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
assert s.execution_type == StageExecutionType.LLM_GENERATION
|
||||
assert s.final_output_type == "audio"
|
||||
assert s.custom_process_input_func is not None
|
||||
@@ -924,14 +1024,18 @@ class TestQwen3OmniPipeline:
|
||||
|
||||
class TestQwen2_5OmniPipeline:
|
||||
def test_registered(self):
|
||||
p = _PIPELINE_REGISTRY.get("qwen2_5_omni")
|
||||
assert p is not None
|
||||
p = StageConfigFactory.resolve_pipeline_config("qwen2_5_omni")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "Qwen2_5OmniForConditionalGeneration"
|
||||
assert len(p.stages) == 3
|
||||
assert p.validate() == []
|
||||
|
||||
def test_thinker(self):
|
||||
s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(0)
|
||||
p = StageConfigFactory.resolve_pipeline_config("qwen2_5_omni")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
|
||||
s = p.get_stage(0)
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
assert s.model_stage == "thinker"
|
||||
assert s.execution_type == StageExecutionType.LLM_AR
|
||||
assert s.owns_tokenizer is True
|
||||
@@ -939,13 +1043,21 @@ class TestQwen2_5OmniPipeline:
|
||||
assert s.requires_multimodal_data is True
|
||||
|
||||
def test_talker(self):
|
||||
s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(1)
|
||||
p = StageConfigFactory.resolve_pipeline_config("qwen2_5_omni")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
|
||||
s = p.get_stage(1)
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
assert s.input_sources == (0,)
|
||||
assert s.sampling_constraints["stop_token_ids"] == [8294]
|
||||
assert s.custom_process_input_func is not None
|
||||
|
||||
def test_code2wav(self):
|
||||
s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(2)
|
||||
p = StageConfigFactory.resolve_pipeline_config("qwen2_5_omni")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
|
||||
s = p.get_stage(2)
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
assert s.execution_type == StageExecutionType.LLM_GENERATION
|
||||
assert s.final_output_type == "audio"
|
||||
assert s.engine_output_type == "audio"
|
||||
@@ -953,14 +1065,19 @@ class TestQwen2_5OmniPipeline:
|
||||
|
||||
class TestQwen3TTSPipeline:
|
||||
def test_registered(self):
|
||||
p = _PIPELINE_REGISTRY.get("qwen3_tts")
|
||||
p = StageConfigFactory.resolve_pipeline_config("qwen3_tts")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p is not None
|
||||
assert p.model_arch == "Qwen3TTSTalkerForConditionalGeneration"
|
||||
assert len(p.stages) == 2
|
||||
assert p.validate() == []
|
||||
|
||||
def test_talker_stage(self):
|
||||
s = _PIPELINE_REGISTRY["qwen3_tts"].get_stage(0)
|
||||
p = StageConfigFactory.resolve_pipeline_config("qwen3_tts")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
|
||||
s = p.get_stage(0)
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
assert s.model_stage == "qwen3_tts"
|
||||
assert s.execution_type == StageExecutionType.LLM_AR
|
||||
assert s.owns_tokenizer is True
|
||||
@@ -970,7 +1087,11 @@ class TestQwen3TTSPipeline:
|
||||
assert s.model_arch is None
|
||||
|
||||
def test_code2wav_stage_has_per_stage_model_arch(self):
|
||||
s = _PIPELINE_REGISTRY["qwen3_tts"].get_stage(1)
|
||||
p = StageConfigFactory.resolve_pipeline_config("qwen3_tts")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
|
||||
s = p.get_stage(1)
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
assert s.execution_type == StageExecutionType.LLM_GENERATION
|
||||
assert s.final_output_type == "audio"
|
||||
assert s.engine_output_type == "audio"
|
||||
@@ -986,7 +1107,7 @@ class TestQwen3TTSPipeline:
|
||||
pytest.skip("qwen3_tts deploy yaml not found")
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
|
||||
pipeline = OMNI_PIPELINES["qwen3_tts"]
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
# Stage 0 inherits pipeline-level model_arch
|
||||
@@ -1025,14 +1146,18 @@ class TestQwen3TTSPipeline:
|
||||
|
||||
class TestMingFlashOmniPipeline:
|
||||
def test_registered(self):
|
||||
p = _PIPELINE_REGISTRY.get("ming_flash_omni")
|
||||
assert p is not None
|
||||
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
|
||||
assert len(p.stages) == 2
|
||||
assert p.validate() == []
|
||||
|
||||
def test_thinker_stage(self):
|
||||
s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(0)
|
||||
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
|
||||
s = p.get_stage(0)
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
assert s.model_stage == "thinker"
|
||||
assert s.execution_type == StageExecutionType.LLM_AR
|
||||
assert s.owns_tokenizer is True
|
||||
@@ -1042,7 +1167,11 @@ class TestMingFlashOmniPipeline:
|
||||
assert s.sampling_constraints["detokenize"] is True
|
||||
|
||||
def test_talker_stage(self):
|
||||
s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(1)
|
||||
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
|
||||
s = p.get_stage(1)
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
assert s.model_stage == "ming_tts"
|
||||
assert s.execution_type == StageExecutionType.LLM_GENERATION
|
||||
assert s.input_sources == (0,)
|
||||
@@ -1060,20 +1189,28 @@ class TestMingFlashOmniPipeline:
|
||||
Lazy string references only fail at first inference otherwise — this
|
||||
catches typos in the pipeline declaration at import / registration time.
|
||||
"""
|
||||
s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(1)
|
||||
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
|
||||
s = p.get_stage(1)
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
module_path, _, attr = s.custom_process_input_func.rpartition(".")
|
||||
module = importlib.import_module(module_path)
|
||||
assert callable(getattr(module, attr))
|
||||
|
||||
def test_tts_pipeline_registered(self):
|
||||
p = _PIPELINE_REGISTRY.get("ming_flash_omni_tts")
|
||||
assert p is not None
|
||||
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_tts")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "MingFlashOmniTalkerForConditionalGeneration"
|
||||
assert len(p.stages) == 1
|
||||
assert p.validate() == []
|
||||
|
||||
def test_tts_stage(self):
|
||||
s = _PIPELINE_REGISTRY["ming_flash_omni_tts"].get_stage(0)
|
||||
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_tts")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
|
||||
s = p.get_stage(0)
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
assert s.model_stage == "ming_tts"
|
||||
assert s.execution_type == StageExecutionType.LLM_GENERATION
|
||||
assert s.input_sources == ()
|
||||
@@ -1096,7 +1233,8 @@ class TestMingFlashOmniPipeline:
|
||||
# We won't test stage 0/1 colocation contract here,
|
||||
# as there could exist more variant of custom device setup
|
||||
|
||||
pipeline = _PIPELINE_REGISTRY["ming_flash_omni"]
|
||||
pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
assert len(stages) == 2
|
||||
assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration"
|
||||
@@ -1112,20 +1250,24 @@ class TestMingFlashOmniPipeline:
|
||||
assert len(deploy.stages) == 1
|
||||
assert deploy.pipeline == "ming_flash_omni_tts"
|
||||
|
||||
pipeline = _PIPELINE_REGISTRY["ming_flash_omni_tts"]
|
||||
pipeline = OMNI_PIPELINES["ming_flash_omni_tts"]
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
assert len(stages) == 1
|
||||
assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniTalkerForConditionalGeneration"
|
||||
|
||||
def test_thinker_only_pipeline_registered(self):
|
||||
p = _PIPELINE_REGISTRY.get("ming_flash_omni_thinker_only")
|
||||
assert p is not None
|
||||
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_thinker_only")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
|
||||
assert len(p.stages) == 1
|
||||
assert p.validate() == []
|
||||
|
||||
def test_thinker_only_stage(self):
|
||||
s = _PIPELINE_REGISTRY["ming_flash_omni_thinker_only"].get_stage(0)
|
||||
p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_thinker_only")
|
||||
assert isinstance(p, PipelineConfig)
|
||||
|
||||
s = p.get_stage(0)
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
assert s.model_stage == "thinker"
|
||||
assert s.execution_type == StageExecutionType.LLM_AR
|
||||
assert s.input_sources == ()
|
||||
@@ -1146,20 +1288,21 @@ class TestMingFlashOmniPipeline:
|
||||
assert len(deploy.stages) == 1
|
||||
assert deploy.pipeline == "ming_flash_omni_thinker_only"
|
||||
|
||||
pipeline = _PIPELINE_REGISTRY["ming_flash_omni_thinker_only"]
|
||||
pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_thinker_only")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
assert len(stages) == 1
|
||||
assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration"
|
||||
|
||||
def test_image_pipeline_registered(self):
|
||||
p = _PIPELINE_REGISTRY.get("ming_flash_omni_image")
|
||||
p = OMNI_PIPELINES.get("ming_flash_omni_image")
|
||||
assert p is not None
|
||||
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
|
||||
assert len(p.stages) == 2
|
||||
assert p.validate() == []
|
||||
|
||||
def test_image_thinker_stage(self):
|
||||
s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0)
|
||||
s = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_image").get_stage(0)
|
||||
assert s.model_stage == "thinker"
|
||||
assert s.execution_type == StageExecutionType.LLM_AR
|
||||
assert s.input_sources == ()
|
||||
@@ -1173,7 +1316,7 @@ class TestMingFlashOmniPipeline:
|
||||
assert s.prompt_expand_func is not None
|
||||
|
||||
def test_image_dit_stage(self):
|
||||
s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1)
|
||||
s = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_image").get_stage(1)
|
||||
assert s.model_stage == "dit"
|
||||
assert s.execution_type == StageExecutionType.DIFFUSION
|
||||
assert s.input_sources == (0,)
|
||||
@@ -1185,8 +1328,11 @@ class TestMingFlashOmniPipeline:
|
||||
|
||||
def test_image_processor_wiring_resolves(self):
|
||||
"""The prompt_expand_func and custom_process_input_func strings must point to real callables."""
|
||||
thinker = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0)
|
||||
dit = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1)
|
||||
pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_image")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
|
||||
thinker = pipeline.get_stage(0)
|
||||
dit = pipeline.get_stage(1)
|
||||
for ref in (thinker.prompt_expand_func, dit.custom_process_input_func):
|
||||
module_path, _, attr = ref.rpartition(".")
|
||||
module = importlib.import_module(module_path)
|
||||
@@ -1205,7 +1351,8 @@ class TestMingFlashOmniPipeline:
|
||||
assert deploy.connectors is not None
|
||||
assert "shared_memory_connector" in deploy.connectors
|
||||
|
||||
pipeline = _PIPELINE_REGISTRY["ming_flash_omni_image"]
|
||||
pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_image")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
assert len(stages) == 2
|
||||
# Stage 0 thinker: AR worker that emits latents.
|
||||
@@ -1356,12 +1503,16 @@ class TestCLIOverrideFlow:
|
||||
"""Test --stage-overrides JSON merge into StageConfig."""
|
||||
|
||||
def test_stage_overrides_merge(self):
|
||||
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
|
||||
deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
|
||||
if not deploy_path.exists():
|
||||
pytest.skip("Deploy config not found")
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
pipeline = StageConfigFactory.resolve_pipeline_config(
|
||||
"qwen3_omni_moe",
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
# Simulate --stage-overrides '{"0": {"gpu_memory_utilization": 0.5}}'
|
||||
@@ -1370,12 +1521,16 @@ class TestCLIOverrideFlow:
|
||||
assert stages[0].runtime_overrides["gpu_memory_utilization"] == 0.5
|
||||
|
||||
def test_global_override_applies_to_all(self):
|
||||
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
|
||||
deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
|
||||
if not deploy_path.exists():
|
||||
pytest.skip("Deploy config not found")
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
pipeline = StageConfigFactory.resolve_pipeline_config(
|
||||
"qwen3_omni_moe",
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
overrides = {"enforce_eager": True}
|
||||
@@ -1392,8 +1547,11 @@ class TestAuraOmniDeploy:
|
||||
assert deploy.pipeline == "aura_omni"
|
||||
|
||||
def test_aura_omni_deploy_resolves_four_native_stages(self):
|
||||
pipeline_cfg = StageConfigFactory.resolve_pipeline_config("aura_omni")
|
||||
|
||||
stages = StageConfigFactory._create_from_registry(
|
||||
"qwen3_tts",
|
||||
pipeline_cfg,
|
||||
cli_overrides={},
|
||||
deploy_config_path=str(Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml"),
|
||||
)
|
||||
@@ -1417,8 +1575,14 @@ class TestSentinelDefaultPrecedence:
|
||||
to YAML / dataclass defaults (#3035)."""
|
||||
|
||||
def _stages(self, cli_overrides):
|
||||
model_type = "qwen3_omni_moe"
|
||||
pipeline_cfg = StageConfigFactory.resolve_pipeline_config(
|
||||
model_type,
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
return StageConfigFactory._create_from_registry(
|
||||
"qwen3_omni_moe",
|
||||
pipeline_cfg,
|
||||
cli_overrides=cli_overrides,
|
||||
)
|
||||
|
||||
@@ -1483,7 +1647,8 @@ class TestSentinelDefaultPrecedence:
|
||||
"""A single ``qwen3_tts`` pipeline picks per-chunk vs end-to-end
|
||||
processors based on ``deploy.async_chunk``, without needing a
|
||||
separate variant pipeline registration."""
|
||||
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
|
||||
pipeline = StageConfigFactory.resolve_pipeline_config("qwen3_tts")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
|
||||
# async_chunk=True → stage 0's per-chunk processor wires up, stage 1
|
||||
# has no sync input processor.
|
||||
@@ -1552,9 +1717,8 @@ class TestSentinelDefaultPrecedence:
|
||||
``thinker2talker_token_only`` placeholder (sync_process_input_func).
|
||||
Merge under either async_chunk mode must not re-introduce a
|
||||
stage-0 full-payload hook."""
|
||||
from vllm_omni.config.stage_config import DeployConfig, merge_pipeline_deploy
|
||||
|
||||
pipeline = _PIPELINE_REGISTRY["ming_flash_omni"]
|
||||
pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
|
||||
stage0, stage1 = pipeline.stages
|
||||
assert stage0.custom_process_next_stage_input_func is None, (
|
||||
@@ -1587,15 +1751,30 @@ class TestSamplingConstraintsPrecedence:
|
||||
"""Test that pipeline sampling_constraints override deploy defaults."""
|
||||
|
||||
def test_constraints_win(self):
|
||||
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
|
||||
deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
|
||||
if not deploy_path.exists():
|
||||
pytest.skip("Deploy config not found")
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
pipeline = StageConfigFactory.resolve_pipeline_config(
|
||||
"qwen3_omni_moe",
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
# Pipeline says detokenize=True for thinker, deploy can't override
|
||||
assert stages[0].yaml_extras["default_sampling_params"]["detokenize"] is True
|
||||
# Pipeline says stop_token_ids=[2150] for talker
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [2150]
|
||||
|
||||
|
||||
class TestPipelineConfigResolvers:
|
||||
@pytest.mark.parametrize("resolver", [obj for obj in OMNI_PIPELINES.values() if callable(obj)])
|
||||
def test_all_resolvers_reject_bad_types(self, resolver):
|
||||
"""Ensure that all resolvers registered reject incorrect config types."""
|
||||
|
||||
class NotTheRightHfConfig(PretrainedConfig):
|
||||
pass
|
||||
|
||||
assert resolver(NotTheRightHfConfig()) is None
|
||||
@@ -10,7 +10,7 @@ from collections.abc import Mapping
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.config.stage_config import StageConfigFactory
|
||||
from vllm_omni.config.config_factory import StageConfigFactory
|
||||
from vllm_omni.diffusion.data import (
|
||||
DiffusionParallelConfig,
|
||||
OmniDiffusionConfig,
|
||||
|
||||
@@ -19,7 +19,7 @@ import torch
|
||||
|
||||
class TestHiggsAudioV3Config:
|
||||
def test_default_config_loads(self):
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
|
||||
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
|
||||
@@ -34,7 +34,7 @@ class TestHiggsAudioV3Config:
|
||||
assert config.tie_modality_embeddings is True
|
||||
|
||||
def test_custom_audio_encoder_config(self):
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
|
||||
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
|
||||
@@ -51,7 +51,7 @@ class TestHiggsAudioV3Config:
|
||||
assert config.tie_modality_embeddings is False
|
||||
|
||||
def test_invalid_num_codebooks_rejected(self):
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
|
||||
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
|
||||
@@ -59,7 +59,7 @@ class TestHiggsAudioV3Config:
|
||||
HiggsAudioV3Config(audio_encoder_config={"num_codebooks": 0, "vocab_size": 1026})
|
||||
|
||||
def test_negative_num_codebooks_rejected(self):
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
|
||||
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
|
||||
@@ -67,7 +67,7 @@ class TestHiggsAudioV3Config:
|
||||
HiggsAudioV3Config(audio_encoder_config={"num_codebooks": -1, "vocab_size": 1026})
|
||||
|
||||
def test_special_token_ids_initially_none(self):
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
|
||||
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
|
||||
@@ -77,7 +77,7 @@ class TestHiggsAudioV3Config:
|
||||
assert config.audio_continuation_id is None
|
||||
|
||||
def test_hidden_size_from_text_config(self):
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
|
||||
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
|
||||
@@ -797,12 +797,12 @@ class TestAudioFeedback:
|
||||
class TestCodecStrictness:
|
||||
def test_bundled_missing_quantizer_key_raises(self):
|
||||
"""Bundled codec load must fail when a quantizer codebook key is missing."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
|
||||
HiggsAudioV3Code2Wav,
|
||||
)
|
||||
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
|
||||
config = HiggsAudioV3Config()
|
||||
c2w = HiggsAudioV3Code2Wav(config=config)
|
||||
@@ -821,12 +821,12 @@ class TestCodecStrictness:
|
||||
|
||||
def test_bundled_missing_fc2_raises(self):
|
||||
"""Bundled codec load must fail when fc2 keys are missing."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
|
||||
HiggsAudioV3Code2Wav,
|
||||
)
|
||||
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
|
||||
config = HiggsAudioV3Config()
|
||||
c2w = HiggsAudioV3Code2Wav(config=config)
|
||||
@@ -844,12 +844,12 @@ class TestCodecStrictness:
|
||||
def test_bundled_unknown_decoder_key_rejected(self, monkeypatch):
|
||||
"""Unknown decoder-side key should be rejected as unconsumed."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_code2wav as c2w_mod
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
|
||||
HiggsAudioV3Code2Wav,
|
||||
)
|
||||
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
|
||||
class FakeDAC(torch.nn.Module):
|
||||
def __init__(self):
|
||||
@@ -871,12 +871,12 @@ class TestCodecStrictness:
|
||||
def test_bundled_encoder_side_keys_accepted(self, monkeypatch):
|
||||
"""Known encoder-side keys should not trigger unconsumed-key rejection."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_code2wav as c2w_mod
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
|
||||
HiggsAudioV3Code2Wav,
|
||||
)
|
||||
from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
|
||||
class FakeDAC(torch.nn.Module):
|
||||
def __init__(self):
|
||||
@@ -1199,9 +1199,9 @@ class TestRegistry:
|
||||
assert "HiggsAudioV3Code2WavForConditionalGeneration" in _OMNI_MODELS
|
||||
|
||||
def test_pipeline_registered(self):
|
||||
from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
|
||||
|
||||
assert "higgs_multimodal_qwen3" in _OMNI_PIPELINES
|
||||
assert "higgs_multimodal_qwen3" in OMNI_PIPELINES
|
||||
|
||||
def test_deploy_yaml_exists(self):
|
||||
import os
|
||||
|
||||
@@ -8,11 +8,11 @@ import pytest
|
||||
import yaml
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.config.config_factory import StageConfigFactory
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
|
||||
from vllm_omni.config.stage_config import (
|
||||
_PIPELINE_REGISTRY,
|
||||
DeployConfig,
|
||||
PipelineConfig,
|
||||
StageConfigFactory,
|
||||
StageDeployConfig,
|
||||
StagePipelineConfig,
|
||||
)
|
||||
@@ -44,9 +44,9 @@ _TEST_DEPLOY = DeployConfig(
|
||||
@pytest.fixture()
|
||||
def mock_stages(monkeypatch):
|
||||
"""Register a fake pipeline and mock deploy YAML loading."""
|
||||
monkeypatch.setitem(_PIPELINE_REGISTRY._loaded, _TEST_MODEL, _TEST_PIPELINE)
|
||||
monkeypatch.setitem(OMNI_PIPELINES, _TEST_MODEL, _TEST_PIPELINE)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.config.stage_config.load_deploy_config",
|
||||
"vllm_omni.config.config_factory.load_deploy_config",
|
||||
lambda _path: _TEST_DEPLOY,
|
||||
)
|
||||
return __file__
|
||||
@@ -462,6 +462,7 @@ def test_explicit_cli_arg_reaches_runtime_overrides(mock_stages):
|
||||
explicit_kwargs = ns.get_explicit_kwargs_dict()
|
||||
stages = StageConfigFactory._create_from_registry(
|
||||
_TEST_MODEL,
|
||||
_TEST_PIPELINE,
|
||||
explicit_kwargs,
|
||||
deploy_config_path=mock_stages,
|
||||
)
|
||||
@@ -478,6 +479,7 @@ def test_omitted_default_not_in_runtime_overrides(mock_stages):
|
||||
explicit_kwargs = ns.get_explicit_kwargs_dict()
|
||||
stages = StageConfigFactory._create_from_registry(
|
||||
_TEST_MODEL,
|
||||
_TEST_PIPELINE,
|
||||
explicit_kwargs,
|
||||
deploy_config_path=mock_stages,
|
||||
)
|
||||
@@ -499,6 +501,7 @@ def test_config_file_args_reach_runtime_overrides(mock_stages):
|
||||
explicit_kwargs = ns.get_explicit_kwargs_dict()
|
||||
stages = StageConfigFactory._create_from_registry(
|
||||
_TEST_MODEL,
|
||||
_TEST_PIPELINE,
|
||||
explicit_kwargs,
|
||||
deploy_config_path=mock_stages,
|
||||
)
|
||||
@@ -516,6 +519,7 @@ def test_per_stage_override_routes_correctly(mock_stages):
|
||||
explicit_kwargs = ns.get_explicit_kwargs_dict()
|
||||
stages = StageConfigFactory._create_from_registry(
|
||||
_TEST_MODEL,
|
||||
_TEST_PIPELINE,
|
||||
explicit_kwargs,
|
||||
deploy_config_path=mock_stages,
|
||||
)
|
||||
@@ -537,6 +541,7 @@ def test_explicit_args_omitted_from_yaml(mock_stages):
|
||||
explicit_kwargs = ns.get_explicit_kwargs_dict()
|
||||
stages = StageConfigFactory._create_from_registry(
|
||||
_TEST_MODEL,
|
||||
_TEST_PIPELINE,
|
||||
explicit_kwargs,
|
||||
deploy_config_path=mock_stages,
|
||||
)
|
||||
|
||||
@@ -1,9 +1,15 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, GPUARModelRunner
|
||||
from vllm_omni.outputs import OmniModelRunnerOutput
|
||||
from vllm_omni.worker.gpu_ar_model_runner import (
|
||||
ExecuteModelState,
|
||||
GPUARModelRunner,
|
||||
OmniAsyncGPUModelRunnerOutput,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -52,6 +58,300 @@ def test_sparse_mm_req_ids_requires_sparse_audio_marker():
|
||||
assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"]
|
||||
|
||||
|
||||
def test_omni_async_gpu_model_runner_output_builds_lazily_once():
|
||||
async_output = object.__new__(OmniAsyncGPUModelRunnerOutput)
|
||||
calls = []
|
||||
sync_calls = []
|
||||
|
||||
def builder():
|
||||
calls.append("build")
|
||||
return OmniModelRunnerOutput(req_ids=["r1"], req_id_to_index={"r1": 0})
|
||||
|
||||
async_output._model_runner_output = None
|
||||
async_output._model_runner_output_builder = builder
|
||||
async_output._invalid_req_indices = []
|
||||
async_output.sampled_token_ids_cpu = torch.tensor([[7]], dtype=torch.long)
|
||||
async_output.async_copy_ready_event = SimpleNamespace(synchronize=lambda: sync_calls.append("sync"))
|
||||
async_output._sampled_token_ids = torch.tensor([[7]], dtype=torch.long)
|
||||
async_output._logprobs_tensors = None
|
||||
async_output._logprobs_tensors_cpu = None
|
||||
async_output._routed_experts = None
|
||||
async_output._routed_experts_cpu = None
|
||||
async_output.vocab_size = 10
|
||||
|
||||
output = async_output.get_output()
|
||||
|
||||
assert calls == ["build"]
|
||||
assert sync_calls == ["sync"]
|
||||
assert async_output._model_runner_output_builder is None
|
||||
assert output.req_ids == ["r1"]
|
||||
assert output.sampled_token_ids == [[7]]
|
||||
|
||||
|
||||
def test_omni_async_gpu_model_runner_output_reraises_background_exception():
|
||||
async_output = object.__new__(OmniAsyncGPUModelRunnerOutput)
|
||||
joined = []
|
||||
|
||||
class FakeThread:
|
||||
def join(self):
|
||||
joined.append("join")
|
||||
|
||||
async_output._background_thread = FakeThread()
|
||||
async_output._background_exception = RuntimeError("background failed")
|
||||
|
||||
with pytest.raises(RuntimeError, match="background failed"):
|
||||
async_output.get_output()
|
||||
|
||||
assert joined == ["join"]
|
||||
assert async_output._background_thread is None
|
||||
|
||||
|
||||
def _make_async_output_runner(engine_output_type: str = "audio"):
|
||||
runner = object.__new__(GPUARModelRunner)
|
||||
model_config = SimpleNamespace(
|
||||
engine_output_type=engine_output_type,
|
||||
async_chunk=True,
|
||||
enable_return_routed_experts=False,
|
||||
)
|
||||
runner.vllm_config = SimpleNamespace(model_config=model_config)
|
||||
runner.model_config = model_config
|
||||
runner.omni_prefix_cache = None
|
||||
runner.requests = {"r1": object(), "r2": object()}
|
||||
runner.supports_mm_inputs = False
|
||||
runner.routed_experts_initialized = False
|
||||
runner.model = SimpleNamespace(has_postprocess=False)
|
||||
runner.model_intermediate_buffer = {}
|
||||
runner.input_batch = SimpleNamespace(
|
||||
req_ids=["mutated"],
|
||||
req_id_to_index={"mutated": 0},
|
||||
)
|
||||
return runner
|
||||
|
||||
|
||||
def test_build_omni_output_uses_snapshots_and_connector_after_accumulation(monkeypatch):
|
||||
runner = _make_async_output_runner()
|
||||
events = []
|
||||
|
||||
monkeypatch.setattr(
|
||||
GPUARModelRunner,
|
||||
"_resolve_pooler_payload_req_ids",
|
||||
lambda self, req_ids: ("audio", req_ids),
|
||||
)
|
||||
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: True)
|
||||
monkeypatch.setattr(
|
||||
GPUARModelRunner,
|
||||
"accumulate_full_payload_output",
|
||||
lambda self, rid, payload, request: events.append(f"accumulate:{rid}"),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
GPUARModelRunner,
|
||||
"get_omni_connector_output",
|
||||
lambda self: events.append("connector") or "connector-output",
|
||||
)
|
||||
|
||||
output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
|
||||
runner,
|
||||
scheduler_output=SimpleNamespace(
|
||||
total_num_scheduled_tokens=3,
|
||||
num_scheduled_tokens={"r1": 1, "r2": 2},
|
||||
),
|
||||
hidden_states=torch.tensor([[1.0], [2.0], [3.0]]),
|
||||
staged_hidden_states_cpu=None,
|
||||
multimodal_outputs={"foo": torch.tensor([10.0, 20.0, 30.0])},
|
||||
req_ids_output_copy=["r1", "r2"],
|
||||
req_id_to_index_output_copy={"r1": 0, "r2": 1},
|
||||
valid_sampled_token_ids=[[101], [102]],
|
||||
logprobs_lists=None,
|
||||
prompt_logprobs_dict={},
|
||||
num_nans_in_logits=None,
|
||||
kv_connector_output=None,
|
||||
ec_connector_output=None,
|
||||
cudagraph_stats=None,
|
||||
kv_extracted_req_ids=["r2"],
|
||||
seq_len=3,
|
||||
num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(),
|
||||
query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
|
||||
)
|
||||
|
||||
assert output.req_ids == ["r1", "r2"]
|
||||
assert torch.equal(output.multimodal_outputs[0]["hidden"], torch.tensor([[1.0]]))
|
||||
assert torch.equal(output.multimodal_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]]))
|
||||
assert output.kv_extracted_req_ids == ["r2"]
|
||||
assert output.omni_connector_output == "connector-output"
|
||||
assert events == ["accumulate:r1", "accumulate:r2", "connector"]
|
||||
|
||||
|
||||
def test_build_omni_output_copies_hidden_for_partial_downstream_batch(monkeypatch):
|
||||
runner = _make_async_output_runner(engine_output_type="latent")
|
||||
|
||||
monkeypatch.setattr(
|
||||
GPUARModelRunner,
|
||||
"_resolve_pooler_payload_req_ids",
|
||||
lambda self, req_ids: ("latent", ["r2"]),
|
||||
)
|
||||
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
|
||||
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
|
||||
monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
|
||||
|
||||
output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
|
||||
runner,
|
||||
scheduler_output=SimpleNamespace(
|
||||
total_num_scheduled_tokens=6,
|
||||
num_scheduled_tokens={"r1": 1, "r2": 2, "r3": 3},
|
||||
),
|
||||
hidden_states=torch.tensor([[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]]),
|
||||
staged_hidden_states_cpu=None,
|
||||
multimodal_outputs={},
|
||||
req_ids_output_copy=["r1", "r2", "r3"],
|
||||
req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2},
|
||||
valid_sampled_token_ids=[[], [], []],
|
||||
logprobs_lists=None,
|
||||
prompt_logprobs_dict={},
|
||||
num_nans_in_logits=None,
|
||||
kv_connector_output=None,
|
||||
ec_connector_output=None,
|
||||
cudagraph_stats=None,
|
||||
kv_extracted_req_ids=None,
|
||||
seq_len=6,
|
||||
num_scheduled_tokens_np=np.array([1, 2, 3], dtype=np.int32),
|
||||
query_start_loc_cpu=torch.tensor([0, 1, 3], dtype=torch.long),
|
||||
)
|
||||
|
||||
assert output.multimodal_outputs is not None
|
||||
assert output.multimodal_outputs[0] == {}
|
||||
assert torch.equal(output.multimodal_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]]))
|
||||
assert output.multimodal_outputs[2] == {}
|
||||
|
||||
|
||||
def test_process_additional_information_uses_snapshot_request_order(monkeypatch):
|
||||
runner = _make_async_output_runner()
|
||||
seen = []
|
||||
|
||||
class PostprocessModel:
|
||||
has_postprocess = True
|
||||
|
||||
def postprocess(self, hidden_states, **kwargs):
|
||||
seen.append(hidden_states.clone())
|
||||
return {}
|
||||
|
||||
runner.model = PostprocessModel()
|
||||
runner.model_intermediate_buffer = {"r1": {}, "r2": {}}
|
||||
|
||||
monkeypatch.setattr(
|
||||
GPUARModelRunner,
|
||||
"_resolve_pooler_payload_req_ids",
|
||||
lambda self, req_ids: ("audio", req_ids),
|
||||
)
|
||||
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
|
||||
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
|
||||
monkeypatch.setattr(GPUARModelRunner, "_update_intermediate_buffer", lambda *args, **kwargs: None)
|
||||
|
||||
GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
|
||||
runner,
|
||||
scheduler_output=SimpleNamespace(
|
||||
total_num_scheduled_tokens=3,
|
||||
num_scheduled_tokens={"r1": 1, "r2": 2},
|
||||
),
|
||||
hidden_states=torch.tensor([[1.0], [2.0], [3.0]]),
|
||||
staged_hidden_states_cpu=None,
|
||||
multimodal_outputs={},
|
||||
req_ids_output_copy=["r1", "r2"],
|
||||
req_id_to_index_output_copy={"r1": 0, "r2": 1},
|
||||
valid_sampled_token_ids=[[], []],
|
||||
logprobs_lists=None,
|
||||
prompt_logprobs_dict={},
|
||||
num_nans_in_logits=None,
|
||||
kv_connector_output=None,
|
||||
ec_connector_output=None,
|
||||
cudagraph_stats=None,
|
||||
kv_extracted_req_ids=None,
|
||||
seq_len=3,
|
||||
num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(),
|
||||
query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
|
||||
)
|
||||
|
||||
assert len(seen) == 2
|
||||
assert torch.equal(seen[0], torch.tensor([[1.0]]))
|
||||
assert torch.equal(seen[1], torch.tensor([[2.0], [3.0]]))
|
||||
|
||||
|
||||
def test_async_omni_output_guard_requires_safe_conditions():
|
||||
runner = _make_async_output_runner()
|
||||
runner.use_async_scheduling = True
|
||||
runner.speculative_config = None
|
||||
runner.model.use_async_omni_output = True
|
||||
|
||||
assert GPUARModelRunner._should_use_async_omni_output(runner)
|
||||
|
||||
runner.omni_prefix_cache = object()
|
||||
assert not GPUARModelRunner._should_use_async_omni_output(runner)
|
||||
|
||||
runner.omni_prefix_cache = None
|
||||
runner.model.has_postprocess = True
|
||||
assert not GPUARModelRunner._should_use_async_omni_output(runner)
|
||||
|
||||
runner.model.eager_omni_postprocess_before_async_output = True
|
||||
assert GPUARModelRunner._should_use_async_omni_output(runner)
|
||||
|
||||
|
||||
def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
|
||||
runner = _make_async_output_runner(engine_output_type="latent")
|
||||
runner.model.omni_pooler_payload_include_hidden = False
|
||||
|
||||
monkeypatch.setattr(
|
||||
GPUARModelRunner,
|
||||
"_resolve_pooler_payload_req_ids",
|
||||
lambda self, req_ids: ("latent", req_ids),
|
||||
)
|
||||
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
|
||||
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
|
||||
monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
|
||||
|
||||
output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
|
||||
runner,
|
||||
scheduler_output=SimpleNamespace(
|
||||
total_num_scheduled_tokens=2,
|
||||
num_scheduled_tokens={"r1": 2},
|
||||
),
|
||||
hidden_states=torch.tensor([[1.0], [2.0]]),
|
||||
staged_hidden_states_cpu=None,
|
||||
multimodal_outputs={"codes": {"audio": torch.tensor([[7, 8], [9, 10]], dtype=torch.long)}},
|
||||
req_ids_output_copy=["r1"],
|
||||
req_id_to_index_output_copy={"r1": 0},
|
||||
valid_sampled_token_ids=[[101]],
|
||||
logprobs_lists=None,
|
||||
prompt_logprobs_dict={},
|
||||
num_nans_in_logits=None,
|
||||
kv_connector_output=None,
|
||||
ec_connector_output=None,
|
||||
cudagraph_stats=None,
|
||||
kv_extracted_req_ids=None,
|
||||
seq_len=2,
|
||||
num_scheduled_tokens_np=np.array([2], dtype=np.int32),
|
||||
query_start_loc_cpu=torch.tensor([0], dtype=torch.long),
|
||||
)
|
||||
|
||||
assert output.multimodal_outputs is not None
|
||||
assert len(output.multimodal_outputs) == 1
|
||||
assert "hidden" not in output.multimodal_outputs[0]
|
||||
assert torch.equal(output.multimodal_outputs[0]["codes.audio"], torch.tensor([[7, 8], [9, 10]], dtype=torch.long))
|
||||
|
||||
|
||||
def test_async_snapshot_payload_omits_hidden_when_model_opts_out():
|
||||
runner = _make_async_output_runner()
|
||||
runner.model.omni_pooler_payload_include_hidden = False
|
||||
|
||||
payload = GPUARModelRunner._build_omni_async_snapshot_payload(
|
||||
runner,
|
||||
hidden_states=torch.tensor([[1.0], [2.0]]),
|
||||
staged_hidden_states_cpu=torch.tensor([[3.0]]),
|
||||
multimodal_outputs={"codes": {"audio": torch.tensor([[1]], dtype=torch.long)}},
|
||||
)
|
||||
|
||||
assert set(payload.keys()) == {"multimodal_outputs"}
|
||||
assert payload["multimodal_outputs"]["codes"]["audio"].tolist() == [[1]]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("query_start_loc_attr", ["method", "tensor_attr"])
|
||||
def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monkeypatch, query_start_loc_attr):
|
||||
runner = object.__new__(GPUARModelRunner)
|
||||
@@ -133,3 +433,54 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
|
||||
output.multimodal_outputs[1]["hidden"],
|
||||
torch.tensor([[2.0, 20.0], [3.0, 30.0]]),
|
||||
)
|
||||
|
||||
|
||||
def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch):
|
||||
"""Tail-only prefix-cache models still need per-step mm passthrough (e.g. codes.audio)."""
|
||||
runner = _make_async_output_runner(engine_output_type="latent")
|
||||
runner.omni_prefix_cache = object()
|
||||
|
||||
monkeypatch.setattr(
|
||||
GPUARModelRunner,
|
||||
"_resolve_pooler_payload_req_ids",
|
||||
lambda self, req_ids: ("latent", req_ids),
|
||||
)
|
||||
monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False)
|
||||
monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: {"codes.audio"})
|
||||
monkeypatch.setattr(GPUARModelRunner, "_stage_deferred_prefix_cache_mm_outputs", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(
|
||||
GPUARModelRunner,
|
||||
"_prepare_prefix_cache_pooler_payload_sources",
|
||||
lambda *args, **kwargs: (None, None, None),
|
||||
)
|
||||
monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
|
||||
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
|
||||
|
||||
codes = torch.tensor([[11.0, 12.0], [21.0, 22.0]], dtype=torch.float32)
|
||||
output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
|
||||
runner,
|
||||
scheduler_output=SimpleNamespace(
|
||||
total_num_scheduled_tokens=2,
|
||||
num_scheduled_tokens={"r1": 1, "r2": 1},
|
||||
),
|
||||
hidden_states=torch.tensor([[1.0], [2.0]]),
|
||||
staged_hidden_states_cpu=None,
|
||||
multimodal_outputs={"codes.audio": codes},
|
||||
req_ids_output_copy=["r1", "r2"],
|
||||
req_id_to_index_output_copy={"r1": 0, "r2": 1},
|
||||
valid_sampled_token_ids=[[], []],
|
||||
logprobs_lists=None,
|
||||
prompt_logprobs_dict={},
|
||||
num_nans_in_logits=None,
|
||||
kv_connector_output=None,
|
||||
ec_connector_output=None,
|
||||
cudagraph_stats=None,
|
||||
kv_extracted_req_ids=None,
|
||||
seq_len=2,
|
||||
num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
|
||||
query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
|
||||
)
|
||||
|
||||
assert torch.equal(output.multimodal_outputs[0]["codes.audio"], codes[0:1])
|
||||
assert torch.equal(output.multimodal_outputs[1]["codes.audio"], codes[1:2])
|
||||
Loaded 100 of 279 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user