Compare commits

...
Author SHA1 Message Date
yujuancao07 d44831591a Clean up PR template
Signed-off-by: yujuancao07 <yujuancao07@gmail.com>
2026-06-11 15:28:00 +08:00
Juan Pablo ZuluagaandJuanPZuluaga 0342827dd7 [Bugfix] qwen3-tts prefix cache: drop per-key size cap that corrupted… (#4317)
Signed-off-by: JuanPZuluaga <juanz9312@gmal.com>
Co-authored-by: JuanPZuluaga <juanz9312@gmal.com>
2026-06-10 21:20:08 +08:00
Zeyu Huang | 黃澤宇 362b3ab013 [skip ci] Add width and height args to offline i2i example script (#4031)
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
2026-06-10 21:18:11 +08:00
Guanzhe (Leo) Huang c2b4564a67 [BugFix] moss_tts_nano: eager-init lm + audio_tokenizer in __init__ so load_format: dummy works (#3230)
Signed-off-by: leohuang257 <masugchds@gmail.com>
2026-06-10 20:50:53 +08:00
Clodagh Walsh aa92a3a35d [CI/Build] Voxtral TTS Tests (#3738)
Signed-off-by: Clodagh Walsh <clodaghwalsh17@gmail.com>
2026-06-10 20:32:21 +08:00
akshatvishuandHongsheng Liu 2fa4c7a4a9 [Model] Add Ming-omni-tts dense 0.5B pipeline (#2906)
Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-06-10 20:11:04 +08:00
39ae6e12f8 [Docs] Update README supported models section with TTS and Diffusion categories (#4300)
Signed-off-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>
Co-authored-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
2026-06-10 19:40:50 +08:00
zhumingjue138 2d9e5e0515 [Test] skip oom test case for issue #4285 (#4311)
Signed-off-by: zhumingjue <zhumingjue@huawei.com>
2026-06-10 10:29:33 +00:00
Optance d50bf1aa42 [Bugfix] Fix Fish Speech prefix cache collision from missing cache_salt (#4008)
Signed-off-by: nagelanping <xiao.yi.ming@foxmail.com>
2026-06-10 17:11:01 +08:00
Yuekai ZhangandClaude Opus 4.8 f1bf544208 [Perf][TTS] Optimize cosyvoice TTFP and throughput using TensorRT (#4168)
Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-10 16:42:34 +08:00
wangyuandAlicia 57c8f3abfc [CI] skip unrelated L3 merge tests via diff-aware upload (#4291)
Signed-off-by: wangyu <410167048@qq.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-06-10 14:27:27 +08:00
FED4andGao Han 6ae355b9a9 [BugFix] Remove pydub dependency for Python 3.13 compatibility (#4035)
Signed-off-by: FED4 <feiyyfrank@outlook.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-06-09 20:45:57 -07:00
longguoandnpuichigo 73df8326d7 [Bugfix] purge chunk-transfer zombies on every schedule tick to keep engine-core alive on aborts (fixes #3736) (#3774)
Signed-off-by: abinggo <107740309+abinggo@users.noreply.github.com>
Co-authored-by: npuichigo <11533479+npuichigo@users.noreply.github.com>
2026-06-10 09:00:23 +08:00
Yueqian Lin 95f633f830 [skip ci] fix(config): pin VoxCPM2 KV cache to avoid OOM on small GPUs (#4279)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-06-10 08:48:35 +08:00
e1b0b83238 add moriio transfer engine (#1742)
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Signed-off-by: zejwang  <zejwang@amd.com>
Signed-off-by: junkang1991 <junkangchow@gmail.com>
Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn>
Co-authored-by: zejwang <zejwang@amd.com>
Co-authored-by: junkang1991 <junkangchow@gmail.com>
Co-authored-by: Zejian Wang <zejianwang@sjtu.edu.cn>
Co-authored-by: Zejian Wang <Zejian.Wang@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-10 08:37:33 +08:00
57227dc7f3 [Perf][Bugfix] qwen3-tts hot path: prefix-cache OOM guards + talker/orchestrator micro-opts (#3689)
Signed-off-by: JuanPZuluaga <juanz9312@gmal.com>
Co-authored-by: JuanPZuluaga <juanz9312@gmal.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-06-09 10:36:57 -07:00
Zeng Chuang ca75623ff3 [BugFix] Fix HunyuanImage3 CoT truncation when stop token stripped by detokenizer (#4260)
Signed-off-by: zengchuang <zengchuang3@huawei.com>
2026-06-09 15:30:25 +00:00
3753b21c3d [Feature] LoRA support for SenseNova-U1 (#3971)
Signed-off-by: leohuang257 <masugchds@gmail.com>
Signed-off-by: Guanzhe (Leo) Huang <116847811+leohuang257@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com>
2026-06-09 14:36:35 +00:00
TJian e0fde88a5f [ROCm] [CI] Add group feature and envs feature (#4208)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-06-09 14:17:55 +00:00
Mu GuanLin bf03413c36 [Feat] Add vae-decode-parallel for LTX-2.3 (#4277)
Signed-off-by: mglyn <1203789601@qq.com>
2026-06-09 13:14:16 +00:00
Chendi.Xue f174ee54b0 [XPU] Add sage_attn backend (#3785)
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
2026-06-09 19:28:18 +08:00
Samit 6ddaf6188d Fix README.md typo (#4292) 2026-06-09 18:08:09 +08:00
Mu GuanLin 02eb7a4ecb [Bugfix] Register LTX RMSNorm identity weight as buffer (#4278)
Signed-off-by: mglyn <1203789601@qq.com>
2026-06-09 09:09:39 +00:00
yangjianjuan cb1c5a0c39 [Perf] lance perf optimize (t2i & i2i) (#4214)
Signed-off-by: yangjianjuan <510818155@qq.com>
2026-06-09 08:31:57 +00:00
weizhoublue 1b85df0510 fix: Propagate Quantization Configuration to HunyuanVideo-1.5 I2V Transformer to Enable FP8 Layers (#4245)
Signed-off-by: weizhoublue <weizhou.lan@daocloud.io>
2026-06-09 15:28:32 +08:00
5de3b11822 docs: update README and supported models for v0.22.0 release (#4233)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
2026-06-09 14:39:19 +08:00
Alex Brooks 7dcbeaea03 [Core] Cleanup DiffusersPipelineLoader (#1932)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-06-09 11:55:32 +08:00
b8f68174ad [Refactor] Remove dead legacy pipeline.yaml loader and duplicate diff… (#4023)
Signed-off-by: dph97 <nenbaying@163.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: zhuhonghan <z00806815@china.huawei.com>
2026-06-09 07:34:35 +08:00
SYLAR b5352a5869 [Quant][Perf] Default Blackwell FP8 GEMM to quack CuteDSL fused-bias kernel (#4241)
Signed-off-by: WeiQing Chen <david6666666@users.noreply.github.com>
2026-06-08 19:44:20 +00:00
MaciejBalaNV 2f25195750 Cosmos3 video to video generation (#4266)
Signed-off-by: Maciej Bala <mbala@nvidia.com>
2026-06-09 03:04:05 +08:00
Sy03 9c9d4ed4a8 [Fix][HiggsAudioV3] Fix ramp-down off-by-one crash and buffer state bugs in Stage0 talker (#4219)
Signed-off-by: Sy03 <1370724210@qq.com>
2026-06-08 17:33:20 +00:00
Mu GuanLin 13f9ecaef8 [Perf] Keep LTX2.3 auxiliary modules resident by default (#4144)
Signed-off-by: mglyn <1203789601@qq.com>
2026-06-08 19:15:06 +08:00
Hongsheng LiuandClaude Opus 4.7 72cc496fb6 [Doc] Add precheck-pr Claude Code skill (#4216)
Signed-off-by: hsliu <liuhongsheng4@huawei.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
2026-06-08 16:41:33 +08:00
amy-why-3459 fc7393e485 [Refactor]Refactoring audio_in_video implementation (#3566)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-06-08 14:14:02 +08:00
WeiQing Chen 5f1a876cbd [skip ci] docs: update WeChat QR code (#4242)
Signed-off-by: David Chen <530634352@qq.com>
2026-06-08 12:36:09 +08:00
eb1b782dd1 [BugFix] fix hunyuan image3 offline cot (#4174)
Signed-off-by: zzh <943967662@qq.com>
Signed-off-by: zengchuang <zengchuang3@huawei.com>
Co-authored-by: zengchuang <zengchuang3@huawei.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>
2026-06-08 10:58:44 +08:00
SYLAR 29b18ea298 [Quant] ModelOpt FP8 for Wan2.2 & HunyuanVideo-1.5 video-gen (#3305)
Signed-off-by: WeiQing Chen <david6666666@users.noreply.github.com>
2026-06-08 02:53:06 +08:00
Lokesh 176b00538c [Model] Support languages added by fine-tuned Qwen3-TTS checkpoints (#4210)
Signed-off-by: lokesh <lokesh.v9911@gmail.com>
2026-06-07 10:10:40 -06:00
Optance 4d71a162e6 [Bugfix] Fish Speech Gradio hardcoded default voice causes 400 error (#3941)
Signed-off-by: nagelanping <xiao.yi.ming@foxmail.com>
2026-06-07 15:01:25 +00:00
yangyonggit 923c1d9402 [Doc] Add Stable-Diffusion-3.5 recipe for 1x RTX A6000 48GB (#2645) (#4052)
Signed-off-by: leo.yang <leo.yang.engineer@gmail.com>
2026-06-07 21:49:41 +08:00
Yueqian Lin 82cce633a3 [Perf][Higgs-Audio-V3] Turn on Stage-0 prefix caching by default (#4199)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-06-07 03:09:47 +00:00
Yueqian Lin fce0f51089 [Perf][Higgs-Audio-V3] LRU cache for voice-clone ref-audio encode (#4200)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-06-07 03:09:16 +00:00
Yuekai ZhangandYueqian Lin 82e74ce972 [TTS][New Model] support bosonai/higgs-audio-v3-tts-4b (#4169)
Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
2026-06-06 15:59:12 +00:00
Gao Han 963ba1ab7a [CI] Revert "[Feature] gen-only FP8 Quantization Support for SenseNova-U1" (#4196) 2026-06-06 08:39:10 +08:00
ChenWenjing e25239c2ec [bugfix] fix realtime ci timeout error (#4187)
Signed-off-by: CHEN <116010019@link.cuhk.edu.cn>
2026-06-06 00:05:07 +08:00
Zeyu Huang | 黃澤宇andGao Han 2414b6f94d [CI][bugfix]: Improve Qwen Image accuracy test with diffusers' RMSNorm (#4143)
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: Zeyu Huang | 黃澤宇 <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-06-05 07:07:26 -07:00
Chenguang Zheng b3dd0a4baf [Fix] Update Qwen3 Omni multi-replica perf baselines (#4175)
Signed-off-by: Chenguang ZHENG <645327136@qq.com>
2026-06-05 09:22:53 +00:00
Jianbao Huang 6cefe91370 [Perf] Enable fused RMSNorm for HunyuanImage3 (#3959) 2026-06-05 07:41:56 +00:00
Yaegaki1Erikaandreandeng 84569280b1 [Bugfix]Fix HunyuanImage3 conditional image and prompt kwargs alignment (#4145)
Signed-off-by: reandeng <reandeng@tencent.com>
Signed-off-by: Yaegaki1Erika <70182590+Yaegaki1Erika@users.noreply.github.com>
Co-authored-by: reandeng <reandeng@tencent.com>
2026-06-05 12:47:25 +08:00
amy-why-3459 1e43656ab1 [Bugfix] Update the value of --max-seed-tts-mean-wer in the accuracy test. (#4160)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-06-05 11:39:07 +08:00
Shylin 3875b81c18 [skip ci]cleanup(assets): remove dead vllm_omni/assets/video.py (#4120)
Signed-off-by: Shylin26 <parishachauhan26@gmail.com>
2026-06-05 11:21:38 +08:00
Guanzhe (Leo) Huang b63b36ded3 [Feature] gen-only FP8 Quantization Support for SenseNova-U1 (#3943)
Signed-off-by: leohuang257 <masugchds@gmail.com>
2026-06-05 02:51:11 +00:00
Yuanheng Zhaoandakshatvishu 95d56cf062 [Bugfix] Fix Compatibility of Ming-flash-omni-2.0 on transformers 5.X and vllm 0.22 (#4080)
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
Co-authored-by: akshatvishu <akshatnayak197@gmail.com>
2026-06-05 09:27:42 +08:00
WeiQing Chenandlishunyang12 6e8ef213d8 [BugFix] Support ModelOpt FP8 under batched diffusion serving (#4155)
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
2026-06-05 06:25:01 +08:00
Chendi.XueandMengkejiergeli Ba 31f882935d [XPU][S2V] Optimize Wan2.2 S2V: RoPE refactor + cache_dit enabling (#4062)
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
Co-authored-by: Mengkejiergeli Ba <mengkejiergeli.ba@intel.com>
2026-06-05 06:20:16 +08:00
Zheng Wengang fccd95b12b [BugFix] Fix incorrect GPU device mapping in multi-replica stages (#4132)
Signed-off-by: ZhengWG <zwg0606@gmail.com>
2026-06-04 15:32:58 +00:00
dengyunyang 469c1dc87e [HunyuanImage3.0][Performance][Optimazation]Adjust perf config (#4149)
Signed-off-by: dengyunyang <584797741@qq.com>
2026-06-04 22:54:46 +08:00
amy-why-3459 519cf3ecf1 [Refactor] [Qwen3-Omni]Modify the thinker's sampling parameters to align with transformers. (#4137)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-06-04 22:49:04 +08:00
ZacheryAU bfc7cedcf1 Benchmark data statistics for each stage of omni models (#3628)
Signed-off-by: ZacheryAU <70869627+ZacheryAU@users.noreply.github.com>
2026-06-04 21:43:35 +08:00
Alex Brooks e10aca30ec [CI] Update Bagel Pixels (#4081)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-06-04 10:37:34 +00:00
Alicia c6e8d78eda Add dependency FlagEmbedding (#3980)
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-06-04 18:00:12 +08:00
BruceLoveDecimal e5030c7524 [BugFix] Cast bf16 video frames to float32 before .numpy() in /v1/videos (#4114)
Signed-off-by: BruceLoveDecimal <39156883+BruceLoveDecimal@users.noreply.github.com>
2026-06-04 17:49:07 +08:00
Yuekai Zhang a1cc583665 fix: cosyvoice3 batch>1 inference (#3910)
Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
2026-06-04 08:58:53 +00:00
dengyunyang a57b8f858b [HunyuanImage3][CI] fix ci (#4134)
Signed-off-by: dengyunyang <584797741@qq.com>
2026-06-04 16:15:49 +08:00
herotai214 9a9c8f7503 [Test] Add L4 diffusion feature test for GLM-Image (#3451)
Signed-off-by: herotai214 <herotai214@gmail.com>
2026-06-04 00:11:17 -07:00
Zhou Taichang af02713952 [Bugfix] harden diffusion model prefetch against transformers v5 shard-resolution race (#4076) 2026-06-04 12:14:25 +08:00
tanhaoan333 d8dc1defa7 [NPU][BugFix] Upgrade parts of ModelRunner to v0.22.0 (#4130)
Signed-off-by: tanhaoan333 <tanhaoan@huawei.com>
2026-06-04 03:59:26 +00:00
tanhaoan333 800c2b8251 [bugfix]qwen3tts code2wav (#4123)
Signed-off-by: tanhaoan333 <tanhaoan@huawei.com>
2026-06-04 02:20:00 +00:00
Canlin Guo 221584a489 [CI] Skip online moss test temp (#4122) 2026-06-04 01:56:11 +00:00
amy-why-3459 4e1c13236e [Perf][Qwen3-Omni]Optimize TTFP using initial_codec_chunk_frames (#4054)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-06-04 09:08:13 +08:00
706bad253d Add Cosmos3 action modality (#4102)
Signed-off-by: Maciej Bala <mbala@nvidia.com>
Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: Maciej Bala <mbala@nvidia.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
2026-06-03 19:30:39 +00:00
amy-why-3459 8949fdc440 [BugFix] Fix the issue of vllm failing to start. (#4105)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-06-03 15:44:46 +00:00
Canlin Guo 3818ba4c65 [CI] Remove omni mark for MOSS-TTS and temporarily skipped (#4097)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-06-03 23:12:11 +08:00
c33ca60208 [Perf][VoxCPM2] Optimize VoxCPM2 high-concurrency decode throughput. (#3882)
Signed-off-by: Sy03 <1370724210@qq.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-06-03 23:08:32 +08:00
tanhaoan333 0fc4032a71 Update qwen3_tts_code2wav.py (#4075)
Signed-off-by: tanhaoan333 <tanhaoan@huawei.com>
2026-06-03 13:34:47 +00:00
ChenWenjing 2e108a4109 [ci] add merge/ready ci for audio realtime api (#4069)
Signed-off-by: CHEN <116010019@link.cuhk.edu.cn>
2026-06-03 18:57:16 +08:00
Canlin Guo ce69eb35e9 [Refactor] Refactor HunyuanImage3 SigLIP2 ViT to vLLM layers (#3297)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-06-03 17:51:06 +08:00
07737e6c61 [Test] Add prefix caching + audio output regression test (#3510) (#3604)
Signed-off-by: Ricardo Noriega De Soto <rnoriega@redhat.com>
Signed-off-by: Alex Brooks <albrooks@redhat.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Alex Brooks <albrooks@redhat.com>
2026-06-03 09:17:53 +00:00
d2e65a92f1 [CI][Accuracy] Add HunyuanImage3 pixel accuracy test and nightly CI (#3790)
Signed-off-by: zzh <943967662@qq.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-06-03 16:59:21 +08:00
ChenWenjing 3842970a82 [bugfix] fix streaming input issue after rebase 0.22.0 (#4085)
Signed-off-by: CHEN <116010019@link.cuhk.edu.cn>
2026-06-03 16:07:41 +08:00
299 changed files with 29016 additions and 4539 deletions

No files matched your search

+4 -4
View File
@@ -1,14 +1,14 @@
# Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci
# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs for the nightly
# exception in upload_pipeline_with_skip_ci.sh (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
# exception in upload_pipeline.py (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
# can be skipped together with test-ready/test-merge. Weekly uploads also depend on image-build.
# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
# Document 2: appended after `---`; same file, read by upload_pipeline.py (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
steps:
- label: ":github: Resolve skip-ci & upload pipeline"
key: upload-ci-pipeline
commands:
- "bash .buildkite/scripts/upload_pipeline_with_skip_ci.sh"
- "python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/pipeline.yml"
agents:
queue: "cpu_queue_premerge"
@@ -41,7 +41,7 @@ steps:
key: upload-merge-pipeline
if: __UPLOAD_MERGE_IF__
commands:
- buildkite-agent pipeline upload .buildkite/test-merge.yml
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
agents:
queue: "cpu_queue_premerge"
+1 -1
View File
@@ -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 .buildkite/test-merge.yml
agents:
queue: "cpu_queue_premerge"
+308
View File
@@ -0,0 +1,308 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Render and optionally upload Buildkite pipeline YAML with diff-aware logic.
Bootstrap mode (pipeline.yml with __IMAGE_BUILD_IF__ placeholders):
- Detect docs-only PR/main changes and substitute skip-ci ``if`` expressions.
Test pipeline mode (e.g. test-merge.yml):
- Drop steps or groups whose ``source_file_dependencies`` do not match changed files.
- ``source_file_dependencies`` may be set on a leaf step or on a group step.
- Always strip ``source_file_dependencies`` from uploaded YAML (Buildkite does not
recognize this uploader-only key).
Usage:
python3 upload_pipeline.py [--upload] <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
"""
from __future__ import annotations
import argparse
import os
import subprocess
import sys
from pathlib import Path
from typing import Any
try:
import yaml
except ImportError:
subprocess.check_call(
[sys.executable, "-m", "pip", "install", "-q", "pyyaml"],
)
import yaml
LOG = "upload_pipeline"
ROOT = Path(__file__).resolve().parent.parent.parent
DOC_SEP = "\n---\n"
BOOTSTRAP_MARKER = "__IMAGE_BUILD_IF__"
def _log(message: str) -> None:
print(f"{LOG}: {message}", file=sys.stderr)
def _git(*args: str) -> subprocess.CompletedProcess[str]:
return subprocess.run(
["git", *args],
check=False,
capture_output=True,
text=True,
)
def resolve_changed_files() -> list[str] | None:
"""Return changed file paths, or None when diff cannot be resolved."""
is_pr = os.environ.get("BUILDKITE_PULL_REQUEST", "false") != "false" and os.environ.get(
"BUILDKITE_PULL_REQUEST", ""
)
commit = os.environ.get("BUILDKITE_COMMIT", "")
if is_pr:
base_branch = os.environ.get("BUILDKITE_PULL_REQUEST_BASE_BRANCH", "main")
base_ref = f"origin/{base_branch}"
if _git("rev-parse", "--verify", base_ref).returncode != 0:
_log(f"origin/{base_branch} not found locally; trying fetch")
_git("fetch", "--depth=200", "origin", base_branch)
if _git("rev-parse", "--verify", base_ref).returncode != 0:
if _git("rev-parse", "--verify", base_branch).returncode == 0:
base_ref = base_branch
else:
_log(f"cannot resolve PR base {base_branch}; using safe defaults")
return None
diff_range = f"{base_ref}...{commit}"
elif os.environ.get("BUILDKITE_BRANCH", "") == "main":
if _git("rev-parse", "--verify", f"{commit}^").returncode != 0:
_log("main commit has no parent; using safe defaults")
return None
diff_range = f"{commit}^..{commit}"
else:
_log("not PR/main build; using safe defaults")
return None
result = _git("diff", "--name-only", diff_range)
if result.returncode != 0:
_log(f"git diff failed ({diff_range}); using safe defaults")
return None
files = [line for line in result.stdout.splitlines() if line.strip()]
_log(f"{len(files)} changed file(s)")
return files
def is_docs_only_change(changed_files: list[str]) -> bool:
has_any = False
for file_path in changed_files:
if not file_path:
continue
has_any = True
if file_path.startswith("docs/"):
continue
if file_path.endswith(".md"):
continue
if file_path == "mkdocs.yml":
continue
return False
return has_any
def resolve_skip_ci(changed_files: list[str] | None) -> bool:
if changed_files is None:
_log("skip-ci=0 (could not resolve changed files)")
return False
if is_docs_only_change(changed_files):
_log("docs-only change detected; skip-ci=1")
return True
_log("non-doc changes detected; skip-ci=0")
return False
def render_bootstrap_pipeline(text: str, *, skip_ci: bool) -> str:
if DOC_SEP in text:
_, continuation = text.split(DOC_SEP, 1)
else:
continuation = text
nightly_only = (
'(build.pull_request.labels includes "nightly-test") || (build.branch == "main" && build.env("NIGHTLY") == "1")'
)
if skip_ci:
image_if = f"'{nightly_only}'"
ready_if = "'false'"
merge_if = "'false'"
else:
image_if = "'true'"
ready_if = '\'build.branch != "main" && build.pull_request.labels includes "ready"\''
merge_if = (
'\'(build.branch == "main" && build.env("NIGHTLY") != "1" '
'&& build.env("WEEKLY") != "1") || '
'(build.branch != "main" && build.pull_request.labels includes "merge-test")\''
)
return (
continuation.replace("__IMAGE_BUILD_IF__", image_if)
.replace("__UPLOAD_READY_IF__", ready_if)
.replace("__UPLOAD_MERGE_IF__", merge_if)
)
def _matches_dependencies(changed_files: list[str], prefixes: list[str]) -> bool:
for path in changed_files:
for prefix in prefixes:
normalized = prefix.rstrip("/")
if path == normalized or path.startswith(f"{normalized}/"):
return True
return False
def _strip_source_file_dependencies(step: dict[str, Any]) -> dict[str, Any]:
return {key: value for key, value in step.items() if key != "source_file_dependencies"}
def _step_label(step: dict[str, Any]) -> str:
return str(step.get("group") or step.get("label") or "<step>")
def _filter_steps(steps: list[Any], changed_files: list[str]) -> list[Any]:
filtered: list[Any] = []
for step in steps:
if not isinstance(step, dict):
filtered.append(step)
continue
deps = step.get("source_file_dependencies")
if deps is not None and not isinstance(deps, list):
raise ValueError(
f"source_file_dependencies must be a list in step {_step_label(step)!r}",
)
if deps is not None and not _matches_dependencies(changed_files, deps):
_log(
f"skip {_step_label(step)!r} (no changes under {deps})",
)
continue
nested = step.get("steps")
if nested is not None:
kept_nested = _filter_steps(nested, changed_files)
if not kept_nested:
_log(f"omit empty group {_step_label(step)!r}")
continue
new_step = _strip_source_file_dependencies(step)
new_step["steps"] = kept_nested
filtered.append(new_step)
continue
if deps is not None:
filtered.append(_strip_source_file_dependencies(step))
else:
filtered.append(step)
return filtered
def _strip_uploader_metadata_from_steps(steps: list[Any]) -> list[Any]:
"""Remove uploader-only keys while keeping all steps (no diff filtering)."""
stripped: list[Any] = []
for step in steps:
if not isinstance(step, dict):
stripped.append(step)
continue
deps = step.get("source_file_dependencies")
if deps is not None and not isinstance(deps, list):
raise ValueError(
f"source_file_dependencies must be a list in step {_step_label(step)!r}",
)
nested = step.get("steps")
new_step = _strip_source_file_dependencies(step)
if nested is not None:
new_step["steps"] = _strip_uploader_metadata_from_steps(nested)
stripped.append(new_step)
return stripped
def render_test_pipeline(
doc: dict[str, Any],
changed_files: list[str] | None,
) -> dict[str, Any]:
steps = doc.get("steps")
if not isinstance(steps, list):
return doc
if changed_files is not None:
steps = _filter_steps(steps, changed_files)
else:
steps = _strip_uploader_metadata_from_steps(steps)
return {**doc, "steps": steps}
def resolve_pipeline_path(arg: str) -> Path:
path = Path(arg)
if path.is_absolute():
return path
return ROOT / path
def render_pipeline(path: Path) -> str:
text = path.read_text(encoding="utf-8")
changed_files = resolve_changed_files()
if BOOTSTRAP_MARKER in text:
rendered = render_bootstrap_pipeline(text, skip_ci=resolve_skip_ci(changed_files))
return rendered
doc = yaml.safe_load(text)
if not isinstance(doc, dict):
raise ValueError(f"invalid pipeline YAML: {path}")
doc = render_test_pipeline(doc, changed_files)
return yaml.safe_dump(doc, sort_keys=False)
def upload_to_buildkite(content: str) -> None:
subprocess.run(
["buildkite-agent", "pipeline", "upload"],
input=content,
text=True,
check=True,
)
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"pipeline",
nargs="?",
default=".buildkite/pipeline.yml",
help="Pipeline YAML path (default: .buildkite/pipeline.yml)",
)
parser.add_argument(
"--upload",
action="store_true",
help="Pipe rendered YAML to buildkite-agent pipeline upload",
)
args = parser.parse_args()
path = resolve_pipeline_path(args.pipeline)
if not path.is_file():
_log(f"missing pipeline file: {path}")
return 1
rendered = render_pipeline(path)
if args.upload:
upload_to_buildkite(rendered)
else:
sys.stdout.write(rendered)
return 0
if __name__ == "__main__":
raise SystemExit(main())
@@ -1,146 +0,0 @@
#!/usr/bin/env bash
# Evaluate docs-only skip-ci and upload continuation steps from the given pipeline YAML
# (YAML document after the first `---`). Buildkite `if` is evaluated at upload time.
#
# Usage: upload_pipeline_with_skip_ci.sh [pipeline_yaml]
# pipeline_yaml: path relative to repo root or absolute (default: .buildkite/pipeline.yml)
set -euo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
PIPELINE_ARG="${1:-.buildkite/pipeline.yml}"
if [[ "${PIPELINE_ARG}" = /* ]]; then
PIPELINE_YML="${PIPELINE_ARG}"
else
PIPELINE_YML="${ROOT}/${PIPELINE_ARG}"
fi
# Prints a single digit to stdout: 1 = skip image CI, 0 = run. Logs go to stderr.
is_docs_only_change() {
local file_path
local has_any=0
while IFS= read -r file_path; do
[[ -z "${file_path}" ]] && continue
has_any=1
if [[ "${file_path}" == docs/* ]]; then
continue
fi
if [[ "${file_path}" == *.md ]]; then
continue
fi
if [[ "${file_path}" == "mkdocs.yml" ]]; then
continue
fi
return 1
done
[[ "${has_any}" -eq 1 ]]
}
resolve_skip_ci() {
local is_pr_build=0
local files
local base_branch base_ref
if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" && -n "${BUILDKITE_PULL_REQUEST:-}" ]]; then
is_pr_build=1
fi
if [[ "${is_pr_build}" -eq 1 ]]; then
base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
if ! git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
echo "resolve_skip_ci: origin/${base_branch} not found locally; trying fetch" >&2
git fetch --depth=200 origin "${base_branch}" >/dev/null 2>&1 || true
fi
base_ref=""
if git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
base_ref="origin/${base_branch}"
elif git rev-parse --verify "${base_branch}" >/dev/null 2>&1; then
base_ref="${base_branch}"
else
echo "resolve_skip_ci: cannot resolve PR base ${base_branch}; skip-ci=0" >&2
echo -n 0
return 0
fi
if ! files="$(git diff --name-only "${base_ref}...${BUILDKITE_COMMIT}" 2>/dev/null)"; then
echo "resolve_skip_ci: failed to compute PR changed files; skip-ci=0" >&2
echo -n 0
return 0
fi
elif [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then
if ! git rev-parse --verify "${BUILDKITE_COMMIT}^" >/dev/null 2>&1; then
echo "resolve_skip_ci: commit has no parent on main; skip-ci=0" >&2
echo -n 0
return 0
fi
if ! files="$(git diff --name-only "${BUILDKITE_COMMIT}^..${BUILDKITE_COMMIT}" 2>/dev/null)"; then
echo "resolve_skip_ci: failed to compute main changed files; skip-ci=0" >&2
echo -n 0
return 0
fi
else
echo "resolve_skip_ci: not PR/main build; skip-ci=0" >&2
echo -n 0
return 0
fi
if is_docs_only_change <<< "${files}"; then
echo "resolve_skip_ci: docs-only change detected; skip-ci=1" >&2
echo -n 1
return 0
fi
echo "resolve_skip_ci: non-doc changes detected; skip-ci=0" >&2
echo -n 0
}
SKIP_CI="$(resolve_skip_ci)"
if [[ ! -f "${PIPELINE_YML}" ]]; then
echo "upload_pipeline_with_skip_ci: missing ${PIPELINE_YML}" >&2
exit 1
fi
export ROOT SKIP_CI PIPELINE_YML
python3 <<'PY' | buildkite-agent pipeline upload
import os
import pathlib
path = pathlib.Path(os.environ["PIPELINE_YML"])
text = path.read_text(encoding="utf-8")
sep = "\n---\n"
# Two supported layouts:
# - multi-doc: doc 1 = bootstrap (loaded by Buildkite), doc 2 (after `---`) = uploaded steps.
# - single-doc: whole file is the uploaded steps; caller (UI init step) is the bootstrap.
if sep in text:
_, continuation = text.split(sep, 1)
else:
continuation = text
skip = os.environ.get("SKIP_CI") == "1"
# When docs-only skip-ci: skip default CI image, but still build for L4 nightly (PR label nightly-test or
# main NIGHTLY=1), otherwise upload-nightly (depends_on image-build) would be skipped too.
nightly_only = (
'(build.pull_request.labels includes "nightly-test") '
'|| (build.branch == "main" && build.env("NIGHTLY") == "1")'
)
# Placeholder in pipeline.yml is `if: __IMAGE_BUILD_IF__` (valid YAML); replace value only.
if skip:
rep = f"'{nightly_only}'"
ready_rep = "'false'"
merge_rep = "'false'"
else:
rep = "'true'"
ready_rep = "'build.branch != \"main\" && build.pull_request.labels includes \"ready\"'"
merge_rep = "'(build.branch == \"main\" && build.env(\"NIGHTLY\") != \"1\" && build.env(\"WEEKLY\") != \"1\") || (build.branch != \"main\" && build.pull_request.labels includes \"merge-test\")'"
rendered = (
continuation
.replace("__IMAGE_BUILD_IF__", rep)
.replace("__UPLOAD_READY_IF__", ready_rep)
.replace("__UPLOAD_MERGE_IF__", merge_rep)
)
print(rendered, end="")
PY
+159 -148
View File
@@ -1,54 +1,111 @@
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
HF_HUB_DOWNLOAD_TIMEOUT: 300
HF_HUB_ETAG_TIMEOUT: 60
steps:
- label: "Simple · Diffusion & Model Executor Test"
agent_pool: mi325_1
depends_on: amd-build
timeout_in_minutes: 50
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- group: ":card_index_dividers: Simple Test"
steps:
- label: "Simple · Diffusion & Model Executor Test"
agent_pool: mi325_1
depends_on: amd-build
timeout_in_minutes: 50
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Simple · Other Test"
agent_pool: mi325_1
depends_on: amd-build
timeout_in_minutes: 50
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Simple · Other Test"
agent_pool: mi325_1
depends_on: amd-build
timeout_in_minutes: 50
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Diffusion Images API LoRA E2E"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
- group: ":card_index_dividers: Diffusion Test"
steps:
- label: "Diffusion · LoRA Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
- label: "Diffusion Model CPU offloading Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
- label: "Diffusion Model CPU offloading Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
agent_pool: mi325_4
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
- timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
# merge-only tests
- label: "Diffusion Tensor Parallelism Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
- label: "Diffusion Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
- label: "Diffusion GPU Worker Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
- label: "Diffusion · Wan22 Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
"
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
# - label: "Audio Generation Model Test"
@@ -62,42 +119,6 @@ steps:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
agent_pool: mi325_4
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
- timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
# merge-only tests
- label: "Diffusion Tensor Parallelism Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
- label: "Diffusion Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
- label: "Diffusion GPU Worker Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
- label: "Engine Test"
agent_pool: mi325_1
depends_on: amd-build
@@ -111,50 +132,65 @@ steps:
# This model seems to be important, so I am trying to run it here.
# Since this is from test-nightly.yml,
# I keep this non-blocking.
- label: "Omni Model Test Qwen2-5-Omni"
agent_pool: mi325_4
depends_on: amd-build
mirror_hardwares: [amdproduction]
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
- group: ":card_index_dividers: E2E Test"
steps:
- label: "Omni Model Test Qwen2-5-Omni"
agent_pool: mi325_4
depends_on: amd-build
mirror_hardwares: [amdproduction]
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
- label: "Omni Model Test Qwen3-Omni"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
- label: "Omni Model Test Qwen3-Omni"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
- label: "Qwen3-TTS CustomVoice E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
'
- label: "Qwen3-TTS CustomVoice E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
'
- label: "Qwen3-TTS Base E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 30m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
'
- label: "Qwen3-TTS Base E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 30m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
'
- label: "Voxtral-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
'
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
# Need to debug before reneable numerical changes across large PRs
@@ -195,28 +231,3 @@ steps:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - export VLLM_ROCM_USE_AITER_RMSNORM=0
# - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
- label: "Diffusion · Wan22 Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
"
- label: "Voxtral-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
'
+156 -144
View File
@@ -1,42 +1,33 @@
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
HF_HUB_DOWNLOAD_TIMEOUT: 300
HF_HUB_ETAG_TIMEOUT: 60
steps:
- label: "Simple · Diffusion & Model Executor Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- group: ":card_index_dividers: Simple Test"
steps:
- label: "Simple · Diffusion & Model Executor Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Simple · Other Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Voxtral TTS CUDA Unit Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
- label: "Diffusion · Batch Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
- label: "Simple · Other Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Custom Pipeline Test"
agent_pool: mi325_1
@@ -46,23 +37,92 @@ steps:
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
- label: "Diffusion Model CPU offloading Test"
- label: "Voxtral TTS CUDA Unit Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
- timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
- group: ":card_index_dividers: Diffusion Test"
steps:
- label: "Diffusion · Batch Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
- label: "Diffusion · Offloader Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
- label: "Diffusion · Cache Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
- label: "Diffusion Sequence Parallelism Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
- label: "Diffusion GPU Worker Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
- label: "Diffusion · Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
- label: "Diffusion · Wan22 Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 40m bash -c "
export VLLM_TEST_CLEAN_GPU_MEMORY=1
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
"
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
# - label: "Audio Generation Model Test"
@@ -76,56 +136,6 @@ steps:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
- label: "Diffusion Cache Backend Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
- label: "Diffusion Sequence Parallelism Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
- label: "Diffusion GPU Worker Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
- label: "Diffusion · Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
- label: "Diffusion · Wan22 Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 40m bash -c "
export VLLM_TEST_CLEAN_GPU_MEMORY=1
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
"
- label: "Engine Test"
agent_pool: mi325_1
depends_on: amd-build
@@ -149,43 +159,56 @@ steps:
# - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
# - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
- label: "MiMo-Audio E2E Test with H100"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 30m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
'
- group: ":card_index_dividers: E2E Test"
steps:
- label: "MiMo-Audio E2E Test with H100"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 30m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
'
- label: "Qwen3-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
- label: "Qwen3-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
- label: "Voxtral-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
'
- label: "Voxtral-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
'
- label: "CosyVoice3-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model"
'
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
# Need to debug before reneable numerical changes across large PRs
@@ -225,14 +248,3 @@ steps:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - export VLLM_ROCM_USE_AITER_RMSNORM=0
# - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
- label: "CosyVoice3-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model"
'
+119 -25
View File
@@ -44,6 +44,9 @@ steps:
depends_on: upload-merge-pipeline
steps:
- label: "Diffusion · LoRA Test"
source_file_dependencies:
- tests/diffusion/lora/
- vllm_omni/diffusion/lora/
timeout_in_minutes: 30
commands:
- pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model"
@@ -61,6 +64,9 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Distributed Test"
source_file_dependencies:
- tests/distributed/
- vllm_omni/distributed/
timeout_in_minutes: 60
commands:
- |
@@ -105,6 +111,9 @@ steps:
- group: ":card_index_dividers: Entrypoints Test"
depends_on: upload-merge-pipeline
source_file_dependencies:
- tests/entrypoints/
- vllm_omni/entrypoints/
steps:
- label: "Entrypoint Test with H100"
timeout_in_minutes: 60
@@ -165,6 +174,13 @@ steps:
depends_on: upload-merge-pipeline
steps:
- label: "TTS · Qwen3-TTS CustomVoice Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen3_tts_customvoice.py
- tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
- vllm_omni/model_executor/models/qwen3_tts/
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
- vllm_omni/deploy/qwen3_tts.yaml
- vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml
commands:
- |
timeout 20m bash -c "
@@ -186,6 +202,12 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "TTS · Qwen3-TTS Base Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen3_tts_base.py
- tests/e2e/offline_inference/test_qwen3_tts_base.py
- vllm_omni/model_executor/models/qwen3_tts/
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
- vllm_omni/deploy/qwen3_tts.yaml
commands:
- |
timeout 20m bash -c "
@@ -206,7 +228,34 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "TTS · Ming-Omni-TTS Test"
timeout_in_minutes: 20
commands:
- |
timeout 20m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
pytest -s -v tests/e2e/online_serving/test_ming_tts.py tests/e2e/offline_inference/test_ming_tts.py -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"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "TTS · Voxtral-TTS Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_voxtral_tts.py
- tests/e2e/online_serving/test_voxtral_tts.py
- vllm_omni/model_executor/models/voxtral_tts/
- vllm_omni/model_executor/stage_input_processors/voxtral_tts.py
- vllm_omni/deploy/voxtral_tts.yaml
timeout_in_minutes: 20
commands:
- |
@@ -249,6 +298,12 @@ steps:
type: DirectoryOrCreate
- label: "TTS · GLM-TTS Test"
source_file_dependencies:
- tests/e2e/online_serving/test_glm_tts.py
- tests/e2e/offline_inference/test_glm_tts.py
- vllm_omni/model_executor/models/glm_tts/
- vllm_omni/model_executor/stage_input_processors/glm_tts.py
- vllm_omni/deploy/glm_tts.yaml
timeout_in_minutes: 20
commands:
- |
@@ -291,6 +346,12 @@ steps:
type: DirectoryOrCreate
- label: "Omni · Qwen3-Omni Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_qwen3_omni.py
- tests/e2e/online_serving/test_qwen3_omni.py
- vllm_omni/model_executor/models/qwen3_omni/
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
- vllm_omni/deploy/qwen3_omni_moe.yaml
timeout_in_minutes: 30
commands:
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -329,6 +390,11 @@ steps:
type: DirectoryOrCreate
- label: "Omni · Mimo Audio Test"
source_file_dependencies:
- tests/e2e/online_serving/test_mimo_audio.py
- vllm_omni/model_executor/models/mimo_audio/
- vllm_omni/model_executor/stage_input_processors/mimo_audio.py
- vllm_omni/deploy/mimo_audio.yaml
timeout_in_minutes: 30
commands:
- pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -366,37 +432,19 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Omni · VoxCPM2 Test"
timeout_in_minutes: 20
commands:
- |
timeout 20m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/offline_inference/test_voxcpm2.py -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: "TTS · VoxCPM2 Online Test"
- label: "TTS · VoxCPM2 Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_voxcpm2_tts.py
- tests/e2e/online_serving/test_voxcpm2_tts.py
- vllm_omni/model_executor/models/voxcpm2/
- vllm_omni/deploy/voxcpm2.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/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
"
agents:
queue: "gpu_1_queue"
@@ -413,6 +461,12 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni · Higgs-Audio V2 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_higgs_audio_v2.py
- tests/e2e/offline_inference/test_higgs_audio_v2.py
- vllm_omni/model_executor/models/higgs_audio_v2/
- vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py
- vllm_omni/deploy/higgs_audio_v2.yaml
timeout_in_minutes: 30
commands:
- |
@@ -438,6 +492,9 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion · Qwen Image Edit Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image_edit.py
- vllm_omni/diffusion/models/qwen_image/
timeout_in_minutes: 20
commands:
- pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -476,6 +533,9 @@ steps:
type: DirectoryOrCreate
- label: "Diffusion · Longcat Image Edit Test"
source_file_dependencies:
- tests/e2e/online_serving/test_longcat_image_edit.py
- vllm_omni/diffusion/models/longcat_image/
timeout_in_minutes: 20
commands:
- pytest -s -v tests/e2e/online_serving/test_longcat_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -514,6 +574,9 @@ steps:
type: DirectoryOrCreate
- label: "Diffusion · Longcat Image Test"
source_file_dependencies:
- tests/e2e/online_serving/test_longcat_image.py
- vllm_omni/diffusion/models/longcat_image/
timeout_in_minutes: 20
commands:
- pytest -s -v tests/e2e/online_serving/test_longcat_image.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -552,6 +615,13 @@ steps:
type: DirectoryOrCreate
- label: "Diffusion · Bagel Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_bagel.py
- tests/e2e/online_serving/test_bagel.py
- vllm_omni/model_executor/models/bagel/
- vllm_omni/diffusion/models/bagel/
- vllm_omni/model_executor/stage_input_processors/bagel.py
- vllm_omni/deploy/bagel.yaml
timeout_in_minutes: 20
commands:
- pytest -s -v tests/e2e/offline_inference/test_bagel.py tests/e2e/online_serving/test_bagel.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -591,6 +661,9 @@ steps:
- label: "Diffusion · Qwen Image Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image.py
- vllm_omni/diffusion/models/qwen_image/
commands:
- |
timeout 40m bash -c '
@@ -632,6 +705,9 @@ steps:
type: DirectoryOrCreate
- label: "Diffusion · Qwen Image Layered Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image_layered.py
- vllm_omni/diffusion/models/qwen_image/
commands:
- |
timeout 40m bash -c "
@@ -673,6 +749,10 @@ steps:
type: DirectoryOrCreate
- label: "Diffusion · Wan22 Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_wan22_t2v.py
- tests/e2e/online_serving/test_wan22_t2v.py
- vllm_omni/diffusion/models/wan2_2/
commands:
- |
timeout 40m bash -c "
@@ -714,6 +794,9 @@ steps:
type: DirectoryOrCreate
- label: "Diffusion · HunyuanVideo 1.5 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_hunyuan_video_15.py
- vllm_omni/diffusion/models/hunyuan_video/
timeout_in_minutes: 40
commands:
- |
@@ -755,6 +838,10 @@ steps:
type: DirectoryOrCreate
- label: "Diffusion · Z Image Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_z_image.py
- tests/e2e/online_serving/test_z_image.py
- vllm_omni/diffusion/models/z_image/
timeout_in_minutes: 30
commands:
- pytest -s -v tests/e2e/offline_inference/test_z_image.py tests/e2e/online_serving/test_z_image.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
@@ -772,6 +859,10 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion · FLUX.2 Klein Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_flux2_klein.py
- tests/e2e/online_serving/test_flux2_klein.py
- vllm_omni/diffusion/models/flux2_klein/
timeout_in_minutes: 30
commands:
- pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py tests/e2e/online_serving/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
@@ -789,6 +880,9 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion · SD3 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_sd3.py
- vllm_omni/diffusion/models/sd3/
timeout_in_minutes: 45
commands:
- |
+120 -1
View File
@@ -394,7 +394,7 @@ steps:
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
steps:
- label: ":full_moon: TTS · Function Test"
- label: ":full_moon: TTS · Function Test with L4"
timeout_in_minutes: 90
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
@@ -413,6 +413,45 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
# - label: ":full_moon: TTS · Function Test with H100"
# timeout_in_minutes: 90
# commands:
# - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
# - pytest -s -v tests/e2e/ -m "full_model and H100 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# resources:
# limits:
# nvidia.com/gpu: 1
# volumeMounts:
# - name: devshm
# mountPath: /dev/shm
# - name: hf-cache
# mountPath: /root/.cache/huggingface
# env:
# - name: HF_HOME
# value: /root/.cache/huggingface
# - name: HF_TOKEN
# valueFrom:
# secretKeyRef:
# name: hf-token-secret
# key: token
# nodeSelector:
# node.kubernetes.io/instance-type: gpu-h100-sxm
# volumes:
# - name: devshm
# emptyDir:
# medium: Memory
# - name: hf-cache
# hostPath:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
- label: ":full_moon: TTS · Accuracy Test · VoxCPM2"
timeout_in_minutes: 30
commands:
@@ -552,6 +591,7 @@ steps:
tests/e2e/online_serving/test_longcat_image_edit_expansion.py
tests/e2e/online_serving/test_flux_2_dev_expansion.py
tests/e2e/online_serving/test_bagel_expansion.py
tests/e2e/offline_inference/test_sensenova_u1_expansion.py
-m "full_model and diffusion and H100 and not distributed_cuda"
--run-level "full_model"
agents:
@@ -762,6 +802,45 @@ steps:
type: DirectoryOrCreate
- label: ":full_moon: Diffusion X2I(&A&T) · GLM-Image Function Test with H100"
timeout_in_minutes: 120
commands:
- pytest -sv tests/e2e/online_serving/test_glm_image_expansion.py -m "full_model and diffusion and H100" --run-level "full_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
timeout_in_minutes: 60
commands:
@@ -859,6 +938,46 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: HunyuanImage3-DIT · Accuracy Test"
timeout_in_minutes: 180
commands:
- export HUNYUAN_IMAGE3_MODEL="tencent/HunyuanImage-3.0-Instruct"
- export HUNYUAN_IMAGE3_DEVICES="0,1,2,3"
- pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m full_model --run-level full_model
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 4
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image"
key: nightly-diffusion-x2iat-performance-qwen-image
timeout_in_minutes: 180
+54 -18
View File
@@ -273,37 +273,50 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Omni · VoxCPM2 Test"
timeout_in_minutes: 20
- label: "Omni · Qwen3-Omni Realtime WebSocket"
commands:
- |
timeout 20m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
timeout 30m bash -c "
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/offline_inference/test_voxcpm2.py -m 'core_model' --run-level 'core_model'
pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_streaming_audio_input_pcm_output_async_chunk_smoke -m 'core_model' --run-level 'core_model'
"
agents:
queue: "gpu_1_queue"
queue: "mithril-h100-pool"
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"
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "TTS · VoxCPM2 Online Test"
- label: "TTS · VoxCPM2 Test"
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/online_serving/test_voxcpm2_tts.py -m 'core_model and tts' --run-level 'core_model'
pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model'
"
agents:
queue: "gpu_1_queue"
@@ -363,6 +376,29 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "TTS · Higgs-Audio-V3 Test"
commands:
- |
timeout 20m bash -c "
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_USE_DEEP_GEMM=0
export VLLM_MOE_USE_DEEP_GEMM=0
pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_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
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion · Qwen Image Edit Test"
commands:
- |
+79 -3
View File
@@ -6,6 +6,13 @@
{% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %}
{% set default_working_dir = "/app/vllm-omni" %}
{% if env %}
env:
{% for name in env %}
{{ name }}: "{{ env[name] }}"
{% endfor %}
{% endif %}
steps:
- group: "AMD Tests"
depends_on: ~
steps:
@@ -31,8 +38,15 @@
agents:
queue: amd-cpu
{% for step in steps %}
{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
{% for step_group in steps %}
{% if step_group.group %}
- group: "{{ step_group.group }}"
{% if step_group.depends_on %}
depends_on: {{ step_group.depends_on }}
{% endif %}
steps:
{% for step in step_group.steps %}
{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
- label: "{{ step.agent_pool }}: {{ step.label }}"
depends_on: amd-build
agents:
@@ -46,8 +60,23 @@
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
env:
DOCKER_BUILDKIT: "1"
{% if step.env %}
{% for name in step.env %}
{{ name }}: "{{ step.env[name] }}"
{% endfor %}
{% endif %}
TEST_COMMAND: |-
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
{% if env %}
{% for name in env %}
export {{ name }}="{{ env[name] }}"
{% endfor %}
{% endif %}
{% if step.env %}
{% for name in step.env %}
export {{ name }}="{{ step.env[name] }}"
{% endfor %}
{% endif %}
{% if "mi250" in step.agent_pool %}
python3 -m pip uninstall -y amd-aiter
{% endif %}
@@ -61,5 +90,52 @@
{% else %}
soft_fail: true
{% endif%}
{% endif %}
{% endfor %}
{% elif step_group.mirror_hardwares and mirror_hw in step_group.mirror_hardwares %}
{% set step = step_group %}
- label: "{{ step.agent_pool }}: {{ step.label }}"
depends_on: amd-build
agents:
{% if step.agent_pool %}
queue: amd_{{ step.agent_pool }}
{% else %}
queue: amd_mi325_1
{% endif %}
{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
{% set indented_cmd = cmd_body | replace("\n", "\n ") %}
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
env:
DOCKER_BUILDKIT: "1"
{% if step.env %}
{% for name in step.env %}
{{ name }}: "{{ step.env[name] }}"
{% endfor %}
{% endif %}
TEST_COMMAND: |-
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
{% if env %}
{% for name in env %}
export {{ name }}="{{ env[name] }}"
{% endfor %}
{% endif %}
{% if step.env %}
{% for name in step.env %}
export {{ name }}="{{ step.env[name] }}"
{% endfor %}
{% endif %}
{% if "mi250" in step.agent_pool %}
python3 -m pip uninstall -y amd-aiter
{% endif %}
{{ indented_cmd | safe }}
priority: 100
{% if step.timeout_in_minutes %}
timeout_in_minutes: {{ step.timeout_in_minutes }}
{% endif %}
{% endfor %}
{% if step.grade and step.grade == "Blocking" %}
soft_fail: false
{% else %}
soft_fail: true
{% endif%}
{% endif %}
{% endfor %}
+97
View File
@@ -0,0 +1,97 @@
---
name: precheck-pr
description: Self-check your branch before creating a PR — catch dead code, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub.
---
# PR Pre-Check
Self-review your branch before creating a PR against `vllm-project/vllm-omni`. Two modes: **quick** catches showstoppers, **full** does a thorough maintainer-grade review. Never posts to GitHub; the report is for the contributor's terminal only.
## Mode Selection
| Mode | When | Time |
|------|------|------|
| **Quick** | About to push, final sanity check | ~3 min |
| **Full** | Ready for review, want maintainer-level scan | ~10 min |
Default to quick if unsure. Run full before marking a PR "ready for review."
## Workflow
### Step 1: Detect Base Branch
```bash
BASE_SHA=$(git merge-base HEAD origin/main 2>/dev/null \
|| git merge-base HEAD main 2>/dev/null \
|| echo origin/main)
echo "diffing against ${BASE_SHA}"
git diff --name-only ${BASE_SHA}...HEAD
```
### Step 2: Validate PR Title
Check the most recent commit message (or branch name if no commit yet) against the project convention. Valid prefixes:
| Prefix | Applies to |
|--------|-----------|
| `[Bugfix]` | Bug fixes |
| `[CI/Build]` | Build or CI improvements |
| `[Doc]` | Documentation changes |
| `[Model]` | New/improved models (include model name) |
| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
| `[Kernel]` | CUDA/kernel changes |
| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
| `[Misc]` | Other changes (use sparingly) |
✗ if: missing prefix, wrong case (`[bugfix]`), or WIP/Draft in title.
⚠ if: `[Model]` prefix without the model identifier (e.g., `[Model] Add new model` — should be `[Model] Add <ModelName> ...`).
### Step 3: Categorize the PR
| Diff contains | PR type |
|---------------|---------|
| New files under `vllm_omni/model_executor/models/<name>/` | **New Model** |
| Changes to `vllm_omni/diffusion/` | **Diffusion Model** |
| `[Bugfix]` prefix or single-file fix | **Bug Fix** |
| Perf/benchmark/throughput claims in commit msg or diff | **Performance** |
| Everything else | **General** |
If multiple rows apply (e.g., a diffusion model is also a new model), union the checklists.
### Step 4: Run Checklist
Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type from [references/checklists.md](references/checklists.md). Each item produces ✓, ✗, or ⚠.
### Step 5: Print Report
```
Pre-check report for <branch>
Mode: quick | full
Type: <new-model | diffusion-model | bug-fix | perf | general>
Dimension Result
───────────────── ──────
PR title format ✓
PR desc integrity ✓
Registry/config ✓
Dead code ⚠ 2 warnings
Accuracy ✓
Benchmark ✗ missing software versions
Verdict: 1 blocking | 2 warnings | recommend fixing ✗ before PR
```
**Severity:**
| Mark | Meaning |
|------|---------|
| ✗ | Blocking — fix before opening PR |
| ⚠ | Warning — consider fixing |
| ✓ | Pass |
| — | Skipped (not applicable) |
### Stop Here
Do not post comments, open PRs, or modify files. The report is for the contributor's terminal only.
@@ -0,0 +1,135 @@
# Pre-Check Checklists
Each item is a concrete check to run against the diff. ✗ = fix before PR. ⚠ = consider fixing.
---
## All PRs: Title Format (check first)
PR title must follow the project convention documented in `docs/contributing/README.md`:
| Prefix | Applies to |
|--------|-----------|
| `[Bugfix]` | Bug fixes |
| `[CI/Build]` | Build or CI improvements |
| `[Doc]` | Documentation changes |
| `[Model]` | New/improved models (include model name) |
| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
| `[Kernel]` | CUDA/kernel changes |
| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
| `[Misc]` | Other changes (use sparingly) |
**Examples:**
```
✓ [Model] Add MiniCPM-o 4.5 support
✓ [Bugfix] Fix Ovis image text encoder dtype
✓ [Core] Qwen3-Omni performance optimization
✓ [Hardware][Ascend] Fix VoxCPM2 on Ascend
✗ Fix Ovis image text encoder dtype ← missing prefix
✗ [bugfix] Fix Ovis image text encoder ← wrong case
✗ [WIP] Add model ← WIP in title
```
- [ ] **Title has valid prefix** — ✗ if missing prefix, wrong case, or WIP/Draft in title. For `[Model]` PRs, the title MUST also include the model identifier (e.g., `[Model] Add <ModelName> ...`). ⚠ if missing the model name.
---
## Bug Fix PRs
### Quick
- [ ] **PR body has 4 sections:** what broke (error message), repro steps, root cause, fix description
- [ ] **Fix is minimal:** diff touches only files directly related to the bug. No unrelated refactors or style changes.
- [ ] **Repro is runnable:** the PR body includes a copy-paste command that reproduces the bug
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
### Full (adds)
- [ ] **Regression test exists:** `git diff --name-only` includes at least one `tests/` file
- [ ] **Regression test reproduces the original error:** the test would fail on main, pass on this branch
- [ ] **Fix matches root cause exactly:** no "fix the symptom + something else" — if the root cause is dtype, the fix is dtype, not dtype + reformatting
- [ ] **No silent failure risk:** no bare `except: pass`, no `try: ... except: return None`, no empty fallback added
- [ ] **Upstream pattern match:** the fix follows the same pattern as existing code for similar cases (grep for analogous `from_pretrained` calls, etc.)
- [ ] **Environment documented:** torch/transformers/vllm versions listed if the bug is version-dependent
---
## Performance PRs
### Quick
- [ ] **Before/after numbers in PR body:** at minimum, one metric (RTF, RPS, latency, VRAM) with main vs PR values
- [ ] **Hardware specified:** GPU model, count, VRAM. "L4" or "H20" is sufficient. "tested on GPU" is not.
- [ ] **Concurrency scaling:** if the optimization claims better throughput, numbers at 2+ concurrency levels
- [ ] **No unexplained regressions:** if any metric goes down at any concurrency level, explain why
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
### Full (adds)
- [ ] **Software versions:** torch, CUDA, vllm, vllm-omni versions in PR body
- [ ] **Warmup stated:** "1 warmup + 3 measured" or equivalent
- [ ] **Stddev or range:** not just a single best run. "average of 2 rounds" is minimum.
- [ ] **VRAM measured:** peak GPU memory before and after
- [ ] **Benchmark script checked in:** under `tests/` or `examples/`, runnable by reviewer
- [ ] **Exact command line in PR:** not pseudocode, not "run the benchmark"
- [ ] **Pytest summary included:** `36 passed, 3 skipped, 17 warnings in 2897.66s` — proves the whole suite passes
- [ ] **Config changes documented:** if `max_num_seqs`, `gpu_memory_utilization`, or deploy YAML changed, old and new values stated with reasoning
---
## New Model PRs
### Quick
- [ ] **PR body file list matches `git diff --name-only`:** every file path in the description exists in the diff
- [ ] **Registry entries match claimed count:** count `_OMNI_MODELS` or `_DIFFUSION_MODELS` entries added vs claimed in body
- [ ] **Pipeline config `custom_process_input_func` paths resolve:** each referenced module.function exists
- [ ] **`__init__.py` exports match registry:** each registered class appears in `__init__.py` and `__all__`
- [ ] **Shell script modes match README:** `run_curl.sh` supports every mode the README documents
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
### Full (adds)
- [ ] **Dead code scan (5 patterns):**
1. Dead `forward()` — defined but never called through the pipeline path
2. Dead factory/builder functions — instantiated nowhere
3. Dead wrappers — thin wrappers with no callers
4. Dead branch guards — `if version > X` where X is always satisfied
5. Unused parameters — `__init__` args never accessed
- [ ] **Copy-paste detection:** no string constant defined in 3+ files, no cross-module validation duplication, no near-identical shape coercion fns
- [ ] **Import hygiene:** no re-export-only imports, no module-level side effects without docs, no `import os` in function body when already at top. Function-body imports are acceptable when the import is only reachable from a closure or factory return value (i.e., `import` would otherwise execute at module load even though the user may never trigger that path). Either lift to top-of-file or annotate with `# noqa: PLC0415`.
- [ ] **Accuracy:** at least one test compares output against reference; audio = valid WAV at correct sample rate; determinism verified if seed claimed. ⚠ if missing and not explicitly deferred to a follow-up PR or RFC section; — if explicitly deferred (the deferral target SHOULD be linked in the PR body).
- [ ] **Performance:** hardware, software versions, warmup, and metrics (RTF, VRAM, RPS, latency) stated. ⚠ if missing and not explicitly deferred; — if explicitly deferred.
- [ ] **Benchmark settings:** model config (TP, PP, max_model_len, enforce_eager, quant), runtime config (batch, input/output spec), environment versions
- [ ] **Benchmark script checked in + exact command line + pytest summary**
---
## Diffusion Model PRs
Diffusion models live under `vllm_omni/diffusion/`. In addition to the New Model checklist above, verify:
- [ ] **Transformer adapter follows contract:** inherits from `nn.Module`, implements `load_weights()`, uses `vllm_omni.diffusion.attention.layer.Attention`
- [ ] **Pipeline has both offline and online paths:** `forward()` for batch generation, streaming path for online serving
- [ ] **Cache-DiT integration:** at least one acceleration (TeaCache, FBCache, etc.) is wired up. ✗ for new **top-level** pipelines. For **subclass** pipelines that inherit `__call__` from a parent where Cache-DiT is already absent, downgrade to ⚠ and recommend filing a follow-up issue against the parent pipeline family.
- [ ] **Parallelism config:** TP/SP/USP/CFG-Parallel options exposed in pipeline config
- [ ] **Docs table updated:** model added to `docs/models/supported_models.md` with correct metadata
- [ ] **E2E test exists:** at least one test exercises the full generate path
---
## General PRs
### Quick
- [ ] **PR body matches diff:** description doesn't claim changes not in the diff
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
- [ ] **CI gates passing:** DCO, pre-commit, build — check `gh pr view --json statusCheckRollup` if a PR already exists
### Full (adds)
- [ ] **No new dead code in changed files:** scan for unreachable paths, unused imports
- [ ] **Test coverage:** if core code changed (`engine/`, `stages/`, `connectors/`), tests are added or existing tests cover the path
- [ ] **Import hygiene:** no redundant imports added
- [ ] **No unrelated changes:** diff scope matches PR description
+5
View File
@@ -25,7 +25,12 @@ include:
- `add-tts-model`: covers integration of new TTS models and related serving
workflows
- `generate-release-note`: helps prepare release notes for repository changes
- `precheck-pr`: self-check a branch before creating a PR — validates PR title
format, catches dead code, verifies accuracy/perf claims, and confirms merge
readiness
- `review-pr`: provides a structured workflow for reviewing pull requests
- `vllm-omni-npu-model-runner-upgrade`: upgrades NPU model runners to align with the
latest vllm-ascend NPUModelRunner
## Maintenance Guidelines
+7 -13
View File
@@ -1,21 +1,15 @@
<!-- markdownlint-disable -->
PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
PLEASE FILL IN THE PR DESCRIPTION HERE.
## Purpose
## Test Plan
**vLLM Version:**
**vLLM-Omni Commit:**
## Test Result
---
<details>
<summary> Essential Elements of an Effective PR Description Checklist </summary>
- [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
- [ ] The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
- [ ] The test results. Please paste the results comparison before and after, or the e2e results.
- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.**
- [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
</details>
**BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)
**BEFORE SUBMITTING:** read [CONTRIBUTING.md](https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md) and run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
(anything written below this line will be removed by GitHub Actions)
+2
View File
@@ -1,3 +1,5 @@
# Contributing to vLLM-Omni
You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/)
Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
+4 -2
View File
@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
---
*Latest News* 🔥
- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
@@ -56,8 +57,9 @@ vLLM-Omni is flexible and easy to use with:
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
- Omni-modality models (e.g. Qwen-Omni)
- Multi-modality generation models (e.g. Qwen-Image)
- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos, HunyuanImage, BAGEL)
- **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
- **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
## Getting Started
+9
View File
@@ -151,6 +151,9 @@ class VllmOmniImageClient:
seed: int | None = None,
negative_prompt: str | None = None,
output_compression: int | None = None,
bot_task: str | None = None,
sys_type: str | None = None,
system_prompt: str | None = None,
) -> Image.Image:
if not isinstance(images, list):
images = [images]
@@ -170,6 +173,12 @@ class VllmOmniImageClient:
data["negative_prompt"] = negative_prompt
if output_compression is not None:
data["output_compression"] = str(output_compression)
if bot_task is not None:
data["bot_task"] = bot_task
if sys_type is not None:
data["sys_type"] = sys_type
if system_prompt is not None:
data["system_prompt"] = system_prompt
files = [
(
@@ -64,6 +64,47 @@ def resolve_model_name(*, model_name: str | None, model: str | None = None, outp
raise ValueError("model-name is required when it cannot be inferred from model or output-root")
def _coerce_score_int(value: Any) -> int:
if isinstance(value, bool):
raise TypeError(f"Expected numeric score, got bool: {value!r}")
if isinstance(value, (int, float)):
return int(value)
if isinstance(value, str):
stripped = value.strip()
if stripped.lstrip("-").isdigit():
return int(stripped)
raise TypeError(f"Expected numeric score, got {type(value).__name__}: {value!r}")
def _expand_score_entry(entry: Any) -> list[int]:
if isinstance(entry, dict):
if "score" in entry:
return [_coerce_score_int(entry["score"])]
if "value" in entry:
return [_coerce_score_int(entry["value"])]
numeric_values = [_coerce_score_int(value) for value in entry.values()]
if numeric_values:
return numeric_values
raise TypeError(f"Score dict has no numeric values: {entry!r}")
return [_coerce_score_int(entry)]
def _normalize_score_list(score: Any) -> list[int]:
if isinstance(score, dict):
entries: list[Any] = [score]
elif isinstance(score, list):
entries = score
else:
entries = [score]
normalized: list[int] = []
for entry in entries:
normalized.extend(_expand_score_entry(entry))
if not normalized:
raise ValueError(f"Empty score list after normalization: {score!r}")
return normalized
def parse_score_payload(raw_text: str) -> dict[str, Any]:
try:
parsed = extract_json_object(raw_text)
@@ -73,10 +114,7 @@ def parse_score_payload(raw_text: str) -> dict[str, Any]:
parsed = {"score": json.loads(stripped), "reasoning": ""}
else:
raise
score = parsed.get("score", [])
if not isinstance(score, list):
score = [score]
parsed["score"] = [int(value) for value in score]
parsed["score"] = _normalize_score_list(parsed.get("score", []))
return parsed
@@ -391,6 +429,9 @@ class GEditBenchRunner:
num_inference_steps: int = 20,
guidance_scale: float | None = None,
seed: int | None = 42,
bot_task: str | None = None,
sys_type: str | None = None,
system_prompt: str | None = None,
):
self.dataset_ref = dataset_ref
self.output_root = output_root
@@ -400,6 +441,9 @@ class GEditBenchRunner:
self.num_inference_steps = num_inference_steps
self.guidance_scale = guidance_scale
self.seed = seed
self.bot_task = bot_task
self.sys_type = sys_type
self.system_prompt = system_prompt
self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
def generate(
@@ -478,6 +522,9 @@ class GEditBenchRunner:
num_inference_steps=self.num_inference_steps,
guidance_scale=self.guidance_scale,
seed=self.seed,
bot_task=self.bot_task,
sys_type=self.sys_type,
system_prompt=self.system_prompt,
)
edited_image.save(output_path)
return {
@@ -633,6 +680,25 @@ def build_parser() -> argparse.ArgumentParser:
generate.add_argument("--num-inference-steps", type=int, default=20)
generate.add_argument("--guidance-scale", type=float, default=None)
generate.add_argument("--seed", type=int, default=42)
generate.add_argument(
"--bot-task",
type=str,
default=None,
choices=["think", "recaption", "think_recaption", "vanilla"],
help="HunyuanImage-3 prompting mode for /v1/images/edits (optional).",
)
generate.add_argument(
"--sys-type",
type=str,
default=None,
help="Hunyuan system-prompt type (e.g. en_unified). Defaults from --bot-task when omitted.",
)
generate.add_argument(
"--system-prompt",
type=str,
default=None,
help="Custom Hunyuan system prompt body (sys_type=custom).",
)
generate.add_argument("--workers", type=int, default=1)
generate.add_argument("--max-samples", type=int, default=None)
generate.add_argument("--samples-per-group", type=int, default=None)
@@ -675,6 +741,9 @@ def main(argv: list[str] | None = None) -> int:
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
bot_task=args.bot_task,
sys_type=args.sys_type,
system_prompt=args.system_prompt,
)
records = runner.generate(
model_name=model_name,
-1
View File
@@ -92,7 +92,6 @@ Core scheduling and caching components.
Configuration classes.
- [vllm_omni.config.model.OmniModelConfig][]
- [vllm_omni.config.stage_config.ModelPipeline][]
- [vllm_omni.config.stage_config.StageConfig][]
- [vllm_omni.config.stage_config.StageConfigFactory][]
- [vllm_omni.config.stage_config.StageType][]
Binary file not shown.

Before

Width:  |  Height:  |  Size: 408 KiB

After

Width:  |  Height:  |  Size: 403 KiB

+134 -26
View File
@@ -2,29 +2,29 @@
The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
## Notes
We currently only support using the "openai-chat-omni" backend.
We currently only support using the "openai-chat-omni" and "openai-image-edits-omni" backend.
## Basic Parameter Description
You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
- `--omni`
- `--omni`
Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio.
- `--backend`
- `--backend`
Specify the backend adapter as openai-chat-omni, using OpenAI Chat compatible API behavior as the protocol. Currently only openai-chat-omni is supported.
- `--model`
- `--model`
The model identifier to load, filled according to the models supported by vLLM-Omni.
- `--endpoint`
- `--endpoint`
The API endpoint exposed externally, to which clients send their requests.
- `--dataset-name`
- `--dataset-name`
The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio).
- `--num-prompts`
- `--num-prompts`
The total number of requests to send, an integer.
- `--max-concurrency`
- `--max-concurrency`
"Maximum number of concurrent requests. This can be used "
"to help simulate an environment where a higher level component "
"is enforcing a maximum number of concurrent requests. While the "
@@ -34,69 +34,73 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
"actual request rate may be lower than specified with --request-rate, "
"if the server is not processing requests fast enough to keep up."
- `--request-rate`
- `--request-rate`
"Number of requests per second. If this is inf, "
"then all the requests are sent at time 0. "
"Otherwise, we use Poisson process or gamma distribution "
"to synthesize the request arrival times."
- `--ignore-eos`
- `--ignore-eos`
"Set ignore_eos flag when sending the benchmark request."
- `--metric-percentiles`
- `--metric-percentiles`
Comma-separated list of percentiles for selected metrics. "
"To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". "
"Default value is \"99\"."
"Use \"--percentile-metrics\" to select metrics.
- `--percentile-metrics`
- `--percentile-metrics`
"Comma-separated list of selected metrics to report percentiles."
"This argument specifies the metrics to report percentiles."
'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
'Allowed metric names are "ttft", "tpot", "itl", "ttfc", "tpoc", "icl", '
'"tpop", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
- `--save-result`
- `-print-stage`
Print per-stage benchmark metrics for --omni serving when stage metrics are returned by the server. Disabled by default.
- `--save-result`
Specify to save benchmark results to a json file
- `--save-detailed`
- `--save-detailed`
"When saving the results, whether to include per request "
"information such as response, error, ttfs, tpots, etc."
- `--result-dir`
- `--result-dir`
"Specify directory to save benchmark json results."
"If not specified, results are saved in the current directory."
- `--result-filename`
- `--result-filename`
"Specify the filename to save benchmark json results."
"If not specified, results will be saved in "
"{label}-{args.request_rate}qps-{base_model_id}-{current_dt}.json"
- `--random-prefix-len`
- `--random-prefix-len`
Number of fixed prefix tokens before the random context in a request.
The total input length is the sum of random-prefix-len and a random
context length sampled from [input_len * (1 - range_ratio),
input_len * (1 + range_ratio)].Only the random and random-mm modes
support this parameter.
- `--random-input-len`
- `--random-input-len`
Number of input tokens per request.Only the random and random-mm modes support this parameter.
- `--random-output-len`
- `--random-output-len`
Number of output tokens per request.Only the random and random-mm modes support this parameter.
- `--random-range-ratio`
- `--random-range-ratio`
Range ratio for sampling input/output length,
used only for random sampling. Must be in the range [0, 1) to define
a symmetric sampling range
[length * (1 - range_ratio), length * (1 + range_ratio)].
Only the random and random-mm modes support this parameter.
- `--random-mm-base-items-per-request`
- `--random-mm-base-items-per-request`
Base number of multimodal items per request for random-mm.
Actual per-request count is sampled around this base using
--random-mm-num-mm-items-range-ratio.
Only the random-mm mode supports this parameter.
- `--random-mm-limit-mm-per-prompt`
- `--random-mm-limit-mm-per-prompt`
Per-modality hard caps for items attached per request, e.g.
'{"image": 3, "video": 1, "audio": 1}'. The sampled per-request item
count is clamped to the sum of these limits. When a modality
@@ -104,7 +108,7 @@ Specify to save benchmark results to a json file
renormalized.
Only the random-mm mode supports this parameter.
- `--random-mm-num-mm-items-range-ratio`
- `--random-mm-num-mm-items-range-ratio`
Range ratio r in [0, 1] for sampling items per request.
We sample uniformly from the closed integer range
[floor(n*(1-r)), ceil(n*(1+r))]
@@ -115,7 +119,7 @@ Specify to save benchmark results to a json file
An error is raised if the computed min exceeds the max.
Only the random-mm mode supports this parameter.
- `--random-mm-bucket-config`
- `--random-mm-bucket-config`
The bucket config is a dictionary mapping a multimodal item
sampling configuration to a probability.
Currently allows for 3 modalities: audio, images and videos.
@@ -276,7 +280,9 @@ Generate synthetic image、video、audio inputs alongside random text prompts to
Notes:
- Works only with online benchmark via the OpenAI backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
- Works only with online benchmark via:
- the OpenAI chat backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
- the OpenAI edit image backend (`--backend openai-image-edits-omni`) and endpoint `/v1/images/edits`.
Start the server (example):
@@ -357,3 +363,105 @@ How sampling works:
This should be seen as an edge case, and if this behavior can be avoided by setting `--random-mm-limit-mm-per-prompt` to a large number. Note that this might result in errors due to engine config `--limit-mm-per-prompt`.
- The resulting request contains synthetic image data in `multi_modal_data` (OpenAI Chat format). When `random-mm` is used with the OpenAI Chat backend, prompts remain text and MM content is attached via `multi_modal_data`.
</details>
### Multi-Stage Benchmark
<details class="admonition abstract" markdown="1">
<summary>Show more</summary>
Start the server (change the port and model):
```bash
vllm serve --omni --port 29999 --model ~/Qwen3-Omni-30B-A3B-Instruct
```
Then run the benchmarking script (remember to add the flag `--print-stage`):
```bash
vllm bench serve --omni \
--dataset-name random \
--port 29999 \
--max-concurrency 1 \
--num-warmups 2 \
--model ~/Qwen3-Omni-30B-A3B-Instruct \
--endpoint /v1/chat/completions \
--backend openai-chat-omni \
--num-prompts 3 \
--random-input-len 2500 \
--ignore-eos \
--percentile-metrics ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,ttfc,tpoc,icl \
--random-output-len 900 \
--extra_body '{"modalities": ["text","audio"]}' \
--print-stage
```
Besides the "Serving Benchmark Result", there will be "Stage Benchmark Result" below:
```text
============= Stage Benchmark Result =============
=============== Stage 0 (thinker) ================
-------------------Stage Timing-------------------
Mean stage_gen_time (ms): 18243.16
Median stage_gen_time (ms): 18220.83
P99 stage_gen_time (ms): 18296.75
================== Text Result ===================
Stage generated tokens: 2700
-----------Serving Time to First Token------------
Mean Serving TTFT (ms): 816.31
Median Serving TTFT (ms): 818.94
P99 Serving TTFT (ms): 829.76
-----Time per Output Token (excl. 1st token)------
Mean TPOP (ms): 19.41
Median TPOP (ms): 19.40
P99 TPOP (ms): 19.47
---------------Inter-token Latency----------------
Mean ITL (ms): 19.41
Median ITL (ms): 19.12
P99 ITL (ms): 23.58
================ Stage 1 (talker) ================
-------------------Stage Timing-------------------
Mean stage_gen_time (ms): 18849.64
Median stage_gen_time (ms): 13641.15
P99 stage_gen_time (ms): 33655.58
============= Internal Stream Result =============
-----------Serving Time to First Chunk------------
Mean Serving TTFC (ms): 1029.14
Median Serving TTFC (ms): 1027.98
P99 Serving TTFC (ms): 1032.59
-----Time per Output Chunk (excl. 1st chunk)------
Mean TPOP (ms): 17.53
Median TPOP (ms): 18.70
P99 TPOP (ms): 18.88
---------------Inter-chunk Latency----------------
Mean ICL (ms): 16.27
Median ICL (ms): 18.19
P99 ICL (ms): 25.86
=============== Stage 2 (code2wav) ===============
-------------------Stage Timing-------------------
Mean stage_gen_time (ms): 18860.74
Median stage_gen_time (ms): 13655.84
P99 stage_gen_time (ms): 33659.79
================== Audio Result ==================
Stage audio duration generated(s): 260.11
Stage audio frames generated: 6242730
-----------Serving Time to First Packet-----------
Mean Serving AUDIO_TTFP (ms): 1366.16
Median Serving AUDIO_TTFP (ms): 1367.20
P99 Serving AUDIO_TTFP (ms): 1368.26
-----------------Real Time Factor-----------------
Mean AUDIO_RTF: 0.24
Median AUDIO_RTF: 0.26
P99 AUDIO_RTF: 0.27
```
Explanation:
- stage_gen_time: Time from submitting a request to a specific stage to that stage finishing generation.
- Serving TTFC (Time to First Chunk): Time from the HTTP request being accepted by the serving frontend to
the stage producing its first non-empty output chunk. Keep the name TTFT in text stage, TTFP in audio stage
for easier mapping.
- TPOP (Time per Output Chunk): Average time from the first output chunk to stage completion, divided by the
number of remaining output chunks. The TPOP abbreviation follows Qwen3.5-Omni Technical Report.
- ICL (Inter-Chunk Latency): Time between two consecutive output chunks produced by the same stage. Keep the
name ITL in text stage for easier mapping.
</details>
+10
View File
@@ -116,6 +116,16 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
!!! note
If the PR spans more than one category, please include all relevant prefixes.
### Pre-Check Before Submitting
Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-review against project conventions:
The skill offers two modes:
- **Quick (~3 min):** catches showstoppers — PR title format, missing benchmark claims, rebase status
- **Full (~10 min):** thorough maintainer-grade review — dead code scan, copy-paste detection, import hygiene
The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Model, and General. Each type has a tailored checklist that validates evidence quality (repro steps, A/B benchmarks, registry entries, etc.). See the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) for the full checklist.
### Local Test
Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
@@ -8,6 +8,7 @@ Backend-specific setup lives in separate docs:
- [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
- [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
- [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
- [MoriTransferEngineConnector](omni_connectors/mori_transfer_engine_connector.md)
- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
- [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md)
@@ -23,6 +24,7 @@ Current connectors operate in D2H2D (device to host to device) mode.
| Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
| Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
| Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
| Multi node (Mori RDMA) | MoriTransferEngineConnector | RDMA direct transfer via Mori IOEngine. |
| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
| Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. |
@@ -0,0 +1,76 @@
# MoriTransferEngineConnector
## When to Use
Currently supports intra-node deployment with Mori.
As noted in #1742, inter-node support will be added back in a future
refactor.
## Mechanism
Uses Mori's `IOEngine` / `MemoryDesc` API for zero-copy RDMA transfers.
- Data Plane: RDMA (InfiniBand/RoCE) with managed memory pool.
- Control Plane: ZMQ for pull-request handshake and async completion.
## Installation
See the [Mori repository](https://github.com/ROCm/mori) for installation instructions.
## Configuration
Mori is configured through the new deploy-config schema (see
[`docs/configuration/stage_configs.md`](../../../configuration/stage_configs.md)).
Define the connector at the top level of the deploy YAML and reference it
by name from each stage's `input_connectors` / `output_connectors`:
```yaml
connectors:
mori_connector:
name: MoriTransferEngineConnector
extra:
host: "auto"
zmq_port: 50051
device_name: ""
memory_pool_size: 536870912
memory_pool_device: "cuda"
stages:
- stage_id: 0
output_connectors:
to_stage_1: mori_connector
- stage_id: 1
input_connectors:
from_stage_0: mori_connector
```
A ready-to-run intra-node example for Qwen3-Omni-MoE on AMD MI300X lives
at `vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml` and can be loaded
with:
```bash
vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --log-stats \
--deploy-config vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml
```
The yaml wires `MoriTransferEngineConnector` (with `backend_type: xgmi`) to
the chunk_transfer_adapter path (`async_chunk: true`) so stage-to-stage
hidden-state and codec-frame streams ship GPU-to-GPU over AMD Infinity
Fabric instead of SHM. Qwen2.5-Omni + Mori is not yet functional on
the chunk path: it needs `thinker2talker_async_chunk` /
`talker2code2wav_async_chunk` input processors that do not exist yet
(the orchestrator-level path the upstream PR originally targeted was
lost during the entrypoints → engine refactor and removed in #1742).
Parameters:
- host: local RDMA IP (`"auto"` for auto-detect).
- zmq_port: ZMQ base port for control-plane communication.
- device_name: RDMA device (e.g., `"mlx5_0"`), empty for auto-detect.
- memory_pool_size: RDMA memory pool size in bytes.
- memory_pool_device: `"cpu"` (pinned) or `"cuda"` (GPUDirect / XGMI RDMA).
For more details, refer to the
[Mori repository](https://github.com/ROCm/mori).
+2 -2
View File
@@ -305,8 +305,8 @@ class GPUWorker:
)
# Load model
model_loader = DiffusersPipelineLoader(load_config)
self.pipeline = model_loader.load_model(od_config, load_device=f"cuda:{rank}")
model_loader = DiffusersPipelineLoader(load_config, self.od_config)
self.pipeline = model_loader.load_model(load_device=f"cuda:{rank}")
# Setup cache backend
from vllm_omni.diffusion.cache.selector import get_cache_backend
@@ -45,7 +45,7 @@ If you do not need to modify source code of vLLM, you can directly install the s
uv pip install vllm==0.22.0 --torch-backend=auto
```
The 0.22.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
The 0.22.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -70,8 +70,8 @@ git checkout v0.22.0
```
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
```bash
#For CUDA 13.0 (the default for v0.22.0; the wheel filename has no `+cu130` suffix)
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_35_x86_64.whl
#For CUDA 12.9 (the default for v0.22.0; the wheel filename has no `+cu129` suffix)
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_28_x86_64.whl
```
Install vllm with command below (If you have no existing PyTorch).
```bash
+1 -1
View File
@@ -22,7 +22,7 @@ source .venv/bin/activate
uv pip install vllm==0.22.0 --torch-backend=auto
# On ROCm
uv pip install vllm==0.22.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm721
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
+10 -1
View File
@@ -33,7 +33,7 @@ th {
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, T2V with sound | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
@@ -50,6 +50,8 @@ th {
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | |
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
@@ -61,7 +63,12 @@ th {
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | |
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
| `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` | ✅︎ | ✅︎ | | |
@@ -69,11 +76,13 @@ th {
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | |
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
+1 -1
View File
@@ -116,7 +116,7 @@ Content-Type: application/json
| `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
**Supported languages:** Only applicable to Qwen3-TTS. Derived from the model configuration (`talker_config.codec_language_id` in the checkpoint's `config.json`), plus `Auto`, which is always accepted. Official Qwen3-TTS checkpoints support: Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian.
#### Voice Clone Parameters (Base task)
+44 -3
View File
@@ -67,8 +67,9 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input_reference` | file | null | Uploaded reference image for image-to-video requests |
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` |
| `input_reference` | file | null | Uploaded reference image or video for image-to-video/video-to-video requests |
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` or `video_reference` |
| `video_reference` | string | null | JSON-encoded reference video payload; do not combine with `input_reference` or `image_reference` |
| `width` | integer | model default | Output video width |
| `height` | integer | model default | Output video height |
| `num_frames` | integer | model default | Number of generated frames |
@@ -80,6 +81,8 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
| `flow_shift` | number | null | Scheduler flow-shift value |
| `true_cfg_scale` | number | null | True CFG scale when supported by the model |
| `seed` | integer | null | Random seed for reproducibility |
| `generate_sound` | boolean | false | Request model-generated audio for video models that support sound generation |
| `sound_duration` | number | null | Duration in seconds for generated audio; defaults to generated video duration |
| `negative_prompt` | string | null | Text describing what to avoid in the generated video |
| `enable_frame_interpolation` | boolean | null | Enable post-generation frame interpolation |
| `frame_interpolation_exp` | integer | null | Interpolation exponent; `1=2x`, `2=4x`, and so on |
@@ -123,6 +126,41 @@ curl -s http://localhost:8091/v1/videos \
-F "fps=16"
```
### Video-to-Video
For models that support video conditioning, upload the reference video with
`input_reference`:
```bash
curl -s http://localhost:8091/v1/videos \
-F "prompt=continue this motion with consistent subjects and lighting" \
-F "input_reference=@input.mp4;type=video/mp4" \
-F "width=1280" \
-F "height=720" \
-F "num_frames=80" \
-F "fps=16"
```
You can also pass a JSON-safe video URL or `data:video/...;base64,...` payload
through `video_reference`. Do not send `video_reference` together with
`input_reference` or `image_reference`.
```bash
curl -s http://localhost:8091/v1/videos \
-F "prompt=continue this motion with consistent subjects and lighting" \
-F 'video_reference={"video_url":"https://example.com/input.mp4"}' \
-F "width=1280" \
-F "height=720" \
-F "num_frames=80" \
-F "fps=16"
```
JSON references currently support `image_url`/`video_url`; `file_id` references
are not implemented yet. Models may expose additional V2V controls through
`extra_params`. For example, Cosmos3 supports
`condition_frame_indexes_vision` and `condition_video_keep` to select which
decoded reference frames are used as clean conditioning.
### Synchronous Generation
```bash
@@ -146,7 +184,10 @@ export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
## Model-Specific Examples
For complete text-to-video and image-to-video walkthroughs, see:
For complete text-to-video, image-to-video, and model-specific video-to-video
walkthroughs, see:
- [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md)
- [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md)
- [Cosmos3 recipes](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md)
for model-specific video-to-video examples and conditioning controls
+2 -1
View File
@@ -145,9 +145,10 @@ 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.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+44
View File
@@ -28,6 +28,50 @@ in deep DiT blocks.
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
guide. FP8 on Ampere may use a weight-only path where available.
### Faster FP8 GEMM on Blackwell (quack)
On Blackwell (SM 100+), vLLM runs FP8 linears through the FlashInfer kernel, which
applies the bias as a separate kernel after the GEMM. On the small GEMMs in video
DiTs this bias add is a significant overhead. Installing the optional `quack` kernel
lets vLLM-Omni fuse `alpha * (A @ B) + bias` into a single CuteDSL GEMM, recovering
that overhead (e.g. HunyuanVideo-1.5 FP8 goes from slower-than-BF16 to faster).
```bash
# CUDA 12.9
pip install vllm-omni[quack]
# CUDA 13.x
pip install 'quack-kernels[cu13]' --extra-index-url https://download.pytorch.org/whl/cu130
```
It is enabled automatically once installed (no flag needed) and is **Blackwell-only**:
on Hopper/Ada the CUTLASS FP8 kernel already fuses bias, so quack is not used there.
Set `VLLM_OMNI_USE_QUACK_FP8=0` to force the FlashInfer path. If `quack-kernels` is
not installed, FP8 still works — it just keeps the unfused FlashInfer path.
#### Compile cache and warmup
quack JIT-compiles its kernel once per distinct GEMM shape (tens of seconds, longer
the first time across all autotuned configs). The compiled `.o` files are cached on
disk and reused on later runs, so this is a one-time cost — **not** per request.
vLLM-Omni points that cache at `~/.cache/vllm_omni/quack` (override with
`QUACK_CACHE_DIR`) instead of quack's default under `/tmp`, so it survives restarts.
In containers, set `QUACK_CACHE_DIR` to a mounted/persistent path — or bake it into
the image — so the first cold start does not recompile. The engine's startup dummy
run already exercises the kernels, so with a warm cache the first real request is fast.
To pre-warm specific shapes (e.g. at image build time):
```python
from vllm_omni.quantization.quack_fp8 import warmup_quack_fp8
# (M, K, N) per linear; M = number of tokens for your resolution/frame count
warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
```
> The PyPI package is `quack-kernels` (imported as `quack`); plain `pip install
> quack` is an unrelated statistics library. Requires CUDA 12.9+ and Python 3.12.
## Model Type Support
### Diffusion Model (Qwen-Image, Wan2.2)
@@ -0,0 +1,14 @@
{
"prompt": "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. The robotic arm is in the process of pouring a transparent liquid into the glass. To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. The lighting is bright, casting soft shadows on the table. The robotic arm's movements are smooth and controlled, demonstrating precision in its task. As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. The other robotic arm on the right side also stays stationary throughout the video. The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper.",
"negative_prompt": "",
"vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4",
"height": 720,
"width": 1280,
"num_frames": 189,
"num_inference_steps": 35,
"guidance_scale": 6.0,
"flow_shift": 10.0,
"fps": 24,
"condition_frame_indexes_vision": [0, 1],
"condition_video_keep": "first"
}
@@ -143,6 +143,20 @@ def parse_args() -> argparse.Namespace:
default=None,
required=False,
)
parser.add_argument(
"--width",
type=int,
default=None,
metavar="W",
help="Output image width in pixels. Default: None (pipeline's default).",
)
parser.add_argument(
"--height",
type=int,
default=None,
metavar="H",
help="Output image height in pixels. Default: None (pipeline's default).",
)
parser.add_argument(
"--seed",
type=int,
@@ -236,8 +250,8 @@ def parse_args() -> argparse.Namespace:
parser.add_argument(
"--resolution",
type=int,
default=640,
help="Bucket in (640, 1024) to determine the condition and output resolution",
default=None,
help="Bucket in (640, 1024) to determine the condition and output resolution. If width and height are not provided, this will be set to default 640.",
)
parser.add_argument(
@@ -384,6 +398,15 @@ def parse_args() -> argparse.Namespace:
def main():
args = parse_args()
if args.resolution and (args.width or args.height):
raise ValueError("--resolution and --width/--height cannot be specified together")
if args.width is not None and args.width <= 0:
raise ValueError("--width must be a positive integer")
if args.height is not None and args.height <= 0:
raise ValueError("--height must be a positive integer")
if not args.width and not args.height and not args.resolution:
args.resolution = 640
# Validate input images exist and load them
input_images = []
for image_path in args.image:
@@ -464,6 +487,7 @@ def main():
print(
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}"
)
print(f" Output image size: {args.width or 'default'}x{args.height or 'default'}")
print(f"{'=' * 60}\n")
generation_start = time.perf_counter()
@@ -473,22 +497,26 @@ def main():
omni.start_profile()
# Generate edited image
sampling_params = OmniDiffusionSamplingParams(
width=args.width,
height=args.height,
generator=generator,
true_cfg_scale=args.cfg_scale,
guidance_scale=args.guidance_scale,
guidance_scale_2=args.guidance_scale_2,
num_inference_steps=args.num_inference_steps,
num_outputs_per_prompt=args.num_outputs_per_prompt,
layers=args.layers,
)
if args.resolution:
sampling_params.resolution = args.resolution
outputs = omni.generate(
{
"prompt": args.prompt,
"negative_prompt": args.negative_prompt,
"multi_modal_data": {"image": input_image},
},
OmniDiffusionSamplingParams(
generator=generator,
true_cfg_scale=args.cfg_scale,
guidance_scale=args.guidance_scale,
guidance_scale_2=args.guidance_scale_2,
num_inference_steps=args.num_inference_steps,
num_outputs_per_prompt=args.num_outputs_per_prompt,
layers=args.layers,
resolution=args.resolution,
),
sampling_params,
)
generation_end = time.perf_counter()
generation_time = generation_end - generation_start
@@ -16,9 +16,10 @@ Usage:
"""
import argparse
import os
import json
import time
from pathlib import Path
from typing import Any
import torch
@@ -29,6 +30,16 @@ from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
def parse_profiler_config(value: str) -> dict[str, Any]:
try:
config = json.loads(value)
except json.JSONDecodeError as e:
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
if not isinstance(config, dict):
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
return config
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Generate a talking-head video from a reference image and audio (Wan2.2 S2V)."
@@ -87,6 +98,12 @@ def parse_args() -> argparse.Namespace:
default=3.0,
help="Scheduler flow shift (default: 3.0).",
)
parser.add_argument(
"--boundary-ratio",
type=float,
default=None,
help="Boundary split ratio for low/high DiT stages (S2V). Default varies by model.",
)
parser.add_argument(
"--output",
type=str,
@@ -178,30 +195,10 @@ def parse_args() -> argparse.Namespace:
help="Enable cache-dit summary logging after diffusion forward passes.",
)
parser.add_argument(
"--profile-dir",
type=str,
"--profiler-config",
type=parse_profiler_config,
default=None,
help="Enable torch profiler and save traces to this directory.",
)
parser.add_argument(
"--profile-record-shapes",
action="store_true",
help="Record tensor shapes in profiler (increases trace size).",
)
parser.add_argument(
"--profile-with-stack",
action="store_true",
help="Record stack traces in profiler (increases overhead).",
)
parser.add_argument(
"--profile-with-memory",
action="store_true",
help="Profile memory usage.",
)
parser.add_argument(
"--profile-with-flops",
action="store_true",
help="Estimate FLOPs for operations.",
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
return parser.parse_args()
@@ -239,23 +236,9 @@ def main():
ulysses_degree=args.ulysses_degree,
)
# Check if profiling is requested via CLI args or environment variable
profile_dir = args.profile_dir or os.getenv("VLLM_TORCH_PROFILER_DIR")
profiler_enabled = bool(profile_dir)
profiler_config = None
if profiler_enabled:
from vllm.config import ProfilerConfig
profiler_enabled = args.profiler_config is not None
profiler_config = ProfilerConfig(
profiler="torch",
torch_profiler_dir=profile_dir,
torch_profiler_record_shapes=args.profile_record_shapes,
torch_profiler_with_stack=args.profile_with_stack,
torch_profiler_with_memory=args.profile_with_memory,
torch_profiler_with_flops=args.profile_with_flops,
)
omni = Omni(
omni_kwargs = dict(
model=args.model,
model_class_name="WanS2VPipeline",
flow_shift=args.flow_shift,
@@ -269,8 +252,12 @@ def main():
cache_backend=args.cache_backend,
cache_config=cache_config,
enable_cache_dit_summary=args.enable_cache_dit_summary,
profiler_config=profiler_config,
profiler_config=args.profiler_config,
)
if args.boundary_ratio is not None:
omni_kwargs["boundary_ratio"] = args.boundary_ratio
omni = Omni(**omni_kwargs)
# Print generation configuration
print(f"\n{'=' * 60}")
@@ -325,13 +312,16 @@ def main():
if profiler_enabled:
print("\n[Profiler] Stopping profiler and collecting results...")
profile_results = omni.stop_profile()
if profile_results and isinstance(profile_results, list):
if profile_results and isinstance(profile_results, dict):
traces = profile_results.get("traces", [])
print("\n" + "=" * 60)
print("PROFILING RESULTS:")
for i, result in enumerate(profile_results):
print(f"\nStage {i}:")
if result:
print(f" {result}")
for rank, trace in enumerate(traces):
print(f"\nRank {rank}:")
if trace:
print(f" • Trace: {trace}")
if not traces:
print(" No traces collected.")
print("=" * 60)
else:
print("[Profiler] No valid profiling data returned.")
@@ -17,6 +17,7 @@ list of supported architectures across all modalities, see
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz |
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
| MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | single (AR + codec) | ✓ (required) | ✓ | voice_clone, continuation | 48 kHz |
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
@@ -159,6 +160,46 @@ Streaming requires `async_chunk: true` in the stage config.
---
## Ming-omni-tts
Dense 0.5B two-stage TTS pipeline (`AR + flow` + audio VAE) at 44.1 kHz. The example covers style, IP voice, music-only generation, text-to-audio events, emotion, dialect, zero-shot cloning, podcast, speech+BGM, and speech+environment-sound cases.
### Quick start
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case style \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
### Voice cloning
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case zero_shot \
--ref-audio /path/to/reference.wav \
--ref-text "在此奉劝大家别乱打美白针。" \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
### Streaming
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case basic \
--ref-audio /path/to/reference.wav \
--streaming \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
### Notes
- `style`, `ip`, `bgm`, and `tta` do not require reference audio.
- Reference-audio cases use `--ref-audio`; `zero_shot` also requires `--ref-text`.
- `podcast` uses multiple references via `--ref-audio-paths`.
- Full case details live in [`ming_tts/README.md`](ming_tts/README.md).
---
## Ming-flash-omni-TTS
Standalone talker-only deployment of Ming-flash-omni-2.0 at 44.1 kHz. Voice is controlled through caption fields (`风格` / `IP` / `语速`/`基频`/`音量`) rather than reference audio.
@@ -0,0 +1,204 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline higgs-audio v3 inference example.
Runs Stage 0 (Qwen3 talker) + Stage 1 (HiggsAudio codec) end-to-end
through the vLLM-Omni engine without going through the HTTP server, and
saves a 24 kHz mono WAV per prompt.
Example (plain TTS):
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
--texts "Hello world." \\
--output-dir results/higgs_v3_wavs
Example (voice clone):
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
--texts "Hello world." \\
--ref-audio path/to/reference.wav \\
--ref-text "Transcript of the reference clip." \\
--output-dir results/higgs_v3_wavs
"""
from __future__ import annotations
import os
os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
import re
import time
from pathlib import Path
import numpy as np
import soundfile as sf
import torch
from vllm_omni import Omni
SAMPLE_RATE = 24_000
DEFAULT_TEXTS = (
"Hello world.",
"The quick brown fox jumps over the lazy dog.",
"Today is a beautiful day for a walk in the park.",
)
def parse_args():
import argparse
parser = argparse.ArgumentParser(description="Offline higgs-audio v3 inference")
parser.add_argument(
"--model",
type=str,
default="bosonai/higgs-audio-v3-tts-4b",
help="Stage-0 talker model id or path.",
)
parser.add_argument(
"--texts",
type=str,
nargs="+",
default=list(DEFAULT_TEXTS),
help="One or more plain-text prompts to synthesize.",
)
parser.add_argument(
"--output-dir",
type=str,
default="results/higgs_v3_wavs",
help="Directory to write per-prompt WAV files.",
)
parser.add_argument(
"--deploy-config",
type=str,
default=None,
help="Override the deploy config path.",
)
parser.add_argument(
"--max-new-tokens",
type=int,
default=2048,
help="Cap on Stage-0 codec frames per request.",
)
parser.add_argument(
"--ref-audio",
type=str,
default=None,
help="Reference audio for voice clone (WAV/FLAC/MP3 path).",
)
parser.add_argument(
"--ref-text",
type=str,
default=None,
help="Transcript of the reference audio. Optional but improves fidelity.",
)
return parser.parse_args()
def _slugify(text: str) -> str:
slug = re.sub(r"\s+", "_", text.strip().lower())
slug = re.sub(r"[^a-z0-9_]+", "", slug)
return slug[:48] or "prompt"
def _extract_pcm(multimodal_output: dict) -> torch.Tensor:
audio = multimodal_output.get("model_outputs")
if audio is None:
audio = multimodal_output.get("audio")
if audio is None:
raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
if isinstance(audio, list):
valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
if not valid:
raise ValueError("audio list is empty")
return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
return torch.as_tensor(audio).float().cpu().reshape(-1)
def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray:
arr = pcm.numpy()
if arr.dtype.kind == "f":
arr = np.clip(arr, -1.0, 1.0)
arr = (arr * 32767.0).astype(np.int16)
else:
arr = arr.astype(np.int16)
return arr
def main():
args = parse_args()
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
from transformers import AutoTokenizer
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
HiggsAudioV3TokenizerAdapter,
apply_delay_pattern,
encode_reference_audio,
)
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
adapter = HiggsAudioV3TokenizerAdapter(tokenizer)
# Load and encode reference audio once (if voice cloning)
ref_codes_delayed: torch.Tensor | None = None
if args.ref_audio is not None:
ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
if ref_wav.ndim == 2:
ref_wav = ref_wav.mean(axis=1)
ref_codes_raw = encode_reference_audio(ref_wav, int(ref_sr))
ref_codes_delayed = apply_delay_pattern(ref_codes_raw)
print(f"Reference : {args.ref_audio}")
print(f"Ref codes : {ref_codes_raw.shape[0]} frames -> {ref_codes_delayed.shape[0]} delayed")
if args.ref_text:
print(f"Ref text : {args.ref_text}")
print(f"Model : {args.model}")
print(f"Prompts : {len(args.texts)}")
print(f"Output dir : {output_dir}")
print(f"Voice clone : {'yes' if ref_codes_delayed is not None else 'no'}")
total_elapsed = 0.0
total_dur = 0.0
for text in args.texts:
if ref_codes_delayed is not None:
prompt_ids = adapter.build_prompt(
text,
num_ref_tokens=int(ref_codes_delayed.shape[0]),
reference_text=args.ref_text,
)
prompt = {
"prompt_token_ids": prompt_ids,
"additional_information": {
"audio_input_ids": ref_codes_delayed.to(torch.long),
"audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool),
},
}
else:
prompt_ids = adapter.build_prompt(text)
prompt = {"prompt_token_ids": prompt_ids}
t_start = time.perf_counter()
outputs = engine.generate([prompt])
elapsed = time.perf_counter() - t_start
total_elapsed += elapsed
mm = outputs[0].outputs[0].multimodal_output
pcm = _extract_pcm(mm)
slug = _slugify(text)
out_path = output_dir / f"{slug}.wav"
sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16")
dur = pcm.numel() / SAMPLE_RATE
total_dur += dur
print(f" {slug:<50} dur={dur:5.2f}s -> {out_path}")
rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
print(f"Total infer : {total_elapsed:.2f}s total audio: {total_dur:.2f}s RTF: {rtf:.3f}")
if __name__ == "__main__":
main()
@@ -0,0 +1,248 @@
# Ming-omni-tts Offline Inference
`end2end.py` runs Ming dense 0.5B end to end with vLLM-Omni. It uses the in-repo Ming prompt assembly helper directly, so the example request shape matches the real integration instead of a simplified wrapper.
## Files
| File | Purpose |
|---|---|
| `end2end.py` | Driver: CLI, case loading, prompt construction, orchestration (~150 lines) |
| `cases.yaml` | All 11 built-in case definitions (prompt, text, instruction, ref-audio flags, flow controls) |
| `runner.py` | Engine management and audio output (streaming + blocking paths) |
## Model Overview
Ming dense 0.5B is exposed here as a two-stage offline pipeline:
- **Stage 0**: Qwen2-based AR generation with Ming prompt formatting and inline flow controls
- **Stage 1**: audio VAE decode to mono 44.1 kHz waveform
`config_ming_tts.py` adapts the checkpoint's HuggingFace config fields
(LLM, DiT, aggregator, AudioVAE, token ids). `vllm_omni/deploy/ming_tts.yaml`
selects the vLLM-Omni pipeline and stage runtime topology, including
connectors, async chunking, memory limits, and sampling defaults.
The example supports both:
- **Blocking eager** via `vllm_omni/deploy/ming_tts.yaml`
- **Async chunk eager** via `vllm_omni/deploy/ming_tts.yaml` (default `async_chunk: true`)
## Setup
Install vLLM-Omni with the platform requirements for your accelerator:
```bash
uv pip install -e .
```
The Ming offline example does not require a separate upstream Ming package.
Reference-audio cases use the repo dependencies for audio loading,
resampling, and CampPlus speaker extraction, including `soundfile`,
`torchaudio`, and `onnxruntime-rocm` in the documented ROCm environment.
The tested ROCm environment is summarized in the
[Ming recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/inclusionAI/Ming-omni-tts-0.5B.md).
## Supported Cases
These cases cover the upstream dense 0.5B cookbook surface that maps cleanly onto the current vLLM-Omni example:
- `style`: zero-speaker style-conditioned speech
- `ip`: zero-speaker IP voice generation
- `bgm`: music-only generation
- `tta`: text-to-audio event generation with FlowLoss controls
- `emotion`: reference-audio speech with emotion control
- `basic`: reference-audio speech with speed / pitch / volume control
- `dialect`: reference-audio speech with dialect control
- `zero_shot`: reference-audio cloning with explicit transcript
- `podcast`: multi-reference dialogue generation with automatic speaker embedding extraction
- `speech_bgm`: speech with background music conditioning
- `speech_sound`: speech with environmental sound conditioning
## Quick Start
Run the zero-speaker style example:
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case style \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
Run zero-shot cloning with a transcript:
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case zero_shot \
--ref-audio /path/to/10002287-00000094.wav \
--ref-text "在此奉劝大家别乱打美白针。" \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
Run emotion-controlled speech:
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case emotion \
--ref-audio /path/to/emotion_prompt.wav \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
Run podcast generation with two reference clips:
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case podcast \
--ref-audio-paths /path/to/CTS-CN-F2F-2019-11-11-423-012-A.wav /path/to/CTS-CN-F2F-2019-11-11-423-012-B.wav \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
The script automatically extracts one 192-d speaker embedding per reference WAV using the Ming model's `campplus.onnx`.
If you already have precomputed multi-speaker embeddings, you can override extraction with:
```bash
--speaker-embedding /path/to/podcast_speaker_embeddings.json
```
where the JSON is a list of speaker embeddings, one 192-d vector per speaker.
Run text-to-audio event generation:
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case tta \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
Use async_chunk streaming:
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case basic \
--ref-audio /path/to/10002287-00000095.wav \
--streaming \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
`--streaming` uses `AsyncOmni` and the async_chunk stage config. It currently
supports one prompt per process invocation; use blocking mode for
`--num-prompts > 1`.
Collect runtime stats and a manifest:
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case style \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager \
--enable-stats \
--stats-log-file output_audio/ming_style_pipeline.log \
--metadata-json output_audio/ming_style_manifest.json
```
## Reference Fixtures
The upstream Ming cookbook uses these public audio fixtures from `inclusionAI/Ming-omni-tts/data/wavs`:
- `10002287-00000094.wav` for zero-shot cloning
- `10002287-00000095.wav` for `basic`
- `emotion_prompt.wav` for `emotion`
- `yue_prompt.wav` for `dialect`
- `00000309-00000300.wav` for `speech_bgm` and `speech_sound`
- `CTS-CN-F2F-2019-11-11-423-012-A.wav` and `CTS-CN-F2F-2019-11-11-423-012-B.wav` for `podcast`
## Validation Matrix
The repo-facing example is intended to cover the same dense TTS workflows used
by the local Ming validation script:
| Case | Blocking `deploy/ming_tts.yaml` | Async chunk `deploy/ming_tts.yaml` | Extra inputs |
|---|---:|---:|---|
| `style` | Yes | Optional smoke test | none |
| `ip` | Yes | Optional smoke test | none |
| `bgm` | Yes | Optional smoke test | none |
| `tta` | Yes | Optional smoke test | none |
| `emotion` | Yes | Yes | `--ref-audio emotion_prompt.wav` |
| `basic` | Yes | Yes | `--ref-audio 10002287-00000095.wav` |
| `dialect` | Yes | Yes | `--ref-audio yue_prompt.wav` |
| `zero_shot` | Yes | Yes | `--ref-audio 10002287-00000094.wav --ref-text ...` |
| `podcast` | Yes | Yes | two `--ref-audio-paths` |
| `speech_bgm` | Yes | Yes | `--ref-audio 00000309-00000300.wav` |
| `speech_sound` | Yes | Yes | `--ref-audio 00000309-00000300.wav` |
## Validated Outputs
The following measurements are retained from an earlier L4 CUDA validation;
they are not ROCm benchmark results. Default async_chunk matched blocking
output frame counts and Stage-1 patch counts for every case:
| Case | Blocking frames / patches / sec | Async chunk frames / patches / sec |
|---|---:|---:|
| `style` | 409248 / 29 / 9.28 | 409248 / 29 / 9.28 |
| `ip` | 183456 / 13 / 4.16 | 183456 / 13 / 4.16 |
| `bgm` | 1326528 / 94 / 30.08 | 1326528 / 94 / 30.08 |
| `tta` | 465696 / 33 / 10.56 | 465696 / 33 / 10.56 |
| `emotion` | 324576 / 23 / 7.36 | 324576 / 23 / 7.36 |
| `basic` | 211680 / 15 / 4.80 | 211680 / 15 / 4.80 |
| `dialect` | 239904 / 17 / 5.44 | 239904 / 17 / 5.44 |
| `zero_shot` | 409248 / 29 / 9.28 | 409248 / 29 / 9.28 |
| `podcast` | 437472 / 31 / 9.92 | 437472 / 31 / 9.92 |
| `speech_bgm` | 296352 / 21 / 6.72 | 296352 / 21 / 6.72 |
| `speech_sound` | 352800 / 25 / 8.00 | 352800 / 25 / 8.00 |
## Key Arguments
| Argument | Description |
|---|---|
| `--model` | Hugging Face repo or local Ming checkpoint path |
| `--deploy-config` | Deploy config YAML. Use `vllm_omni/deploy/ming_tts.yaml` |
| `--case` | Built-in demo case |
| `--ref-audio` | Single reference wav path for cloning-style cases |
| `--ref-audio-paths` | Multiple reference wav paths, used by `podcast` |
| `--ref-text` | Reference transcript. Required for `zero_shot` |
| `--instructions` | Free-form Ming instruction string |
| `--instruction-json` | Structured Ming instruction JSON |
| `--speaker-embedding` | JSON file containing a 192-d speaker embedding |
| `--extract-speaker-embeddings` | Force CampPlus speaker extraction from the provided reference audio paths |
| `--max-decode-steps` | Override `ming_max_decode_steps` |
| `--num-prompts` | Repeat the same case N times. Outputs are indexed when `N > 1` |
| `--streaming` | Use `AsyncOmni` and async_chunk transport |
| `--enforce-eager` | Recommended for Ming dense; non-eager is out of scope |
| `--enable-stats` / `--log-stats` | Enable vLLM-Omni per-request stats logging |
| `--stats-log-file` | Optional path for the stats log |
| `--metadata-json` | Optional path for the run manifest JSON |
| `--stage-init-timeout` | Per-stage initialization timeout in seconds |
| `--init-timeout` | Total initialization timeout in seconds |
| `--batch-timeout` | Batch timeout in seconds |
| `--worker-backend` | `multi_process` or `ray` |
| `--ray-address` | Ray cluster address when using `--worker-backend ray` |
## Output
- The script writes one mono 44.1 kHz WAV file per run
- Default output directory: `output_audio/`
- Default filename: `ming_<case>.wav`
- When `--num-prompts > 1`, outputs are indexed as `ming_<case>_00000.wav`, `..._00001.wav`, etc.
- When stats are enabled, the script can also write:
- a stats log file such as `ming_style_pipeline.log`
- a manifest JSON with per-output metadata, stage durations, peak memory info,
and streaming client latency metrics when `--streaming` is used
## Notes
- `style` and `ip` are zero-speaker paths and do not require a reference clip
- `emotion`, `basic`, `dialect`, `speech_bgm`, and `speech_sound` require one reference clip
- `zero_shot` requires both `--ref-audio` and `--ref-text`
- `podcast` requires at least two reference clips via `--ref-audio-paths`
- `podcast` automatically extracts one speaker embedding per reference clip
- `--speaker-embedding` may contain either one 192-d vector or a list of 192-d vectors
- `--enforce-eager` was used for the validated runs
- The earlier L4 validation used SDPA for the Ming audio VAE instead of
FlashAttention2, which is the preferred default when available.
@@ -0,0 +1,110 @@
style:
prompt: "Please generate speech based on the following description.\n"
text: "我会一直在这里陪着你,直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗?"
instruction:
风格: >-
这是一种ASMR耳语,属于一种旨在引发特殊感官体验的创意风格。这个女性使用轻柔的普通话进行耳语,声音气音成分重。音量极低,紧贴麦克风,语速极慢,旨在制造触发听者颅内快感的声学刺激。
use_zero_spk_emb: true
max_decode_steps: 200
ip:
prompt: "Please generate speech based on the following description.\n"
text: "这款产品的名字,叫变态坑爹牛肉丸。"
instruction:
IP: "灵小甄"
use_zero_spk_emb: true
max_decode_steps: 200
bgm:
prompt: "Please generate music based on the following description.\n"
text: " Genre: 电子舞曲. Mood: 自信 / 坚定. Instrument: 架子鼓. Theme: 节日. Duration: 30s."
instruction: null
use_zero_spk_emb: false
max_decode_steps: 400
tta:
prompt: "Please generate audio events based on given text.\n"
text: "Thunder and a gentle rain"
instruction: null
use_zero_spk_emb: false
max_decode_steps: 200
cfg: 4.5
sigma: 0.3
temperature: 2.5
emotion:
prompt: "Please generate speech based on the following description.\n"
text: "我竟然抢到了陈奕迅的演唱会门票!太棒了!终于可以现场听一听他的歌声了!"
instruction:
情感: "高兴"
requires_ref_audio: true
auto_extract_speaker_embeddings: true
max_decode_steps: 200
basic:
prompt: "Please generate speech based on the following description.\n"
text: "简单地说,这相当于惠普把消费领域市场拱手相让了。"
instruction:
语速: "快速"
基频: "中"
音量: "中"
requires_ref_audio: true
auto_extract_speaker_embeddings: true
max_decode_steps: 200
dialect:
prompt: "Please generate speech based on the following description.\n"
text: "我觉得社会企业同个人都有责任"
instruction:
方言: "广粤话"
requires_ref_audio: true
auto_extract_speaker_embeddings: true
max_decode_steps: 200
zero_shot:
prompt: "Please generate speech based on the following description.\n"
text: "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。"
instruction: null
requires_ref_audio: true
requires_ref_text: true
auto_extract_speaker_embeddings: true
max_decode_steps: 200
podcast:
prompt: "Please generate speech based on the following description.\n"
text: " speaker_1:你可以说一下,就大概说一下,可能虽然我也不知道,我看过那部电影没有。\n speaker_2:就是那个叫什么,变相一节课的嘛。\n speaker_1:嗯。\n speaker_2:一部搞笑的电影。\n speaker_1:一部搞笑的。\n"
instruction: null
prompt_text: " speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试,做个,当服务员还要去笔试了\n speaker_2:对啊,这真的很奇怪,就是 单纯的因,单纯自己工资不高,只是因为可能人家那个店比较出名一点,就对你苛刻要求\n"
requires_ref_audio_count: 2
auto_extract_speaker_embeddings: true
max_decode_steps: 200
speech_bgm:
prompt: "Please generate speech based on the following description.\n"
text: "此次业绩下滑原因,可归结为企业停止服务某些品牌,而带来的负面影响。"
instruction:
BGM:
Genre: "当代古典音乐."
Mood: "温暖 / 友善."
Instrument: "电吉他"
Theme: "节日."
SNR: 10.0
ENV: null
requires_ref_audio: true
auto_extract_speaker_embeddings: true
max_decode_steps: 200
speech_sound:
prompt: "Please generate speech based on the following description.\n"
text: "此次业绩下滑原因,可归结为企业停止服务某些品牌,而带来的负面影响。"
instruction:
BGM:
ENV: "Birds chirping"
SNR: 10.0
Genre: null
Mood: null
Instrument: null
Theme: null
requires_ref_audio: true
auto_extract_speaker_embeddings: true
max_decode_steps: 200
@@ -0,0 +1,200 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline inference demo for Ming-omni-tts via vLLM Omni."""
import json
import os
from pathlib import Path
import soundfile as sf
import torch
import torchaudio
import yaml
from transformers import AutoTokenizer
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
KEY_CFG,
KEY_MAX_DECODE_STEPS,
KEY_SIGMA,
KEY_TEMPERATURE,
SAMPLE_RATE,
)
from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt
from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor
try:
from .runner import (
build_manifest,
build_sampling_params,
resolve_metadata_json,
resolve_stats_log_file,
run_generation,
)
except ImportError:
from runner import (
build_manifest,
build_sampling_params,
resolve_metadata_json,
resolve_stats_log_file,
run_generation,
)
_DEFAULT_MODEL = "inclusionAI/Ming-omni-tts-0.5B"
_REPO_ROOT = Path(__file__).resolve().parents[4]
_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni/deploy/ming_tts.yaml")
_CASES_FILE = Path(__file__).with_name("cases.yaml")
CASE_DEFAULTS = yaml.safe_load(_CASES_FILE.read_text(encoding="utf-8")) or {}
if not CASE_DEFAULTS:
raise RuntimeError(f"Empty or missing case definitions in {_CASES_FILE}")
def _build_parser() -> FlexibleArgumentParser:
p = FlexibleArgumentParser(description="Offline Ming-omni-tts example")
p.add_argument("--model", default=_DEFAULT_MODEL, help="Model name or local path")
p.add_argument("--deploy-config", default=None, help="Deploy config YAML; auto-selected when omitted")
p.add_argument("--case", choices=sorted(CASE_DEFAULTS), default="style", help="Built-in demo case")
p.add_argument("--text", default=None, help="Override case text")
p.add_argument("--prompt", default=None, help="Override the system prompt prefix")
p.add_argument("--instructions", default=None, help="Free-form Ming instruction string")
p.add_argument(
"--instruction-json", default=None, help='Structured Ming instruction JSON, e.g. \'{"方言":"广粤话"}\''
)
p.add_argument("--ref-audio", default=None, help="Single reference audio path for cloning-style cases")
p.add_argument("--ref-audio-paths", nargs="+", default=None, help="Multiple reference audio paths (podcast)")
p.add_argument("--ref-text", default=None, help="Reference transcript; required for zero_shot")
p.add_argument("--speaker-embedding", default=None, help="Path to a JSON speaker embedding file")
p.add_argument(
"--extract-speaker-embeddings",
action="store_true",
help="Extract speaker embeddings from ref audio via campplus.onnx",
)
p.add_argument("--max-decode-steps", type=int, default=None, help="Override ming_max_decode_steps")
p.add_argument("--output-dir", default="output_audio", help="Directory for output wav files")
p.add_argument("--output-name", default=None, help="Output wav filename")
p.add_argument("--num-prompts", type=int, default=1, help="Repeat the same prompt N times")
p.add_argument("--streaming", action="store_true", help="Use AsyncOmni with async_chunk streaming")
p.add_argument("--trust-remote-code", action="store_true")
p.add_argument("--enforce-eager", action="store_true")
p.add_argument(
"--log-stats", "--enable-stats", dest="log_stats", action="store_true", help="Enable Omni stats logging"
)
p.add_argument("--stats-log-file", default=None, help="Path for the Omni stats log file")
p.add_argument("--metadata-json", default=None, help="Path for the run manifest JSON")
p.add_argument("--stage-init-timeout", type=int, default=300, help="Per-stage init timeout (s)")
p.add_argument("--init-timeout", type=int, default=600, help="Total init timeout (s)")
p.add_argument("--batch-timeout", type=int, default=5, help="Batch timeout (s)")
p.add_argument("--shm-threshold-bytes", type=int, default=65536)
p.add_argument("--worker-backend", default="multi_process", choices=["multi_process", "ray"])
p.add_argument("--ray-address", default=None, help="Ray cluster address (--worker-backend ray)")
return p
def _finalize_args(args) -> None:
if args.instructions is not None and args.instruction_json is not None:
raise RuntimeError("Use either --instructions or --instruction-json, not both")
if args.num_prompts < 1:
raise RuntimeError("--num-prompts must be at least 1")
if args.streaming and args.num_prompts != 1:
raise RuntimeError("--streaming currently supports exactly one prompt")
if args.deploy_config is None:
args.deploy_config = _DEFAULT_DEPLOY_CONFIG
def _load_waveform(path: str) -> torch.Tensor:
samples, sr = sf.read(path, dtype="float32")
wav = torch.as_tensor(samples, dtype=torch.float32)
if wav.ndim == 2:
wav = wav.mean(dim=1)
wav = wav.reshape(1, -1)
if int(sr) != SAMPLE_RATE:
wav = torchaudio.functional.resample(wav, int(sr), SAMPLE_RATE)
return wav
def _load_speaker_embedding(path: str) -> torch.Tensor:
return torch.as_tensor(json.loads(Path(path).read_text(encoding="utf-8")), dtype=torch.float32)
def _ref_audio_paths(args) -> list[str]:
if args.ref_audio is not None and args.ref_audio_paths is not None:
raise RuntimeError("Use either --ref-audio or --ref-audio-paths, not both")
if args.ref_audio_paths is not None:
return list(args.ref_audio_paths)
return [args.ref_audio] if args.ref_audio else []
def _resolve_reference_inputs(args, case: dict, paths: list[str]):
required = int(case.get("requires_ref_audio_count", 0))
if required > 0 and len(paths) < required:
raise RuntimeError(f"Case '{args.case}' needs {required} ref audio paths via --ref-audio-paths")
if required <= 0 and case.get("requires_ref_audio") and not paths:
raise RuntimeError(f"--ref-audio required for case '{args.case}'")
if not paths:
return None
wavs = [_load_waveform(p) for p in paths]
return wavs[0] if len(wavs) == 1 else wavs
def _resolve_speaker_embedding(args, case: dict, paths: list[str]):
if args.speaker_embedding:
return _load_speaker_embedding(args.speaker_embedding)
if not (case.get("auto_extract_speaker_embeddings") or args.extract_speaker_embeddings) or not paths:
return None
embs = MingSpeakerEmbeddingExtractor(args.model).extract_many(paths)
if not embs:
raise RuntimeError("Speaker extraction produced no embeddings")
return embs[0] if len(embs) == 1 else torch.stack(embs, dim=0)
def _build_prompt_payload(tokenizer, args):
case = CASE_DEFAULTS[args.case]
paths = _ref_audio_paths(args)
prompt_text = args.ref_text if args.ref_text is not None else case.get("prompt_text")
if case.get("requires_ref_text") and not prompt_text:
raise RuntimeError(f"--ref-text required for case '{args.case}'")
prompt_waveform = _resolve_reference_inputs(args, case, paths) if prompt_text is not None else None
speaker_embedding = _resolve_speaker_embedding(args, case, paths)
use_zero_spk_emb = bool(case.get("use_zero_spk_emb")) and prompt_waveform is None and speaker_embedding is None
runtime_controls = {KEY_MAX_DECODE_STEPS: args.max_decode_steps or case["max_decode_steps"]}
for key, field in [(KEY_CFG, "cfg"), (KEY_SIGMA, "sigma"), (KEY_TEMPERATURE, "temperature")]:
if field in case:
runtime_controls[key] = case[field]
instruction = (
json.loads(args.instruction_json) if args.instruction_json else (args.instructions or case.get("instruction"))
)
return build_ming_dense_prompt(
tokenizer,
prompt=args.prompt or case["prompt"],
text=args.text or case["text"],
runtime_controls=runtime_controls,
instruction=instruction,
prompt_text=prompt_text,
prompt_waveform=prompt_waveform,
speaker_embedding=speaker_embedding,
use_zero_spk_emb=use_zero_spk_emb,
)
def main():
args = _build_parser().parse_args()
_finalize_args(args)
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=False)
prompt_payload = _build_prompt_payload(tokenizer, args)
case = CASE_DEFAULTS[args.case]
sampling_params_list = build_sampling_params(args.max_decode_steps or case["max_decode_steps"])
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
stats_log_file = resolve_stats_log_file(args)
summaries = run_generation(args, prompt_payload, sampling_params_list, output_dir, stats_log_file)
metadata_json = resolve_metadata_json(args)
manifest = build_manifest(args, prompt_payload, stats_log_file, summaries)
if metadata_json is not None:
Path(metadata_json).write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
print(f"Saved run manifest to {metadata_json}")
if __name__ == "__main__":
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
main()
@@ -0,0 +1,277 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
import json
import time
import uuid
import wave
from pathlib import Path
import torch
from vllm import SamplingParams
from vllm_omni import AsyncOmni, Omni
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
KEY_SPEAKER_EMBEDDING,
SAMPLE_RATE,
TEXT_EOS_TOKEN_ID,
)
def coerce_audio_tensor(audio, *, async_chunk: bool) -> torch.Tensor:
if isinstance(audio, list):
parts = []
for item in audio:
tensor = torch.as_tensor(item, dtype=torch.float32).reshape(-1)
if tensor.numel() > 0:
parts.append(tensor)
if not parts:
return torch.zeros((0,), dtype=torch.float32)
return torch.cat(parts, dim=0)
return torch.as_tensor(audio, dtype=torch.float32).reshape(-1)
def resolve_sr(sr) -> int:
if isinstance(sr, list):
sr = sr[-1]
if hasattr(sr, "item"):
return int(sr.item())
return int(sr)
def extract_sample_rate(multimodal_output: dict) -> int:
sr = multimodal_output.get("sr")
if sr is None:
raise RuntimeError("Expected multimodal_output['sr']")
return resolve_sr(sr)
def write_wav(path: str, audio: torch.Tensor, sample_rate: int) -> None:
audio = audio.clamp(-1.0, 1.0)
pcm16 = (audio * 32767.0).round().to(torch.int16).cpu().numpy()
with wave.open(path, "wb") as wav_file:
wav_file.setnchannels(1)
wav_file.setsampwidth(2)
wav_file.setframerate(int(sample_rate))
wav_file.writeframes(pcm16.tobytes())
def request_index(request_id: str | None, fallback: int) -> int:
try:
return int(request_id)
except (TypeError, ValueError):
if isinstance(request_id, str):
head = request_id.split("_", 1)[0]
if head.isdigit():
return int(head)
return fallback
def audio_summary(audio: torch.Tensor, sample_rate: int) -> dict:
waveform = audio.detach().cpu().reshape(-1).to(torch.float32)
return {
"sample_rate": int(sample_rate),
"num_samples": int(waveform.numel()),
"duration_seconds": float(waveform.numel()) / float(sample_rate),
"max_abs_amplitude": float(waveform.abs().max().item()) if waveform.numel() > 0 else 0.0,
}
def resolve_output_name(output_name: str | None, case: str, index: int, total: int) -> str:
if total == 1:
return output_name or f"ming_{case}.wav"
base = Path(output_name or f"ming_{case}.wav")
return f"{base.stem}_{index:05d}{base.suffix or '.wav'}"
def resolve_stats_log_file(args) -> str | None:
if not args.log_stats:
return None
if args.stats_log_file:
return args.stats_log_file
base = Path(args.output_name or f"ming_{args.case}.wav").stem
return str(Path(args.output_dir) / f"{base}_pipeline.log")
def resolve_metadata_json(args) -> str | None:
if args.metadata_json:
return args.metadata_json
if args.log_stats:
base = Path(args.output_name or f"ming_{args.case}.wav").stem
return str(Path(args.output_dir) / f"{base}_manifest.json")
return None
def build_manifest(args, prompt_payload, stats_log_file: str | None, outputs: list[dict]) -> dict:
additional_information = {}
if isinstance(prompt_payload, dict):
additional_information = dict(prompt_payload.get("additional_information", {}))
return {
"model": args.model,
"case": args.case,
"streaming": bool(args.streaming),
"deploy_config": args.deploy_config,
"enforce_eager": bool(args.enforce_eager),
"num_prompts": int(args.num_prompts),
"log_stats": bool(args.log_stats),
"stats_log_file": stats_log_file,
"prompt_text": additional_information.get("prompt_text"),
"instruction": additional_information.get("instruction"),
"speaker_embedding_shape": (
list(additional_information[KEY_SPEAKER_EMBEDDING].shape)
if KEY_SPEAKER_EMBEDDING in additional_information
and hasattr(additional_information[KEY_SPEAKER_EMBEDDING], "shape")
else None
),
"outputs": outputs,
"generated_at_unix": time.time(),
}
def build_engine_kwargs(args, stats_log_file: str | None) -> dict:
kwargs = {
"model": args.model,
"deploy_config": args.deploy_config,
"enforce_eager": args.enforce_eager,
"trust_remote_code": args.trust_remote_code,
"log_stats": args.log_stats,
"stage_init_timeout": args.stage_init_timeout,
"init_timeout": args.init_timeout,
"batch_timeout": args.batch_timeout,
"shm_threshold_bytes": args.shm_threshold_bytes,
"worker_backend": args.worker_backend,
}
if stats_log_file is not None:
kwargs["log_file"] = stats_log_file
if args.ray_address is not None:
kwargs["ray_address"] = args.ray_address
return kwargs
def build_sampling_params(max_decode_steps: int) -> list[SamplingParams]:
return [
SamplingParams(
temperature=0.0,
max_tokens=max_decode_steps + 1,
stop_token_ids=[int(TEXT_EOS_TOKEN_ID)],
),
SamplingParams(temperature=0.0, max_tokens=1),
]
async def run_streaming(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
engine = AsyncOmni(**build_engine_kwargs(args, stats_log_file))
try:
all_audio_chunks = []
accumulated_samples = 0
chunk_idx = 0
start_time = time.time()
chunk_times = []
ttfp_seconds = None
final_stage_output = None
async for stage_output in engine.generate(
prompt=prompt_payload,
request_id=str(uuid.uuid4()),
sampling_params_list=sampling_params_list,
):
final_stage_output = stage_output
multimodal_output = stage_output.multimodal_output or {}
audio = multimodal_output.get("audio")
if audio is None:
continue
finished = stage_output.finished
if isinstance(audio, torch.Tensor):
if finished:
audio_chunk = audio[accumulated_samples:].float().detach().cpu()
else:
audio_chunk = audio.float().detach().cpu()
elif isinstance(audio, list):
audio_chunk = torch.as_tensor(audio[chunk_idx], dtype=torch.float32).reshape(-1).cpu()
else:
audio_chunk = torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()
accumulated_samples += int(audio_chunk.numel())
chunk_idx += 1
if audio_chunk.numel() > 0:
now = time.time()
if ttfp_seconds is None:
ttfp_seconds = now - start_time
chunk_times.append(now)
all_audio_chunks.append(audio_chunk)
if not all_audio_chunks:
raise RuntimeError("Streaming Ming example produced no audio chunks")
waveform = torch.cat(all_audio_chunks, dim=0)
output_name = resolve_output_name(args.output_name, args.case, 0, 1)
output_path = str(output_dir / output_name)
write_wav(output_path, waveform, SAMPLE_RATE)
summary = {
"request_id": getattr(final_stage_output, "request_id", None),
"stage_id": getattr(final_stage_output, "stage_id", None),
"output_path": output_path,
"stage_durations": getattr(final_stage_output, "stage_durations", {}),
"peak_memory_mb": getattr(final_stage_output, "peak_memory_mb", 0.0),
"ttfp_seconds": ttfp_seconds,
"mean_inter_chunk_seconds": (
sum(t1 - t0 for t0, t1 in zip(chunk_times, chunk_times[1:])) / (len(chunk_times) - 1)
if len(chunk_times) > 1
else None
),
}
summary.update(audio_summary(waveform, SAMPLE_RATE))
print(f"Saved streaming output to {output_path}")
print(json.dumps(summary, ensure_ascii=False, indent=2))
return [summary]
finally:
engine.shutdown()
def run_non_streaming(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
engine = Omni(**build_engine_kwargs(args, stats_log_file))
try:
outputs = engine.generate(
prompts=[prompt_payload for _ in range(args.num_prompts)],
sampling_params_list=sampling_params_list,
py_generator=False,
)
summaries = []
for fallback_index, output in enumerate(outputs):
if output.final_output_type != "audio":
continue
multimodal_output = output.multimodal_output or {}
waveform = coerce_audio_tensor(multimodal_output.get("audio"), async_chunk=False)
sample_rate = extract_sample_rate(multimodal_output)
output_name = resolve_output_name(
args.output_name,
args.case,
request_index(output.request_id, fallback_index),
args.num_prompts,
)
output_path = str(output_dir / output_name)
write_wav(output_path, waveform, sample_rate)
summary = {
"request_id": output.request_id,
"stage_id": output.stage_id,
"output_path": output_path,
"stage_durations": output.stage_durations,
"peak_memory_mb": output.peak_memory_mb,
}
summary.update(audio_summary(waveform, sample_rate))
summaries.append(summary)
print(f"Saved output to {output_path}")
print(json.dumps(summary, ensure_ascii=False, indent=2))
if not summaries:
raise RuntimeError("Non-streaming Ming example produced no audio outputs")
return summaries
finally:
engine.close()
def run_generation(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
if args.streaming:
return asyncio.run(run_streaming(args, prompt_payload, sampling_params_list, output_dir, stats_log_file))
return run_non_streaming(args, prompt_payload, sampling_params_list, output_dir, stats_log_file)
@@ -212,7 +212,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
output_audio_dur = 0.0
for batch_idx, o in enumerate(outputs):
audio_tensor = torch.cat(o.multimodal_output["audio"])
audio_tensor = torch.cat((o.multimodal_output["audio"],))
audio_array = audio_tensor.tolist()
output_audio_dur += float(len(audio_array)) / 24000
if args.write_audio:
@@ -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) | — | ✓ |
| 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) | — | ✓ |
| OmniVoice | `k2-fsa/OmniVoice` | ✓ | — | — | — |
@@ -202,6 +203,51 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server
---
## Ming-omni-tts
Dense 0.5B two-stage TTS served through `/v1/audio/speech`. Ming uses the standard speech endpoint plus structured controls in `instructions`, `voice`, `language`, `ref_audio`, `ref_text`, and `speaker_embedding`.
### Launch
```bash
bash examples/online_serving/text_to_speech/ming_tts/run_server.sh
```
Equivalent manual command:
```bash
vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--host 0.0.0.0 --port 8091 \
--enforce-eager --omni
```
### Sending requests
```bash
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
--text "你好,这是 Ming 在线语音合成测试。"
```
Structured dialect control:
```bash
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
--text "我觉得社会企业同个人都有责任" \
--instruction-json '{"方言":"广粤话"}' \
--ref-audio /path/to/yue_prompt.wav
```
Zero-shot cloning:
```bash
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
--text "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。" \
--ref-audio /path/to/10002287-00000094.wav \
--ref-text "在此奉劝大家别乱打美白针。"
```
### Notes
- `run_curl.sh` keeps a small sanity subset; use the Ming README for the broader request cookbook.
- Online serving is speech-shaped today; music-only `bgm` and text-to-audio `tta` remain offline examples.
- Full request details live in [`ming_tts/README.md`](ming_tts/README.md).
---
## Ming-flash-omni-TTS
Standalone talker-only deployment of Ming-flash-omni-2.0. Voice is controlled through caption text passed via `instructions`.
@@ -51,13 +51,18 @@ def build_payload(
response_format: str = "wav",
stream: bool = False,
) -> dict:
"""Build the /v1/audio/speech request payload."""
"""Build the /v1/audio/speech request payload.
Fish Speech has no built-in speakers; omit ``voice`` so the model
generates with a random timbre (controllable via text tags like
``[cool]``, ``[warm]``, etc.). When ``ref_audio`` is supplied we
enter voice-cloning mode instead.
"""
if not text or not text.strip():
raise gr.Error("Please enter text to synthesize.")
payload: dict = {
"input": text.strip(),
"voice": "default",
"response_format": "pcm" if stream else response_format,
"stream": stream,
}
@@ -193,7 +198,8 @@ def build_interface(api_base: str, stream_chunk_seconds: float = 0.5):
with gr.Accordion("Voice Cloning (optional)", open=False):
gr.Markdown(
"Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice."
"Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice. "
"Leave empty for random timbre; use text tags like `[cool]`, `[warm]`, `[happy]` to steer the output."
)
ref_audio = gr.Audio(
label="Reference Audio",
@@ -39,7 +39,6 @@ def run_tts(args) -> None:
payload = {
"model": args.model,
"input": args.text,
"voice": "default",
"response_format": args.response_format,
}
@@ -0,0 +1,40 @@
# Higgs-Audio V3 Online Serving
## Start the server
```bash
# Default: GPU 0, port 8095
./examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh
# Custom GPU / port
PORT=8096 GPUS=0,1 ./examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh
```
## Plain text TTS
```bash
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
--base-url http://localhost:8095 \
--output-dir /tmp/higgs_v3_batch \
--prompts "Hello world." "The quick brown fox jumps over the lazy dog."
```
## Voice clone
```bash
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
--base-url http://localhost:8095 \
--output-dir /tmp/higgs_v3_clone \
--ref-audio path/to/reference.wav \
--ref-text "transcript of the reference clip" \
--prompts "Text to synthesize in the cloned voice."
```
## curl example
```bash
curl -X POST http://localhost:8095/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{"model": "higgs_audio_v3", "input": "Hello world."}' \
--output hello.wav
```
@@ -0,0 +1,115 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Batch client for the higgs-audio v3 online server.
Sends prompts to ``/v1/audio/speech`` and saves the returned WAV files.
Usage (plain text -> speech):
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
--base-url http://localhost:8095 \
--output-dir /tmp/higgs_v3_batch \
--prompts "Hello world." "The quick brown fox jumps over the lazy dog."
Usage (voice clone):
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
--base-url http://localhost:8095 \
--output-dir /tmp/higgs_v3_clone \
--ref-audio path/to/reference.wav \
--ref-text "the transcript of the reference clip" \
--prompts "Hello world."
"""
from __future__ import annotations
import argparse
import base64
import sys
from pathlib import Path
DEFAULT_PROMPTS = (
"Hello world.",
"The quick brown fox jumps over the lazy dog.",
"Today is a beautiful day for a walk in the park.",
"Innovation distinguishes between a leader and a follower.",
)
def _slug(text: str) -> str:
import re
s = re.sub(r"\s+", "_", text.strip().lower())
return re.sub(r"[^a-z0-9_]+", "", s)[:32] or "prompt"
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--base-url", default="http://localhost:8095")
parser.add_argument("--model", default="higgs_audio_v3")
parser.add_argument("--prompts", nargs="+", default=list(DEFAULT_PROMPTS))
parser.add_argument("--output-dir", type=Path, default=Path("/tmp/higgs_v3_batch"))
parser.add_argument("--format", choices=("wav", "pcm"), default="wav")
parser.add_argument("--max-new-tokens", type=int, default=2048)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--timeout-s", type=float, default=120.0)
parser.add_argument(
"--ref-audio",
type=Path,
default=None,
help="Reference clip for voice clone (WAV/FLAC/MP3 path). Pair with --ref-text.",
)
parser.add_argument(
"--ref-text",
type=str,
default=None,
help="Transcript of the reference clip. Optional but improves fidelity.",
)
args = parser.parse_args()
ref_audio_data_url: str | None = None
if args.ref_audio is not None:
if not args.ref_audio.exists():
print(f"ref-audio file not found: {args.ref_audio}", file=sys.stderr)
return 2
mime = "audio/wav" if args.ref_audio.suffix.lower() == ".wav" else "audio/mpeg"
ref_b64 = base64.b64encode(args.ref_audio.read_bytes()).decode("ascii")
ref_audio_data_url = f"data:{mime};base64,{ref_b64}"
try:
import httpx
except ImportError:
print("this client needs `httpx`. Install with `pip install httpx`.", file=sys.stderr)
return 2
args.output_dir.mkdir(parents=True, exist_ok=True)
url = args.base_url.rstrip("/") + "/v1/audio/speech"
failures = 0
with httpx.Client(timeout=args.timeout_s) as client:
for prompt in args.prompts:
payload = {
"model": args.model,
"input": prompt,
"response_format": args.format,
"max_new_tokens": args.max_new_tokens,
"seed": args.seed,
}
if ref_audio_data_url is not None:
payload["ref_audio"] = ref_audio_data_url
if args.ref_text:
payload["ref_text"] = args.ref_text
resp = client.post(url, json=payload)
if resp.status_code != 200:
print(f"[FAIL] {prompt!r} -> {resp.status_code}: {resp.text[:200]}", file=sys.stderr)
failures += 1
continue
suffix = ".wav" if args.format == "wav" else ".pcm"
out = args.output_dir / f"{_slug(prompt)}{suffix}"
out.write_bytes(resp.content)
print(f"[ ok ] {prompt!r} -> {out} ({len(resp.content)} bytes)")
return 1 if failures else 0
if __name__ == "__main__":
sys.exit(main())
@@ -0,0 +1,32 @@
#!/bin/bash
# Launch vLLM-Omni server for higgs-audio v3.
#
# Supports plain text TTS and voice cloning via /v1/audio/speech.
#
# Usage:
# ./run_server.sh # default port 8095, GPU 0
# PORT=8096 GPUS=0,1 ./run_server.sh
# MODEL=/path/to/local/checkpoint ./run_server.sh
set -e
MODEL="${MODEL:-bosonai/higgs-audio-v3-tts-4b}"
PORT="${PORT:-8095}"
GPUS="${GPUS:-0}"
GPU_UTIL="${GPU_UTIL:-0.6}"
echo "Starting higgs-audio v3 server"
echo " MODEL=$MODEL"
echo " PORT=$PORT"
echo " CUDA_VISIBLE_DEVICES=$GPUS"
CUDA_VISIBLE_DEVICES="$GPUS" \
VLLM_USE_DEEP_GEMM=0 \
VLLM_MOE_USE_DEEP_GEMM=0 \
vllm-omni serve "$MODEL" \
--deploy-config vllm_omni/deploy/higgs_multimodal_qwen3.yaml \
--host 0.0.0.0 \
--port "$PORT" \
--gpu-memory-utilization "$GPU_UTIL" \
--trust-remote-code \
--omni
@@ -0,0 +1,113 @@
# Ming-omni-tts Online Serving
Serve the dense `inclusionAI/Ming-omni-tts-0.5B` two-stage TTS model through
the OpenAI-compatible `/v1/audio/speech` endpoint.
## Start Server
```bash
vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--omni \
--port 8091 \
--enforce-eager
```
Or:
```bash
cd examples/online_serving/text_to_speech/ming_tts
./run_server.sh
```
The tested ROCm environment is summarized in the
[Ming recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/inclusionAI/Ming-omni-tts-0.5B.md).
## Send Requests
The Python client targets `http://localhost:8091/v1` with `api_key=EMPTY`; it
does not call OpenAI's hosted API.
```bash
python openai_speech_client.py \
--text "你好,这是 Ming 在线语音合成测试。" \
--max-new-tokens 200
```
Style or dialect controls can be plain text or Ming JSON. The upstream
dialect example also uses `yue_prompt.wav` for speaker conditioning:
```bash
python openai_speech_client.py \
--text "我觉得社会企业同个人都有责任" \
--instruction-json '{"方言":"广粤话"}' \
--ref-audio /path/to/yue_prompt.wav \
--max-new-tokens 200
```
When `--ref-audio` is supplied without `--ref-text`, the server extracts the
Ming speaker embedding, matching upstream `use_spk_emb=True`, without using the
audio as a zero-shot prompt.
Reference-audio cloning:
```bash
python openai_speech_client.py \
--text "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。" \
--ref-audio /path/to/10002287-00000094.wav \
--ref-text "在此奉劝大家别乱打美白针。" \
--max-new-tokens 200
```
Podcast-style multi-speaker prompt:
```bash
python openai_speech_client.py \
--text " speaker_1:你可以说一下,就大概说一下,可能虽然我也不知道,我看过那部电影没有。
speaker_2:就是那个叫什么,变相一节课的嘛。
speaker_1:嗯。
speaker_2:一部搞笑的电影。
speaker_1:一部搞笑的。" \
--ref-audio /path/to/CTS-CN-F2F-2019-11-11-423-012-A.wav \
--ref-audio /path/to/CTS-CN-F2F-2019-11-11-423-012-B.wav \
--ref-text " speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试,做个,当服务员还要去笔试了
speaker_2:对啊,这真的很奇怪,就是 单纯的因,单纯自己工资不高,只是因为可能人家那个店比较出名一点,就对你苛刻要求"
```
Streaming PCM:
```bash
python openai_speech_client.py \
--text "你好,这是流式输出测试。" \
--stream \
--output ming_output.pcm
```
`run_curl.sh` keeps small smoke checks:
```bash
./run_curl.sh basic
REF_AUDIO=/path/to/reference.wav REF_TEXT="在此奉劝大家别乱打美白针。" ./run_curl.sh zero_shot
./run_curl.sh stream
```
## Request Fields
| Field | Ming meaning |
|-------|--------------|
| `input` | target text |
| `instructions` | plain style text, or JSON object for structured Ming controls |
| `voice` | Ming IP voice label unless it resolves to an uploaded speaker |
| `language` | Ming `方言` control |
| `ref_audio` | speaker reference; with `ref_text`, also supplies the prompt waveform |
| `ref_text` | transcript enabling zero-shot or podcast prompt-latent conditioning |
| `speaker_embedding` | 192-d Ming speaker embedding |
| `max_new_tokens` | Ming `max_decode_steps` |
## Notes
- `ref_audio` accepts local paths through the client, remote URLs, `file://`,
or `data:` URLs.
- Non-streaming responses return WAV bytes; streaming responses return PCM.
- Music-only `bgm` generation is offline-only until the API exposes Ming
prompt-mode selection.
@@ -0,0 +1,225 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""OpenAI-compatible client for Ming-omni-tts via /v1/audio/speech.
Examples:
python openai_speech_client.py --text "你好,世界"
python openai_speech_client.py --text "我觉得社会企业同个人都有责任" \
--instruction-json '{"方言":"广粤话"}' --ref-audio yue_prompt.wav --max-new-tokens 200
python openai_speech_client.py \
--text "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。" \
--ref-audio 10002287-00000094.wav --ref-text "在此奉劝大家别乱打美白针。" \
--max-new-tokens 200
python openai_speech_client.py --text "speaker_1:你好。 speaker_2:你好。" \
--ref-audio speaker_1.wav --ref-audio speaker_2.wav --ref-text "speaker_1:你好。 speaker_2:你好。"
python openai_speech_client.py --text "你好,这是流式输出测试。" \
--stream --output ming_output.pcm
"""
import argparse
import base64
import json
import os
import httpx
DEFAULT_API_BASE = "http://localhost:8091"
DEFAULT_API_KEY = "EMPTY"
DEFAULT_MODEL = "inclusionAI/Ming-omni-tts-0.5B"
EXPECTED_SPEAKER_EMBEDDING_DIM = 192
def encode_audio_to_base64(audio_path: str) -> str:
"""Encode a local audio file to a base64 data URL."""
if not os.path.exists(audio_path):
raise FileNotFoundError(f"Audio file not found: {audio_path}")
ext = audio_path.lower().rsplit(".", 1)[-1]
mime_map = {
"wav": "audio/wav",
"mp3": "audio/mpeg",
"flac": "audio/flac",
"ogg": "audio/ogg",
"aac": "audio/aac",
}
mime_type = mime_map.get(ext, "audio/wav")
with open(audio_path, "rb") as f:
audio_b64 = base64.b64encode(f.read()).decode("utf-8")
return f"data:{mime_type};base64,{audio_b64}"
def load_speaker_embedding(path: str) -> list[float]:
"""Load and validate a 192-d Ming speaker embedding JSON file."""
with open(path, encoding="utf-8") as f:
data = json.load(f)
if not isinstance(data, list):
raise ValueError("speaker_embedding file must contain a JSON list")
if len(data) != EXPECTED_SPEAKER_EMBEDDING_DIM:
raise ValueError(
f"Ming dense speaker_embedding must have {EXPECTED_SPEAKER_EMBEDDING_DIM} values, got {len(data)}"
)
values = []
for index, value in enumerate(data):
try:
values.append(float(value))
except (TypeError, ValueError) as exc:
raise ValueError(f"speaker_embedding[{index}] must be a number, got {value!r}") from exc
return values
def build_instruction_payload(args) -> str | None:
"""Return a string payload for the API `instructions` field."""
if args.instructions and args.instruction_json:
raise ValueError("Use either --instructions or --instruction-json, not both")
if args.instruction_json:
parsed = json.loads(args.instruction_json)
return json.dumps(parsed, ensure_ascii=False)
return args.instructions
def validate_args(args) -> None:
"""Fail fast on invalid combinations before hitting the server."""
if args.ref_text and not args.ref_audio:
raise ValueError("--ref-audio is required when --ref-text is provided")
if args.speaker_embedding and args.ref_audio and len(args.ref_audio) > 1:
raise ValueError("--speaker-embedding cannot be combined with multiple --ref-audio values")
def run_tts(args) -> None:
"""Generate speech via the OpenAI-compatible /v1/audio/speech API."""
validate_args(args)
payload = {
"model": args.model,
"input": args.text,
"response_format": args.response_format,
}
if args.voice:
payload["voice"] = args.voice
if args.task_type:
payload["task_type"] = args.task_type
if args.dialect:
payload["language"] = args.dialect
instructions = build_instruction_payload(args)
if instructions:
payload["instructions"] = instructions
if args.ref_audio:
ref_audio = []
for audio in args.ref_audio:
if audio.startswith(("http://", "https://", "data:")):
ref_audio.append(audio)
else:
ref_audio.append(encode_audio_to_base64(audio))
payload["ref_audio"] = ref_audio[0] if len(ref_audio) == 1 else ref_audio
if args.ref_text:
payload["ref_text"] = args.ref_text
if args.speaker_embedding:
payload["speaker_embedding"] = load_speaker_embedding(args.speaker_embedding)
if args.max_new_tokens:
payload["max_new_tokens"] = args.max_new_tokens
if args.stream:
payload["stream"] = True
payload["response_format"] = "pcm"
api_url = f"{args.api_base}/v1/audio/speech"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {args.api_key}",
}
print(f"Model: {args.model}")
print(f"Text: {args.text}")
print(f"Payload keys: {sorted(payload)}")
if args.stream:
output_path = args.output or "ming_output.pcm"
with httpx.Client(timeout=300.0) as client:
with client.stream("POST", api_url, json=payload, headers=headers) as response:
if response.status_code != 200:
print(f"Error: {response.status_code}")
print(response.read().decode())
return
with open(output_path, "wb") as f:
for chunk in response.iter_bytes():
f.write(chunk)
print(f"Streamed PCM audio to: {output_path}")
return
with httpx.Client(timeout=300.0) as client:
response = client.post(api_url, json=payload, headers=headers)
if response.status_code != 200:
print(f"Error: {response.status_code}")
print(response.text)
return
try:
text = response.content.decode("utf-8")
if text.startswith('{"error"'):
print(f"Error: {text}")
return
except UnicodeDecodeError:
pass
output_path = args.output or "ming_output.wav"
with open(output_path, "wb") as f:
f.write(response.content)
print(f"Audio saved to: {output_path}")
def main():
parser = argparse.ArgumentParser(description="OpenAI-compatible client for Ming-omni-tts via /v1/audio/speech")
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name or path")
parser.add_argument("--text", required=True, help="Text to synthesize")
parser.add_argument(
"--task-type",
default=None,
choices=["CustomVoice", "VoiceDesign", "Base"],
help="Optional compatibility task type. Ming accepts the same field but primarily uses prompt metadata.",
)
parser.add_argument(
"--voice",
default=None,
help="Maps to Ming `IP` when using built-in character voices, or to an uploaded voice sample name",
)
parser.add_argument("--dialect", default=None, help="Maps to Ming `方言`")
parser.add_argument("--instructions", default=None, help="Free-form Ming instruction string")
parser.add_argument(
"--instruction-json",
default=None,
help='Structured Ming instruction JSON, for example \'{"情感":"高兴"}\'',
)
parser.add_argument(
"--ref-audio",
action="append",
default=None,
help="Reference audio path, URL, or data URL. Repeat for podcast-style multi-speaker prompts.",
)
parser.add_argument("--ref-text", default=None, help="Reference transcript for cloning")
parser.add_argument(
"--speaker-embedding", default=None, help="Path to a JSON file containing a 192-d speaker embedding"
)
parser.add_argument("--max-new-tokens", type=int, default=None, help="Override ming_max_decode_steps")
parser.add_argument("--stream", action="store_true", help="Enable streaming PCM output")
parser.add_argument(
"--response-format",
default="wav",
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
help="Audio format when not streaming",
)
parser.add_argument("--output", "-o", default=None, help="Output file path")
args = parser.parse_args()
try:
run_tts(args)
except (FileNotFoundError, ValueError, json.JSONDecodeError) as exc:
raise SystemExit(f"Error: {exc}") from exc
if __name__ == "__main__":
main()
@@ -0,0 +1,76 @@
#!/bin/bash
set -euo pipefail
MODE="${1:-basic}"
HOST="${HOST:-localhost}"
PORT="${PORT:-8091}"
MODEL="${MODEL:-inclusionAI/Ming-omni-tts-0.5B}"
API_URL="http://${HOST}:${PORT}/v1/audio/speech"
TEXT="${TEXT:-你好,这是 Ming 在线语音合成测试。}"
OUTPUT="${OUTPUT:-ming_output.wav}"
STREAM_OUTPUT="${STREAM_OUTPUT:-ming_output.pcm}"
REF_AUDIO="${REF_AUDIO:-}"
REF_TEXT="${REF_TEXT:-}"
post_json() {
local payload="$1"
local output_path="$2"
curl -X POST "$API_URL" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer EMPTY" \
-d "$payload" \
--output "$output_path"
}
case "$MODE" in
basic)
post_json "{
\"model\": \"${MODEL}\",
\"input\": \"${TEXT}\",
\"response_format\": \"wav\"
}" "$OUTPUT"
;;
zero_shot)
if [ -z "$REF_AUDIO" ] || [ -z "$REF_TEXT" ]; then
echo "zero_shot requires REF_AUDIO and REF_TEXT" >&2
exit 1
fi
python - <<'PY' > /tmp/ming_zero_shot_payload.json
import base64
import json
import mimetypes
import os
from pathlib import Path
path = Path(os.environ["REF_AUDIO"])
mime_type = mimetypes.guess_type(path.name)[0] or "audio/wav"
payload = {
"model": os.environ["MODEL"],
"input": os.environ["TEXT"],
"ref_audio": f"data:{mime_type};base64,{base64.b64encode(path.read_bytes()).decode('utf-8')}",
"ref_text": os.environ["REF_TEXT"],
"response_format": "wav",
}
print(json.dumps(payload, ensure_ascii=False))
PY
curl -X POST "$API_URL" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer EMPTY" \
--data-binary @/tmp/ming_zero_shot_payload.json \
--output "$OUTPUT"
rm -f /tmp/ming_zero_shot_payload.json
;;
stream)
post_json "{
\"model\": \"${MODEL}\",
\"input\": \"${TEXT}\",
\"stream\": true,
\"response_format\": \"pcm\"
}" "$STREAM_OUTPUT"
;;
*)
echo "Unknown mode: $MODE" >&2
echo "Supported sanity checks: basic, zero_shot, stream" >&2
exit 1
;;
esac
@@ -0,0 +1,25 @@
#!/bin/bash
# Launch vLLM-Omni server for Ming-omni-tts.
#
# Usage:
# ./run_server.sh
# PORT=8000 ./run_server.sh
set -e
DIR="$(cd "$(dirname "$0")" && pwd)"
ROOT="$(cd "$DIR/../../../.." && pwd)"
MODEL="${MODEL:-inclusionAI/Ming-omni-tts-0.5B}"
PORT="${PORT:-8091}"
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$ROOT/vllm_omni/deploy/ming_tts.yaml}"
echo "Starting Ming-omni-tts server with model: $MODEL"
echo "Deploy config: $DEPLOY_CONFIG"
vllm-omni serve "$MODEL" \
--deploy-config "$DEPLOY_CONFIG" \
--host 0.0.0.0 \
--port "$PORT" \
--enforce-eager \
--omni
@@ -0,0 +1,247 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Verify a ModelOpt FP8 diffusers checkpoint exported by
quantize_wan2_2_modelopt_fp8.py / quantize_hunyuanvideo_15_modelopt_fp8.py.
Three checks:
A. transformer/config.json has a sane quantization_config block.
B. transformer/*.safetensors contains FP8 (float8_e4m3fn) quantized tensors.
C. transformer disk size is materially smaller than a BF16 baseline.
Example:
python examples/quantization/check_modelopt_fp8_export.py \\
--output ./hv15-480p-modelopt-fp8
# Optional: compare disk size against a local or HF BF16 baseline.
python examples/quantization/check_modelopt_fp8_export.py \\
--output ./hv15-480p-modelopt-fp8 \\
--baseline hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
"""
from __future__ import annotations
import argparse
import json
import sys
from collections import Counter
from pathlib import Path
SUPPORTED_ALGOS = ("FP8",)
def _check_config(transformer_dir: Path) -> tuple[int, str | None]:
"""Returns (status, quant_algo). status: 0 pass, 1 fail, 2 warn."""
cfg_path = transformer_dir / "config.json"
if not cfg_path.exists():
print(f"[FAIL] {cfg_path} missing.")
return 1, None
with cfg_path.open(encoding="utf-8") as f:
cfg = json.load(f)
qc = cfg.get("quantization_config")
if not isinstance(qc, dict):
print(f"[FAIL] No `quantization_config` block in {cfg_path}.")
return 1, None
print(f"[A] quantization_config from {cfg_path}:")
print(json.dumps(qc, indent=2))
quant_algo = qc.get("quant_algo")
issues = []
if qc.get("quant_method") != "modelopt":
issues.append(f"quant_method={qc.get('quant_method')!r} (expected 'modelopt')")
if quant_algo not in SUPPORTED_ALGOS:
issues.append(
f"quant_algo={quant_algo!r} (expected one of {SUPPORTED_ALGOS} — vllm-omni adapter may not auto-detect)"
)
if issues:
print("[A] WARN — config looks incomplete:")
for issue in issues:
print(f" - {issue}")
return 2, quant_algo
print(f"[A] PASS — config looks correct (quant_algo={quant_algo}).")
return 0, quant_algo
def _read_safetensors_header(path: Path) -> dict:
"""Read the JSON header of a safetensors file. Bypass-safe — doesn't materialize tensors.
Returns {tensor_name: {'dtype': 'F8_E4M3', 'shape': [...], 'data_offsets': [...]}}.
Header dtype strings: F8_E4M3, F8_E5M2, BF16, F16, F32, F64, I8, I16, I32, I64, BOOL, U8, ...
"""
import struct
with open(path, "rb") as f:
header_len = struct.unpack("<Q", f.read(8))[0]
header = json.loads(f.read(header_len))
header.pop("__metadata__", None)
return header
def _classify_weight_scale_granularity(weight_scale_shapes: list[list[int]]) -> str:
"""Infer per-tensor vs per-channel vs per-block from sample weight_scale shapes.
ModelOpt block-wise produces shapes like `[16, 1, 16, 1]` (broadcasting dims of 1
interleaved with block-count dims). We count "meaningful" dims — ones with size > 1 —
and classify: 0 meaningful dims = per-tensor (scalar), 1 = per-channel, 2+ = per-block.
"""
if not weight_scale_shapes:
return "no weight_scale tensors found"
def meaningful_dims(shape: list[int]) -> int:
return sum(1 for d in shape if d > 1)
per_tensor = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 0)
per_channel = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 1)
per_block = sum(1 for s in weight_scale_shapes if meaningful_dims(s) >= 2)
total = len(weight_scale_shapes)
if per_tensor == total:
return "per-tensor (all scalar scales)"
if per_channel == total:
return "per-channel (1 meaningful dim)"
if per_block == total:
return "per-block (2+ meaningful dims — e.g. [M//bm, 1, N//bn, 1] for tiles)"
return f"mixed: per-tensor={per_tensor}, per-channel={per_channel}, per-block={per_block} of {total}"
def _check_safetensors(transformer_dir: Path) -> int:
"""Returns 0 on pass, 1 on fail. Reads on-disk dtype from the safetensors header."""
files = sorted(transformer_dir.glob("*.safetensors"))
if not files:
print(f"[FAIL] No *.safetensors in {transformer_dir}.")
return 1
header_dtype_counts: Counter[str] = Counter()
sample_quant_weight_keys: list[str] = []
sample_scale_keys: list[str] = []
weight_scale_shapes: list[list[int]] = []
sample_weight_scale_entries: list[tuple[str, list[int]]] = []
for f in files:
try:
header = _read_safetensors_header(f)
except Exception as exc:
print(f"[B] WARN — could not parse header of {f}: {exc}")
continue
for k, info in header.items():
dtype = info.get("dtype", "?")
header_dtype_counts[dtype] += 1
# FP8 stores weights as F8_E4M3 directly.
if dtype.startswith("F8") and k.endswith(".weight") and len(sample_quant_weight_keys) < 5:
sample_quant_weight_keys.append(k)
if k.endswith(("_scale", ".weight_scale", ".input_scale", "_scale_inv")) and len(sample_scale_keys) < 5:
sample_scale_keys.append(k)
if k.endswith(".weight_scale"):
weight_scale_shapes.append(info.get("shape", []))
if len(sample_weight_scale_entries) < 5:
sample_weight_scale_entries.append((k, info.get("shape", [])))
print(f"\n[B] On-disk dtype counts across {len(files)} safetensors file(s) (from header, not get_tensor):")
for dtype, count in sorted(header_dtype_counts.items(), key=lambda kv: -kv[1]):
marker = " <-- FP8" if dtype.startswith("F8") else ""
print(f" {dtype:10s} {count:>6d}{marker}")
quant_count = sum(c for d, c in header_dtype_counts.items() if d.startswith("F8"))
if quant_count == 0:
print("[B] FAIL — no FP8 tensors on disk. Calibration likely did not actually quantize the weights.")
return 1
print(f"[B] PASS — {quant_count} FP8 tensors stored on disk.")
if sample_quant_weight_keys:
print(f" sample quantized weight tensors: {sample_quant_weight_keys[:3]}")
if sample_scale_keys:
print(f" sample scale tensors: {sample_scale_keys[:3]}")
print(" (Note: torch's get_tensor() may return these as bf16 views on some versions —")
print(" irrelevant; vLLM's loader uses native FP8 ops.)")
# Weight-scale granularity — per-tensor (scalar) vs per-channel (1-D) vs per-block (N-D).
print(f"\n weight_scale granularity: {_classify_weight_scale_granularity(weight_scale_shapes)}")
for key, shape in sample_weight_scale_entries[:3]:
print(f" {key}: shape {shape}")
return 0
def _disk_size_gib(p: Path) -> float:
return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) / (1024**3)
def _check_size_vs_baseline(transformer_dir: Path, baseline: str | None) -> int:
"""Returns 0 always (informational only)."""
quant_size = _disk_size_gib(transformer_dir)
print(f"\n[C] FP8 transformer disk size: {quant_size:.2f} GiB")
if baseline is None:
print("[C] SKIP — pass --baseline <path or HF id> to compare against BF16.")
return 0
baseline_path = Path(baseline)
if not baseline_path.exists():
# Try HF download.
try:
from huggingface_hub import snapshot_download
except ImportError:
print("[C] SKIP — huggingface_hub not installed and baseline not a local path.")
return 0
print(f" Downloading baseline transformer from HF: {baseline}")
baseline_path = Path(snapshot_download(baseline, allow_patterns=["transformer/*"]))
bf16_dir = baseline_path / "transformer" if (baseline_path / "transformer").exists() else baseline_path
bf16_size = _disk_size_gib(bf16_dir)
if bf16_size == 0:
print(f"[C] WARN — baseline transformer dir empty: {bf16_dir}")
return 0
# Expected reduction for FP8: ~50%.
min_reduction = 30
reduction = (1 - quant_size / bf16_size) * 100
print(f"[C] BF16 baseline transformer disk size: {bf16_size:.2f} GiB ({bf16_dir})")
print(f"[C] Disk reduction: {reduction:.1f}% (FP8 transformer is {quant_size / bf16_size:.0%} of BF16)")
if reduction < min_reduction:
print(
f"[C] WARN — FP8 should typically reduce disk by ~40-50%; <{min_reduction}% suggests partial quantization."
)
return 0
def main() -> None:
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--output", required=True, help="Path to the exported ModelOpt FP8 checkpoint root.")
p.add_argument(
"--baseline",
default=None,
help="Optional BF16 baseline (local diffusers dir or HF id) for disk-size comparison.",
)
args = p.parse_args()
out_root = Path(args.output).expanduser().resolve()
transformer_dir = out_root / "transformer"
if not transformer_dir.exists():
print(f"[FAIL] {transformer_dir} does not exist.")
sys.exit(1)
print(f"Checking: {out_root}\n")
fail = 0
config_status, _quant_algo = _check_config(transformer_dir)
fail |= config_status
fail |= _check_safetensors(transformer_dir)
_check_size_vs_baseline(transformer_dir, args.baseline)
print()
if fail == 0:
print("=" * 60)
print("ALL CHECKS PASSED — checkpoint looks ready for vllm-omni serving.")
elif fail == 1:
print("=" * 60)
print("FAILURES detected — calibration may need to be re-run.")
sys.exit(1)
else:
print("=" * 60)
print("WARNINGS only — checkpoint may serve but with caveats. See [A] above.")
if __name__ == "__main__":
main()
@@ -0,0 +1,448 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Quantize HunyuanVideo-1.5 (480p T2V) to a ModelOpt FP8 Hugging Face checkpoint.
Calibrates the DiT transformer using a small video prompt set and exports a
diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
Layers kept full precision match the #2728 / #2795 pattern: modulation,
AdaLayerNorm, entry/exit projections, embeddings, the token refiner path,
and final proj_out. MHA quantizers are off by default; HV-1.5 self-attention
empirically degrades under FP8 (see #2920 ablation).
Example:
python examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py \\
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
--output ./hv15-480p-modelopt-fp8 \\
--overwrite
"""
from __future__ import annotations
import argparse
import copy
import json
import re
import shutil
import sys
from pathlib import Path
from typing import Any
import torch
from diffusers import DiffusionPipeline
DEFAULT_PROMPTS = [
"A dog running across a field of golden wheat.",
"An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
"A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
"A crackling campfire at night under a starry sky, sparks rising into the dark.",
"An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
"A close-up of a blooming rose covered in morning dew, soft natural light.",
"A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
"A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
]
def _build_parser() -> argparse.ArgumentParser:
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--model", required=True, help="Input HV-1.5 diffusers directory or HF id.")
p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
p.add_argument("--height", type=int, default=480)
p.add_argument("--width", type=int, default=832)
p.add_argument(
"--num-frames",
type=int,
default=33,
help="Frames per calibration sample. 33 matches the typical short benchmark.",
)
p.add_argument("--guidance-scale", type=float, default=6.0)
p.add_argument(
"--calib-steps",
type=int,
default=10,
help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
)
p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
p.add_argument("--seed", type=int, default=42)
p.add_argument(
"--prompt",
action="append",
default=[],
help="Custom calibration prompt. Repeat to provide multiple.",
)
p.add_argument(
"--quantize-mha",
action="store_true",
help="Enable FP8 attention K/V/softmax quantizers. Off by default — empirically degrades HV-1.5 video output.",
)
p.add_argument(
"--weight-block-size",
type=str,
default=None,
help="Per-block weight quantization as 'M,N' (e.g. '128,128' for 128x128 tiles). "
"Default: per-tensor (one scale per linear). Block-wise typically gives tighter quality at "
"negligible memory cost. Static FP8 is exempt from upstream vLLM's online block-wise gate.",
)
p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
return p
def _parse_block_size(spec: str | None) -> list[int] | None:
if spec is None:
return None
parts = [int(x) for x in spec.split(",") if x.strip()]
if len(parts) != 2:
raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
return parts
def _require_modelopt() -> Any:
try:
import modelopt.torch.quantization as mtq
except ModuleNotFoundError as exc:
raise SystemExit(
"NVIDIA ModelOpt is not installed. Install with:\n"
" pip install 'nvidia-modelopt[all]'\n"
f"Original error: {exc}"
) from exc
return mtq
def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
model_path = args.model
output_dir = Path(args.output).expanduser().resolve()
if output_dir.exists():
if not args.overwrite:
raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
shutil.rmtree(output_dir)
return model_path, output_dir
def _select_dtype(name: str) -> torch.dtype:
return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
def _build_prompts(args: argparse.Namespace) -> list[str]:
prompts = args.prompt or DEFAULT_PROMPTS
if args.calib_size <= 0:
raise SystemExit("--calib-size must be positive.")
if len(prompts) < args.calib_size:
repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
prompts = (prompts * repeats)[: args.calib_size]
return prompts[: args.calib_size]
# Layers to KEEP at full precision (mirror of the #2920 wiring + #2728/#2795 skip pattern).
# - x_embedder, image_embedder, context_embedder*, time_embed*, cond_type_embed: entry/embedding
# - norm_out, norm1*.linear, norm1_context*.linear, norm2*, norm2_context*: AdaLayerNorm modulation
# - proj_out: final output projection
# - token_refiner*: text-encoder refinement uses diffusers raw nn.Linear
def _filter_func_hv15(name: str) -> bool:
pattern = re.compile(
r"(proj_out.*|"
r".*(x_embedder|image_embedder|context_embedder|context_embedder_2|"
r"time_embed|cond_type_embed|"
r"norm_out|norm1\.linear|norm1_context\.linear|norm2|norm2_context|"
r"token_refiner).*)"
)
return pattern.match(name) is not None
def _mha_filter_func(name: str) -> bool:
pattern = re.compile(
r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
)
return pattern.match(name) is not None
def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
if not hasattr(mtq, "disable_quantizer"):
return
mtq.disable_quantizer(backbone, _filter_func_hv15)
if not quantize_mha:
mtq.disable_quantizer(backbone, _mha_filter_func)
def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
if hasattr(pipe, "set_progress_bar_config"):
pipe.set_progress_bar_config(disable=True)
pipe.to("cuda")
return pipe
def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
generator = torch.Generator(device="cuda")
# Try to set guidance on the pipeline's guider object up front (modern
# diffusers HV-1.5 uses a Guider abstraction, not a per-call kwarg). Falls
# back silently — calibration uses whatever default the pipeline ships with.
guider = getattr(pipe, "guider", None)
if guider is not None and hasattr(guider, "guidance_scale"):
try:
guider.guidance_scale = args.guidance_scale
except Exception:
pass
base_kwargs = dict(
height=args.height,
width=args.width,
num_frames=args.num_frames,
num_inference_steps=args.calib_steps,
output_type="latent",
)
def forward_loop(*_unused_args, **_unused_kwargs) -> None:
with torch.inference_mode():
for idx, prompt in enumerate(prompts):
generator.manual_seed(args.seed + idx)
# Try with guidance_scale first; fall back without on TypeError
# for pipelines (like HV-1.5) that take CFG via guider config.
try:
pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
except TypeError as exc:
if "guidance_scale" not in str(exc):
raise
pipe(prompt=prompt, generator=generator, **base_kwargs)
return forward_loop
def _summarize_export(output_dir: Path) -> None:
cfg_path = output_dir / "transformer" / "config.json"
if not cfg_path.exists():
print(f"[warn] {cfg_path} missing.", file=sys.stderr)
return
with cfg_path.open(encoding="utf-8") as f:
cfg = json.load(f)
qc = cfg.get("quantization_config")
if not isinstance(qc, dict):
print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
return
print("Export summary:")
print(f" quant_method: {qc.get('quant_method')}")
print(f" quant_algo: {qc.get('quant_algo')}")
producer = qc.get("producer")
if isinstance(producer, dict):
print(f" producer: {producer.get('name')} {producer.get('version')}")
print(f" config path: {cfg_path}")
def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
"""Convert in-memory weights of quantized modules to actual FP8 storage.
`export_hf_checkpoint` skips this step for unknown model types (HV-1.5 isn't
in ModelOpt's recognized-model registry), so we must call the per-weight
export helper ourselves. Same workaround as the HunyuanImage-3 calibration
helper.
"""
from modelopt.torch.export.quant_utils import (
QUANTIZATION_NONE,
get_quantization_format,
quantizer_attr_names,
weight_attr_names,
)
from modelopt.torch.export.unified_export_hf import _export_quantized_weight
exported = 0
for name, module in backbone.named_modules():
try:
quantization_format = get_quantization_format(module)
except Exception as exc:
print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
continue
if quantization_format == QUANTIZATION_NONE:
continue
for weight_name in weight_attr_names(module):
quantizer_attrs = quantizer_attr_names(weight_name)
weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
continue
_export_quantized_weight(module, dtype, weight_name)
exported += 1
return exported
def _hv15_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
"""Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913).
Same shape as the HunyuanImage-3 author's _hunyuan_quant_config(). When
`weight_block_size` is given, advertise block-wise weight quantization in
the saved metadata (so consumers know to expect multi-element scale tensors).
"""
weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
if weight_block_size is not None:
weights_cfg["strategy"] = "block"
weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
return {
"config_groups": {
"group_0": {
"input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
"weights": weights_cfg,
"targets": ["Linear"],
}
},
"ignore": [
"context_embedder*",
"context_embedder_2*",
"cond_type_embed*",
"image_embedder*",
"norm1.linear*",
"norm1_context.linear*",
"norm2*",
"norm2_context*",
"norm_out*",
"proj_out*",
"time_embed*",
"token_refiner*",
"x_embedder*",
],
"producer": {"name": "modelopt"},
"quant_algo": "FP8",
"quant_method": "modelopt",
}
def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
"""Inject quant_algo: FP8 + config_groups into transformer/config.json so
vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8."""
cfg_path = output_dir / "transformer" / "config.json"
with cfg_path.open(encoding="utf-8") as f:
cfg = json.load(f)
new_qc = _hv15_quant_config_block(weight_block_size=weight_block_size)
existing = cfg.get("quantization_config")
if isinstance(existing, dict):
producer = existing.get("producer")
if isinstance(producer, dict):
new_qc["producer"] = producer
cfg["quantization_config"] = new_qc
with cfg_path.open("w", encoding="utf-8") as f:
json.dump(cfg, f, indent=2)
def _save_pipeline_with_fp8_transformer(
pipe: DiffusionPipeline,
model_path: str,
output_dir: Path,
max_shard_size: str = "5GB",
) -> None:
"""Save the pipeline with the (now FP8) transformer.
Copies the source directory verbatim except for `transformer/`, then
saves the transformer with quantizers hidden so the state dict contains
only the FP8 weights + scale tensors.
"""
from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
src = Path(model_path)
if not src.exists():
from huggingface_hub import snapshot_download
src = Path(snapshot_download(model_path))
if output_dir.exists():
shutil.rmtree(output_dir)
shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer"))
transformer_out = output_dir / "transformer"
# `hide_quantizers_from_state_dict` walks named_modules(); pass the actual
# nn.Module (transformer), not the diffusers Pipeline wrapper.
with hide_quantizers_from_state_dict(pipe.transformer):
pipe.transformer.save_pretrained(
str(transformer_out),
safe_serialization=True,
max_shard_size=max_shard_size,
)
def main() -> None:
args = _build_parser().parse_args()
if not torch.cuda.is_available():
raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
mtq = _require_modelopt()
model_path, output_dir = _ensure_paths(args)
dtype = _select_dtype(args.dtype)
prompts = _build_prompts(args)
print("Quantization plan:")
weight_block_size = _parse_block_size(args.weight_block_size)
print(f" input: {args.model}")
print(f" output: {output_dir}")
print(f" dtype: {dtype}")
print(f" height/width: {args.height}x{args.width}")
print(f" num_frames: {args.num_frames}")
print(f" calib_size: {len(prompts)}")
print(f" calib_steps: {args.calib_steps}")
print(f" quantize_mha: {args.quantize_mha}")
print(
f" weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
)
pipe = _load_pipeline(model_path, dtype)
backbone = pipe.transformer
quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
if weight_block_size is not None:
# Switch from per-tensor (default) to block-wise weight quantization.
# ModelOpt's wildcard "*weight_quantizer" matches every linear's weight quantizer.
quant_config["quant_cfg"]["*weight_quantizer"] = {
"num_bits": (4, 3), # E4M3 (FP8 weights, same as default)
"block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
}
print(
f" -> overriding weight quantizer with block_sizes={weight_block_size} "
f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
)
forward_loop = _build_forward_loop(pipe, args, prompts)
quantized = mtq.quantize(backbone, quant_config, forward_loop)
if quantized is not None:
pipe.transformer = quantized
backbone = quantized
_disable_known_problematic_quantizers(mtq, backbone, quantize_mha=args.quantize_mha)
print("\nForcing FP8 weight serialization (HV-1.5 isn't in ModelOpt's recognized-model registry,")
print("so we have to call the per-weight export helper ourselves)...")
exported = _force_export_quantized_weights(backbone, dtype)
print(f" -> {exported} weights converted to FP8 in memory")
if exported == 0:
raise SystemExit(
"No quantized weights were exported. Calibration may have skipped every layer "
"(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
"weight quantizers."
)
print("\nSaving pipeline with FP8 transformer...")
_save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
_patch_quant_config(output_dir, weight_block_size=weight_block_size)
print(f"Saved to: {output_dir}")
_summarize_export(output_dir)
print("\nNext: validate the checkpoint with vllm-omni:")
print(
" python examples/offline_inference/text_to_video/text_to_video.py \\\n"
f" --model {output_dir} \\\n"
" --quantization fp8 \\\n"
" --prompt 'A dog running across a field of golden wheat.' \\\n"
f" --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
" --num-inference-steps 30 --guidance-scale 6.0 --seed 42 \\\n"
" --output outputs/hv15_modelopt_fp8.mp4 \\\n"
" --enforce-eager"
)
print(
"\n (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
"checkpoint's config.json has modelopt metadata.)"
)
if __name__ == "__main__":
main()
@@ -0,0 +1,471 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Quantize Wan2.2 (TI2V-5B, 704x1280 T2V) to a ModelOpt FP8 Hugging Face checkpoint.
Calibrates the DiT transformer using a small video prompt set and exports a
diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
Layers kept full precision match the #2728 / #2795 pattern: condition embedder
(time/text/image), patch embedding, modulation (scale_shift_table), final
norm + proj_out, and sequence-parallel helpers. All attention + FFN linears
are quantized — static calibration handles the numerics that online FP8
couldn't (see #2920 ablation).
Default target is `Wan-AI/Wan2.2-TI2V-5B-Diffusers`, the dense 5B variant that
fits 80GB BF16. The A14B MoE variants need 2+ GPUs and are out of scope here.
Example:
python examples/quantization/quantize_wan2_2_modelopt_fp8.py \\
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \\
--output ./wan22-ti2v-modelopt-fp8 \\
--overwrite
"""
from __future__ import annotations
import argparse
import copy
import json
import re
import shutil
import sys
from pathlib import Path
from typing import Any
import torch
from diffusers import DiffusionPipeline
DEFAULT_PROMPTS = [
"A dog running across a field of golden wheat.",
"An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
"A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
"A crackling campfire at night under a starry sky, sparks rising into the dark.",
"An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
"A close-up of a blooming rose covered in morning dew, soft natural light.",
"A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
"A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
]
def _build_parser() -> argparse.ArgumentParser:
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--model", required=True, help="Input Wan2.2 diffusers directory or HF id.")
p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
p.add_argument("--height", type=int, default=704, help="Calibration video height (Wan2.2 TI2V-5B native: 704).")
p.add_argument("--width", type=int, default=1280, help="Calibration video width (Wan2.2 TI2V-5B native: 1280).")
p.add_argument(
"--num-frames",
type=int,
default=49,
help="Frames per calibration sample. 49 matches the typical short benchmark; "
"use 17 to reduce memory pressure during calibration.",
)
p.add_argument("--guidance-scale", type=float, default=5.0)
p.add_argument(
"--calib-steps",
type=int,
default=10,
help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
)
p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
p.add_argument("--seed", type=int, default=42)
p.add_argument(
"--prompt",
action="append",
default=[],
help="Custom calibration prompt. Repeat to provide multiple.",
)
p.add_argument(
"--quantize-mha",
action="store_true",
help="Enable FP8 attention K/V/softmax quantizers. Off by default — Wan2.2's long attention "
"sequences amplified FP8 drift in the online ablation (see #2920).",
)
p.add_argument(
"--weight-block-size",
type=str,
default=None,
help="Per-block weight quantization as 'M,N' (e.g. '128,128'). Default per-tensor. "
"Note: vllm-omni's ModelOpt adapter may not yet dispatch block-wise scales — check #2924 "
"for the HV-1.5 investigation status before relying on this.",
)
p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
return p
def _parse_block_size(spec: str | None) -> list[int] | None:
if spec is None:
return None
parts = [int(x) for x in spec.split(",") if x.strip()]
if len(parts) != 2:
raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
return parts
def _require_modelopt() -> Any:
try:
import modelopt.torch.quantization as mtq
except ModuleNotFoundError as exc:
raise SystemExit(
"NVIDIA ModelOpt is not installed. Install with:\n"
" pip install 'nvidia-modelopt[all]'\n"
f"Original error: {exc}"
) from exc
return mtq
def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
model_path = args.model
output_dir = Path(args.output).expanduser().resolve()
if output_dir.exists():
if not args.overwrite:
raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
shutil.rmtree(output_dir)
return model_path, output_dir
def _select_dtype(name: str) -> torch.dtype:
return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
def _build_prompts(args: argparse.Namespace) -> list[str]:
prompts = args.prompt or DEFAULT_PROMPTS
if args.calib_size <= 0:
raise SystemExit("--calib-size must be positive.")
if len(prompts) < args.calib_size:
repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
prompts = (prompts * repeats)[: args.calib_size]
return prompts[: args.calib_size]
# Layers to KEEP at full precision. Wan2.2's module naming:
# - condition_embedder: time_embedder, time_proj, text_embedder, image_embedder (I2V)
# - patch_embedding: Conv3dLayer (already not Linear, belt-and-suspenders skip)
# - scale_shift_table: nn.Parameter modulation (not Linear, but pattern guard)
# - norm_out: AdaLayerNorm final
# - proj_out: final nn.Linear
# - timestep_proj_prepare / output_scale_shift_prepare: SP helpers
def _filter_func_wan22(name: str) -> bool:
pattern = re.compile(
r"(proj_out.*|"
r".*(condition_embedder|patch_embedding|"
r"norm_out|scale_shift_table|"
r"timestep_proj_prepare|output_scale_shift_prepare).*)"
)
return pattern.match(name) is not None
def _mha_filter_func(name: str) -> bool:
pattern = re.compile(
r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
)
return pattern.match(name) is not None
def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
if not hasattr(mtq, "disable_quantizer"):
return
mtq.disable_quantizer(backbone, _filter_func_wan22)
if not quantize_mha:
mtq.disable_quantizer(backbone, _mha_filter_func)
def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
if hasattr(pipe, "set_progress_bar_config"):
pipe.set_progress_bar_config(disable=True)
pipe.to("cuda")
return pipe
def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
generator = torch.Generator(device="cuda")
# Try setting guidance on the pipeline's guider if present (newer diffusers APIs).
guider = getattr(pipe, "guider", None)
if guider is not None and hasattr(guider, "guidance_scale"):
try:
guider.guidance_scale = args.guidance_scale
except Exception:
pass
base_kwargs = dict(
height=args.height,
width=args.width,
num_frames=args.num_frames,
num_inference_steps=args.calib_steps,
output_type="latent",
)
def forward_loop(*_unused_args, **_unused_kwargs) -> None:
with torch.inference_mode():
for idx, prompt in enumerate(prompts):
generator.manual_seed(args.seed + idx)
# Try with guidance_scale first; fall back without on TypeError
# for pipelines that take CFG via guider config only.
try:
pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
except TypeError as exc:
if "guidance_scale" not in str(exc):
raise
pipe(prompt=prompt, generator=generator, **base_kwargs)
return forward_loop
def _summarize_export(output_dir: Path) -> None:
cfg_path = output_dir / "transformer" / "config.json"
if not cfg_path.exists():
print(f"[warn] {cfg_path} missing.", file=sys.stderr)
return
with cfg_path.open(encoding="utf-8") as f:
cfg = json.load(f)
qc = cfg.get("quantization_config")
if not isinstance(qc, dict):
print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
return
print("Export summary:")
print(f" quant_method: {qc.get('quant_method')}")
print(f" quant_algo: {qc.get('quant_algo')}")
producer = qc.get("producer")
if isinstance(producer, dict):
print(f" producer: {producer.get('name')} {producer.get('version')}")
print(f" config path: {cfg_path}")
def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
"""Convert in-memory weights of quantized modules to actual FP8 storage.
`export_hf_checkpoint` skips this step for unknown model types (Wan2.2 isn't
in ModelOpt's recognized-model registry), so we must call the per-weight
export helper ourselves. Same workaround as the HunyuanVideo-1.5 / HunyuanImage-3
calibration helpers.
"""
from modelopt.torch.export.quant_utils import (
QUANTIZATION_NONE,
get_quantization_format,
quantizer_attr_names,
weight_attr_names,
)
from modelopt.torch.export.unified_export_hf import _export_quantized_weight
exported = 0
for name, module in backbone.named_modules():
try:
quantization_format = get_quantization_format(module)
except Exception as exc:
print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
continue
if quantization_format == QUANTIZATION_NONE:
continue
for weight_name in weight_attr_names(module):
quantizer_attrs = quantizer_attr_names(weight_name)
weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
continue
_export_quantized_weight(module, dtype, weight_name)
exported += 1
return exported
def _wan22_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
"""Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913)."""
weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
if weight_block_size is not None:
weights_cfg["strategy"] = "block"
weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
return {
"config_groups": {
"group_0": {
"input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
"weights": weights_cfg,
"targets": ["Linear"],
}
},
"ignore": [
"condition_embedder*",
"norm_out*",
"output_scale_shift_prepare*",
"patch_embedding*",
"proj_out*",
"scale_shift_table*",
"timestep_proj_prepare*",
],
"producer": {"name": "modelopt"},
"quant_algo": "FP8",
"quant_method": "modelopt",
}
def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
"""Inject quant_algo: FP8 + config_groups into each transformer's config.json
so vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8.
Patches both ``transformer`` and ``transformer_2`` (A14B) when present."""
for sub in ("transformer", "transformer_2"):
cfg_path = output_dir / sub / "config.json"
if not cfg_path.exists():
continue
with cfg_path.open(encoding="utf-8") as f:
cfg = json.load(f)
new_qc = _wan22_quant_config_block(weight_block_size=weight_block_size)
existing = cfg.get("quantization_config")
if isinstance(existing, dict):
producer = existing.get("producer")
if isinstance(producer, dict):
new_qc["producer"] = producer
cfg["quantization_config"] = new_qc
with cfg_path.open("w", encoding="utf-8") as f:
json.dump(cfg, f, indent=2)
def _list_transformers(pipe: DiffusionPipeline) -> list[tuple[str, torch.nn.Module]]:
"""Return [(attr_name, module)] for every DiT on the pipeline.
Wan2.2 A14B has both ``transformer`` (low-noise) and ``transformer_2``
(high-noise). TI2V-5B has only ``transformer``.
"""
out = []
for attr in ("transformer", "transformer_2"):
mod = getattr(pipe, attr, None)
if mod is not None:
out.append((attr, mod))
return out
def _save_pipeline_with_fp8_transformer(
pipe: DiffusionPipeline,
model_path: str,
output_dir: Path,
max_shard_size: str = "5GB",
) -> None:
"""Copy source dir verbatim minus transformer dirs, then save each quantized DiT.
Handles A14B's dual transformer (``transformer`` + ``transformer_2``) as well
as the single-DiT TI2V-5B.
"""
from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
src = Path(model_path)
if not src.exists():
from huggingface_hub import snapshot_download
src = Path(snapshot_download(model_path))
if output_dir.exists():
shutil.rmtree(output_dir)
shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer", "transformer_2"))
for attr, backbone in _list_transformers(pipe):
transformer_out = output_dir / attr
# Pass the nn.Module (transformer), not the Pipeline wrapper.
with hide_quantizers_from_state_dict(backbone):
backbone.save_pretrained(
str(transformer_out),
safe_serialization=True,
max_shard_size=max_shard_size,
)
def main() -> None:
args = _build_parser().parse_args()
if not torch.cuda.is_available():
raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
mtq = _require_modelopt()
model_path, output_dir = _ensure_paths(args)
dtype = _select_dtype(args.dtype)
prompts = _build_prompts(args)
weight_block_size = _parse_block_size(args.weight_block_size)
print("Quantization plan:")
print(f" input: {args.model}")
print(f" output: {output_dir}")
print(f" dtype: {dtype}")
print(f" height/width: {args.height}x{args.width}")
print(f" num_frames: {args.num_frames}")
print(f" calib_size: {len(prompts)}")
print(f" calib_steps: {args.calib_steps}")
print(f" quantize_mha: {args.quantize_mha}")
print(
f" weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
)
pipe = _load_pipeline(model_path, dtype)
transformers = _list_transformers(pipe)
if not transformers:
raise SystemExit("Pipeline has no transformer or transformer_2 attribute.")
print(f" found {len(transformers)} transformer(s): {', '.join(a for a, _ in transformers)}")
quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
if weight_block_size is not None:
quant_config["quant_cfg"]["*weight_quantizer"] = {
"num_bits": (4, 3),
"block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
}
print(
f" -> overriding weight quantizer with block_sizes={weight_block_size} "
f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
)
forward_loop = _build_forward_loop(pipe, args, prompts)
# Quantize each DiT. The forward_loop runs the full pipeline, so quantizers
# on every attached transformer (A14B: transformer + transformer_2) get
# calibrated during the same passes.
for attr, backbone in transformers:
print(f"\nQuantizing {attr}...")
quantized = mtq.quantize(backbone, quant_config, forward_loop)
if quantized is not None:
setattr(pipe, attr, quantized)
_disable_known_problematic_quantizers(mtq, getattr(pipe, attr), quantize_mha=args.quantize_mha)
# Export weights to FP8 only AFTER all transformers are calibrated. Exporting
# inside the loop would convert transformer's weights to real FP8, and the
# next transformer's forward_loop (full pipeline) would then fake-quantize an
# already-FP8 weight -> "fake_e4m3fy not implemented for Float8_e4m3fn".
print("\nForcing FP8 weight serialization (Wan2.2 isn't in ModelOpt's recognized-model registry,")
print("so we have to call the per-weight export helper ourselves)...")
total_exported = 0
for attr, backbone in _list_transformers(pipe):
exported = _force_export_quantized_weights(backbone, dtype)
print(f" -> {exported} weights converted to FP8 in {attr}")
total_exported += exported
if total_exported == 0:
raise SystemExit(
"No quantized weights were exported. Calibration may have skipped every layer "
"(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
"weight quantizers."
)
print("\nSaving pipeline with FP8 transformer(s)...")
_save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
_patch_quant_config(output_dir, weight_block_size=weight_block_size)
print(f"Saved to: {output_dir}")
_summarize_export(output_dir)
print("\nNext: validate the checkpoint with vllm-omni:")
print(
" python examples/offline_inference/text_to_video/text_to_video.py \\\n"
f" --model {output_dir} \\\n"
" --quantization fp8 \\\n"
" --prompt 'A dog running across a field of golden wheat.' \\\n"
f" --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
" --num-inference-steps 30 --guidance-scale 5.0 --seed 42 \\\n"
" --output outputs/wan22_modelopt_fp8.mp4"
)
print(
"\n (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
"checkpoint's config.json has modelopt metadata.)"
)
if __name__ == "__main__":
main()
+9 -1
View File
@@ -64,7 +64,9 @@ dev = [
"lpips>=0.1.4", # for quantization fp8
"orjson>=3.11.7", # to make v1/embeddings API fast
"voxcpm>=2.0.2", # for VoxCPM2 path-locating in tests
"torchdiffeq>=0.2.5" # for covo_audio Test
"torchdiffeq>=0.2.5", # for covo_audio Test
"FlagEmbedding", # for hunyuan_image3 tests
"kernels<0.15.0" # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
]
demo = [
@@ -88,6 +90,12 @@ minicpmo = [
"stepaudio2-minicpmo",
]
# Optional Blackwell-only fused-bias FP8 GEMM; package is `quack-kernels` (imported
# as `quack`), not the unrelated `quack`. See docs/user_guide/quantization/fp8.md.
quack = [
"quack-kernels>=0.3.11",
]
docs = [
"mkdocs>=1.5.0",
"mkdocs-api-autonav",
+110
View File
@@ -0,0 +1,110 @@
# Higgs-Audio V3 TTS
> Multilingual text-to-speech with voice cloning on 1×H100
## Summary
- Vendor: Boson AI
- Model: `bosonai/higgs-audio-v3-tts-4b`
- Task: Text-to-speech synthesis with optional voice cloning (100+ languages)
- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API; offline batch inference
- Maintainer: @yuekaiz
## When to use this recipe
Use this recipe to serve `bosonai/higgs-audio-v3-tts-4b` for high-quality
multilingual TTS. The model generates 24 kHz speech, supports zero-shot voice
cloning from a reference clip, and handles 100+ languages with inline control
tokens for emotion, style, and prosody. The architecture is a ~4B Qwen3
backbone with fused multi-codebook embedding/head (8 codebooks × 1026 vocab,
MusicGen-style delay pattern).
## References
- Model card: [bosonai/higgs-audio-v3-tts-4b](https://huggingface.co/bosonai/higgs-audio-v3-tts-4b)
- Offline example: [`examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py`](../../examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py)
- Online example: [`examples/online_serving/text_to_speech/higgs_audio_v3/`](../../examples/online_serving/text_to_speech/higgs_audio_v3/)
- Benchmark results: see Performance section below
## Hardware Support
## GPU
### 1×H100 80GB
#### Environment
- OS: Linux
- Python: 3.12+
- CUDA: 12.x
- vLLM version: 0.22.0
- vLLM-Omni version or commit: `36e048fd` (branch `higgs-v3`)
#### Command
**Online serving:**
```bash
vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
--host 0.0.0.0 --port 8095 \
--trust-remote-code --omni
```
The default deploy config `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` is
loaded automatically by model registry (HF `model_type=higgs_multimodal_qwen3`).
**Offline batch inference:**
```bash
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \
--texts "Hello world." "The quick brown fox jumps over the lazy dog." \
--output-dir results/higgs_v3_wavs
```
**Offline voice clone:**
```bash
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \
--texts "Text to synthesize in the cloned voice." \
--ref-audio path/to/reference.wav \
--ref-text "Transcript of the reference clip." \
--output-dir results/higgs_v3_clone
```
#### Verification
Basic TTS via curl:
```bash
curl -X POST http://localhost:8095/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"model": "bosonai/higgs-audio-v3-tts-4b",
"input": "Hello, how are you?"
}' --output hello.wav
```
Voice clone via Python client:
```bash
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
--base-url http://localhost:8095 \
--model bosonai/higgs-audio-v3-tts-4b \
--ref-audio path/to/reference.wav \
--ref-text "Transcript of the reference." \
--prompts "Text to clone."
```
#### Notes
- Memory usage: Stage 0 (talker, ~4B) uses ~60% GPU memory; Stage 1 (codec decoder) uses ~25%.
- Key flags: `--trust-remote-code` and `--omni` are required.
- Output: 24 kHz mono WAV.
- Voice cloning: `ref_audio` accepts WAV/FLAC/MP3; `ref_text` is optional but improves fidelity.
- Deploy config: `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` (auto-discovered from `model_type`).
- `max_num_seqs=16` for both stages.
- `enforce_eager=true` for both stages (CUDA graph available for stage 0 but no throughput gain at batch>1).
- Known limitations:
- CUDA graph for stage 0 helps batch=1 (-13% RTF) but not batch>1 due to model-owned sampler running outside graph.
- Stage 1 (code2wav) must use `enforce_eager=true` (`@torch.inference_mode` incompatible with graph capture).
- Async streaming (chunk-based) not yet implemented; sync-only pipeline.
+4 -2
View File
@@ -33,11 +33,12 @@ recipes/
| [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
| [`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-0.5B.md`](./inclusionAI/Ming-omni-tts-0.5B.md) | Offline + online dense Ming TTS/audio generation | 1x AMD `gfx942` (ROCm 7.2) |
| [`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 |
| [`nvidia/Cosmos3-Nano.md`](./nvidia/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video generation, text to video with sound | 1x H200 141GB / B300 |
| [`nvidia/Cosmos3-Super.md`](./nvidia/Cosmos3-Super.md) | 64B T2I / T2V / I2V generation (+ optional audio) | 8x H200/H100/A100 / 2x H200 / B300 |
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
| [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
@@ -50,6 +51,7 @@ recipes/
| [`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 |
| [`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 |
|
+139
View File
@@ -0,0 +1,139 @@
# Stable Diffusion 3.5
> Text-to-image serving with Stability AI's SD 3.5 family
## Summary
- Vendor: Stability AI
- Model: `stabilityai/stable-diffusion-3.5-medium`, `stabilityai/stable-diffusion-3.5-large`
- Task: Text-to-image generation
- Mode: Online serving with the OpenAI-compatible images API
- Maintainer: Community
## When to use this recipe
Use this recipe when you want a lightweight, fast text-to-image model that
fits on a single consumer or workstation GPU. SD 3.5-medium requires only
~17 GiB VRAM and generates 1024×1024 images in ~4 seconds, making it one of
the most accessible diffusion models in vLLM-Omni. SD 3.5-large produces
noticeably sharper images at the cost of ~34 GiB VRAM and ~21 seconds per
image.
## References
- Model (medium): <https://huggingface.co/stabilityai/stable-diffusion-3.5-medium>
- Model (large): <https://huggingface.co/stabilityai/stable-diffusion-3.5-large>
- Related example under `examples/`:
[`examples/online_serving/text_to_image/README.md`](../../examples/online_serving/text_to_image/README.md)
- User guide:
[`docs/user_guide/examples/online_serving/text_to_image.md`](../../docs/user_guide/examples/online_serving/text_to_image.md)
## Hardware Support
This recipe documents two CUDA GPU configurations on a single RTX A6000 48 GB,
one for each model variant. Extend it with more hardware sections as community
validation lands.
## GPU
### 1x NVIDIA RTX A6000 48GB — stable-diffusion-3.5-medium
#### Environment
- OS: Ubuntu 22.04
- Python: 3.12
- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0
- vLLM version: 0.22.0
- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d
#### Command
```bash
vllm serve stabilityai/stable-diffusion-3.5-medium --omni --port 8091
```
#### Verification
```bash
curl -X POST http://localhost:8091/v1/images/generations \
-H "Content-Type: application/json" \
-d '{
"model": "stabilityai/stable-diffusion-3.5-medium",
"prompt": "a cute cat sitting on a windowsill, highly detailed, sharp focus",
"negative_prompt": "blurry, out of focus",
"guidance_scale": 4.5,
"n": 1,
"size": "1024x1024"
}' | python3 -c "
import json, base64, sys
data = json.load(sys.stdin)
with open('output.png', 'wb') as f:
f.write(base64.b64decode(data['data'][0]['b64_json']))
print('saved output.png')
"
```
Or use the example client:
```bash
python examples/online_serving/text_to_image/openai_chat_client.py \
--server http://localhost:8091 \
--prompt "a cute cat sitting on a windowsill" \
--output output.png
```
#### Notes
- Memory usage: Model weights ~15.6 GiB, peak VRAM ~17.0 GiB (no CFG), ~20.9 GiB with `guidance_scale > 1` (CFG enabled).
- Generation time: ~4 seconds without CFG, ~8 seconds with `guidance_scale=4.5`.
- Model loading time: ~55 seconds (including weight download on first run).
- `negative_prompt` requires `guidance_scale > 1` to take effect; SD3 defaults to 1.0 (no CFG).
### 1x NVIDIA RTX A6000 48GB — stable-diffusion-3.5-large
#### Environment
- OS: Ubuntu 22.04
- Python: 3.12
- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0
- vLLM version: 0.22.0
- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d
#### Command
```bash
vllm serve stabilityai/stable-diffusion-3.5-large --omni --port 8091
```
> **Note**: `stable-diffusion-3.5-large` is a gated model. Accept the license
> on [Hugging Face](https://huggingface.co/stabilityai/stable-diffusion-3.5-large)
> and set `HF_TOKEN` before running.
#### Verification
```bash
curl -X POST http://localhost:8091/v1/images/generations \
-H "Content-Type: application/json" \
-d '{
"model": "stabilityai/stable-diffusion-3.5-large",
"prompt": "a cute cat sitting on a windowsill, highly detailed, sharp focus",
"negative_prompt": "blurry, out of focus",
"guidance_scale": 4.5,
"n": 1,
"size": "1024x1024"
}' | python3 -c "
import json, base64, sys
data = json.load(sys.stdin)
with open('output.png', 'wb') as f:
f.write(base64.b64decode(data['data'][0]['b64_json']))
print('saved output.png')
"
```
#### Notes
- Memory usage: Peak VRAM ~31.6 GiB at 1024×1024, ~45.5 GiB at 2048×2048 with `guidance_scale=4.5`.
- Generation time: ~19 seconds at 1024×1024, ~102 seconds at 2048×2048.
- Image quality is noticeably better than medium, with sharper details.
- `negative_prompt` requires `guidance_scale > 1` to take effect; SD3 defaults to 1.0 (no CFG).
- 2048×2048 fits within 48 GB but leaves little headroom (~45.5 GiB used out of 49 GiB).
@@ -1,28 +1,55 @@
# Cosmos3-Nano
> Text-to-image, text-to-video, and image-to-video serving
> Text-to-image, text-to-video, image-to-video, and video-to-video serving
## Summary
- Vendor: NVIDIA
- Model: `nvidia/Cosmos3-Nano`
- Task: Text-to-image (T2I), text-to-video (T2V), and image-to-video (I2V) generation, with optional synchronized audio (video + sound)
- Task: Text-to-image (T2I), text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) generation, with optional synchronized audio (video + sound), action policy
- Mode: Online serving with the OpenAI-compatible image/video APIs, plus offline generation via the `Omni` API
- Maintainer: Community
## When to use this recipe
Use this recipe to deploy `nvidia/Cosmos3-Nano` for image and video generation.
A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves all three modes;
the mode is selected per request:
A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves these modes; the
mode is selected per request:
- **T2I** — `POST /v1/images/generations` (or a prompt carrying `modalities=["image"]`).
- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image.
- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image/video.
- **I2V** — `POST /v1/videos/sync` with a reference image (`input_reference` file
upload, or `image_reference` JSON).
- **V2V** — `POST /v1/videos/sync` with a reference video (`input_reference`
video upload, or `video_reference` JSON). Cosmos3 conditions on selected
reference-video latent frames; use `extra_params.condition_frame_indexes_vision`
and `extra_params.condition_video_keep` to choose which prefix/tail frames guide
generation.
- **T2VS / I2VS** — add `generate_sound=true` (and optional `sound_duration`) to a
T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into
the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples.
- **Action** — pass `extra_params={"action_mode": ...}` to drive Physical-AI tasks:
- `forward_dynamics` — given a first frame or video **and** an action trajectory,
roll out the resulting video. Synchronous: `POST /v1/videos/sync`.
- `policy` — given a first frame or video and a language instruction,
**predict** the action trajectory (and a rollout video). Use the async
`POST /v1/videos` endpoint and read the predicted action from the top-level
`action` field.
- `inverse_dynamics` — given a video, **recover** the action trajectory. Use
the async `POST /v1/videos` endpoint and read the recovered action from
the top-level `action` field
(`{data, shape, dtype, raw_action_dim, domain_id}`).
Action requests also take `domain_name` (e.g. `av`, `bridge_orig_lerobot`,
`droid_lerobot`, `agibotworld`, …; or a numeric `domain_id`), `raw_action_dim`,
and `action_chunk_size` (must equal `num_frames` or `num_frames - 1`). For
`forward_dynamics` also pass the `action` array. The dedicated policy checkpoint
**`nvidia/Cosmos3-Nano-Policy-DROID`** is served the same way
(`domain_name=droid_lerobot`).
Action requests can use `input_reference` or `video_reference` for video input.
`policy` and `forward_dynamics` can also use an image reference; `inverse_dynamics`
requires a video reference.
## References
@@ -125,6 +152,36 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
-F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
-o cosmos3_i2v.mp4
# Video-to-video -> /v1/videos/sync with an uploaded reference video.
# By default Cosmos3 conditions on latent indexes [0, 1]. For the default
# temporal VAE stride this decodes only the first 5 input frames.
# The model works best when the prompt describes the actual situation happening in the video.
# Generic prompts may create sub-standard generations.
curl -sS -X POST http://localhost:8000/v1/videos/sync \
-H "Accept: video/mp4" \
-F "model=nvidia/Cosmos3-Nano" \
-F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
-F "negative_prompt=blurry, distorted, low quality, jittery, deformed" \
-F "size=1280x720" -F "num_frames=189" -F "fps=24" \
-F "num_inference_steps=35" -F "guidance_scale=6.0" \
-F "max_sequence_length=4096" -F "flow_shift=10.0" \
-F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true,"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}' \
-F "seed=2222" \
-F "input_reference=@/path/to/reference.mp4;type=video/mp4" \
-o cosmos3_v2v.mp4
# V2V can also use a JSON-safe URL/data-URL video reference. Do not combine
# video_reference with input_reference or image_reference.
curl -sS -X POST http://localhost:8000/v1/videos/sync \
-H "Accept: video/mp4" \
-F "model=nvidia/Cosmos3-Nano" \
-F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
-F "size=1280x720" -F "num_frames=189" -F "fps=24" \
-F "num_inference_steps=35" -F "guidance_scale=6.0" \
-F "max_sequence_length=4096" -F "flow_shift=10.0" \
-F 'extra_params={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"last"}' \
-F 'video_reference={"video_url":"https://example.com/reference.mp4"}' \
-o cosmos3_v2v_from_url.mp4
# Text-to-video-with-sound
curl -sS -X POST http://localhost:8000/v1/videos/sync \
@@ -144,6 +201,50 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
-F "sound_duration=7.875" \
-F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
-o cosmos3_t2v_with_sound.mp4
# Action — forward dynamics (first frame + action trajectory -> rollout video).
# Synchronous; `action` is a JSON array shaped [action_chunk_size, raw_action_dim].
curl -sS -X POST http://localhost:8000/v1/videos/sync \
-H "Accept: video/mp4" \
--form-string "model=nvidia/Cosmos3-Nano" \
--form-string "prompt=You are an autonomous vehicle. This video is captured from a first-person perspective." \
-F "input_reference=@first_frame.jpg;type=image/jpeg" \
-F "size=640x480" -F "num_frames=61" -F "fps=10" \
-F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \
--form-string "extra_params={\"action_mode\":\"forward_dynamics\",\"domain_name\":\"av\",\"raw_action_dim\":9,\"action_chunk_size\":60,\"action\":$(cat action.json)}" \
-F "seed=0" \
-o cosmos3_forward_dynamics.mp4
# Action — policy (first frame + instruction -> predicted action trajectory + video).
# Asynchronous: POST returns a job id; poll, then read the predicted action from
# the top-level `action` field ({data, shape, dtype, raw_action_dim, domain_id}).
VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
-H "Accept: application/json" \
--form-string "model=nvidia/Cosmos3-Nano" \
--form-string "prompt=Pick up the banana and place it in the bowl." \
-F "input_reference=@first_frame.jpg;type=image/jpeg" \
-F "size=640x480" -F "num_frames=17" -F "fps=5" \
-F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \
--form-string 'extra_params={"action_mode":"policy","domain_name":"bridge_orig_lerobot","raw_action_dim":10,"action_chunk_size":16}' \
-F "seed=0" | jq -r '.id')
# poll until status == completed, then:
curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_policy.mp4
# Action — inverse dynamics (video -> recovered action trajectory).
# Asynchronous: use the job metadata to read the recovered action.
VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
-H "Accept: application/json" \
--form-string "model=nvidia/Cosmos3-Nano" \
--form-string "prompt=Recover the robot action trajectory from this clip." \
-F "input_reference=@motion_clip.mp4;type=video/mp4" \
-F "size=640x480" -F "num_frames=17" -F "fps=5" \
-F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \
--form-string 'extra_params={"action_mode":"inverse_dynamics","domain_name":"bridge_orig_lerobot","raw_action_dim":10,"action_chunk_size":16}' \
-F "seed=0" | jq -r '.id')
# poll until status == completed, then:
curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_inverse_dynamics.mp4
```
#### Notes
@@ -152,6 +253,7 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
- T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s**
- T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s**
- I2V 1280×720, 189 frames @ 35 steps → ~**99 s**
- Action 640×480 @ 30 steps — forward-dynamics 61f ~**4 s**, policy 17f ~**1–3 s**.
- Guardrails-on overhead: ~8% on T2I, negligible on video.
- **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 GPU;
full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer)
@@ -159,22 +261,27 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
- **Determinism:** identical seed reproduces identical output on the same
hardware; outputs are not bit-identical across different GPU types.
- **Supported sizes (per model card):** 256p / 480p / 720p at 16:9, 4:3, 1:1,
3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V 1280×720,
189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`.
3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V/V2V
1280×720, 189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`.
- **Key flags / params:** `--no-guardrails` (server) or
`extra_params={"guardrails":false}` (per request) toggles safety. The
per-request flag only takes effect when the server was launched **with**
guardrails enabled (it cannot re-enable them on a `--no-guardrails` server).
`use_resolution_template` / `use_duration_template` are off by default and only
needed when not using upsampled prompts that already encode resolution/duration.
For V2V, `condition_frame_indexes_vision` selects the clean conditioned latent
frame indexes (default `[0, 1]`), and `condition_video_keep` selects whether the
API decodes the first or last needed reference frames (`"first"` by default).
- **Known limitations:**
- Guardrails-on requires `cosmos-guardrail` **and** access to the gated
`nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise
the server fails at pipeline build with a gated-repo / safety-checker error.
- A guardrail-blocked prompt currently returns HTTP 500
(`"Guardrail blocked prompt"`).
- Action (policy / forward- / inverse-dynamics) modalities are not part of
this integration yet.
- Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported
online. Use async `POST /v1/videos` when you need the predicted/recovered
action payload under the top-level `action` field; sync `/v1/videos/sync`
returns raw MP4 bytes and does not expose action metadata in the response body.
### 1x GPU (Offline generation)
@@ -238,8 +345,9 @@ python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image
- Same `Cosmos3OmniDiffusersPipeline` as online; mode is chosen by
`prompt["modalities"]` (`["image"]` → T2I, `["video"]` → T2V) plus
`num_frames`/`fps`, and `multi_modal_data={"image": ...}` for I2V. For video,
frames are returned in `outputs[0].request_output.images` as an
`num_frames`/`fps`, `multi_modal_data={"image": ...}` for I2V, and
`multi_modal_data={"video": [<PIL frames>]}` or a video tensor/array for V2V.
For video, frames are returned in `outputs[0].request_output.images` as an
`(B, F, H, W, 3)` array.
- The offline entry must be guarded by `if __name__ == "__main__":` — the engine
spawns workers with the `spawn` start method.
@@ -1,12 +1,12 @@
# Cosmos3-Super
> Frontier 64B world model: text-to-image, text-to-video, image-to-video (+ optional audio)
> Frontier 64B world model: text-to-image, text-to-video, image-to-video, video-to-video (+ optional audio)
## Summary
- Vendor: NVIDIA
- Model: `nvidia/Cosmos3-Super` (64B; also `Cosmos3-Super-Text2Image`, `Cosmos3-Super-Image2Video`)
- Task: T2I, T2V, I2V generation, with optional synchronized audio (video + sound)
- Task: T2I, T2V, I2V, V2V generation, with optional synchronized audio (video + sound)
- Mode: Online serving with the OpenAI-compatible image/video APIs
- Maintainer: Community
@@ -16,8 +16,8 @@ Use this recipe to deploy the 64B `nvidia/Cosmos3-Super` for the highest-quality
Cosmos3 generation. It shares the same `Cosmos3OmniDiffusersPipeline` and request
formats as [Cosmos3-Nano](./Cosmos3-Nano.md) — only the checkpoint size and the
recommended parallelism differ. Mode is selected per request (T2I →
`/v1/images/generations`; T2V/I2V → `/v1/videos/sync`; add `generate_sound=true`
for audio).
`/v1/images/generations`; T2V/I2V/V2V → `/v1/videos/sync`; add
`generate_sound=true` for audio).
## References
@@ -62,8 +62,9 @@ disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs.
#### Verification
Requests are identical to Nano (see [`Cosmos3-Nano.md`](./Cosmos3-Nano.md) for full
T2I/T2V/I2V/T2VS curls); official params: `size=1280x720, num_frames=189, fps=24,
num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0, max_sequence_length=4096`.
T2I/T2V/I2V/V2V/T2VS curls); official params: `size=1280x720, num_frames=189,
fps=24, num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0,
max_sequence_length=4096`.
```bash
curl http://localhost:8000/v1/models
@@ -84,6 +85,15 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
-F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
-o cosmos3_super_i2v.mp4
# V2V — add an uploaded reference video. condition_video_keep can be "first" or "last".
curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
-F "model=nvidia/Cosmos3-Super" -F "prompt=Continue the same scene with smooth natural motion." \
-F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
-F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
-F 'extra_params={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}' \
-F "seed=2222" -F "input_reference=@/path/to/reference.mp4;type=video/mp4" \
-o cosmos3_super_v2v.mp4
# T2V + sound — add generate_sound/sound_duration (output muxes AAC 48 kHz stereo)
curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
-F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \
@@ -103,6 +113,12 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
- T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo**
- (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
- **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk.
- Same generation defaults, supported sizes, and `generate_sound`/`sound_duration`
semantics as Nano. Action (policy / forward- / inverse-dynamics) modalities are
not part of this integration yet.
- Same generation defaults, supported sizes, V2V reference-video controls
(`condition_frame_indexes_vision`, `condition_video_keep`), and
`generate_sound`/`sound_duration`
semantics as Nano, including the **action** modality: `forward_dynamics`,
`policy`, and `inverse_dynamics` — see the Cosmos3-Nano recipe for the request
shapes. Use async `/v1/videos` when you need predicted/recovered action metadata
under the top-level `action` field. Verified on the 64B Super under
`--cfg-parallel-size 2`: async `policy` returns the predicted action (`[16, 10]`)
and the rollout video reliably.
+9 -6
View File
@@ -57,12 +57,6 @@ Adjust `devices` in the YAML to match your hardware.
#### Command
Thinker only (text output):
```bash
vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091
```
Thinker + talker (text and/or audio output):
```bash
@@ -72,6 +66,15 @@ vllm serve Jonathan1909/Ming-flash-omni-2.0 \
--log-stats
```
Thinker only (text-only output):
```bash
vllm serve Jonathan1909/Ming-flash-omni-2.0 \
--omni \
--deploy-config vllm_omni/deploy/ming_flash_omni_thinker_only.yaml \
--port 8091
```
`--log-stats` is optional but recommended while validating the deployment.
#### Verification
+82
View File
@@ -0,0 +1,82 @@
# Ming-omni-tts 0.5B
> Offline and online dense Ming text-to-speech on ROCm
## Summary
- Vendor: inclusionAI
- Model: `inclusionAI/Ming-omni-tts-0.5B`
- Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls
- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API; offline inference
- Maintainer: Community
## References
- [Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
- [Upstream repository](https://github.com/inclusionAI/Ming-omni-tts)
- [Offline example](../../examples/offline_inference/text_to_speech/ming_tts/)
- [Online example](../../examples/online_serving/text_to_speech/ming_tts/)
## Hardware Support
## ROCm
### 1x AMD `gfx942`
#### Environment
- OS: Ubuntu 22.04.5 LTS, x86_64
- Python: 3.12.13
- ROCm / HIP: 7.2.53211
- vLLM version: 0.22.0
- vLLM-Omni version or commit: 0.1.dev1873 / `99c35c410`
- Docker image: `vllm/vllm-omni-rocm:v0.22.0`
#### Command
From the vLLM-Omni repository root:
```bash
docker run --rm \
--group-add=video \
--ipc=host \
--cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined \
--device /dev/kfd \
--device /dev/dri \
-v ~/.cache/huggingface:/root/.cache/huggingface \
-v "$PWD":/app/vllm-omni \
-w /app/vllm-omni \
-e VLLM_ROCM_USE_AITER=0 \
-p 8091:8091 \
vllm/vllm-omni-rocm:v0.22.0 \
--model inclusionAI/Ming-omni-tts-0.5B \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--omni \
--port 8091 \
--enforce-eager
```
#### Verification
```bash
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
--text "我觉得社会企业同个人都有责任" \
--instruction-json '{"方言":"广粤话"}' \
--ref-audio /path/to/yue_prompt.wav \
--max-new-tokens 200 \
--output dialect.wav
```
`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text`
for the dialect case.
## Notes
- The official ROCm image includes the platform dependencies.
- See the [ROCm installation guide](../../docs/getting_started/installation/gpu.md)
for interactive and source-build workflows.
- The [offline example](../../examples/offline_inference/text_to_speech/ming_tts/)
can be run from an interactive container.
- The tested environment uses `--enforce-eager`.
- Output is mono 44.1 kHz audio.
-1
View File
@@ -22,4 +22,3 @@ prettytable>=3.8.0
aenum==3.1.16
pyzmq>=25.0.0
janus>=1.0.0
pydub
+1
View File
@@ -1,2 +1,3 @@
-r common.txt
onnxruntime>=1.23.2
auto-round-lib
Binary file not shown.

After

Width:  |  Height:  |  Size: 1.1 MiB

@@ -345,6 +345,14 @@ def test_parse_score_payload_handles_raw_json_and_delimited_json():
assert parse_score_payload(wrapped)["score"] == [6]
def test_parse_score_payload_handles_qwen_vl_nested_score_dicts():
qwen_style = '{"score": [{"naturalness": 8, "artifact_free": 7}], "reasoning": "good quality"}'
assert parse_score_payload(qwen_style)["score"] == [8, 7]
nested_values = '{"score": [{"score": 9}, {"score": 6}], "reasoning": "ok"}'
assert parse_score_payload(nested_values)["score"] == [9, 6]
def test_summarize_gedit_generated_records_groups_by_task_and_language():
records = []
for group in GEDIT_GROUPS[:2]:
@@ -61,7 +61,7 @@ def test_stage0_streaming_update_discards_outstanding_async_placeholder_token()
sched._update_request_as_session(session, _make_update([10, 20]))
assert session.discard_latest_async_tokens is True
assert session.async_tokens_to_discard == 1
assert session.num_output_placeholders == 0
assert session.spec_token_ids == []
# The async placeholder makes token 9 unconfirmed, so only 7 and 8 are
@@ -83,7 +83,7 @@ def test_stage0_streaming_update_keeps_all_computed_tokens_without_placeholder()
sched._update_request_as_session(session, _make_update([10, 20]))
assert session.discard_latest_async_tokens is False
assert getattr(session, "async_tokens_to_discard", 0) == 0
assert session.num_output_placeholders == 0
assert session.prompt_token_ids == [1, 2, 3, 7, 8, 9, 10, 20]
assert list(session._all_token_ids) == [1, 2, 3, 7, 8, 9, 10, 20]
@@ -20,6 +20,7 @@ from vllm.sampling_params import SamplingParams
from vllm.v1.engine import EngineCoreEventType
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
# isort: on
@@ -142,3 +143,300 @@ class TestReplaceSessionWithStreamingUpdate:
assert session.events
assert session.events[-1].type == EngineCoreEventType.QUEUED
class _RealignSchedulerStub(OmniSchedulerMixin):
"""Minimal scheduler surface for ``_realign_request_status_to_queues``.
Pins ``self.requests`` / ``self.running`` / ``self.waiting`` to the
same shapes the upstream ``Scheduler.finish_requests`` reads, so the
helper sees realistic state without spinning a real scheduler.
"""
def __init__(
self,
*,
requests: dict,
running: list,
waiting: list,
) -> None:
self.requests = requests
self.running = running
self.waiting = waiting
class TestRealignRequestStatusToQueues:
"""Regression for the residual hang described in
https://github.com/vllm-project/vllm-omni/pull/3774 -- chunk-transfer
adapter's ``requests_origin_status`` table goes stale on the
``waiting → running`` admit transition, and an abort that lands
before the next deque round-trip stomps stale ``WAITING`` onto a
request that actually lives in ``self.running``. After
``max_num_seqs`` such aborts every ``input_batch`` slot is leaked and
new requests hang at ``chunks=0``.
"""
def test_running_with_stale_waiting_status_is_realigned_to_running(
self,
) -> None:
"""The exact race the hang reproduces: a request lives in
``self.running`` but its ``status`` is ``WAITING``. After
realign, status must be ``RUNNING`` so upstream
``Scheduler.finish_requests`` removes it from ``self.running``.
"""
req = _make_request(request_id="req-stale")
req.status = RequestStatus.WAITING # stale -- actually in running
sched = _RealignSchedulerStub(
requests={req.request_id: req},
running=[req],
waiting=[],
)
sched._realign_request_status_to_queues([req.request_id])
assert req.status == RequestStatus.RUNNING
def test_waiting_with_stale_running_status_is_realigned_to_waiting(
self,
) -> None:
"""Symmetric case: request lives in ``self.waiting`` but status
is ``RUNNING``. Upstream's else branch should run, so realign to
``WAITING``.
"""
req = _make_request(request_id="req-stale-2")
req.status = RequestStatus.RUNNING # stale -- actually in waiting
sched = _RealignSchedulerStub(
requests={req.request_id: req},
running=[],
waiting=[req],
)
sched._realign_request_status_to_queues([req.request_id])
assert req.status == RequestStatus.WAITING
def test_already_aligned_status_is_left_unchanged(self) -> None:
"""Healthy case: status matches actual queue. No status mutation
means no spurious side effects.
"""
req_running = _make_request(request_id="req-r")
req_running.status = RequestStatus.RUNNING
req_waiting = _make_request(request_id="req-w")
req_waiting.status = RequestStatus.WAITING
sched = _RealignSchedulerStub(
requests={
req_running.request_id: req_running,
req_waiting.request_id: req_waiting,
},
running=[req_running],
waiting=[req_waiting],
)
sched._realign_request_status_to_queues([req_running.request_id, req_waiting.request_id])
assert req_running.status == RequestStatus.RUNNING
assert req_waiting.status == RequestStatus.WAITING
def test_unknown_request_id_is_skipped_silently(self) -> None:
"""Aligning ids that aren't in ``self.requests`` (already freed
upstream, or never existed) must be a no-op.
"""
sched = _RealignSchedulerStub(requests={}, running=[], waiting=[])
sched._realign_request_status_to_queues(["never-existed"]) # no raise
def test_request_in_neither_queue_is_left_unchanged(self) -> None:
"""If a tracked request is in neither queue (e.g. parked in a
chunk-transfer deque), realign must not invent a status. The
adapter / deque purge owns that surface; the realign helper is
only for the admit-transition staleness.
"""
req = _make_request(request_id="req-parked")
req.status = RequestStatus.WAITING_FOR_CHUNK
sched = _RealignSchedulerStub(
requests={req.request_id: req},
running=[],
waiting=[],
)
sched._realign_request_status_to_queues([req.request_id])
assert req.status == RequestStatus.WAITING_FOR_CHUNK
def test_request_ids_str_is_treated_as_single_id(self) -> None:
"""Match upstream ``Scheduler.finish_requests`` resolution: a
bare string is treated as one id, not iterated as characters.
"""
req = _make_request(request_id="req-s")
req.status = RequestStatus.WAITING
sched = _RealignSchedulerStub(
requests={req.request_id: req},
running=[req],
waiting=[],
)
sched._realign_request_status_to_queues(req.request_id)
assert req.status == RequestStatus.RUNNING
def test_request_ids_none_aligns_every_known_request(self) -> None:
"""``request_ids=None`` matches upstream's "all requests" path.
Realign must walk every entry in ``self.requests`` and fix any
stale status it finds.
"""
stale = _make_request(request_id="req-stale-none")
stale.status = RequestStatus.WAITING # but actually in running
clean = _make_request(request_id="req-clean-none")
clean.status = RequestStatus.RUNNING
sched = _RealignSchedulerStub(
requests={
stale.request_id: stale,
clean.request_id: clean,
},
running=[stale, clean],
waiting=[],
)
sched._realign_request_status_to_queues(None)
assert stale.status == RequestStatus.RUNNING
assert clean.status == RequestStatus.RUNNING
def test_finished_request_is_skipped(self) -> None:
"""Already-finished requests must not be touched -- they may
have legitimate finished statuses (FINISHED_STOPPED etc.) that
a status flip would corrupt.
"""
req = _make_request(request_id="req-finished")
req.status = RequestStatus.FINISHED_STOPPED
sched = _RealignSchedulerStub(
requests={req.request_id: req},
running=[req], # not realistic, but we want to prove the guard fires
waiting=[],
)
sched._realign_request_status_to_queues([req.request_id])
assert req.status == RequestStatus.FINISHED_STOPPED
class TestPurgeFinishedFromRunning:
"""Regression for the residual ``self.running`` slot leak surface
paired with ``_realign_request_status_to_queues``: even after
realign + ``super().finish_requests`` runs, corner cases can leave
already-finished or untracked entries in ``self.running`` -- e.g.
a connector cleanup that pops from ``self.requests`` without
unwinding ``self.running``, or a ``status`` mid-transition when
finish ran. The defensive post-finish purge sweeps those residues
so the worker's ``input_batch`` slot never pins a freed request.
See https://github.com/vllm-project/vllm-omni/pull/3774 discussion
and the residual-hang reproduction in that PR.
"""
def test_finished_request_is_purged_from_running(self) -> None:
"""``is_finished()`` request lingering in ``self.running`` must
be swept so its ``input_batch`` slot is freed."""
finished = _make_request(request_id="req-finished")
finished.status = RequestStatus.FINISHED_STOPPED # is_finished() True
sched = _RealignSchedulerStub(
requests={finished.request_id: finished},
running=[finished],
waiting=[],
)
sched._purge_finished_from_running()
assert sched.running == []
def test_untracked_request_is_purged_from_running(self) -> None:
"""Request lingering in ``self.running`` but no longer present
in ``self.requests`` (already freed by upstream / connector)
must be swept."""
untracked = _make_request(request_id="req-untracked")
untracked.status = RequestStatus.RUNNING
sched = _RealignSchedulerStub(
requests={}, # already deleted from self.requests
running=[untracked],
waiting=[],
)
sched._purge_finished_from_running()
assert sched.running == []
def test_healthy_running_is_left_unchanged(self) -> None:
"""Live tracked running requests must not be touched -- the
purge is defensive, not aggressive."""
alive = _make_request(request_id="req-alive")
alive.status = RequestStatus.RUNNING
sched = _RealignSchedulerStub(
requests={alive.request_id: alive},
running=[alive],
waiting=[],
)
sched._purge_finished_from_running()
assert sched.running == [alive]
def test_empty_running_is_noop(self) -> None:
"""Empty ``self.running`` must short-circuit cleanly."""
sched = _RealignSchedulerStub(requests={}, running=[], waiting=[])
sched._purge_finished_from_running()
assert sched.running == []
def test_mixed_alive_and_dead_keeps_only_alive_in_order(self) -> None:
"""Mixed ``self.running`` -- some alive, some finished, some
untracked. Sweep keeps only the alive ones and preserves their
relative order."""
alive_a = _make_request(request_id="req-alive-a")
alive_a.status = RequestStatus.RUNNING
finished_b = _make_request(request_id="req-finished-b")
finished_b.status = RequestStatus.FINISHED_STOPPED
untracked_c = _make_request(request_id="req-untracked-c")
untracked_c.status = RequestStatus.RUNNING
alive_d = _make_request(request_id="req-alive-d")
alive_d.status = RequestStatus.RUNNING
sched = _RealignSchedulerStub(
requests={
alive_a.request_id: alive_a,
finished_b.request_id: finished_b, # tracked but is_finished
# untracked_c is NOT in self.requests
alive_d.request_id: alive_d,
},
running=[alive_a, finished_b, untracked_c, alive_d],
waiting=[],
)
sched._purge_finished_from_running()
assert sched.running == [alive_a, alive_d]
def test_in_place_mutation_keeps_helper_local_list_identity(self) -> None:
"""The slice-assign form (``self.running[:] = ...``) avoids an
extra rebind inside this helper. Note that across a full
``finish_requests`` call upstream ``Scheduler.finish_requests``
already rebinds ``self.running``, so global identity is not a
contract -- this test only pins the local helper behavior so a
future refactor away from in-place mutation is intentional.
"""
finished = _make_request(request_id="req-finished")
finished.status = RequestStatus.FINISHED_STOPPED
alive = _make_request(request_id="req-alive")
alive.status = RequestStatus.RUNNING
sched = _RealignSchedulerStub(
requests={
finished.request_id: finished,
alive.request_id: alive,
},
running=[finished, alive],
waiting=[],
)
held_ref = sched.running
sched._purge_finished_from_running()
assert held_ref is sched.running
assert held_ref == [alive]
@@ -0,0 +1,158 @@
"""Integration tests for the post-finish defensive purge of
``self.running`` placed in
``OmniARScheduler.finish_requests`` /
``OmniGenerationScheduler.finish_requests``.
Pins *placement* — that ``_purge_finished_from_running`` runs AFTER
``super().finish_requests`` so it catches entries upstream's removal
missed (e.g. status mid-transition, connector cleanup popping
``self.requests`` without unwinding ``self.running``). Direct helper
unit tests live in ``test_omni_scheduler_mixin.py``; these tests
exercise the full ``finish_requests`` path end-to-end.
Regression for the residual self.running slot leak surface paired
with ``_realign_request_status_to_queues`` (vllm-omni#3774).
"""
from __future__ import annotations
import pytest
from vllm.v1.request import RequestStatus
import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _StubRequest:
"""Minimal Request stub with the surface the helpers exercise."""
def __init__(self, request_id: str, status: RequestStatus) -> None:
self.request_id = request_id
self.status = status
def is_finished(self) -> bool:
return RequestStatus.is_finished(self.status)
_SCHEDULER_PARAMS = [
pytest.param(OmniARScheduler, ar_sched_mod, id="ar"),
pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"),
]
def _make_scheduler(scheduler_cls, *, requests, running, waiting):
"""Bare scheduler with just the surface ``finish_requests`` reads."""
scheduler = scheduler_cls.__new__(scheduler_cls)
scheduler.chunk_transfer_adapter = None
scheduler.input_coordinator = None
scheduler.requests = requests
scheduler.running = running
scheduler.waiting = waiting
return scheduler
@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
def test_finish_requests_purges_finished_request_from_running(
monkeypatch: pytest.MonkeyPatch,
scheduler_cls,
scheduler_mod,
) -> None:
"""A request whose ``status`` is a finished value but still lives
in ``self.running`` after ``super().finish_requests`` returns must
be swept by the post-finish purge. Mocks ``super()`` to simulate
the corner case where upstream's status-driven removal missed the
entry.
"""
finished = _StubRequest("req-finished", RequestStatus.FINISHED_STOPPED)
scheduler = _make_scheduler(
scheduler_cls,
requests={finished.request_id: finished},
running=[finished],
waiting=[],
)
def fake_finish_requests(self, request_ids, finished_status):
# super() returning the abort tuple but leaving self.running untouched
return [(finished.request_id, 0)]
monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
result = scheduler_cls.finish_requests(
scheduler,
[finished.request_id],
RequestStatus.FINISHED_STOPPED,
)
assert scheduler.running == []
assert result == [(finished.request_id, 0)]
@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
def test_finish_requests_purges_untracked_request_from_running(
monkeypatch: pytest.MonkeyPatch,
scheduler_cls,
scheduler_mod,
) -> None:
"""A request that has been deleted from ``self.requests`` (e.g.
after upstream's ``_free_request`` ran but a status mismatch
caused the ``self.running`` removal branch to no-op) must be
swept by the post-finish purge.
"""
untracked = _StubRequest("req-untracked", RequestStatus.RUNNING)
scheduler = _make_scheduler(
scheduler_cls,
requests={}, # already deleted from self.requests
running=[untracked],
waiting=[],
)
def fake_finish_requests(self, request_ids, finished_status):
return [(untracked.request_id, 0)]
monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
result = scheduler_cls.finish_requests(
scheduler,
[untracked.request_id],
RequestStatus.FINISHED_STOPPED,
)
assert scheduler.running == []
assert result == [(untracked.request_id, 0)]
@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
def test_finish_requests_leaves_healthy_running_intact(
monkeypatch: pytest.MonkeyPatch,
scheduler_cls,
scheduler_mod,
) -> None:
"""A live ``RUNNING`` request that is not in the finish set must
not be swept by the defensive purge -- this pins that the predicate
is conservative and won't aggressively drop healthy entries.
"""
alive = _StubRequest("req-alive", RequestStatus.RUNNING)
leaving_id = "req-leaving"
scheduler = _make_scheduler(
scheduler_cls,
requests={alive.request_id: alive},
running=[alive], # leaving already removed by super() (typical happy path)
waiting=[],
)
def fake_finish_requests(self, request_ids, finished_status):
return [(leaving_id, 0)]
monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
scheduler_cls.finish_requests(
scheduler,
[leaving_id],
RequestStatus.FINISHED_STOPPED,
)
assert scheduler.running == [alive]
@@ -46,6 +46,9 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids(
scheduler = scheduler_cls.__new__(scheduler_cls)
scheduler.chunk_transfer_adapter = None
scheduler.input_coordinator = coordinator
scheduler.requests = {}
scheduler.running = []
scheduler.waiting = []
def fake_finish_requests(self, request_ids, finished_status):
assert request_ids == ["req-a", "req-b"]
@@ -1,7 +1,7 @@
[
{
"test_name": "test_hunyuan_image3_ti2i_tp4_tp4_inline",
"description": "AR (TP=4) + DiT (TP=4) - using inline dataset and deploy-config",
"test_name": "test_hunyuan_image3_ti2i_inline",
"description": "AR (TP=2) + 2 * DiT (TP=2) - using inline dataset and deploy-config",
"server_type": "vllm-omni",
"benchmark_endpoint": "/v1/images/edits",
"server_params": {
@@ -21,7 +21,7 @@
"zmq_port": 50051,
"protocol": "rdma",
"device_name": "",
"memory_pool_size": 4294967296,
"memory_pool_size": 17179869184,
"memory_pool_device": "cpu"
}
}
@@ -32,13 +32,13 @@
"is_comprehension": true,
"final_output": true,
"final_output_type": "text",
"max_num_seqs": 4,
"max_num_seqs": 16,
"gpu_memory_utilization": 0.95,
"enforce_eager": false,
"enable_prefix_caching": false,
"max_num_batched_tokens": 32768,
"devices": "0,1,2,3",
"tensor_parallel_size": 4,
"devices": "0,1",
"tensor_parallel_size": 2,
"hf_overrides": {
"rope_parameters": {
"mrope_section": [0, 32, 32],
@@ -65,13 +65,14 @@
"stage_id": 1,
"max_num_seqs": 1,
"enforce_eager": true,
"devices": "4,5,6,7",
"devices": "2,3",
"distributed_executor_backend": "mp",
"omni_kv_config": {
"need_recv_cache": true
},
"parallel_config": {
"tensor_parallel_size": 4,
"tensor_parallel_size": 2,
"vae_patch_parallel_size": 2,
"enable_expert_parallel": true
},
"input_connectors": {
@@ -114,8 +115,8 @@
},
"task": "ti2i",
"bot-task": "think_recaption",
"max-concurrency": 4,
"num-prompts": 16,
"max-concurrency": 16,
"num-prompts": 96,
"skip-performance-assertion": true,
"baseline": {
"throughput_qps": 0.1,
@@ -23,10 +23,10 @@
"ignore_eos": true,
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": [1000, 2000, 3000, 4000],
"mean_tpot_ms": [200, 200, 200, 200],
"mean_audio_ttfp_ms": [10000, 15000, 20000, 25000],
"mean_audio_rtf": [1.5, 1.5, 1.5, 1.5]
"mean_ttft_ms": [203, 283, 470, 869],
"mean_tpot_ms": [28, 36, 41, 45],
"mean_audio_ttfp_ms": [861, 1200, 1528, 2100],
"mean_audio_rtf": [0.24, 0.30, 0.35, 0.39]
}
},
{
@@ -49,9 +49,9 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": [2000],
"mean_audio_ttfp_ms": [2000],
"mean_audio_rtf": [0.25]
"mean_ttft_ms": [377],
"mean_audio_ttfp_ms": [684],
"mean_audio_rtf": [0.14]
}
},
{
@@ -76,9 +76,9 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": [6000],
"mean_audio_ttfp_ms": [6000],
"mean_audio_rtf": [0.7]
"mean_ttft_ms": [228],
"mean_audio_ttfp_ms": [686],
"mean_audio_rtf": [0.16]
}
},
{
@@ -105,9 +105,9 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": [12000],
"mean_audio_ttfp_ms": [12000],
"mean_audio_rtf": [1.0]
"mean_ttft_ms": [522],
"mean_audio_ttfp_ms": [1119],
"mean_audio_rtf": [0.21]
}
}
]
@@ -31,11 +31,6 @@
],
"random_input_len": 2500,
"random_output_len": 900,
"temperature": 0.4,
"top_p": 0.9,
"top_k": 1,
"seed": 42,
"repetition_penalty": 1.05,
"ignore_eos": true,
"extra_body": {
"modalities": [
@@ -74,11 +69,6 @@
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
"temperature": 0.4,
"top_p": 0.9,
"top_k": 1,
"seed": 42,
"repetition_penalty": 1.05,
"ignore_eos": true,
"extra_body": {
"modalities": [
@@ -117,11 +107,6 @@
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
"temperature": 0.4,
"top_p": 0.9,
"top_k": 1,
"seed": 42,
"repetition_penalty": 1.05,
"ignore_eos": true,
"extra_body": {
"modalities": [
@@ -162,11 +147,6 @@
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
"temperature": 0.4,
"top_p": 0.9,
"top_k": 1,
"seed": 42,
"repetition_penalty": 1.05,
"ignore_eos": true,
"extra_body": {
"modalities": [
+61
View File
@@ -339,5 +339,66 @@
}
}
]
},
{
"test_name": "test_voxtral_tts",
"server_params": {
"model": "mistralai/Voxtral-4B-TTS-2603"
},
"benchmark_params": [
{
"task": "default_voice",
"eval_phase": "latency",
"dataset_name": "seed-tts-text",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
"num_prompts": [20],
"max_concurrency": [1],
"seed_tts_locale": "en",
"extra_body": {"voice": "casual_male"},
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [150],
"median_audio_rtf": [0.15]
}
},
{
"task": "default_voice",
"eval_phase": "throughput",
"dataset_name": "seed-tts-text",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
"num_prompts": [80],
"max_concurrency": [8],
"seed_tts_locale": "en",
"extra_body": {"voice": "casual_male"},
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [1500],
"median_audio_rtf": [0.30],
"audio_throughput": [30.0]
}
},
{
"task": "default_voice",
"eval_phase": "stress",
"dataset_name": "seed-tts-text",
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
"num_prompts": [100],
"request_rate": [2.0],
"seed_tts_locale": "en",
"extra_body": {"voice": "casual_male"},
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [3000],
"median_audio_rtf": [0.25],
"audio_throughput": [4.0]
}
}
]
}
]
@@ -319,6 +319,7 @@ QWEN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLO
@pytest.mark.slow
@pytest.mark.skip(reason="issue#4285")
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
@@ -387,6 +388,7 @@ def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
@pytest.mark.slow
@pytest.mark.skip(reason="issue#4285")
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, openai_client_function) -> None:
@@ -429,6 +431,7 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi
@pytest.mark.slow
@pytest.mark.skip(reason="issue#4285")
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, openai_client_function) -> None:
@@ -771,7 +774,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
@pytest.mark.slow
@pytest.mark.skip(reason="OOM injection causes the serving instance to exit")
@pytest.mark.skip(reason="issue#4285")
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
def test_reliability_fault_gpu_oom_state_converges_after_fault_removed(
@@ -1,19 +1,46 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Tests for the DiffusersPipelineLoader.
"""
from types import SimpleNamespace
import pytest
import torch
import torch.nn as nn
from huggingface_hub import snapshot_download
from vllm.config.load import LoadConfig
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
from vllm_omni.diffusion.models.helios import HeliosPipeline
from vllm_omni.diffusion.registry import initialize_model
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
model_path = "hf-internal-testing/tiny-helios-modular-pipe"
@pytest.fixture(scope="module")
def prefetch_helios_model():
"""Downloads the tiny helios model prior to running a test."""
snapshot_download(model_path)
@pytest.fixture(scope="function")
def mock_tp_group(mocker):
"""Mocks the tensor parallel group; this is needed to initialize the Helios model."""
mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
mock_group = mocker.MagicMock()
mock_group.world_size = 1
mock_group.rank_in_group = 0
mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
class _DummyPipelineModel(nn.Module):
def __init__(self, *, source_prefix: str):
@@ -42,7 +69,13 @@ class _DummyPipelineModel(nn.Module):
def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader:
loader = object.__new__(DiffusersPipelineLoader)
od_config = SimpleNamespace(
dtype=torch.float32,
parallel_config=SimpleNamespace(use_hsdp=False),
quantization_config=None,
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
loader.counter_before_loading_weights = 0.0
loader.counter_after_loading_weights = 0.0
@@ -150,16 +183,15 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
quantization_config=None,
)
loader = object.__new__(DiffusersPipelineLoader)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
loader.load_weights = lambda model: None # type: ignore[assignment]
loader._process_weights_after_loading = lambda model, target_device: None # type: ignore[assignment]
loader._is_gguf_quantization = lambda _od_config: False # type: ignore[assignment]
loader._is_gguf_quantization = lambda: False # type: ignore[assignment]
monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))
model = loader.load_model(
od_config,
load_device="cpu",
load_format="custom_pipeline",
custom_pipeline_name="tests.dummy.ConfigAwarePipeline",
@@ -168,3 +200,33 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
assert model.captured_config is od_config
assert model.seen_config_during_init is od_config
assert get_current_diffusion_config_or_none() is None
def test_get_all_weights(prefetch_helios_model, mock_tp_group):
"""Ensure that get all weights on a tiny model resolves to nonempty weights."""
od_config = OmniDiffusionConfig(
model_class_name="HeliosPipeline",
model=model_path,
)
loader = DiffusersPipelineLoader(
load_config=LoadConfig(),
od_config=od_config,
)
pipeline = HeliosPipeline(od_config=od_config)
weights = list(loader.get_all_weights(pipeline))
assert len(weights) > 0
def test_load_model(prefetch_helios_model, mock_tp_group):
"""Ensure that load model creates an instance of the expected pipeline class."""
od_config = OmniDiffusionConfig(
model_class_name="HeliosPipeline",
model=model_path,
)
loader = DiffusersPipelineLoader(
load_config=LoadConfig(),
od_config=od_config,
)
model = loader.load_model(load_device="cpu")
assert isinstance(model, HeliosPipeline)
@@ -62,6 +62,7 @@ def _make_loader() -> DiffusersPipelineLoader:
ignore_patterns=["*.tmp"],
)
loader.od_config = None
loader.quant_config = None
loader.counter_before_loading_weights = 0.0
loader.counter_after_loading_weights = 0.0
return loader
@@ -69,53 +70,47 @@ def _make_loader() -> DiffusersPipelineLoader:
def test_is_gguf_quantization_accepts_dict_config():
loader = _make_loader()
od_config = SimpleNamespace(
quantization_config={
"method": "gguf",
"gguf_model": "weights.gguf",
}
)
loader.quant_config = {
"method": "gguf",
"gguf_model": "weights.gguf",
}
assert loader._is_gguf_quantization(od_config) is True
assert loader._is_gguf_quantization() is True
def test_is_gguf_quantization_rejects_non_gguf_dict_config():
loader = _make_loader()
od_config = SimpleNamespace(
quantization_config={
"method": "fp8",
"gguf_model": "weights.gguf",
}
)
loader.quant_config = {
"method": "fp8",
"gguf_model": "weights.gguf",
}
assert loader._is_gguf_quantization(od_config) is False
assert loader._is_gguf_quantization() is False
def test_is_gguf_quantization_requires_gguf_model_for_dict_config():
loader = _make_loader()
od_config = SimpleNamespace(quantization_config={"method": "gguf"})
loader.quant_config = {"method": "gguf"}
with pytest.raises(ValueError, match="gguf_model"):
loader._is_gguf_quantization(od_config)
loader._is_gguf_quantization()
def test_is_gguf_quantization_accepts_object_config():
loader = _make_loader()
od_config = SimpleNamespace(
quantization_config=SimpleNamespace(
get_name=lambda: "gguf",
gguf_model="weights.gguf",
)
loader.quant_config = SimpleNamespace(
get_name=lambda: "gguf",
gguf_model="weights.gguf",
)
assert loader._is_gguf_quantization(od_config) is True
assert loader._is_gguf_quantization() is True
def test_is_gguf_quantization_uses_fallback_object_without_get_name():
loader = _make_loader()
od_config = SimpleNamespace(quantization_config=SimpleNamespace(gguf_model="weights.gguf"))
loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
assert loader._is_gguf_quantization(od_config) is True
assert loader._is_gguf_quantization() is True
def test_get_model_loadable_names_collects_parameters_and_buffers():
@@ -171,16 +166,14 @@ def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.Mon
def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.MonkeyPatch):
loader = _make_loader()
loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
loader.od_config = SimpleNamespace(revision="main")
source = DiffusersPipelineLoader.ComponentSource(
model_or_path="dummy",
subfolder="transformer",
revision=None,
prefix="transformer.",
)
od_config = SimpleNamespace(
revision="main",
quantization_config=SimpleNamespace(gguf_model="weights.gguf"),
)
class _Adapter:
def weights_iterator(self):
@@ -190,7 +183,7 @@ def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.Mon
monkeypatch.setattr(loader, "_resolve_gguf_model_path", lambda gguf_model, revision: "resolved.gguf")
monkeypatch.setattr(loader_module, "get_gguf_adapter", lambda gguf_file, model, source, od_config: _Adapter())
weights = list(loader._get_gguf_weights_iterator(source, object(), od_config))
weights = list(loader._get_gguf_weights_iterator(source, object()))
assert weights[0][0] == "transformer.weight"
assert torch.equal(weights[0][1], torch.ones((2, 2)))
@@ -201,13 +194,12 @@ def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.Mon
def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(monkeypatch: pytest.MonkeyPatch):
loader = _make_loader()
model = _DummyModel()
od_config = SimpleNamespace()
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
monkeypatch.setattr(
loader,
"_get_gguf_weights_iterator",
lambda source, model, od_config: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
lambda source, model: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
)
def _hf_iter(source):
@@ -222,7 +214,7 @@ def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(
lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
)
loaded = loader._load_weights_with_gguf(model, od_config)
loaded = loader._load_weights_with_gguf(model)
assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
assert model.calls[0] == ["transformer.weight"]
@@ -233,13 +225,12 @@ def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(
def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete(monkeypatch: pytest.MonkeyPatch):
loader = _make_loader()
model = _DummyModel()
od_config = SimpleNamespace()
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
monkeypatch.setattr(
loader,
"_get_gguf_weights_iterator",
lambda source, model, od_config: iter(
lambda source, model: iter(
[
("transformer.weight", torch.ones((2, 2), dtype=torch.float32)),
("transformer.bias", torch.zeros(2, dtype=torch.float32)),
@@ -262,7 +253,7 @@ def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete
lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
)
loaded = loader._load_weights_with_gguf(model, od_config)
loaded = loader._load_weights_with_gguf(model)
assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
assert hf_calls == ["vae"]
@@ -99,12 +99,16 @@ class StubCosmos3Transformer(nn.Module):
sound_gen: bool = False,
sound_dim: int = 3,
sound_latent_fps: float = 25.0,
action_gen: bool = False,
action_dim: int = 4,
) -> None:
super().__init__()
self.latent_channel_size = latent_channel_size
self.sound_gen = sound_gen
self.sound_dim = sound_dim
self.sound_latent_fps = sound_latent_fps
self.action_gen = action_gen
self.action_dim = action_dim
self.cached_kv: Any | None = None
self.cached_freqs_gen: Any | None = None
self.calls: list[dict[str, Any]] = []
@@ -139,7 +143,10 @@ class StubCosmos3Transformer(nn.Module):
marker = torch.tensor([token], dtype=torch.float32)
self.cached_kv = [(marker, marker + 100)]
self.cached_freqs_gen = (marker + 200, marker + 300)
action_latents = kwargs.get("action_latents")
outputs: list[torch.Tensor] = [torch.full_like(hidden_states, float(token))]
if action_latents is not None:
outputs.append(torch.full_like(action_latents, float(token + 20)))
if sound_latents is not None:
outputs.append(torch.full_like(sound_latents, float(token + 10)))
return outputs[0] if len(outputs) == 1 else tuple(outputs)
@@ -344,7 +351,7 @@ def test_pipeline_init_passes_tokenizer_attrs_into_transformer(
assert pipeline.transformer.audio_proj_out.out_features == 5
def test_preprocess_i2v_image_input() -> None:
def test_preprocess_i2v_image_and_action_video_inputs() -> None:
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_pre_process_func
preprocess = get_cosmos3_pre_process_func(SimpleNamespace())
@@ -357,6 +364,29 @@ def test_preprocess_i2v_image_input() -> None:
assert (result.sampling_params.height, result.sampling_params.width) == (672, 1344)
assert tuple(result.prompts[0]["additional_information"]["preprocessed_image"].shape[-2:]) == (672, 1344)
frames = [Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue")]
action = SimpleNamespace(
prompts=[{"prompt": "Move.", "multi_modal_data": {"video": frames}}],
sampling_params=SimpleNamespace(height=16, width=32, extra_args={"action_mode": "forward_dynamics"}),
)
additional = preprocess(action).prompts[0]["additional_information"]
assert tuple(additional["preprocessed_image"].shape) == (1, 3, 16, 32)
assert tuple(additional["preprocessed_video"].shape) == (1, 3, 3, 16, 32)
frames = [Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue", "yellow", "purple", "black")]
v2v = SimpleNamespace(
prompts=[{"prompt": "Continue.", "multi_modal_data": {"video": frames}}],
sampling_params=SimpleNamespace(
height=16,
width=32,
extra_args={"condition_frame_indexes_vision": [0, 1], "condition_video_keep": "last"},
),
)
additional = preprocess(v2v).prompts[0]["additional_information"]
assert tuple(additional["preprocessed_video"].shape) == (1, 3, 5, 16, 32)
assert additional["condition_frame_indexes_vision"] == [0, 1]
def test_postprocess_handles_image_video_audio_and_validation() -> None:
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_post_process_func
@@ -434,7 +464,7 @@ def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> No
assert {key: Cosmos3OmniDiffusersPipeline._remap_ckpt_key(key) for key in remaps} == remaps
def test_prepare_latents_for_video_image_and_sound(make_cosmos3_pipeline) -> None:
def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline) -> None:
pipeline = make_cosmos3_pipeline()
latents = pipeline._prepare_latents(16, 24, 5, torch.Generator(device="cpu").manual_seed(0))
assert latents.shape == (1, 2, 2, 2, 3)
@@ -447,6 +477,19 @@ def test_prepare_latents_for_video_image_and_sound(make_cosmos3_pipeline) -> Non
assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]]]]]
assert image_latent.shape == (1, 2, 1, 2, 3)
pipeline._encode_video_tensor = lambda *args, **kwargs: torch.full((1, 2, 3, 2, 3), 6.0)
v2v_latents, v2v_velocity_mask, v2v_condition = pipeline._prepare_latents_v2v(
torch.zeros(1, 3, 5, 16, 24),
16,
24,
9,
torch.Generator(device="cpu").manual_seed(0),
[0, 1],
)
torch.testing.assert_close(v2v_latents[:, :, 0:2], torch.full((1, 2, 2, 2, 3), 6.0))
assert v2v_velocity_mask.tolist() == [[[[[0.0]], [[0.0]], [[1.0]]]]]
assert v2v_condition.shape == (1, 2, 3, 2, 3)
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
pipeline._sound_tokenizer = SimpleNamespace(
sample_rate=10,
@@ -463,8 +506,20 @@ def test_prepare_latents_for_video_image_and_sound(make_cosmos3_pipeline) -> Non
assert (sound_latents.shape, latent_frames) == (torch.Size([1, 3, 6]), 6)
assert pipeline._decode_sound_latents(torch.zeros(1, 3, 6), target_audio_samples=21).shape == (1, 2, 21)
pipeline.transformer = pipeline.transformer.__class__(action_gen=True, action_dim=4)
action, action_mask, clean, raw_dim = pipeline._prepare_action_latents(
mode="forward_dynamics",
action_chunk_size=2,
raw_action_dim=None,
generator=torch.Generator(device="cpu").manual_seed(0),
sp=SimpleNamespace(extra_args={"action": [[1.0, 2.0], [3.0, 4.0]]}),
)
assert raw_dim == 2
assert action_mask.tolist() == [[[0.0], [0.0]]]
torch.testing.assert_close(action, clean)
def test_diffuse_covers_cfg_i2v_and_sound_steps(make_cosmos3_pipeline) -> None:
def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> None:
pipeline = make_cosmos3_pipeline()
latents = torch.zeros(1, 2, 1, 1, 1)
@@ -496,20 +551,22 @@ def test_diffuse_covers_cfg_i2v_and_sound_steps(make_cosmos3_pipeline) -> None:
)
torch.testing.assert_close(i2v[:, :, 0:1], torch.full((1, 2, 1, 1, 1), 7.0))
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
video_result, sound_result = pipeline.diffuse(
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4)
video_result, action_result = pipeline.diffuse(
latents=latents,
sound_latents=torch.zeros(1, 3, 4),
action_latents=torch.zeros(1, 3, 4),
action_velocity_mask=torch.ones(1, 3, 1),
action_condition_latents=torch.zeros(1, 3, 4),
timesteps=torch.tensor([7, 3]),
cond_ids=_ids(2),
cond_mask=_mask(),
uncond_ids=_ids(1),
uncond_mask=_mask(),
guidance_scale=1.0,
shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "action_domain_ids": torch.tensor([0])},
)
torch.testing.assert_close(video_result, torch.full_like(latents, 4.0))
torch.testing.assert_close(sound_result, torch.full((), 24.0).expand_as(sound_result))
torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result))
def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -> None:
@@ -568,6 +625,8 @@ class TestForwardRouting:
def fake_diffuse(**kwargs):
captured["diffuse_calls"].append(kwargs)
outputs = [kwargs["latents"] + len(captured["diffuse_calls"])]
if kwargs.get("action_latents") is not None:
outputs.append(kwargs["action_latents"] + 3.0)
if kwargs.get("sound_latents") is not None:
outputs.append(kwargs["sound_latents"] + 2.0)
return outputs[0] if len(outputs) == 1 else tuple(outputs)
@@ -612,7 +671,7 @@ class TestForwardRouting:
assert captured["flow_shifts"] == expected["flow"]
assert [call[0] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"]
def test_forward_i2v_and_sound_routes(self, make_cosmos3_pipeline) -> None:
def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> None:
pipeline = make_cosmos3_pipeline()
captured = self._install_forward_stubs(pipeline)
image_tensor = torch.zeros(1, 3, 16, 16)
@@ -637,6 +696,34 @@ class TestForwardRouting:
)
assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask
video_tensor = torch.zeros(1, 3, 5, 16, 16)
v2v_condition = torch.full((1, 2, 2, 1, 1), 4.0)
v2v_mask = torch.tensor([[[[[0.0]], [[1.0]]]]])
pipeline._prepare_latents_v2v = lambda *args, **kwargs: (
torch.zeros(1, 2, 2, 1, 1),
v2v_mask,
v2v_condition,
)
pipeline.forward(
SimpleNamespace(
prompts=[
{
"prompt": "continue",
"modalities": ["video"],
"additional_information": {
"preprocessed_video": video_tensor,
"condition_frame_indexes_vision": [0],
},
}
],
sampling_params=make_sampling_params(height=16, width=16, num_frames=5),
)
)
assert captured["flow_shifts"][-1] == 10.0
assert captured["format"]["negative_prompt"] == ""
assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is v2v_mask
assert captured["diffuse_calls"][-1]["condition_latents"] is v2v_condition
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
sound_latents = torch.zeros(1, 3, 4)
pipeline._resolve_sound_target_samples = lambda *args: (20, 2.0, 10)
@@ -651,6 +738,31 @@ class TestForwardRouting:
assert captured["diffuse_calls"][-1]["sound_latents"] is sound_latents
assert output.output["audio_sample_rate"] == 10
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4)
output = pipeline.forward(
SimpleNamespace(
prompts=[
{
"prompt": "Pick the block.",
"modalities": ["video"],
"additional_information": {"preprocessed_image": image_tensor},
}
],
sampling_params=make_sampling_params(
height=16,
width=16,
extra_args={
"action_mode": "policy",
"action_chunk_size": 2,
"raw_action_dim": 2,
"domain_name": "bridge_orig_lerobot",
},
),
)
)
assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
assert output.custom_output["action"].shape == (1, 2, 2)
@pytest.mark.parametrize(
("prompt", "sampling_params", "message"),
[
@@ -30,8 +30,9 @@ def _tiny_cosmos3_config(**overrides):
return config
def test_mrope_position_ids_cover_text_video_and_sound() -> None:
def test_mrope_position_ids_cover_text_video_sound_and_action() -> None:
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
compute_mrope_position_ids_action,
compute_mrope_position_ids_sound,
compute_mrope_position_ids_text,
compute_mrope_position_ids_vision,
@@ -62,6 +63,10 @@ def test_mrope_position_ids_cover_text_video_and_sound() -> None:
torch.testing.assert_close(sound_ids[0], torch.tensor([10.0, 10.96, 11.92]))
assert sound_offset == 12
action_ids, action_offset = compute_mrope_position_ids_action(3, temporal_offset=10, action_fps=None)
assert action_ids.tolist() == [[11, 12, 13], [0, 0, 0], [0, 0, 0]]
assert action_offset == 14
@pytest.mark.parametrize(
("key", "value"),
@@ -126,7 +131,7 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
assert tuple(output.shape) == (1, 2, 1, 2, 2)
def test_sound_modules_follow_injected_sound_dim() -> None:
def test_sound_and_action_modules_follow_config() -> None:
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
tiny = _tiny_cosmos3_config()
@@ -137,12 +142,22 @@ def test_sound_modules_follow_injected_sound_dim() -> None:
sound_dim=5,
sound_latent_fps=40.0,
)
with_action = Cosmos3VFMTransformer(
SimpleNamespace(
tf_model_config={**tiny, "action_gen": True, "max_action_dim": 6, "num_embodiment_domains": 9},
dtype=torch.float32,
)
)
assert no_modal.sound_gen is False
assert no_modal.action_gen is False
assert not hasattr(no_modal, "audio_proj_in")
assert not hasattr(no_modal, "action_proj_in")
assert with_sound.sound_dim == 5
assert with_sound.sound_latent_fps == 40.0
assert with_sound.audio_proj_in.in_features == 5
assert with_action.action_dim == 6
assert with_action.action_proj_in.num_domains == 9
@pytest.mark.parametrize(
@@ -163,33 +178,56 @@ def test_transformer_requires_sound_dim_and_fps_when_sound_gen_true(kwargs: dict
)
def test_sound_pack_unpack_validate_shapes() -> None:
def test_sound_and_action_pack_unpack_validate_shapes() -> None:
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
model = object.__new__(Cosmos3VFMTransformer)
nn.Module.__init__(model)
model.sound_dim = 3
model.action_dim = 3
sound = torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(2, 3, 4)
action = torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(2, 4, 3)
torch.testing.assert_close(model.unpack_sound(model.pack_sound(sound)), sound)
torch.testing.assert_close(model.unpack_action(model.pack_action(action)), action)
with pytest.raises(ValueError, match="channel mismatch"):
model.pack_sound(torch.zeros(1, 4, 2))
with pytest.raises(ValueError, match="dimension mismatch"):
model.pack_action(torch.zeros(1, 2, 4))
def test_forward_returns_video_and_sound_predictions(monkeypatch: pytest.MonkeyPatch) -> None:
@pytest.mark.parametrize(
("config", "transformer_kwargs", "extra_kwargs", "expected_shapes"),
[
(
_tiny_cosmos3_config(),
{"sound_gen": True, "sound_dim": 3, "sound_latent_fps": 24.0},
{"sound_latents": torch.zeros(1, 3, 4)},
[(1, 2, 1, 2, 2), (1, 3, 4)],
),
(
_tiny_cosmos3_config(action_gen=True, max_action_dim=3, num_embodiment_domains=4),
{},
{"action_latents": torch.zeros(1, 5, 3), "action_domain_ids": torch.tensor([2])},
[(1, 2, 1, 2, 2), (1, 5, 3)],
),
],
)
def test_forward_returns_video_plus_optional_modality_predictions(
monkeypatch: pytest.MonkeyPatch,
config,
transformer_kwargs,
extra_kwargs,
expected_shapes,
) -> None:
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
output = transformer_cosmos3.Cosmos3VFMTransformer(
SimpleNamespace(
tf_model_config=_tiny_cosmos3_config(),
dtype=torch.float32,
),
sound_gen=True,
sound_dim=3,
sound_latent_fps=40.0,
SimpleNamespace(tf_model_config=config, dtype=torch.float32),
**transformer_kwargs,
)(
hidden_states=torch.zeros(1, 2, 1, 2, 2),
timestep=torch.tensor([1.0]),
@@ -197,11 +235,12 @@ def test_forward_returns_video_and_sound_predictions(monkeypatch: pytest.MonkeyP
text_mask=torch.ones(1, 2, dtype=torch.long),
video_shape=(1, 2, 2),
fps=24.0,
sound_latents=torch.zeros(1, 3, 4),
action_noisy_mask=torch.ones(1, 5, 1),
**extra_kwargs,
)
assert isinstance(output, tuple)
assert [tuple(tensor.shape) for tensor in output] == [(1, 2, 1, 2, 2), (1, 3, 4)]
assert [tuple(tensor.shape) for tensor in output] == expected_shapes
def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -227,7 +266,7 @@ def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch
)
def test_compute_rope_freqs_places_text_video_and_sound_positions() -> None:
def test_compute_rope_freqs_places_text_video_action_and_sound_positions() -> None:
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
class FakeRotary:
@@ -275,9 +314,11 @@ def test_compute_rope_freqs_places_text_video_and_sound_positions() -> None:
fps=24.0,
device=torch.device("cpu"),
dtype=torch.float32,
t_action=2,
action_start_frame_offset=1,
t_sound=1,
)
_, gen_pos = rotary.position_ids
assert gen_pos.shape == (3, 1, 3)
assert gen_pos[0, 0].tolist() == [102, 103, 102]
assert gen_pos.shape == (3, 1, 5)
assert gen_pos[0, 0].tolist() == [102, 103, 103, 104, 102]
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -0,0 +1,160 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for HunyuanVideo-1.5 quant_config propagation through transformer creation.
Tests cover:
- HunyuanVideo15Pipeline passes quant_config to HunyuanVideo15Transformer3DModel
- HunyuanVideo15I2VPipeline passes quant_config to HunyuanVideo15Transformer3DModel
"""
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
import torch
import vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 as t2v_module
import vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v as i2v_module
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
class TestHunyuanVideoQuantConfigPropagation:
"""Verify quant_config is propagated to the transformer model in HunyuanVideo-1.5 pipelines."""
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.get_local_device",
return_value="cpu",
)
@patch("vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.prefetch_subfolders")
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.from_pretrained_with_prefetch",
return_value=MagicMock(),
)
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.Qwen2Tokenizer.from_pretrained",
return_value=MagicMock(),
)
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.ByT5Tokenizer.from_pretrained",
return_value=MagicMock(),
)
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.AutoConfig.from_pretrained",
return_value=MagicMock(),
)
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.T5EncoderModel",
return_value=MagicMock(),
)
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.FlowMatchEulerDiscreteScheduler.from_pretrained",
return_value=MagicMock(),
)
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.get_transformer_config_kwargs",
return_value={},
)
def test_t2v_quant_config_passed_through(
self,
mock_get_kwargs,
mock_scheduler,
mock_t5_encoder,
mock_auto_config,
mock_byt5_tok,
mock_qwen2_tok,
mock_from_pretrained,
mock_prefetch,
mock_get_device,
):
captured_kwargs = {}
class FakeTransformer:
def __init__(self, **kwargs):
captured_kwargs.update(kwargs)
with patch.object(t2v_module, "HunyuanVideo15Transformer3DModel", FakeTransformer):
fake_qc = MagicMock()
od_config = SimpleNamespace(
model="fake-t2v-model",
dtype=torch.bfloat16,
flow_shift=None,
tf_model_config=SimpleNamespace(),
quantization_config=fake_qc,
enable_diffusion_pipeline_profiler=False,
)
t2v_module.HunyuanVideo15Pipeline(od_config=od_config)
assert captured_kwargs.get("quant_config") is fake_qc
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.get_local_device",
return_value="cpu",
)
@patch("vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.prefetch_subfolders")
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.from_pretrained_with_prefetch",
return_value=MagicMock(),
)
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.Qwen2Tokenizer.from_pretrained",
return_value=MagicMock(),
)
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.ByT5Tokenizer.from_pretrained",
return_value=MagicMock(),
)
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.AutoConfig.from_pretrained",
return_value=MagicMock(),
)
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.T5EncoderModel",
return_value=MagicMock(),
)
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.SiglipImageProcessor.from_pretrained",
return_value=MagicMock(),
)
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.FlowMatchEulerDiscreteScheduler.from_pretrained",
return_value=MagicMock(),
)
@patch(
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.get_transformer_config_kwargs",
return_value={},
)
def test_i2v_quant_config_passed_through(
self,
mock_get_kwargs,
mock_scheduler,
mock_siglip_processor,
mock_t5_encoder,
mock_auto_config,
mock_byt5_tok,
mock_qwen2_tok,
mock_from_pretrained,
mock_prefetch,
mock_get_device,
):
captured_kwargs = {}
class FakeTransformer:
def __init__(self, **kwargs):
captured_kwargs.update(kwargs)
with patch.object(i2v_module, "HunyuanVideo15Transformer3DModel", FakeTransformer):
fake_qc = MagicMock()
od_config = SimpleNamespace(
model="fake-i2v-model",
dtype=torch.bfloat16,
flow_shift=None,
tf_model_config=SimpleNamespace(),
quantization_config=fake_qc,
enable_diffusion_pipeline_profiler=False,
)
i2v_module.HunyuanVideo15I2VPipeline(od_config=od_config)
assert captured_kwargs.get("quant_config") is fake_qc
@@ -72,6 +72,285 @@ class TestPipelineIndependence:
assert issubclass(LTX23Pipeline, CFGParallelMixin)
def test_ltx23_pipeline_declares_offload_components(self):
"""LTX23Pipeline must expose LTX-2.3-specific modules to offload discovery."""
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
pipe = object.__new__(LTX23Pipeline)
torch.nn.Module.__init__(pipe)
pipe.transformer = torch.nn.Linear(1, 1)
pipe.text_encoder = torch.nn.Linear(1, 1)
pipe.connectors = torch.nn.Linear(1, 1)
pipe.vae = torch.nn.Linear(1, 1)
pipe.audio_vae = torch.nn.Linear(1, 1)
pipe.vocoder = torch.nn.Linear(1, 1)
modules = ModuleDiscovery.discover(pipe)
assert LTX23Pipeline._dit_modules == ["transformer"]
assert LTX23Pipeline._encoder_modules == ["text_encoder", "connectors"]
assert LTX23Pipeline._vae_modules == ["vae", "audio_vae"]
assert LTX23Pipeline._resident_modules == ["vocoder"]
assert modules.dit_names == ["transformer"]
assert modules.encoder_names == ["text_encoder", "connectors"]
assert modules.resident_names == ["vocoder"]
assert len(modules.vaes) == 2
def test_ltx23_pipeline_has_diffusion_pipeline_profiler_mixin(self):
"""LTX23Pipeline must support lightweight stage timing."""
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
assert issubclass(LTX23Pipeline, DiffusionPipelineProfilerMixin)
class TestLTX23VaeDecodeParallel:
"""Test LTX-2.3 video VAE tiled parallel helpers without loading weights."""
def test_ltx23_video_vae_is_distributed_tile_only_class(self):
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
DistributedAutoencoderKLLTX2Video,
)
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin
assert issubclass(DistributedAutoencoderKLLTX2Video, DistributedVaeMixin)
assert not hasattr(DistributedAutoencoderKLLTX2Video, "patch_split")
def test_ltx23_video_vae_tile_split_uses_native_ltx23_tile_geometry(self):
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
DistributedAutoencoderKLLTX2Video,
)
vae = SimpleNamespace(
spatial_compression_ratio=32,
tile_sample_min_height=512,
tile_sample_min_width=512,
tile_sample_stride_height=448,
tile_sample_stride_width=448,
temporal_compression_ratio=8,
dtype=torch.float32,
)
z = torch.zeros(1, 2, 5, 16, 24)
tasks, grid_spec = DistributedAutoencoderKLLTX2Video.tile_split(vae, z)
assert grid_spec.grid_shape == (2, 2)
assert grid_spec.split_dims == (3, 4)
assert grid_spec.tile_spec["sample_height"] == 512
assert grid_spec.tile_spec["sample_width"] == 768
assert grid_spec.tile_spec["blend_height"] == 64
assert grid_spec.tile_spec["blend_width"] == 64
assert grid_spec.tile_spec["max_tile_output_shape"] == (1, 3, 33, 512, 512)
assert grid_spec.tile_spec["tile_output_shapes"] == {
0: (1, 3, 33, 512, 512),
1: (1, 3, 33, 512, 320),
2: (1, 3, 33, 64, 512),
3: (1, 3, 33, 64, 320),
}
assert [task.grid_coord for task in tasks] == [(0, 0), (0, 1), (1, 0), (1, 1)]
assert [tuple(task.tensor.shape) for task in tasks] == [
(1, 2, 5, 16, 16),
(1, 2, 5, 16, 10),
(1, 2, 5, 2, 16),
(1, 2, 5, 2, 10),
]
assert [task.workload for task in tasks] == [5 * 16 * 16, 5 * 16 * 10, 5 * 2 * 16, 5 * 2 * 10]
def test_ltx23_video_vae_tile_merge_blends_and_crops_like_tiled_decode(self):
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
DistributedAutoencoderKLLTX2Video,
)
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import GridSpec
class FakeVae:
def __init__(self):
self.blend_calls = []
def clear_cache(self):
pass
def blend_v(self, _previous, current, blend_height):
self.blend_calls.append(("v", blend_height))
return current
def blend_h(self, _previous, current, blend_width):
self.blend_calls.append(("h", blend_width))
return current
fake_vae = FakeVae()
grid_spec = GridSpec(
split_dims=(3, 4),
grid_shape=(2, 2),
tile_spec={
"sample_height": 10,
"sample_width": 10,
"blend_height": 1,
"blend_width": 2,
"tile_sample_stride_height": 5,
"tile_sample_stride_width": 5,
},
)
tiles = {
(0, 0): torch.full((1, 3, 2, 6, 6), 1.0),
(0, 1): torch.full((1, 3, 2, 6, 6), 2.0),
(1, 0): torch.full((1, 3, 2, 6, 6), 3.0),
(1, 1): torch.full((1, 3, 2, 6, 6), 4.0),
}
merged = DistributedAutoencoderKLLTX2Video.tile_merge(fake_vae, tiles, grid_spec)
assert merged.shape == (1, 3, 2, 10, 10)
assert fake_vae.blend_calls == [("h", 2), ("v", 1), ("v", 1), ("h", 2)]
torch.testing.assert_close(merged[:, :, :, :5, :5], torch.ones(1, 3, 2, 5, 5))
torch.testing.assert_close(merged[:, :, :, :5, 5:], torch.full((1, 3, 2, 5, 5), 2.0))
torch.testing.assert_close(merged[:, :, :, 5:, :5], torch.full((1, 3, 2, 5, 5), 3.0))
torch.testing.assert_close(merged[:, :, :, 5:, 5:], torch.full((1, 3, 2, 5, 5), 4.0))
def test_ltx23_video_vae_tiled_decode_dispatches_to_tile_operator(self):
from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_ltx2
z = torch.zeros(1, 2, 1, 16, 24)
expected = torch.ones(1, 3, 1, 512, 768)
seen = {}
class FakeExecutor:
def execute(self, tensor, operator, broadcast_result=True):
seen["tensor"] = tensor
seen["operator"] = operator
seen["broadcast_result"] = broadcast_result
return expected
vae = SimpleNamespace(distributed_executor=FakeExecutor(), is_distributed_enabled=lambda: True)
vae.tile_split = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_split.__get__(vae)
vae.tile_exec = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_exec.__get__(vae)
vae.tile_merge = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_merge.__get__(vae)
output = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tiled_decode(
vae,
z,
temb=torch.tensor(0.5),
return_dict=False,
)
assert len(output) == 1
assert output[0] is expected
assert seen["tensor"] is z
assert seen["broadcast_result"] is False
assert seen["operator"].split.__name__ == "tile_split"
assert seen["operator"].merge.__name__ == "tile_merge"
def test_ltx23_vae_executor_gathers_known_tile_shapes_and_returns_empty_on_non_rank0(self):
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import LTX2VaeExecutor
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
DistributedOperator,
GridSpec,
TileTask,
)
z = torch.zeros(1, 1, 1, 1, 1)
tile_output_shapes = {
0: (1, 1, 1, 2, 2),
1: (1, 1, 1, 2, 1),
2: (1, 1, 1, 1, 2),
3: (1, 1, 1, 1, 1),
}
tasks = [
TileTask(0, (0, 0), z, workload=4),
TileTask(1, (0, 1), z, workload=2),
TileTask(2, (1, 0), z, workload=2),
TileTask(3, (1, 1), z, workload=1),
]
grid_spec = GridSpec(
split_dims=(3, 4),
grid_shape=(2, 2),
tile_spec={
"max_tile_output_shape": (1, 1, 1, 2, 2),
"tile_output_shapes": tile_output_shapes,
},
output_dtype=torch.float32,
)
seen = {}
def exec_tile(task):
return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1))
def merge_tiles(coord_tensor_map, passed_grid_spec):
seen["merged_shapes"] = {coord: tuple(tile.shape) for coord, tile in coord_tensor_map.items()}
assert passed_grid_spec is grid_spec
return torch.stack(
[
coord_tensor_map[(0, 0)].flatten()[0],
coord_tensor_map[(0, 1)].flatten()[0],
coord_tensor_map[(1, 0)].flatten()[0],
coord_tensor_map[(1, 1)].flatten()[0],
]
)
operator = DistributedOperator(split=lambda _z: (tasks, grid_spec), exec=exec_tile, merge=merge_tiles)
rank0_executor = object.__new__(LTX2VaeExecutor)
rank0_executor.parallel_size = 2
rank0_executor.world_size = 2
rank0_executor.rank = 0
def gather_rank0(local_tile_tensor):
assigned = rank0_executor._balance_tasks(tasks, 2)
rank1_results = [(task.tile_id, exec_tile(task)) for task in assigned[1]]
rank1_tile_tensor = rank0_executor._pack_local_tiles_without_meta(
rank1_results,
list(local_tile_tensor.shape),
z.device,
torch.float32,
)
seen["rank0_gather_shape"] = tuple(local_tile_tensor.shape)
return [local_tile_tensor, rank1_tile_tensor]
def fail_final_sync(*_args, **_kwargs):
raise AssertionError("broadcast_result=False should not sync the final result")
rank0_executor.gather_tensors = gather_rank0
rank0_executor._sync_final_result = fail_final_sync
rank0_result = rank0_executor.execute(z, operator, broadcast_result=False)
torch.testing.assert_close(rank0_result, torch.tensor([1.0, 2.0, 3.0, 4.0]))
assert seen["rank0_gather_shape"] == (2, 1, 1, 1, 2, 2)
assert seen["merged_shapes"] == {
(0, 0): (1, 1, 1, 2, 2),
(0, 1): (1, 1, 1, 2, 1),
(1, 0): (1, 1, 1, 1, 2),
(1, 1): (1, 1, 1, 1, 1),
}
non_rank0_executor = object.__new__(LTX2VaeExecutor)
non_rank0_executor.parallel_size = 2
non_rank0_executor.world_size = 2
non_rank0_executor.rank = 1
def gather_rank1(local_tile_tensor):
seen["rank1_gather_shape"] = tuple(local_tile_tensor.shape)
return None
def fail_non_rank0_merge(*_args, **_kwargs):
raise AssertionError("non-rank0 should not merge gathered tiles")
non_rank0_executor.gather_tensors = gather_rank1
non_rank0_executor._sync_final_result = fail_final_sync
empty_result = non_rank0_executor.execute(
z,
DistributedOperator(
split=lambda _z: (tasks, grid_spec),
exec=exec_tile,
merge=fail_non_rank0_merge,
),
broadcast_result=False,
)
assert tuple(empty_result.shape) == (0,)
assert seen["rank1_gather_shape"] == (2, 1, 1, 1, 2, 2)
class TestCFGParallelHelpers:
"""Test LTX-2.3 CFG helper math without loading model weights."""
@@ -0,0 +1,22 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from torch import nn
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import _make_rms_norm
def test_ltx_rms_norm_no_affine_identity_weight_is_non_persistent_buffer():
norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)
assert "weight" not in dict(norm.named_parameters())
assert "weight" in dict(norm.named_buffers())
assert "weight" not in norm.state_dict()
def test_ltx_rms_norm_affine_weight_remains_parameter():
norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=True)
assert isinstance(dict(norm.named_parameters())["weight"], nn.Parameter)
assert "weight" not in dict(norm.named_buffers())
assert "weight" in norm.state_dict()
+1 -1
View File
@@ -41,7 +41,7 @@ def mock_od_config(mocker: MockerFixture):
config.cache_backend = None
config.cache_config = None
config.model = "test-model"
config.diffusion_load_format = None
config.diffusion_load_format = "default"
config.dtype = "float32"
config.max_cpu_loras = 0
config.lora_path = None
@@ -32,8 +32,8 @@ BAGEL_MOONCAKE_CI_DEPLOY = get_deploy_config_path("ci/bagel_mooncake.yaml")
# Reference pixel data extracted from the known-good output image
# Each entry contains (x, y) position and expected (R, G, B) values
# "Generated with seed=52, num_inference_steps=15,
# prompt='A futuristic city skyline at twilight, cyberpunk style'"
# "Generated with seed=52, num_inference_steps=14,
# prompt='A cute cat'"
REFERENCE_PIXELS = [
{"position": (100, 100), "rgb": (115, 113, 94)},
{"position": (400, 50), "rgb": (159, 160, 144)},
@@ -63,7 +63,7 @@ def _find_free_port() -> int:
return port
def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
def _configure_sampling_params(omni: Omni, num_inference_steps: int = 14) -> list:
"""Configure sampling parameters for Bagel text2img generation.
Args:
@@ -269,7 +269,6 @@ def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str:
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.skip(reason="Skip failed CI issue 3977: https://github.com/vllm-project/vllm-omni/issues/3977")
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_bagel_text2img_mooncake_connector(run_level):
"""Test Bagel text2img with Mooncake connector for inter-stage communication."""
@@ -24,14 +24,13 @@ from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.usefixtures("clean_gpu_memory_between_tests")]
BAGEL_CI_DEPLOY = get_deploy_config_path("ci/bagel.yaml")
# Reference pixel data extracted from the known-good output image
# Generated with seed=52, num_inference_steps=15,
# Generated with seed=52, num_inference_steps=14,
# prompt='Change the grass color to red',
# input image: 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg
REFERENCE_PIXELS = [
@@ -47,23 +46,10 @@ REFERENCE_PIXELS = [
{"position": (256, 256), "rgb": (181, 201, 221)},
]
if current_omni_platform.is_rocm():
REFERENCE_PIXELS = [
{"position": (100, 100), "rgb": (156, 172, 217)},
{"position": (400, 50), "rgb": (105, 144, 217)},
{"position": (700, 100), "rgb": (118, 159, 232)},
{"position": (150, 400), "rgb": (180, 22, 52)},
{"position": (512, 336), "rgb": (221, 211, 194)},
{"position": (700, 400), "rgb": (192, 10, 46)},
{"position": (100, 600), "rgb": (102, 12, 22)},
{"position": (400, 600), "rgb": (161, 28, 47)},
{"position": (700, 600), "rgb": (100, 87, 94)},
{"position": (256, 256), "rgb": (181, 201, 221)},
]
# text2img reference pixels (aligned with offline `bagel/end2end.py` text2img, 15 steps)
# "Generated with seed=52, num_inference_steps=15,
# prompt='A futuristic city skyline at twilight, cyberpunk style'"
# "Generated with seed=52, num_inference_steps=14,
# prompt='A cute cat'"
TEXT2IMG_REFERENCE_PIXELS = [
{"position": (100, 100), "rgb": (115, 113, 94)},
{"position": (400, 50), "rgb": (159, 160, 144)},
@@ -77,19 +63,6 @@ TEXT2IMG_REFERENCE_PIXELS = [
{"position": (256, 256), "rgb": (167, 156, 150)},
]
if current_omni_platform.is_rocm():
TEXT2IMG_REFERENCE_PIXELS = [
{"position": (100, 100), "rgb": (115, 113, 94)},
{"position": (400, 50), "rgb": (159, 160, 144)},
{"position": (700, 100), "rgb": (164, 151, 123)},
{"position": (150, 400), "rgb": (120, 121, 107)},
{"position": (512, 512), "rgb": (165, 133, 127)},
{"position": (700, 400), "rgb": (217, 130, 66)},
{"position": (100, 700), "rgb": (191, 168, 152)},
{"position": (400, 700), "rgb": (130, 96, 77)},
{"position": (700, 700), "rgb": (247, 203, 140)},
{"position": (256, 256), "rgb": (167, 156, 150)},
]
PIXEL_TOLERANCE = 10
TEXT2IMG_PIXEL_TOLERANCE = 5
@@ -115,7 +88,7 @@ def _find_free_port() -> int:
return port
def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
def _configure_sampling_params(omni: Omni, num_inference_steps: int = 14) -> list:
"""Configure sampling parameters for Bagel img2img generation.
Args:
@@ -258,7 +231,6 @@ def _resolve_deploy_config(config_path: str, run_level: str) -> str:
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.skip(reason="Skip failed CI issue 3977: https://github.com/vllm-project/vllm-omni/issues/3977")
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
def test_bagel_img2img_shared_memory_connector(run_level):
"""Test Bagel img2img with shared memory connector."""
@@ -276,7 +248,6 @@ def test_bagel_img2img_shared_memory_connector(run_level):
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.skip(reason="Skip failed CI issue 3977: https://github.com/vllm-project/vllm-omni/issues/3977")
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
def test_bagel_text2img_shared_memory_connector(run_level):
"""Test Bagel text2img with shared memory connector."""
@@ -387,14 +387,15 @@ def test_process_and_restore_queues(build_adapter):
running_req = _req("r1", RequestStatus.RUNNING)
waiting_queue = DummyWaitingQueue([waiting_req])
running_queue = [running_req]
scheduler_requests = {waiting_req.request_id: waiting_req, running_req.request_id: running_req}
adapter.process_pending_chunks(waiting_queue, running_queue)
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert waiting_req.status == RequestStatus.WAITING_FOR_CHUNK
assert running_req.status == RequestStatus.WAITING_FOR_CHUNK
assert waiting_queue == []
assert running_queue == []
adapter.restore_queues(waiting_queue, running_queue)
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert waiting_queue == [waiting_req]
assert running_queue == [running_req]
assert adapter.waiting_for_chunk_waiting_requests == deque()
@@ -897,6 +898,8 @@ def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture):
sched = OmniARScheduler.__new__(OmniARScheduler)
sched.chunk_transfer_adapter = adapter
sched.requests = {}
sched.running = []
sched.waiting = []
with patch.object(VLLMScheduler, "finish_requests", _super_finish):
OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED)
@@ -1170,3 +1173,164 @@ def test_deferred_finish_not_finished_still_emits_output(mocker: MockerFixture):
assert eco.outputs[0].finish_reason == "stop"
assert eco.outputs[0].kv_transfer_params is None
assert scheduler._pending_finish_reqs == []
# ---------------------------------------------------------------
# Zombie purge tests (regression for vllm-project/vllm-omni#3736)
# ---------------------------------------------------------------
def test_process_pending_chunks_purges_zombies_in_running_deque(
build_adapter,
):
"""A request aborted while parked in ``waiting_for_chunk_running_requests``
is no longer tracked by ``scheduler.requests`` once
``Scheduler._free_request`` runs. ``process_pending_chunks`` must drop
that zombie *before* ``restore_queues`` would re-inject it onto the
running queue.
Regression for https://github.com/vllm-project/vllm-omni/issues/3736
(engine-core ``KeyError`` on aborts under chunk-transfer pressure).
"""
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
live_req = _req("live-1", RequestStatus.WAITING_FOR_CHUNK)
zombie_req = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK)
adapter.waiting_for_chunk_running_requests.append(live_req)
adapter.waiting_for_chunk_running_requests.append(zombie_req)
# Mirror state the adapter would carry for an in-flight request so we can
# later assert ``cleanup_receiver`` actually fired against the zombie.
adapter.requests_origin_status[zombie_req.request_id] = RequestStatus.RUNNING
adapter.requests_origin_status[live_req.request_id] = RequestStatus.RUNNING
waiting_queue = DummyWaitingQueue()
running_queue: list = []
# Simulate the post-abort scheduler: only ``live-1`` is still tracked.
scheduler_requests = {live_req.request_id: live_req}
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
# 1. Zombie was popped from the deque.
assert zombie_req not in adapter.waiting_for_chunk_running_requests
# 2. Live request is still in the deque.
assert live_req in adapter.waiting_for_chunk_running_requests
# 3. ``cleanup_receiver`` ran for the zombie (drops origin-status mapping
# and registers the id as cancelled so a late load/poll is dropped too).
assert zombie_req.request_id not in adapter.requests_origin_status
assert zombie_req.request_id in adapter._cancelled_load_reqs
# 4. Live request's bookkeeping is untouched.
assert adapter.requests_origin_status[live_req.request_id] == RequestStatus.RUNNING
assert live_req.request_id not in adapter._cancelled_load_reqs
# 5. ``restore_queues`` (which the scheduler runs in its ``finally``
# clause) now only re-injects the live request -- the zombie is
# gone, so the worker's ``_update_states`` cannot crash on it.
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert running_queue == [live_req]
def test_process_pending_chunks_purges_zombies_in_waiting_deque(build_adapter):
"""Zombie purge applies symmetrically to the waiting-side deque.
Regression for https://github.com/vllm-project/vllm-omni/issues/3736 --
aborted requests that landed in ``waiting_for_chunk_waiting_requests``
must also be removed before ``restore_queues`` re-injects them, since
``restore_queues`` uses ``waiting_queue.add_request`` rather than
``running_queue.extend`` for that path.
"""
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
live_req = _req("live-w", RequestStatus.WAITING_FOR_CHUNK)
zombie_req = _req("zombie-w", RequestStatus.WAITING_FOR_CHUNK)
adapter.waiting_for_chunk_waiting_requests.append(live_req)
adapter.waiting_for_chunk_waiting_requests.append(zombie_req)
waiting_queue = DummyWaitingQueue()
running_queue: list = []
scheduler_requests = {live_req.request_id: live_req}
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert live_req in adapter.waiting_for_chunk_waiting_requests
assert zombie_req not in adapter.waiting_for_chunk_waiting_requests
assert zombie_req.request_id in adapter._cancelled_load_reqs
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert waiting_queue == [live_req]
def test_purge_preserves_live_order_with_interleaved_zombies(build_adapter):
"""Interleaved live/zombie ordering -- pin the ``popleft`` + append-live
in-place filter at chunk_transfer_adapter._purge_untracked_chunk_requests:
survivor order must match insertion order, and ``cleanup_receiver`` must
run exactly once per zombie.
"""
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
live1 = _req("live-1", RequestStatus.WAITING_FOR_CHUNK)
zombie1 = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK)
live2 = _req("live-2", RequestStatus.WAITING_FOR_CHUNK)
zombie2 = _req("zombie-2", RequestStatus.WAITING_FOR_CHUNK)
for req in (live1, zombie1, live2, zombie2):
adapter.waiting_for_chunk_running_requests.append(req)
waiting_queue = DummyWaitingQueue()
running_queue: list = []
scheduler_requests = {live1.request_id: live1, live2.request_id: live2}
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert list(adapter.waiting_for_chunk_running_requests) == [live1, live2]
assert zombie1.request_id in adapter._cancelled_load_reqs
assert zombie2.request_id in adapter._cancelled_load_reqs
def test_restore_queues_purges_late_aborts_after_process_pending_chunks(
build_adapter,
):
"""Race window guard: an abort can fire between ``process_pending_chunks``
and the scheduler's ``finally``-clause ``restore_queues`` call. The
second purge inside ``restore_queues`` must drop the now-untracked
request so it does not get re-injected onto ``running_queue``.
"""
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
req = _req("late-abort", RequestStatus.WAITING_FOR_CHUNK)
adapter.waiting_for_chunk_running_requests.append(req)
waiting_queue = DummyWaitingQueue()
running_queue: list = []
# Tick N: process_pending_chunks runs while req is still tracked.
scheduler_requests: dict = {req.request_id: req}
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert req in adapter.waiting_for_chunk_running_requests
# Mid-tick: abort fires and the scheduler's free path deletes the entry.
del scheduler_requests[req.request_id]
# finally: restore_queues sees the now-untracked req and must drop it
# instead of blindly extending it onto running_queue.
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
assert running_queue == []
assert req.request_id in adapter._cancelled_load_reqs
def test_purge_is_noop_on_empty_deques(build_adapter):
"""Empty deques short-circuit -- guards against any accidental
``IndexError`` from ``popleft`` on an empty deque if a future caller
shadows the empty check.
"""
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
assert len(adapter.waiting_for_chunk_waiting_requests) == 0
assert len(adapter.waiting_for_chunk_running_requests) == 0
waiting_queue = DummyWaitingQueue()
running_queue: list = []
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests={})
assert len(adapter.waiting_for_chunk_waiting_requests) == 0
assert len(adapter.waiting_for_chunk_running_requests) == 0
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests={})
assert running_queue == []
assert waiting_queue == []
@@ -8,7 +8,7 @@ L4-style smoke runs against an already-running server. The pytest wrappers in
H100 / MI325) because they launch the live Omni server inside the test.
1. **Daily-Omni** — MCQ accuracy fields in the saved JSON (``daily_omni_accuracy``, …); by default the
run **fails** if accuracy is strictly below **0.67** (``--min-daily-omni-accuracy`` / ``ACC_BENCH_MIN_DAILY_OMNI_ACCURACY``).
run **fails** if accuracy is strictly below **0.69** (``--min-daily-omni-accuracy`` / ``ACC_BENCH_MIN_DAILY_OMNI_ACCURACY``).
2. **Seed-TTS** — ``seed-tts-eval``-style metrics when ``--seed-tts-wer-eval`` is used
(WER / SIM / UTMOS keys from :func:`compute_seed_tts_wer_metrics`).
@@ -308,9 +308,9 @@ def build_arg_parser() -> argparse.ArgumentParser:
p.add_argument(
"--min-daily-omni-accuracy",
type=float,
default=float((os.environ.get("ACC_BENCH_MIN_DAILY_OMNI_ACCURACY") or "0.67").strip() or "0.67"),
default=float((os.environ.get("ACC_BENCH_MIN_DAILY_OMNI_ACCURACY") or "0.69").strip() or "0.69"),
help="Fail when daily_omni_accuracy is strictly below this threshold (0–1). "
"Default baseline 0.67; override with env ACC_BENCH_MIN_DAILY_OMNI_ACCURACY or pass 0 to disable the floor.",
"Default baseline 0.69; override with env ACC_BENCH_MIN_DAILY_OMNI_ACCURACY or pass 0 to disable the floor.",
)
p.add_argument(
@@ -340,7 +340,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
p.add_argument(
"--max-seed-tts-mean-wer",
type=float,
default=0.5,
default=0.35,
help="If set, fail when seed_tts_content_error_mean is strictly above this value.",
)
p.add_argument(
@@ -218,7 +218,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> tuple[Image.I
torch_dtype=torch.bfloat16,
trust_remote_code=True,
).to("cuda")
pipe.transformer.set_attention_backend("_flash_3_hub")
_diffusers_dummy_run(pipe)
generator = torch.Generator(device="cuda").manual_seed(SEED)
+45 -1
View File
@@ -29,7 +29,46 @@ from tests.e2e.accuracy.helpers import (
)
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner, OmniServer
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
build_prompt_tokens,
resolve_stop_token_ids,
)
# CoT structural markers that MUST appear in AR output text.
# The AR trajectory under think_recaption mode is:
# <think> ... </think> <recaption> ... <relation_1> ... </relation_1>
# <relation_2> ... </relation_2> </recaption>
# Test deploy config sets include_stop_str_in_output=True and
# skip_special_tokens=False, so ALL opening and closing markers
# should be present in the detokenized text.
_THINK_CLOSE = "</" + "think>" # avoid Python slash-t-as-TAB pitfall
_COT_REQUIRED_MARKERS = (
"<recaption>",
"</recaption>",
"<relation_1>",
"</relation_1>",
"<relation_2>",
"</relation_2>",
_THINK_CLOSE,
)
def _assert_cot_structural_markers(cot_text: str, label: str) -> None:
"""Assert that the AR CoT output contains all required structural tags.
Guards against regressions where detokenizer stop-token stripping
or truncation logic loses CoT content. Under the test deploy config
(include_stop_str_in_output=True, skip_special_tokens=False),
ALL markers should be present in the detokenized text.
"""
missing = [m for m in _COT_REQUIRED_MARKERS if m not in cot_text]
assert not missing, (
f"[{label}] CoT text missing required structural markers: {missing}. "
f"This typically indicates that stop-token stripping or "
f"_truncate_at_cot_end truncation removed CoT content. "
f"CoT text length={len(cot_text)}, first 100 chars: {repr(cot_text[:100])}"
)
os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
@@ -129,6 +168,7 @@ _DEPLOY_CONFIG = {
"stop_token_ids": [128025],
"detokenize": True,
"skip_special_tokens": False,
"include_stop_str_in_output": True,
},
},
{
@@ -400,6 +440,8 @@ def test_image_to_image_alignment_online(accuracy_artifact_root: Path, accuracy_
]
print("[ONLINE] " + tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid"))
_assert_cot_structural_markers(online_cot, "ONLINE")
assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], (
f"[ONLINE] COT semantic similarity {cot_results['cot_semantic_sim']:.4f} below threshold {THRESHOLDS['cot_semantic_sim']}"
)
@@ -508,6 +550,8 @@ def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_
print(tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid"))
_assert_cot_structural_markers(omni_cot, "OFFLINE")
assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], (
f"COT semantic similarity {cot_results['cot_semantic_sim']:.4f} is below threshold {THRESHOLDS['cot_semantic_sim']}"
)
@@ -0,0 +1,218 @@
from __future__ import annotations
import base64
import copy
import io
import os
import tempfile
from pathlib import Path
import pytest
import requests
import yaml
from PIL import Image
from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServer
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
SEED = 42
NUM_INFERENCE_STEPS = 50
GUIDANCE_SCALE = 2.5
HEIGHT = 1024
WIDTH = 1024
PROMPT = "A brown and white dog is running on the grass."
MEAN_THRESHOLD = 3e-2
P99_THRESHOLD = 3e-1
SSIM_THRESHOLD = 0.97
PSNR_THRESHOLD = 30.0
_REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
BASELINE_PATH = _REPO_ROOT / "tests" / "assets" / "hunyuan" / "hunyuan_baseline.png"
_OFFLINE_SCRIPT = _REPO_ROOT / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py"
# DiT-only deploy config with trust_remote_code (based on hunyuan_image3_dit.yaml).
_DEPLOY_CONFIG = {
"pipeline": "hunyuan_image3_dit",
"async_chunk": False,
"trust_remote_code": True,
"stages": [
{
"stage_id": 0,
"max_num_seqs": 1,
"gpu_memory_utilization": 0.9,
"enforce_eager": True,
"trust_remote_code": True,
"devices": "0,1,2,3", # set dynamically by _write_deploy_config
"vae_use_slicing": False,
"vae_use_tiling": False,
"parallel_config": {
"pipeline_parallel_size": 1,
"data_parallel_size": 1,
"tensor_parallel_size": 4,
"enable_expert_parallel": True,
"sequence_parallel_size": 1,
"ulysses_degree": 1,
"ring_degree": 1,
"cfg_parallel_size": 1,
"vae_patch_parallel_size": 1,
"use_hsdp": False,
"hsdp_shard_size": -1,
"hsdp_replicate_size": 1,
},
"default_sampling_params": {
"seed": SEED,
},
},
],
}
def _model_name() -> str:
return os.environ.get("HUNYUAN_IMAGE3_MODEL", MODEL_NAME)
def _devices() -> str:
return os.environ.get("HUNYUAN_IMAGE3_DEVICES", "0,1,2,3")
def _write_deploy_config(path: Path) -> None:
config = copy.deepcopy(_DEPLOY_CONFIG)
devices = _devices()
config["stages"][0]["devices"] = devices
config["stages"][0]["parallel_config"]["tensor_parallel_size"] = len(devices.split(","))
path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
def _run_vllm_omni_hunyuan_image3_online(*, model: str, deploy_config: str, output_path: Path) -> Image.Image:
server_args = [
"--deploy-config",
deploy_config,
"--stage-init-timeout",
"300",
"--init-timeout",
"900",
"--enforce-eager",
"--trust-remote-code",
]
with OmniServer(model, server_args, use_omni=True) as omni_server:
response = requests.post(
f"http://{omni_server.host}:{omni_server.port}/v1/images/generations",
json={
"model": omni_server.model,
"prompt": PROMPT,
"size": f"{WIDTH}x{HEIGHT}",
"n": 1,
"response_format": "b64_json",
"num_inference_steps": NUM_INFERENCE_STEPS,
"guidance_scale": GUIDANCE_SCALE,
"seed": SEED,
"bot_task": "none",
"use_system_prompt": "en_unified",
},
timeout=600,
)
response.raise_for_status()
payload = response.json()
assert len(payload["data"]) == 1
image_bytes = base64.b64decode(payload["data"][0]["b64_json"])
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
image.load()
image.save(output_path)
return image
def _run_vllm_omni_hunyuan_image3_offline(*, model: str, deploy_config: str, output_path: Path) -> Image.Image:
import subprocess
import sys
output_dir = str(output_path.parent)
subprocess.run(
[
sys.executable,
str(_OFFLINE_SCRIPT),
"--modality",
"text2img",
"--deploy-config",
deploy_config,
"--prompts",
PROMPT,
"--output",
output_dir,
"--steps",
str(NUM_INFERENCE_STEPS),
"--guidance-scale",
str(GUIDANCE_SCALE),
"--seed",
str(SEED),
"--height",
str(HEIGHT),
"--width",
str(WIDTH),
"--bot-task",
"none",
"--sys-type",
"en_unified",
"--model",
model,
"--enforce-eager",
],
check=True,
)
images = sorted(Path(output_dir).glob("output_*.png"))
assert images, f"No output image found in {output_dir}"
image = Image.open(images[0]).convert("RGB")
image.load()
image.save(output_path)
return image
def _assert_against_baseline(image: Image.Image, label: str) -> None:
assert BASELINE_PATH.exists(), f"Baseline image not found at {BASELINE_PATH}"
baseline_image = Image.open(BASELINE_PATH).convert("RGB")
assert_images_pixel_close(
model_name=f"{MODEL_NAME} ({label} vs baseline)",
vllm_image=image,
diffusers_image=baseline_image,
mean_threshold=MEAN_THRESHOLD,
p99_threshold=P99_THRESHOLD,
)
assert_similarity(
model_name=f"{MODEL_NAME} ({label} vs baseline)",
vllm_image=image,
diffusers_image=baseline_image,
ssim_threshold=SSIM_THRESHOLD,
psnr_threshold=PSNR_THRESHOLD,
)
@hardware_test(res={"cuda": "H100"}, num_cards=4)
def test_hunyuan_image3_pixel_accuracy_online(accuracy_artifact_root: Path) -> None:
model = _model_name()
output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME)
with tempfile.TemporaryDirectory() as tmpdir:
deploy_config_path = Path(tmpdir) / "deploy.yaml"
_write_deploy_config(deploy_config_path)
image = _run_vllm_omni_hunyuan_image3_online(
model=model, deploy_config=str(deploy_config_path), output_path=output_dir / "vllm_omni_online.png"
)
_assert_against_baseline(image, "online")
@hardware_test(res={"cuda": "H100"}, num_cards=4)
def test_hunyuan_image3_pixel_accuracy_offline(accuracy_artifact_root: Path) -> None:
model = _model_name()
output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME)
with tempfile.TemporaryDirectory() as tmpdir:
deploy_config_path = Path(tmpdir) / "deploy.yaml"
_write_deploy_config(deploy_config_path)
image = _run_vllm_omni_hunyuan_image3_offline(
model=model, deploy_config=str(deploy_config_path), output_path=output_dir / "vllm_omni_offline.png"
)
_assert_against_baseline(image, "offline")
+2
View File
@@ -106,6 +106,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image:
trust_remote_code=True,
local_files_only=_local_files_only(model),
).to("cuda")
pipe.transformer.set_attention_backend("_flash_3_hub")
generator = torch.Generator(device="cuda").manual_seed(SEED)
result = pipe( # pyright: ignore[reportCallIssue]
prompt=PROMPT,
@@ -167,6 +168,7 @@ def _run_diffusers_qwen_image_2512(*, model: str, output_path: Path) -> Image.Im
trust_remote_code=True,
local_files_only=_local_files_only(model),
).to("cuda")
pipe.transformer.set_attention_backend("_flash_3_hub")
generator = torch.Generator(device="cuda").manual_seed(SEED_2512)
result = pipe( # pyright: ignore[reportCallIssue]
prompt=PROMPT_2512,
@@ -0,0 +1,273 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end ModelOpt FP8 quality accuracy for video-gen DiTs.
Modeled on ``test_hunyuan_image3.py``'s ``test_quantized_dit_matches_bf16_accuracy``:
for each calibrated FP8 checkpoint the test generates one video with BF16 and one
with the FP8 DiT under the same seed, scores prompt-faithfulness with CLIP on the
middle frame, then asserts both ``CLIP >= absolute_floor`` and
``CLIP_drop <= drop_threshold``. The test is gated by an opt-in env var (model
paths are local) so it never runs in generic CI.
To run::
VIDEOGEN_RUN_QUANT_ACCURACY=1 \\
WAN22_A14B_BF16_MODEL=Wan-AI/Wan2.2-T2V-A14B-Diffusers \\
WAN22_A14B_FP8_MODEL=/path/to/wan22-a14b-modelopt-fp8 \\
HV15_BF16_MODEL=hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
HV15_FP8_MODEL=/path/to/hv15-modelopt-fp8 \\
pytest -s -v tests/e2e/accuracy/test_videogen_modelopt_quant.py
"""
from __future__ import annotations
import gc
import json
import os
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import numpy as np
import pytest
import torch
from PIL import Image
from tests.e2e.accuracy.helpers import CLIPScorer, model_output_dir
pytestmark = [pytest.mark.local_model, pytest.mark.diffusion]
# ---------------------------------------------------------------------------
# Tunables (override via env to suit local hardware / calibration budget)
# ---------------------------------------------------------------------------
SEED = 42
HEIGHT = int(os.environ.get("VIDEOGEN_QUANT_HEIGHT", "480"))
WIDTH = int(os.environ.get("VIDEOGEN_QUANT_WIDTH", "832"))
NUM_FRAMES = int(os.environ.get("VIDEOGEN_QUANT_NUM_FRAMES", "17"))
NUM_INFERENCE_STEPS = int(os.environ.get("VIDEOGEN_QUANT_STEPS", "20"))
WAN_PROMPT = "A red fox running through a snowy forest at sunrise, cinematic."
HV_PROMPT = "A red fox running through a snowy forest at sunrise, cinematic."
# CLIP gates: absolute floor + max drop from the BF16 reference.
CLIP_ABSOLUTE_FLOOR = float(os.environ.get("VIDEOGEN_QUANT_CLIP_FLOOR", "20.0"))
CLIP_DROP_FP8 = float(os.environ.get("VIDEOGEN_QUANT_CLIP_DROP_FP8", "7.0"))
RUN_ENV = "VIDEOGEN_RUN_QUANT_ACCURACY"
_TRUE = {"1", "true", "yes", "on"}
# ---------------------------------------------------------------------------
@dataclass(frozen=True)
class _Spec:
"""A model under test (BF16 + its FP8 variant)."""
family: str
bf16_env: str
fp8_env: str
prompt: str
guidance_scale: float
guidance_scale_2: float | None = None
boundary_ratio: float | None = None
WAN_SPEC = _Spec(
family="wan22_a14b_t2v",
bf16_env="WAN22_A14B_BF16_MODEL",
fp8_env="WAN22_A14B_FP8_MODEL",
prompt=WAN_PROMPT,
guidance_scale=3.5,
guidance_scale_2=4.0,
boundary_ratio=0.875,
)
HV_SPEC = _Spec(
family="hunyuanvideo15_t2v",
bf16_env="HV15_BF16_MODEL",
fp8_env="HV15_FP8_MODEL",
prompt=HV_PROMPT,
guidance_scale=6.0,
)
@dataclass(frozen=True)
class _QuantCase:
name: str
spec: _Spec
clip_drop: float
def _cases() -> list[pytest.ParameterSet]:
cases = [
_QuantCase("wan_fp8", WAN_SPEC, CLIP_DROP_FP8),
_QuantCase("hv_fp8", HV_SPEC, CLIP_DROP_FP8),
]
params: list[pytest.ParameterSet] = []
opted_in = os.environ.get(RUN_ENV, "").lower() in _TRUE
for case in cases:
marks: list[Any] = []
if not opted_in:
marks.append(pytest.mark.skip(reason=f"Set {RUN_ENV}=1 to run video-gen quant accuracy."))
if not os.environ.get(case.spec.bf16_env):
marks.append(pytest.mark.skip(reason=f"Set {case.spec.bf16_env} to run {case.name}."))
if not os.environ.get(case.spec.fp8_env):
marks.append(pytest.mark.skip(reason=f"Set {case.spec.fp8_env} to run {case.name}."))
params.append(pytest.param(case, id=case.name, marks=marks))
return params
# ---------------------------------------------------------------------------
def _middle_frame(frames: list[Image.Image] | list[np.ndarray] | np.ndarray) -> Image.Image:
if isinstance(frames, np.ndarray):
idx = frames.shape[0] // 2
arr = frames[idx]
if arr.dtype != np.uint8:
arr = (arr * 255).clip(0, 255).astype(np.uint8)
return Image.fromarray(arr).convert("RGB")
if not frames:
raise AssertionError("Empty video output")
mid = frames[len(frames) // 2]
if isinstance(mid, np.ndarray):
if mid.dtype != np.uint8:
mid = (mid * 255).clip(0, 255).astype(np.uint8)
return Image.fromarray(mid).convert("RGB")
return mid.convert("RGB")
def _extract_frames(result) -> list[Image.Image]:
"""Pull a frame list out of an Omni text-to-video result."""
from vllm_omni.entrypoints.omni import OmniRequestOutput
if isinstance(result, list):
result = result[0] if result else None
if isinstance(result, OmniRequestOutput) and result.is_pipeline_output and result.request_output is not None:
result = result.request_output
images = getattr(result, "images", None)
if not images:
raise AssertionError("Pipeline output had no images")
if isinstance(images, list) and images and isinstance(images[0], tuple):
# (frames, fps)
images = images[0][0]
return list(images)
def _generate_video(
*,
model: str,
spec: _Spec,
quantization: str | None,
output_path: Path,
) -> tuple[Image.Image, float]:
"""Generate one video and return its middle frame + elapsed wall time."""
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
kwargs: dict[str, Any] = dict(
model=model,
parallel_config=DiffusionParallelConfig(),
enforce_eager=True,
vae_use_tiling=True,
)
if quantization:
kwargs["quantization"] = quantization
if spec.boundary_ratio is not None:
kwargs["boundary_ratio"] = spec.boundary_ratio
omni = Omni(**kwargs)
try:
gen = torch.Generator(device=current_omni_platform.device_type).manual_seed(SEED)
sampling = dict(
height=HEIGHT,
width=WIDTH,
generator=gen,
guidance_scale=spec.guidance_scale,
num_inference_steps=NUM_INFERENCE_STEPS,
num_frames=NUM_FRAMES,
seed=SEED,
)
if spec.guidance_scale_2 is not None:
sampling["guidance_scale_2"] = spec.guidance_scale_2
import time as _time
t0 = _time.perf_counter()
out = omni.generate({"prompt": spec.prompt}, OmniDiffusionSamplingParams(**sampling))
elapsed = _time.perf_counter() - t0
frames = _extract_frames(out)
mid = _middle_frame(frames)
mid.save(output_path)
return mid, elapsed
finally:
del omni
gc.collect()
if torch.cuda.is_available():
torch.accelerator.empty_cache()
@pytest.mark.parametrize("case", _cases())
@pytest.mark.skipif(
not torch.cuda.is_available() or torch.accelerator.device_count() < 1,
reason="Needs at least 1 GPU.",
)
def test_quantized_videogen_dit_matches_bf16_clip(case: _QuantCase, accuracy_artifact_root: Path) -> None:
"""FP8 DiT should preserve prompt-aligned video quality vs BF16."""
out_dir = model_output_dir(accuracy_artifact_root, case.spec.family + "-quant")
with tempfile.TemporaryDirectory():
bf16_model = os.environ[case.spec.bf16_env]
quant_model = os.environ[case.spec.fp8_env]
bf16_frame, bf16_time = _generate_video(
model=bf16_model,
spec=case.spec,
quantization=None,
output_path=out_dir / "bf16.png",
)
quant_frame, quant_time = _generate_video(
model=quant_model,
spec=case.spec,
quantization="fp8",
output_path=out_dir / f"{case.name}.png",
)
clip = CLIPScorer()
bf16_clip = clip.score(bf16_frame, case.spec.prompt)
quant_clip = clip.score(quant_frame, case.spec.prompt)
clip_drop = bf16_clip - quant_clip
speedup = (bf16_time / quant_time) if quant_time else float("nan")
metrics = {
"case": case.name,
"family": case.spec.family,
"quantization": "fp8",
"bf16_model": bf16_model,
"quant_model": quant_model,
"prompt": case.spec.prompt,
"height": HEIGHT,
"width": WIDTH,
"num_frames": NUM_FRAMES,
"num_inference_steps": NUM_INFERENCE_STEPS,
"bf16_elapsed_s": bf16_time,
"quant_elapsed_s": quant_time,
"speedup_vs_bf16": speedup,
"bf16_clip_score": bf16_clip,
"quant_clip_score": quant_clip,
"clip_score_drop": clip_drop,
}
metrics_path = out_dir / f"{case.name}_metrics.json"
metrics_path.write_text(json.dumps(metrics, indent=2, sort_keys=True), encoding="utf-8")
print(f"\n[{case.name}] BF16 {bf16_time:.2f}s | FP8 {quant_time:.2f}s (x{speedup:.2f})")
print(f" CLIP BF16={bf16_clip:.2f} FP8={quant_clip:.2f} drop={clip_drop:.2f}")
print(f" metrics={metrics_path}")
assert quant_clip >= CLIP_ABSOLUTE_FLOOR, (
f"{case.name} CLIP below floor: got {quant_clip:.2f}, expected >= {CLIP_ABSOLUTE_FLOOR:.2f}"
)
assert clip_drop <= case.clip_drop, (
f"{case.name} CLIP drop too large: got {clip_drop:.2f}, expected <= {case.clip_drop:.2f}"
)
@@ -51,7 +51,7 @@ test_params = [(m, c) for m in models for c in [get_eager_config()]]
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
def test_text_to_text(omni_runner, omni_runner_handler) -> None:
"""
@@ -67,7 +67,7 @@ def test_text_to_text(omni_runner, omni_runner_handler) -> None:
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
def test_image_to_text(omni_runner, omni_runner_handler) -> None:
"""
@@ -84,7 +84,7 @@ def test_image_to_text(omni_runner, omni_runner_handler) -> None:
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
def test_audio_to_text(omni_runner, omni_runner_handler) -> None:
"""
@@ -103,7 +103,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler) -> None:
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
def test_video_to_text(omni_runner, omni_runner_handler) -> None:
"""
@@ -120,7 +120,7 @@ def test_video_to_text(omni_runner, omni_runner_handler) -> None:
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
def test_mixed_to_text(omni_runner, omni_runner_handler) -> None:
"""
@@ -140,7 +140,7 @@ def test_mixed_to_text(omni_runner, omni_runner_handler) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
def test_text_to_audio(omni_runner, omni_runner_handler) -> None:
"""
@@ -156,7 +156,7 @@ def test_text_to_audio(omni_runner, omni_runner_handler) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
def test_image_to_audio(omni_runner, omni_runner_handler) -> None:
"""
@@ -0,0 +1,295 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end offline inference tests for Ming-omni-tts."""
import asyncio
import uuid
import numpy as np
import pytest
import torch
from transformers import AutoTokenizer
from vllm import SamplingParams
from vllm.inputs import tokens_input
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 AsyncOmni
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
KEY_MAX_DECODE_STEPS,
SAMPLE_RATE,
TEXT_EOS_TOKEN_ID,
)
from vllm_omni.model_executor.models.ming_tts.prompt_assembly import DEFAULT_PROMPT, build_ming_dense_prompt
MODEL = "inclusionAI/Ming-omni-tts-0.5B"
DEPLOY_CONFIG = get_deploy_config_path("ming_tts.yaml")
TEST_TEXT = "我会一直在这里陪着你,直到你慢慢地沉入那个最温柔的梦里。"
TEST_INSTRUCTION = "轻柔的ASMR耳语,慢速,贴近麦克风"
MIN_AUDIO_SAMPLES = 1000
@pytest.fixture(scope="module")
def ming_tokenizer():
return AutoTokenizer.from_pretrained(MODEL, trust_remote_code=False)
@pytest.fixture
def ming_engine():
with OmniRunner(
MODEL,
deploy_config=DEPLOY_CONFIG,
stage_init_timeout=300,
enforce_eager=True,
) as runner:
yield runner.omni
@pytest.fixture
def async_omni_engine():
engine = AsyncOmni(
model=MODEL,
deploy_config=DEPLOY_CONFIG,
stage_init_timeout=300,
enforce_eager=True,
)
yield engine
engine.shutdown()
def _build_prompt(
tokenizer,
*,
text: str = TEST_TEXT,
instruction=TEST_INSTRUCTION,
use_zero_spk_emb: bool = True,
) -> dict:
prompt_dict = build_ming_dense_prompt(
tokenizer,
prompt=DEFAULT_PROMPT,
text=text,
instruction=instruction,
runtime_controls={KEY_MAX_DECODE_STEPS: 200},
use_zero_spk_emb=use_zero_spk_emb,
)
prompt = tokens_input(prompt_token_ids=prompt_dict["prompt_token_ids"])
prompt["prompt"] = prompt_dict["prompt"]
prompt["text"] = prompt_dict["text"]
prompt["additional_information"] = prompt_dict["additional_information"]
return prompt
def _sampling_params_list() -> list[SamplingParams]:
return [
SamplingParams(
temperature=0.0,
max_tokens=201,
stop_token_ids=[int(TEXT_EOS_TOKEN_ID)],
),
SamplingParams(temperature=0.0, max_tokens=1),
]
def _flatten_audio(audio) -> torch.Tensor:
if isinstance(audio, list):
parts = [torch.as_tensor(item, dtype=torch.float32).reshape(-1).cpu() for item in audio]
parts = [item for item in parts if item.numel() > 0]
if not parts:
return torch.zeros((0,), dtype=torch.float32)
return torch.cat(parts, dim=0)
return torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()
def _extract_audio(multimodal_output: dict) -> torch.Tensor:
assert isinstance(multimodal_output, dict), f"Expected dict, got {type(multimodal_output)}"
audio = multimodal_output.get("audio", multimodal_output.get("model_outputs"))
assert audio is not None, f"No audio output found, keys={list(multimodal_output.keys())}"
waveform = _flatten_audio(audio)
if waveform.numel() == 0:
raise RuntimeError("Generated audio waveform is empty")
return waveform
def _extract_sample_rate(multimodal_output: dict) -> int:
sample_rate = multimodal_output.get("sr")
if sample_rate is None:
raise RuntimeError("Expected multimodal_output['sr']")
if isinstance(sample_rate, list):
sample_rate = sample_rate[-1]
if hasattr(sample_rate, "item"):
sample_rate = sample_rate.item()
return int(sample_rate)
def _extract_final_audio_outputs(outputs):
final_outputs = []
for item in outputs:
if getattr(item, "final_output_type", None) == "audio":
final_outputs.append(item)
continue
request_output = getattr(item, "request_output", None)
if request_output is None:
continue
multimodal_output = getattr(request_output, "multimodal_output", None)
if isinstance(multimodal_output, dict):
final_outputs.append(item)
continue
completions = getattr(request_output, "outputs", None) or []
if any(isinstance(getattr(completion, "multimodal_output", None), dict) for completion in completions):
final_outputs.append(item)
return final_outputs
def _extract_multimodal_output(output) -> dict:
multimodal_output = getattr(output, "multimodal_output", None)
if isinstance(multimodal_output, dict):
return multimodal_output
request_output = getattr(output, "request_output", None)
if request_output is not None:
multimodal_output = getattr(request_output, "multimodal_output", None)
if isinstance(multimodal_output, dict):
return multimodal_output
completions = getattr(request_output, "outputs", None) or []
for completion in completions:
multimodal_output = getattr(completion, "multimodal_output", None)
if isinstance(multimodal_output, dict):
return multimodal_output
raise AssertionError("No multimodal audio output found in Ming generate results")
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_ming_tts_offline_basic(ming_engine, ming_tokenizer) -> None:
"""Test blocking Ming generation through Omni."""
outputs = ming_engine.generate(
prompts=[_build_prompt(ming_tokenizer)],
sampling_params_list=_sampling_params_list(),
py_generator=False,
)
final_outputs = _extract_final_audio_outputs(outputs)
assert len(final_outputs) == 1, f"Expected one final audio output, got {len(final_outputs)}"
multimodal_output = _extract_multimodal_output(final_outputs[0])
waveform = _extract_audio(multimodal_output)
sample_rate = _extract_sample_rate(multimodal_output)
assert waveform.ndim == 1
assert waveform.shape[0] == waveform.numel()
assert waveform.numel() > MIN_AUDIO_SAMPLES
assert np.max(np.abs(waveform.numpy())) > 0.01, "Audio appears silent"
assert sample_rate == SAMPLE_RATE, f"Expected Ming output sample rate {SAMPLE_RATE}, got {sample_rate}"
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_ming_tts_speaker_conditioning_differs(ming_engine, ming_tokenizer) -> None:
"""Test that different Ming speaker controls produce different waveform outputs."""
style_outputs = ming_engine.generate(
prompts=[_build_prompt(ming_tokenizer)],
sampling_params_list=_sampling_params_list(),
py_generator=False,
)
ip_outputs = ming_engine.generate(
prompts=[_build_prompt(ming_tokenizer, text=TEST_TEXT, instruction={"IP": "灵小甄"}, use_zero_spk_emb=True)],
sampling_params_list=_sampling_params_list(),
py_generator=False,
)
style_final_outputs = _extract_final_audio_outputs(style_outputs)
ip_final_outputs = _extract_final_audio_outputs(ip_outputs)
assert len(style_final_outputs) == 1, "No style audio output produced"
assert len(ip_final_outputs) == 1, "No IP audio output produced"
style_waveform = _extract_audio(_extract_multimodal_output(style_final_outputs[0]))
ip_waveform = _extract_audio(_extract_multimodal_output(ip_final_outputs[0]))
assert style_waveform.numel() > MIN_AUDIO_SAMPLES
assert ip_waveform.numel() > MIN_AUDIO_SAMPLES
assert np.max(np.abs(style_waveform.numpy())) > 0.01, "Style audio appears silent"
assert np.max(np.abs(ip_waveform.numpy())) > 0.01, "IP audio appears silent"
overlap = min(int(style_waveform.numel()), int(ip_waveform.numel()))
mean_abs_diff = torch.mean(torch.abs(style_waveform[:overlap] - ip_waveform[:overlap])).item()
assert style_waveform.shape != ip_waveform.shape or mean_abs_diff > 1e-4, (
"Speaker-conditioned outputs should differ, but style and IP waveforms were effectively identical"
)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_ming_tts_multiple_prompts_queued(ming_engine, ming_tokenizer) -> None:
"""Regression: supported max_num_seqs=1 config must still drain queued prompts."""
prompts = [
_build_prompt(
ming_tokenizer,
text="第一条语音用于验证 Ming dense 队列中的长请求可以完成。",
instruction="平静自然的旁白,语速中等",
),
_build_prompt(
ming_tokenizer,
text="第二条短语音也必须生成完成。",
instruction="清晰明亮的女声",
),
]
outputs = ming_engine.generate(
prompts=prompts,
sampling_params_list=_sampling_params_list(),
py_generator=False,
)
final_outputs = _extract_final_audio_outputs(outputs)
assert len(final_outputs) == len(prompts), f"Expected {len(prompts)} audio outputs, got {len(final_outputs)}"
for i, output in enumerate(final_outputs):
waveform = _extract_audio(_extract_multimodal_output(output))
duration_s = waveform.numel() / SAMPLE_RATE
assert 0.1 < duration_s < 30.0, f"Request {i} audio duration out of range: {duration_s:.2f}s"
assert np.max(np.abs(waveform.numpy())) > 0.01, f"Request {i} audio appears silent"
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_ming_tts_offline_streaming(async_omni_engine, ming_tokenizer) -> None:
"""Test async_chunk streaming Ming generation through AsyncOmni."""
async def _run() -> None:
all_audio_chunks = []
accumulated_samples = 0
chunk_idx = 0
sample_rate = None
async for stage_output in async_omni_engine.generate(
prompt=_build_prompt(ming_tokenizer),
request_id=str(uuid.uuid4()),
sampling_params_list=_sampling_params_list(),
):
multimodal_output = stage_output.multimodal_output or {}
audio = multimodal_output.get("audio")
if "sr" in multimodal_output:
sample_rate = _extract_sample_rate(multimodal_output)
if audio is None:
continue
finished = stage_output.finished
if isinstance(audio, torch.Tensor):
if finished:
audio_chunk = audio[accumulated_samples:].float().detach().cpu()
else:
audio_chunk = audio.float().detach().cpu()
elif isinstance(audio, list):
audio_chunk = _flatten_audio(audio)
else:
audio_chunk = torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()
accumulated_samples += int(audio_chunk.numel())
chunk_idx += 1
if audio_chunk.numel() > 0:
all_audio_chunks.append(audio_chunk)
assert all_audio_chunks, "No streaming audio chunks received"
waveform = torch.cat(all_audio_chunks, dim=0)
assert waveform.numel() > MIN_AUDIO_SAMPLES
assert np.max(np.abs(waveform.numpy())) > 0.01, "Audio appears silent"
assert sample_rate is not None, "Streaming path did not return a sample rate"
assert sample_rate == SAMPLE_RATE, f"Expected Ming output sample rate {SAMPLE_RATE}, got {sample_rate}"
asyncio.run(_run())
+2 -7
View File
@@ -32,7 +32,8 @@ from vllm_omni import Omni
# Selected by the nightly "TTS · Function Test" (-m "full_model and L4 and tts").
# Without the tts + run-level marks these tests are collected but never run.
pytestmark = [pytest.mark.full_model, pytest.mark.tts]
# TODO: Fix this test
# pytestmark = [pytest.mark.full_model, pytest.mark.tts]
SAMPLE_RATE = 24_000 # All MOSS-TTS full variants output 24 kHz
@@ -128,7 +129,6 @@ def delay_engine():
return Omni(model=_DELAY_MODEL, stage_init_timeout=300)
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_delay_english(delay_engine, ref_audio_path):
"""MossTTSDelayModel: English voice_clone produces non-empty 24 kHz audio."""
@@ -140,7 +140,6 @@ def test_moss_tts_delay_english(delay_engine, ref_audio_path):
assert not torch.all(audio == 0), "Audio is silence"
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_delay_chinese(delay_engine, ref_audio_path):
"""MossTTSDelayModel: Chinese input produces non-empty audio."""
@@ -152,7 +151,6 @@ def test_moss_tts_delay_chinese(delay_engine, ref_audio_path):
assert not torch.all(audio == 0)
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_delay_deterministic(delay_engine, ref_audio_path):
"""MossTTSDelayModel: same seed yields identical waveforms."""
@@ -164,7 +162,6 @@ def test_moss_tts_delay_deterministic(delay_engine, ref_audio_path):
assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms differ across identical seeds"
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_delay_batch(delay_engine, ref_audio_path):
"""MossTTSDelayModel: batch of two requests each returns non-empty audio."""
@@ -206,7 +203,6 @@ def realtime_engine():
return Omni(model=_REALTIME_MODEL, stage_init_timeout=300)
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_realtime_english(realtime_engine, ref_audio_path):
"""MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
@@ -218,7 +214,6 @@ def test_moss_tts_realtime_english(realtime_engine, ref_audio_path):
assert not torch.all(audio == 0), "Audio is silence"
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_realtime_deterministic(realtime_engine, ref_audio_path):
"""MossTTSRealtime: same seed yields identical waveforms."""
@@ -120,7 +120,6 @@ def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_nano_english(omni_runner: OmniRunner, ref_audio_path) -> None:
"""English TTS produces non-empty 48 kHz stereo audio."""
@@ -133,7 +132,6 @@ def test_moss_tts_nano_english(omni_runner: OmniRunner, ref_audio_path) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_nano_chinese(omni_runner: OmniRunner, ref_audio_path) -> None:
"""Chinese TTS produces non-empty audio."""
@@ -146,7 +144,6 @@ def test_moss_tts_nano_chinese(omni_runner: OmniRunner, ref_audio_path) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_nano_deterministic(omni_runner: OmniRunner, ref_audio_path) -> None:
"""Same seed produces identical waveforms."""
@@ -159,7 +156,6 @@ def test_moss_tts_nano_deterministic(omni_runner: OmniRunner, ref_audio_path) ->
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"})
def test_moss_tts_nano_batch(omni_runner: OmniRunner, ref_audio_path) -> None:
"""Batch of two requests returns audio for each."""
@@ -24,10 +24,11 @@ from vllm import SamplingParams
from tests.helpers.mark import hardware_test
from vllm_omni import Omni
# H100-gated (8B). Picked up by the nightly H100 omni job
# (-m "full_model and H100 and omni"); the tts mark keeps it consistent with
# the other MOSS-TTS tests and any tts-scoped sweeps.
pytestmark = [pytest.mark.full_model, pytest.mark.tts]
# H100-gated (8B). Picked up by the nightly "TTS · Function Test with H100"
# job (-m "full_model and H100 and tts"); the L4 MOSS-TTS tests land in the
# sibling L4 tts lane instead.
# TODO: Fix this test
# pytestmark = [pytest.mark.full_model, pytest.mark.tts]
SAMPLE_RATE = 24_000
REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav"
@@ -101,7 +102,6 @@ def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
raise AssertionError("No stage outputs received")
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"})
def test_moss_tts_v15_voice_clone(v15_engine, ref_audio_path):
"""MOSS-TTS-v1.5 (8B): voice_clone produces non-empty 24 kHz audio."""
@@ -0,0 +1,160 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Expanded end-to-end test for SenseNova-U1 in offline mode.
This test file primarily covers end-to-end tests for LoRA support.
Validates that LoRA adapters are correctly loaded, applied with controllable
scale, and cleanly deactivated. Uses a synthetic rank-4 adapter targeting the
first decoder layer's gen-path (``_mot_gen``) output projection, so the
perturbation is exercised during the denoising loop.
Assertions:
(a) LoRA at scale=1.0 visibly changes the output (diff > 0.5)
(b) scale=2.0 produces a different delta than scale=1.0 (scale sensitivity)
(c) The delta is bounded (not corrupted)
(d) Deactivating LoRA exactly restores the baseline (diff == 0)
Observed on H100 (seed=42, 8 steps, 512x512): diff_1x~62.2, diff_2x~76.9,
diff_restored=0.0. Thresholds match test_bagel_expansion.py and hold with margin.
"""
import json
import os
from pathlib import Path
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
import numpy as np
import pytest
import torch
from PIL import Image
from safetensors.torch import save_file
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.lora.request import LoRARequest
from vllm_omni.lora.utils import stable_lora_int_id
MODEL = "SenseNova/SenseNova-U1-8B-MoT"
DEFAULT_PROMPT = "a photo of a cat sitting on a laptop keyboard"
SIZE = (512, 512) # (width, height); small for fast E2E
# SenseNova-U1-8B-MoT hidden_size (Qwen3-based). o_proj is square (hidden x hidden)
# regardless of GQA, so this is the robust LoRA target. Verify against config.json.
_LORA_DIM = 4096
_LORA_MODULE = "language_model.model.layers.0.self_attn.o_proj_mot_gen"
_LORA_RANK = 4
# ---------------------------------------------------------------------------
# Helpers (reused from test_sensenova_u1_text2img.py patterns)
# ---------------------------------------------------------------------------
def _build_sampling_params(
lora_request: LoRARequest | None = None,
lora_scale: float = 1.0,
) -> OmniDiffusionSamplingParams:
return OmniDiffusionSamplingParams(
height=SIZE[1],
width=SIZE[0],
seed=42,
num_inference_steps=8,
extra_args={
"cfg_scale": 4.0,
"cfg_norm": "none",
"timestep_shift": 3.0,
"cfg_interval": (0.0, 1.0),
"batch_size": 1,
"think": False,
"t_eps": 0.02,
},
lora_request=lora_request,
lora_scale=lora_scale,
)
def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
for req_output in omni_outputs:
if images := getattr(req_output, "images", None):
return images[0]
return None
def _generate(omni: Omni, lora_request: LoRARequest | None = None, lora_scale: float = 1.0) -> Image.Image:
outputs = list(
omni.generate(
prompts={"prompt": DEFAULT_PROMPT, "modalities": ["image"]},
sampling_params_list=_build_sampling_params(lora_request, lora_scale),
)
)
img = _extract_generated_image(outputs)
assert img is not None, "No image generated"
return img
def _make_file_lora_request(adapter_dir: Path) -> LoRARequest:
"""Write synthetic adapter to disk and return a file-backed LoRARequest."""
adapter_dir.mkdir(parents=True, exist_ok=True)
gen = torch.Generator().manual_seed(42)
lora_a = torch.randn((_LORA_RANK, _LORA_DIM), dtype=torch.float32, generator=gen) * 0.1
lora_b = torch.randn((_LORA_DIM, _LORA_RANK), dtype=torch.float32, generator=gen) * 0.5
save_file(
{
f"base_model.model.{_LORA_MODULE}.lora_A.weight": lora_a,
f"base_model.model.{_LORA_MODULE}.lora_B.weight": lora_b,
},
str(adapter_dir / "adapter_model.safetensors"),
)
(adapter_dir / "adapter_config.json").write_text(
json.dumps({"r": _LORA_RANK, "lora_alpha": _LORA_RANK, "target_modules": [_LORA_MODULE]}),
encoding="utf-8",
)
lora_dir = str(adapter_dir)
return LoRARequest(lora_name="test_file", lora_int_id=stable_lora_int_id(lora_dir), lora_path=lora_dir)
# ---------------------------------------------------------------------------
# Test
# ---------------------------------------------------------------------------
@pytest.mark.full_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"})
def test_sensenova_u1_lora_scale_and_deactivation(tmp_path) -> None:
"""Validate LoRA effect, scale sensitivity, bounded perturbation, and clean deactivation."""
lora_request = _make_file_lora_request(tmp_path / "sensenova_lora")
with OmniRunner(MODEL, stage_configs_path=None) as runner:
omni = runner.omni
baseline = _generate(omni)
img_1x = _generate(omni, lora_request, lora_scale=1.0)
img_2x = _generate(omni, lora_request, lora_scale=2.0)
restored = _generate(omni)
baseline_arr = np.array(baseline, dtype=np.int16)
diff_1x = np.abs(baseline_arr - np.array(img_1x, dtype=np.int16)).mean()
diff_2x = np.abs(baseline_arr - np.array(img_2x, dtype=np.int16)).mean()
diff_restored = np.abs(baseline_arr - np.array(restored, dtype=np.int16)).mean()
# (a) Adapter has visible effect at both scales
assert diff_1x > 0.5, f"LoRA scale=1.0 had no visible effect: diff={diff_1x}"
assert diff_2x > 0.5, f"LoRA scale=2.0 had no visible effect: diff={diff_2x}"
# (b) Different scales produce different outputs
assert not np.isclose(diff_1x, diff_2x, atol=1.0), (
f"LoRA scale has no effect: diff_1x={diff_1x:.2f}, diff_2x={diff_2x:.2f}"
)
# (c) Output is not corrupted
assert diff_1x < 80, f"LoRA output looks corrupted: diff_1x={diff_1x}"
assert diff_2x < 120, f"LoRA output looks corrupted: diff_2x={diff_2x}"
# (d) Deactivation fully restores base model
assert diff_restored == 0.0, f"Base model not restored after LoRA deactivation: diff={diff_restored}"
Loaded 100 of 299 files, more files were not shown because too many files have changed in this diff. Show more