Compare commits

...
Author SHA1 Message Date
tzhouamandCursor d2c149f25b fix: harden omni weight snapshot downloads
Serialize full-repo Hugging Face snapshot materialization with a repo-wide
node lock and retry incomplete snapshots so CI does not race on shared caches.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-04 07:51:54 +00:00
Gao Han f0758f5cbe Merge branch 'main' into feat/diffusion-prefetch-protection 2026-06-03 14:59:14 +08:00
amy-why-3459 9bb4974198 [BugFix] Fix the issue of dataset names not being resolved (#4094) 2026-06-03 06:17:04 +00:00
ChenWenjing 0e2d55f46d [ci] add Voxcpm2 accuracy tests (#4055)
Signed-off-by: CHEN <116010019@link.cuhk.edu.cn>
2026-06-03 14:11:04 +08:00
40b2959183 Add Cosmos3 sound generation (#4073)
Signed-off-by: Maciej Bala <mbala@nvidia.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>
Co-authored-by: Bartosz Stefaniak <bstefaniak@nvidia.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
2026-06-03 00:51:44 +00:00
Honghan ZhuandClaude Opus 4.8 3329f1b674 [Refactor] Migrate dynin_omni to pipeline registry, drop legacy stage… (#4078)
Signed-off-by: dph97 <nenbaying@163.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-02 21:28:12 +00:00
Nick CaoandClaude 23484d8a11 [Bugfix] Pass media paths to use_mixed_modalities in example script (#3355)
Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-06-02 21:04:16 +00:00
bdfc77147d [Diffusion] DreamZero world model integration with CFG parallel + OpenPI serving (#2162)
Signed-off-by: Meng <meng_chen99@163.com>
Signed-off-by: Yangshen Deng <yangshen.d@outlook.com>
Signed-off-by: Suli Wang <50828270+wsl2000@users.noreply.github.com>
Co-authored-by: Meng <meng_chen99@163.com>
Co-authored-by: Suli Wang <50828270+wsl2000@users.noreply.github.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-06-02 19:50:20 +00:00
e35cd213d7 [NPU] [Perf] Adjust flash_attn mask shape for hunyuanvideo1.5 on npu (#3178)
Signed-off-by: liqian <65649795+vasede@users.noreply.github.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-06-02 23:23:18 +08:00
Gao Han 69fd466fba Merge branch 'main' into feat/diffusion-prefetch-protection 2026-06-02 23:22:08 +08:00
NumberWan 35ee3c7543 [BugFix]Qwen-Image performance regression by using torch RMSNorm(RMSNorm backend) (#4074)
Signed-off-by: NumberWan <wantszkin2003@gmail.com>
2026-06-02 23:19:17 +08:00
rein yangandHongsheng Liu 3a7c7f1422 [Bugfix] fix qwen3-omni performance regression (#3575)
Signed-off-by: rein yang <ruiruyang2@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-06-02 23:02:07 +08:00
Alex Brooks 12e75c6618 [Core] Integrate TrackingArgumentParser (#3369)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-06-02 21:56:24 +08:00
lsyyyyy bd37f3c27d Support VAE parallel for Bagel (#3982)
Signed-off-by: siyuan.lei <siyuanlei37@gmail.com>
2026-06-02 13:14:01 +00:00
Boopathy Kannappan 7c729e19d7 [Bugfix] Fix SD3 T5 truncation check device mismatch on long prompts (#3949)
Signed-off-by: bkannappan <bkannappan@digitalocean.com>
2026-06-02 11:04:08 +00:00
tzhouam cb08af976b fix: correct logging message format in hub_prefetch.py for better clarity during prefetch attempts 2026-06-02 09:01:47 +00:00
tzhouam 617c6ca55f feat: enhance model loading with prefetching and retry logic
- Introduced `from_pretrained_with_prefetch` to handle racy cache scenarios by re-prefetching and retrying on failures.
- Updated various model pipelines to utilize the new prefetching mechanism, ensuring robust loading of model components.
- Prefetch logic added to multiple models, including HunyuanVideo, LongCatImage, and QwenImage, to mitigate issues with incomplete cache states.
2026-06-02 09:01:47 +00:00
Zhang JianandCanlin Guo b550709bba [TTS][Model] support MOSS-TTS series (#3420)
Signed-off-by: Zhang <jianmusings@gmail.com>
Signed-off-by: Zhang Jian <jianmusings@gmail.com>
Signed-off-by: jian <jianmusings@gmail.com>
Signed-off-by: Canlin Guo <961750412@qq.com>
Co-authored-by: Canlin Guo <961750412@qq.com>
2026-06-02 15:46:01 +08:00
tc-mb da9c45e1de add MiniCPM-o 4.5 recipe under recipes/OpenBMB (#4067)
Signed-off-by: tc-mb <tianchi_cai@icloud.com>
2026-06-02 14:33:46 +08:00
NATURE 1fb423e5f6 [Perf][Bagel] Avoid per-step device syncs in Bagel img2img (#3987)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
2026-06-02 10:58:32 +08:00
Chendi.Xue 2a2033df02 [XPU][Rebase v0.22] Fix for 0.22 rebase (#4059)
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
2026-06-02 10:16:44 +08:00
bc794e625f Add Cosmos3 model support (#3454)
Signed-off-by: Maciej Bala <mbala@nvidia.com>
Signed-off-by: MaciejBalaNV <mbala@nvidia.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
2026-06-02 04:04:31 +08:00
2b7249fbb3 [Model]Support MiniCPM-o 4.5 (#3642)
Signed-off-by: tc-mb <tianchi_cai@icloud.com>
Co-authored-by: GKangaroo <1095103651@qq.com>
Co-authored-by: GKangaroo <gqx24@mails.tsinghua.edu.cn>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-06-01 22:12:24 +08:00
TJian 29343540f7 [ROCm] [CI] Bugfix Existing CI cases (#3946)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-06-01 22:02:15 +08:00
Yueqian LinandHongsheng Liu c067485030 [Perf][TTS] Bounded-K active-stream window for Stage 1 (RFC #3535) (#3592)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-06-01 21:59:10 +08:00
8d83fb59e9 [BugFix] Fix two stop reason for multimodal output (#3374)
Signed-off-by: Mike_Qiu <qiudayu.qdy@antgroup.com>
Signed-off-by: Mike Qiu <qdy220091330@gmail.com>
Co-authored-by: Mike_Qiu <qiudayu.qdy@antgroup.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-06-01 21:51:29 +08:00
LHXuuuandvraiti 41a795b5d6 [Metrics] Add audio SLOs + cross-stage transfer families + per-(stage, replica) wrap for upstream vllm:* (#3576)
Signed-off-by: vraiti <vraiti@redhat.com>
Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
Co-authored-by: vraiti <vraiti@redhat.com>
2026-06-01 20:20:48 +08:00
TaffyOfficialandTaffyOfficial ee33601571 [Perf] Deduplicate AR prefix cache hidden-state CPU staging (#3734)
Signed-off-by: TaffyOfficial <2324465096@qq.com>
Co-authored-by: TaffyOfficial <2324465096@qq.com>
2026-06-01 09:35:07 +00:00
b76291d6a1 [New Model] Add Lance (ByteDance) (#3710)
Signed-off-by: lishunyang <lishunyang03@gmail.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Signed-off-by: WeiQing Chen <david6666666@users.noreply.github.com>
Co-authored-by: lishunyang <lishunyang03@gmail.com>
Co-authored-by: Haco <923390377@qq.com>
2026-06-01 16:01:25 +08:00
boatmanandsphinxkkkbc 2dcfcc0519 [Feat]Support Nonverbal Tags in OmniVoice (#3968)
Signed-off-by: sphinxkkkbc <binchengkang8@gmail.com>
Co-authored-by: sphinxkkkbc <binchengkang8@gmail.com>
2026-06-01 15:51:05 +08:00
634af60fba [Entrypoint] Add realtime OpenPI robot serving API (#3673)
Signed-off-by: Yangshen Deng <yangshen.d@outlook.com>
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: Meng <meng_chen99@163.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-06-01 14:47:23 +08:00
王金旭andCursor ac4106900b [skip ci][Recipe] OpenBMB/VoxCPM2 (#3850)
Signed-off-by: wjinxu <1299461899@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-01 14:39:54 +08:00
9afdab4939 [Feat]upgrade vllm version [skip-ci] (#4022)
Signed-off-by: rongfu.leng <lenronfu@gmail.com>
Co-authored-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-01 12:08:10 +08:00
Yi LiuandCopilot ef71d74c81 [Bugfix] Fix FLUX W4A16/AutoRound quant_config propagation (#3587)
Signed-off-by: yiliu30 <yi4.liu@intel.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-06-01 04:06:02 +00:00
+1 f92d84fcf9 [Rebase] Rebase to vllm releases/v0.22.0 (#3891)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: ZhengWG <zwg0606@gmail.com>
Signed-off-by: david6666666 <530634352@qq.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Signed-off-by: nainiu258 <cperfect02@163.com>
Signed-off-by: nainiu258 <101917677+nainiu258@users.noreply.github.com>
Signed-off-by: JasonJ2021 <jasonj@zju.edu.cn>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: skf <54565339+skf-1999@users.noreply.github.com>
Co-authored-by: Zheng Wengang <zwg0606@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: nainiu258 <101917677+nainiu258@users.noreply.github.com>
Co-authored-by: Jason <72191212+JasonJ2021@users.noreply.github.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-06-01 11:08:12 +08:00
d4c139508f [Feature] Add precomputed custom voices and Qwen3-TTS ref-context cache (#3492)
Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Dan250124 <416947747@qq.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-05-31 17:51:24 +00:00
Qing AnandSamit 57e5f0fe69 [model, omni] feat: Qwen3-Omni Thinker LoRA for RL training (#3915)
Signed-off-by: Qing An <qa4@rice.edu>
Co-authored-by: Samit <285365963@qq.com>
2026-06-01 00:10:14 +08:00
ChenWenjing 5dfdf584dd [Feat]audio streaming input for async chunk (#3614)
Signed-off-by: CHEN <116010019@link.cuhk.edu.cn>
2026-05-31 17:39:25 +08:00
NumberWan 6f4bd3e45c [BugFix]Qwen-Image performance regression by using omni RMSNorm(RMSNorm backend) (#3933)
Signed-off-by: NumberWan <wantszkin2003@gmail.com>
2026-05-31 08:06:18 +08:00
Canlin Guo 412bf42a75 [Refactor] Unify _talker_mtp_forward across GPU and NPU model runners (#3476)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-05-31 07:58:00 +08:00
zzh ffd9771c54 Fix hunyuan resolve stop token ids (#3896)
Signed-off-by: zzh <943967662@qq.com>
2026-05-31 07:47:36 +08:00
DMAN 11c4fcedd3 [Recipe] add mistralai voxtral tts recipe (#3498)
Signed-off-by: Dmaner <2663515256@qq.com>
2026-05-31 00:18:52 +08:00
wuhang b0f3672aec [Bugfix] Fix DiffusionWorker crash on SIGINT: ensure NCCL/ZMQ cleanup on shutdown (#3872)
Signed-off-by: WU Hang <whlbx@hotmail.com>
2026-05-31 00:13:23 +08:00
WeiQing Chen 2289478e3a [CI] Temporarily skip failing Bagel connector tests (#4005)
Signed-off-by: David Chen <530634352@qq.com>
2026-05-30 17:46:29 +08:00
Alex BrooksandHongsheng Liu 8d13b54383 [CI/Bugfix] Async Request ID Aliasing (#3953)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-05-30 16:59:42 +08:00
wuhangandAlicia 802a6242f0 [Bugfix][Tests] Remove unnecessary device map in tests init (#3958)
Signed-off-by: WU Hang <whlbx@hotmail.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-05-30 16:18:48 +08:00
Yuanheng ZhaoandZheng Wengang ea6d811b58 [Config Refactor] Migrate Ming-flash-omni-2.0 Image-Gen deploy configs (#3975)
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
Co-authored-by: Zheng Wengang <zwg0606@gmail.com>
2026-05-30 08:13:21 +00:00
WeiQing Chen 5994a0bbf5 docs: update WeChat QR code (#4003)
Signed-off-by: David Chen <530634352@qq.com>
2026-05-30 16:12:11 +08:00
Viacheslav KlimkovandYueqian Lin 914e92cf80 [Refactor][Qwen3-TTS] Extract reusable prompt-embeds builder and make tts_pad_embed a persistent buffer (#3992)
Signed-off-by: Viacheslav Klimkov <vklimkov@nvidia.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-05-30 16:11:00 +08:00
dengyunyang 3b317285c4 [HunyuanImage][Perf] adapt to deploy config changes (#3996)
Signed-off-by: dengyunyang <584797741@qq.com>
2026-05-30 14:48:50 +08:00
WeiQing Chen a3456752c5 [Diffusion] Support USP and VAE patch parallel for HunyuanVideo 1.5 (#3979)
Signed-off-by: david6666666 <530634352@qq.com>
2026-05-30 14:40:42 +08:00
Chendi.Xue f43bb124bc [CI][XPU]Fix sage_attn hard-code import for cuda (#3994)
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
2026-05-30 09:43:14 +08:00
fywc 272e626100 [Feat] support cache-dit for SenseNova-U1 (#3906)
Signed-off-by: hanzheli <hanzheli@kuaishou.com>
2026-05-29 17:40:54 +00:00
knlnguyen1802andSamit 5377108438 [minor, fix] Allow passing class interface as custom pipeline argument (#2973)
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Samit <285365963@qq.com>
2026-05-29 23:39:38 +08:00
nperraud 9fdb3bcb89 [Bugfix][Model] Qwen3-TTS: don't collapse 2D ref_code list when estimating prompt length (#3940)
Signed-off-by: nperraud <nathanael.perraudin@gmail.com>
2026-05-29 09:18:25 -06:00
Chen ChengandClaude Opus 4.7 f1c2caa3c7 [Perf][Qwen3-TTS] Restore Code2Wav cross-request batching (RFC #3163 P0) (#3322)
Signed-off-by: ischencheng <cheng21@seas.upenn.edu>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-29 08:48:58 -06:00
b690a10e3a [Feat] Support MagCache (#1287)
Signed-off-by: Lancer <maruixiang6688@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-05-29 22:26:23 +08:00
Aleksi VesantoandHongsheng Liu eba5617c95 [ROCm] Enable AITER backend with ring attention (#3511)
Signed-off-by: Aleksi Vesanto <avesanto@amd.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-05-29 22:20:29 +08:00
amy-why-3459andHongsheng Liu 2696751003 [Perf] Qwen3-Omni performance optimization (#3878)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-05-29 22:15:23 +08:00
Alicia 8c4a42b16c Temp skip TEST - Entrypoint Test with H100 (#3989)
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-05-29 22:37:09 +12:00
Huanyu Yang 3457758114 [Bugfix] Fix IndexError in DistributedVaeExecutor when vae_patch_parallel_size < world_size (#3928)
Signed-off-by: QingZhou-YangHY <3868850350@qq.com>
2026-05-29 03:17:34 -07:00
WeiQing Chen 67bb9b4a12 [CI][Bugfix] Fix LTX audio-video warmup output typing (#3964)
Signed-off-by: david6666666 <530634352@qq.com>
2026-05-29 07:46:02 +00:00
boatmanandsphinxkkkbc 4d9262a6eb [Feat]Support voice clone for omnivoice in online serving & add seed parameter for reproducible (#3668)
Signed-off-by: sphinxkkkbc <binchengkang8@gmail.com>
Co-authored-by: sphinxkkkbc <binchengkang8@gmail.com>
2026-05-29 07:42:08 +00:00
c69ed4037d [feat] Support prompt embedding caching for diffusion model (#2962)
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Co-authored-by: Samit <285365963@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-05-29 15:36:03 +08:00
Alex Brooks 68a1d23e4d [Bugfix/Feature] Remove Hardcoded Flash Attention in Bagel & Support GQA in SDPA Backend (#3728)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-05-29 10:36:15 +08:00
Chen-Yo Sun 593c4c06aa [BugFix] Avoid Voxtral TTS loading error msg (#3951)
Signed-off-by: Chen-Yo Sun <chenyo.sun@mistral.ai>
2026-05-29 02:35:01 +00:00
NATURE 4987bec425 [bugfix] Solve Nightly / CI failed - tests/e2e/online_serving/test_bagel_expansion.py #3918 (#3936)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
2026-05-29 02:05:07 +00:00
amy-why-3459 5dd1d0a772 [BugFix] Fix prefix-caching issue (#3726)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-05-29 01:51:54 +00:00
ChenWenjingandGao Han 834abdaf8b [ci][nightly] Voxcpm2 performance benchmark (#3864)
Signed-off-by: CHEN <116010019@link.cuhk.edu.cn>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-05-29 09:44:29 +08:00
Jason 5dcf18a984 [Docs] Add recipe for Helios (#3114)
Signed-off-by: JasonJ2021 <jasonj@zju.edu.cn>
2026-05-29 09:34:37 +08:00
Yuxiang Liang 900b4187a1 [XPU] Add CosyVoice3Model support on Intel XPU (#2325)
Signed-off-by: Yuxiang Liang <yuliang@habana.ai>
Signed-off-by: Yuxiang Liang <yuxiang.liang@intel.com>
2026-05-29 00:59:45 +00:00
wuli666andwuli666 8935439428 [Bugfix] Fix Triton SnakeBeta kernel for bf16/fp16 inputs (#3472)
Signed-off-by: wuli666 <djjpro975@gmail.com>
Co-authored-by: wuli666 <djjpro975@gmail.com>
2026-05-28 17:10:36 +00:00
Dakai An bb9f21d0db Fix Qwen3-omni accuracy degradation from deepstack inputs under torch.compile (#3885)
Signed-off-by: Dakai An <dakaian108@gmail.com>
2026-05-28 23:57:56 +08:00
knlnguyen1802andSamit 5eac50f22c [bugfix, rl] Fix sleep do not release full memory in custom pipeline (#3818)
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Co-authored-by: Samit <285365963@qq.com>
2026-05-28 22:27:27 +08:00
126fb4b4a9 [Blackwell][1/N] Add SageAttention3 diffusion backend on blackwell(GB200/B200/RTX5090/PRO6000/DGX Spark available) (#3015)
Signed-off-by: lishunyang <lishunyang12@163.com>
Signed-off-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
Co-authored-by: lishunyang <lishunyang12@163.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-05-28 14:08:45 +00:00
zhumingjue138 8d9a4afed9 [Test] Add scenarios for L5 reliability test (#3729)
Signed-off-by: zhumingjue <zhumingjue@huawei.com>
2026-05-28 21:41:33 +08:00
WeiQing ChenandWeiQing Chen 1acd27c205 [CI][BugFix] Fix and Validate FP8 Z-Image quality gate (#3929)
Signed-off-by: WeiQing Chen <david6666666@users.noreply.github.com>
Co-authored-by: WeiQing Chen <david6666666@users.noreply.github.com>
2026-05-28 12:52:58 +00:00
nainiu258andHongsheng Liu 740ec608f4 [Docs]Add recipe for GLM-Image on 2x A800 GPUs and 1x A800 GPU (#2950)
Signed-off-by: nainiu258 <cperfect02@163.com>
Signed-off-by: nainiu258 <101917677+nainiu258@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-05-28 19:50:37 +08:00
Zeyu Huang | 黃澤宇 74258b18b7 [CI] Rectify L2~L4 Qwen Image Edit series tests (#3901) 2026-05-28 18:33:31 +08:00
431 changed files with 53056 additions and 5495 deletions

No files matched your search

+1 -1
View File
@@ -10,7 +10,7 @@ steps:
DOCKER_BUILDKIT: "1"
# Buildkite will automatically replace this with the actual commit hash
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
VLLM_VERSION: "v0.21.0"
VLLM_VERSION: "v0.22.0"
priority: 100
timeout_in_minutes: 60
soft_fail: false
+16 -13
View File
@@ -1,15 +1,28 @@
steps:
- label: "Simple Unit Test"
- 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 60m pytest -v -s -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"
- "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
@@ -49,16 +62,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 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 (Need 4 GPUs)"
agent_pool: mi325_4
depends_on: amd-build
@@ -123,7 +126,7 @@ steps:
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 tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
- 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
+15 -4
View File
@@ -1,6 +1,6 @@
steps:
- label: "Simple Unit Test"
- label: "Simple · Diffusion & Model Executor Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
@@ -9,7 +9,18 @@ steps:
- 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 60m pytest -v -s -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"
- "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
@@ -19,13 +30,13 @@ steps:
commands:
- timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
- label: "Diffusion Batching Test"
- label: "Diffusion · Batch Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/batching/test_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
- timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
- label: "Custom Pipeline Test"
agent_pool: mi325_1
+24 -1
View File
@@ -107,7 +107,7 @@ steps:
depends_on: upload-merge-pipeline
steps:
- label: "Entrypoint Test with H100"
timeout_in_minutes: 35
timeout_in_minutes: 60
commands:
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
agents:
@@ -389,6 +389,29 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "TTS · VoxCPM2 Online 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 '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: "Omni · Higgs-Audio V2 Test"
timeout_in_minutes: 30
commands:
+119 -61
View File
@@ -398,7 +398,7 @@ steps:
timeout_in_minutes: 90
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
agents:
queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
@@ -413,6 +413,30 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: ":full_moon: TTS · Accuracy Test · VoxCPM2"
timeout_in_minutes: 30
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_voxcpm2_tts_expansion.py -m 'full_model and cuda and tts' --run-level 'full_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: ":full_moon: TTS · Perf Test"
key: nightly-tts-performance
timeout_in_minutes: 180
@@ -459,6 +483,52 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: TTS · VoxCPM2 · Perf Test"
key: nightly-tts-performance-voxcpm2
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_voxcpm2.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
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
# Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below.
- group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
key: nightly-diffusion-x2iat-group
@@ -624,7 +694,49 @@ steps:
- 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: 3
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 H100 · Multi-GPU DreamZero"
timeout_in_minutes: 120
commands:
- >-
pytest -sv
tests/e2e/online_serving/test_dreamzero_expansion.py
-m "full_model and diffusion and H100 and distributed_cuda"
--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
@@ -843,8 +955,8 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2509"
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2509
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
@@ -852,7 +964,7 @@ steps:
- export CACHE_DIT_VERSION=1.3.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
@@ -987,60 +1099,6 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion X2I(&A&T) · HunyuanImage3 · DiT Perf Test"
key: nightly-hunyuan-image3-performance
soft_fail: true
timeout_in_minutes: 120
if: build.env("RUN_HUNYUAN_IMAGE3_PERF") == "1"
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp4_fp8.json
EXIT1=$$?
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_sp2.json
EXIT2=$$?
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_cfgp2.json
EXIT3=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$((EXIT1 | EXIT2 | EXIT3))
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
# Diffusion x2v only (Wan, HunyuanVideo, …). x2i/x2a/x2t live in the X2I group above, not here.
- group: ":card_index_dividers: Diffusion X2V Model Test"
key: nightly-diffusion-x2v-group
@@ -1288,7 +1346,7 @@ steps:
- nightly-tts-performance
- nightly-diffusion-x2iat-performance-qwen-image
- nightly-diffusion-x2iat-performance-qwen-image-edit
- nightly-diffusion-x2iat-performance-qwen-image-edit-2509
- nightly-diffusion-x2iat-performance-qwen-image-edit-2511
- nightly-diffusion-x2iat-performance-qwen-image-layered
- nightly-diffusion-x2iat-performance-bagel
- nightly-diffusion-x2v-performance
@@ -1305,7 +1363,7 @@ steps:
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2509
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
+24 -1
View File
@@ -183,7 +183,7 @@ steps:
depends_on: upload-ready-pipeline
commands:
- |
timeout 15m bash -c "
timeout 20m bash -c "
pytest -sv tests/distributed/ -m 'core_model and cuda' --run-level 'core_model'
"
agents:
@@ -296,6 +296,29 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "TTS · VoxCPM2 Online 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'
"
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 · Qwen3-TTS CustomVoice Test"
commands:
- |
+3
View File
@@ -53,6 +53,9 @@
{% endif %}
{{ indented_cmd | safe }}
priority: 100
{% if step.timeout_in_minutes %}
timeout_in_minutes: {{ step.timeout_in_minutes }}
{% endif %}
{% if step.grade and step.grade == "Blocking" %}
soft_fail: false
{% else %}
+3
View File
@@ -270,3 +270,6 @@ vllm_omni/_version.py
.worktrees/
# CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
tests/.ci_generated/
# Local Lance precision-matching dump scripts (not part of the repo)
lance_compare/
@@ -0,0 +1,10 @@
se001|ocean waves crashing on a rocky beach|10.0
se002|heavy rain on a metal roof at night|8.0
se003|a dog barking in the distance|6.0
se004|a crackling campfire in a quiet forest|10.0
se005|footsteps on dry leaves|6.0
se006|thunder rolling across a wide valley|8.0
se007|a busy coffee shop with chatter and espresso machines|10.0
se008|wind howling through tall pine trees|8.0
se009|a small stream flowing over pebbles|10.0
se010|a creaking wooden door slowly opening|6.0
@@ -0,0 +1,10 @@
ttsd001|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Hello and welcome to the show today. [S2] Thanks for having me, it's a pleasure to be here. [S1] Let's dive right in.
ttsd002|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you finish the report? [S2] Almost — I'm just polishing the conclusion. [S1] Great, send it over when you're done.
ttsd003|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have you tried that new restaurant downtown? [S2] Yes, the pasta is amazing. [S1] I'll book a table for Friday.
ttsd004|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] What did you think of yesterday's meeting? [S2] Productive, but a bit long. [S1] Agreed, we should set an hour cap next time.
ttsd005|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] The weather is beautiful today. [S2] Perfect for a hike. [S1] Want to drive out to the trails this afternoon?
ttsd006|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] How did the demo go? [S2] Better than expected — the client signed off on the prototype. [S1] Fantastic news!
ttsd007|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Coffee or tea this morning? [S2] Coffee please, with a splash of milk. [S1] Coming right up.
ttsd008|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you read the email about the deadline change? [S2] I did, we now have until next Friday. [S1] That should give us enough buffer.
ttsd009|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Are you joining the team dinner tonight? [S2] I'll be there around seven. [S1] See you then.
ttsd010|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have we backed up the database? [S2] Last backup ran two hours ago. [S1] Good, I'll trigger a fresh snapshot before the migration.
View File
Whitespace-only changes.
@@ -0,0 +1,108 @@
#!/usr/bin/env bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# 2x2 timing matrix (backend x compile mode) for SAGE_ATTN_3 vs TORCH_SDPA,
# on HunyuanVideo-1.5 480p T2V and Wan 2.2 TI2V-5B. Run from the repo root:
#
# bash benchmarks/diffusion/bench_sage_attn3_matrix.sh
#
# Requires: sageattn3 installed (for SAGE_ATTN_3 rows) and a Blackwell GPU.
set -u
STEPS=${STEPS:-50}
SEED=${SEED:-42}
PROMPT=${PROMPT:-"A dog running across a field of golden wheat."}
OUT_DIR=${OUT_DIR:-/tmp/bench_sage3}
mkdir -p "$OUT_DIR"
declare -A TOTAL PER_STEP TEXT_ENC TRANSFORMER VAE
run() {
local name="$1"; shift
local model="$1"; shift
local h="$1"; shift
local w="$1"; shift
local f="$1"; shift
local g="$1"; shift
local log="$OUT_DIR/${name}.log"
echo "[$(date +%T)] Running $name ..."
env "$@" python examples/offline_inference/text_to_video/text_to_video.py \
--model "$model" \
--prompt "$PROMPT" \
--height "$h" --width "$w" --num-frames "$f" \
--num-inference-steps "$STEPS" --seed "$SEED" --guidance-scale "$g" \
--enable-diffusion-pipeline-profiler \
--output "$OUT_DIR/${name}.mp4" > "$log" 2>&1 || {
echo " FAILED — see $log"
TOTAL[$name]="FAIL"; PER_STEP[$name]="FAIL"
TEXT_ENC[$name]="-"; TRANSFORMER[$name]="-"; VAE[$name]="-"
return
}
local wait_ms
wait_ms=$(grep "add_req_and_wait=" "$log" | tail -1 | sed -nE 's/.*add_req_and_wait=([0-9.]+) ms.*/\1/p')
local total_s
total_s=$(grep -oE "Total generation time: [0-9.]+ seconds" "$log" | tail -1 | sed -nE 's/.*: ([0-9.]+) seconds.*/\1/p')
local text_enc_s
text_enc_s=$(grep -oE "text_encoder.forward took [0-9.]+s" "$log" | awk '{sub(/s$/,"",$NF); sum+=$NF} END {printf "%.3f", sum+0}')
local vae_s
vae_s=$(grep -oE "vae.decode took [0-9.]+s" "$log" | awk '{sub(/s$/,"",$NF); sum+=$NF} END {printf "%.3f", sum+0}')
local pipeline_s
pipeline_s=$(grep -oE "Pipeline\.forward took [0-9.]+s" "$log" | tail -1 | sed -nE 's/.*took ([0-9.]+)s.*/\1/p')
local transformer_s
transformer_s=$(awk "BEGIN {printf \"%.3f\", ${pipeline_s:-0} - ${vae_s:-0} - ${text_enc_s:-0}}")
TOTAL[$name]="${total_s:-0}"
PER_STEP[$name]=$(awk "BEGIN {printf \"%.3f\", ${wait_ms:-0} / 1000 / $STEPS}")
TEXT_ENC[$name]="$text_enc_s"
TRANSFORMER[$name]="$transformer_s"
VAE[$name]="$vae_s"
echo " total=${TOTAL[$name]}s per-step=${PER_STEP[$name]}s/it transformer=${TRANSFORMER[$name]}s vae=${VAE[$name]}s"
}
# ========== HunyuanVideo-1.5 480p T2V (480x832x33, CFG=6.0) ==========
HV_MODEL=hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
run hv15_sdpa_eager "$HV_MODEL" 480 832 33 6.0 TORCH_COMPILE_DISABLE=1
run hv15_sdpa_compiled "$HV_MODEL" 480 832 33 6.0
run hv15_sage3_eager "$HV_MODEL" 480 832 33 6.0 TORCH_COMPILE_DISABLE=1 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
run hv15_sage3_compiled "$HV_MODEL" 480 832 33 6.0 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
# ========== Wan 2.2 TI2V-5B (704x1280x33, CFG=5.0) ==========
WAN_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers
run wan22_sdpa_eager "$WAN_MODEL" 704 1280 33 5.0 TORCH_COMPILE_DISABLE=1
run wan22_sdpa_compiled "$WAN_MODEL" 704 1280 33 5.0
run wan22_sage3_eager "$WAN_MODEL" 704 1280 33 5.0 TORCH_COMPILE_DISABLE=1 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
run wan22_sage3_compiled "$WAN_MODEL" 704 1280 33 5.0 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
print_table() {
local title="$1"; shift
echo
echo "$title (${STEPS} steps, seed=${SEED})"
echo
printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
"Config" "Total (s)" "Per-step (s/it)" "Text Enc (s)" "Transformer (s)" "VAE Decode (s)"
printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
"-----------------------" "---------" "---------------" "------------" "---------------" "--------------"
for row in "$@"; do
case "$row" in
*_sdpa_eager) label="SDPA + Eager" ;;
*_sdpa_compiled) label="SDPA + Compiled" ;;
*_sage3_eager) label="SAGE_ATTN_3 + Eager" ;;
*_sage3_compiled) label="SAGE_ATTN_3 + Compiled" ;;
esac
printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
"$label" "${TOTAL[$row]}" "${PER_STEP[$row]}" "${TEXT_ENC[$row]}" "${TRANSFORMER[$row]}" "${VAE[$row]}"
done
}
print_table "HunyuanVideo-1.5 480p T2V, 832x480x33" \
hv15_sdpa_eager hv15_sdpa_compiled hv15_sage3_eager hv15_sage3_compiled
print_table "Wan 2.2 TI2V-5B, 1280x704x33" \
wan22_sdpa_eager wan22_sdpa_compiled wan22_sage3_eager wan22_sage3_compiled
echo
echo "Logs: $OUT_DIR/*.log"
+19
View File
@@ -53,6 +53,8 @@ _TASK_TO_DATASET: dict[str, str] = {
"voice_clone": "seed-tts",
"default_voice": "seed-tts-text",
"voice_design": "seed-tts-design",
"dialogue": "ttsd",
"sound_effect": "sound-effect",
}
# Default design dataset path (bundled with the repo)
@@ -81,6 +83,7 @@ def build_bench_args(
output_dir: str | None,
result_filename: str | None,
extra_cli_args: list[str],
output_len: int | None = None,
) -> list[str]:
"""Build the ``vllm bench serve --omni`` command for one (task, concurrency) run."""
dataset_name = _TASK_TO_DATASET[task]
@@ -130,12 +133,18 @@ def build_bench_args(
if task_extra_body:
cmd += ["--extra-body", json.dumps(task_extra_body, separators=(",", ":"))]
if model_cfg.get("trust_remote_code"):
cmd.append("--trust-remote-code")
if concurrency is not None:
cmd += ["--max-concurrency", str(concurrency), "--request-rate", "inf"]
if wer_eval:
cmd.append("--seed-tts-wer-eval")
if output_len is not None:
cmd += ["--hf-output-len", str(output_len)]
if output_dir or result_filename:
out_dir = output_dir or "."
os.makedirs(out_dir, exist_ok=True)
@@ -235,6 +244,15 @@ def main() -> None:
parser.add_argument("--host", default="localhost")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--model-configs", default=str(_DEFAULT_MODEL_CONFIGS), help="Path to model_configs.yaml")
parser.add_argument(
"--output-len",
type=int,
default=None,
help="Override per-request max_new_tokens forwarded to the server "
"(seed-tts datasets default to 2048; the MOSS-TTS full talker takes "
"this as 'max_new_frames' so smaller values cap audio length per "
"request and keep bench runtime tractable).",
)
parser.add_argument("extra", nargs=argparse.REMAINDER, help="Extra args passed directly to vllm bench serve")
args = parser.parse_args()
@@ -289,6 +307,7 @@ def main() -> None:
output_dir=args.output_dir,
result_filename=result_filename,
extra_cli_args=args.extra or [],
output_len=args.output_len,
)
result = run_one_benchmark(cmd)
if result is not None:
+64
View File
@@ -14,6 +14,7 @@ models:
supported_tasks: [default_voice, voice_design]
backend: openai-audio-speech
endpoint: /v1/audio/speech
trust_remote_code: true
task_extra_body:
default_voice:
voice: Vivian
@@ -27,13 +28,76 @@ models:
supported_tasks: [voice_clone]
backend: openai-audio-speech
endpoint: /v1/audio/speech
trust_remote_code: true
task_extra_body:
voice_clone:
task_type: Base
openbmb/VoxCPM2:
supported_tasks: [voice_clone, default_voice]
trust_remote_code: true
backend: openai-audio-speech
endpoint: /v1/audio/speech
trust_remote_code: true
task_extra_body:
voice_clone: {}
default_voice:
voice: default
# MOSS-TTS family. After the serving_speech.py variant-detection refactor
# all five variants now go through OpenAI /v1/audio/speech with the
# variant-specific request fields (ambient_sound, ref_audio_2, instructions).
OpenMOSS-Team/MOSS-TTS:
supported_tasks: [voice_clone]
backend: openai-audio-speech
endpoint: /v1/audio/speech
task_extra_body:
voice_clone: {}
# MOSS-TTS-v1.5: continued-training upgrade of MOSS-TTS 1.0 — same
# MossTTSDelay-8B architecture and API, so identical wiring. Adds 31-language
# synthesis via the `language` field and inline `[pause Xs]` markers.
OpenMOSS-Team/MOSS-TTS-v1.5:
supported_tasks: [voice_clone]
backend: openai-audio-speech
endpoint: /v1/audio/speech
task_extra_body:
voice_clone: {}
OpenMOSS-Team/MOSS-TTS-Realtime:
supported_tasks: [voice_clone]
backend: openai-audio-speech
endpoint: /v1/audio/speech
task_extra_body:
voice_clone: {}
# MOSS-TTSD synthesizes [S1]/[S2] dialogue. The seed-tts voice_clone dataset
# has single-speaker rows; bench piggybacks on it (both speakers share the
# same ref_audio). For real dialogue evaluation use the ttsd dataset.
OpenMOSS-Team/MOSS-TTSD-v1.0:
supported_tasks: [voice_clone, dialogue]
backend: openai-audio-speech
endpoint: /v1/audio/speech
task_extra_body:
voice_clone: {}
dialogue: {}
# MOSS-SoundEffect ignores text/ref_audio; expects ambient_sound +
# duration_seconds. The sound_effect dataset module emits those.
OpenMOSS-Team/MOSS-SoundEffect:
supported_tasks: [sound_effect]
backend: openai-audio-speech
endpoint: /v1/audio/speech
task_extra_body:
sound_effect:
duration_seconds: 10
# MOSS-VoiceGenerator: instructions-driven (no ref_audio). seed-tts-design
# dataset already emits {instructions, target_text}, so voice_design works
# out of the box once the serving validator accepts instructions.
OpenMOSS-Team/MOSS-VoiceGenerator:
supported_tasks: [voice_design]
backend: openai-audio-speech
endpoint: /v1/audio/speech
task_extra_body:
voice_design: {}
+4 -1
View File
@@ -1,5 +1,8 @@
# Target vllm version: v0.21.0 (releases/v0.21.0 branch).
# The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must
# be updated on every rebase to match the new upstream vllm commit/tag.
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
ARG VLLM_BASE_TAG=v0.21.0
ARG VLLM_BASE_TAG=v0.22.0
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
ARG APP_DIR=/workspace/vllm-omni
WORKDIR ${APP_DIR}
+1 -1
View File
@@ -1,4 +1,4 @@
ARG BASE_IMAGE=vllm/vllm-openai:v0.21.0
ARG BASE_IMAGE=vllm/vllm-openai:v0.22.0
FROM ${BASE_IMAGE}
ARG COMMON_WORKDIR=/app
+1 -1
View File
@@ -1,4 +1,4 @@
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.22.0
FROM ${BASE_IMAGE} AS base
# Declare a variable to know if we want to use the nightly build or the stable build.
+1 -1
View File
@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
ENV UV_INDEX_STRATEGY="unsafe-best-match"
ENV UV_LINK_MODE="copy"
ARG VLLM_VERSION=v0.21.0
ARG VLLM_VERSION=v0.22.0
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
WORKDIR /workspace/vllm
Binary file not shown.

Before

Width:  |  Height:  |  Size: 402 KiB

After

Width:  |  Height:  |  Size: 408 KiB

+4
View File
@@ -13,6 +13,10 @@ This section contains design documents and architecture specifications for vLLM-
- [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
- [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
## Infrastructure Design Documents
- [Prometheus Metrics](metrics.md)
## Module Design Documents
- [AR Module](module/ar_module.md)
+208
View File
@@ -0,0 +1,208 @@
# Prometheus Metrics Design
This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. This iteration covers the text and audio surface; diffusion / image / video families land in a follow-up.
## Objectives
- Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope).
- Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically.
- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss.
- Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable.
- Keep the metrics collection overhead low enough that it does not regress TTFA or throughput.
## Background
### Upstream vLLM Metrics
Upstream vLLM defines ~37 Prometheus metric families under the `vllm:` prefix. These are registered by `PrometheusStatLogger` and cover engine-level state: KV cache usage, running/waiting request counts, token throughput, TTFT, inter-token latency, e2e latency, and so on. They are served via the `/metrics` HTTP endpoint provided by `prometheus_fastapi_instrumentator` and the default `prometheus_client` WSGI handler.
vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` collector whose `_name` attribute contains the substring `"vllm"`. This runs during engine initialization to clean up stale collectors from prior instantiations within the same process.
### The Problem
vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm:omni_` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below).
## Architecture
### Component Overview
```text
+------------------------+
| API Server (FastAPI) |
| GET /metrics |
+-----------+------------+
|
prometheus_client default registry
|
+--------+--------+--------+--------+--------+
| |
vllm:omni_* vllm:*
collectors collectors
| |
+----+--------+ +-----------+ +----------+ +-----------+
| OmniPromet- | | OmniMod- | | OmniTra- | | OmniProm- |
| heusMetrics | | alityMet- | | nsferMe- | | etheusSt- |
| | | rics | | trics | | atLogger |
+----+--------+ +-----+-----+ +----+-----+ +----+------+
| | | |
OmniBase OmniBase Orchestrator Orchestrator
(request life- (finalize + (record_trans- (per-(stage,
cycle, success/ streaming fer_tx/rx replica)
fail counter) hooks via hooks via scheduler/
observe_* emit hook in iteration
APIs) OrchestratorAg- stats)
gregator)
```
### Data Flow
There are four independent paths for metric collection.
**Path 1: Pipeline-level metrics (`vllm:omni_*`)**
`OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress:
- `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`.
- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm:omni_requests_success_total{finished_reason}`.
- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` so a single Counter family covers both natural and aborted completion.
**Path 2: Audio modality metrics (`vllm:omni_audio_*`)**
`OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites:
- `observe_modality_at_finalize(...)` — called from `omni_base._process_single_result` inside the existing `e2e_done` finalize guard. For `output_type == "audio"` it emits `audio_frames_total`, `audio_duration_s`, `audio_rtf` (or `audio_skipped_requests_total{reason="no_audio_data"}` when no audio was produced). Sample rate is resolved from `engine_outputs.multimodal_output` via `definitions.resolve_audio_sample_rate(...)` (fallback chain mirrors `serving_chat.py`'s audio response path).
- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry.
- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`.
**Path 3: Cross-stage transfer metrics (`vllm:omni_transfer_*`)**
`OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution.
The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transfer_rx`. After the existing `TransferEdgeStats` accumulation, the aggregator calls `_emit_transfer_tx` / `_emit_transfer_rx`. Those:
1. Resolve `from_replica` / `to_replica` via a `replica_resolver` callback supplied by `async_omni.py`. The resolver delegates to `stage_pool.get_bound_replica_id(request_id)` — i.e. the orchestrator's existing sticky-routing binding is the source of truth.
2. Convert the underlying `_ms` accumulators to seconds and call the `_s`-suffixed observe methods on `OmniTransferMetrics`.
Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected).
**Path 4: Per-engine metrics (`vllm:*`, stage/replica wrap)**
The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics.
### Shared State Between Threads
The Orchestrator runs in a background thread. The API server (OmniBase) runs in the asyncio event loop thread. `OmniRequestCounter` bridges them — a plain Python object with an `int` field. The Orchestrator increments/decrements it; the entrypoint reads it for gauge updates. No lock is needed because the counter is advisory (a stale read by one Prometheus scrape interval is acceptable). It is created by `AsyncOmniEngine.__init__()` and passed to the Orchestrator at construction time.
### Metric Registration and Lifecycle
All `vllm:omni_*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm:omni_*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
## OmniPrometheusStatLogger Wrap
Upstream `PrometheusStatLogger.__init__` hard-codes `labelnames = ["model_name", "engine"]` as a local variable, references it across ~37 metric-family construction sites, and uses the `engine` label value in five different `.labels()` call shapes (kwarg with int engine, kwarg with str engine, positional with str engine in the middle, plus a `metrics_info["engine"] = str(...)` dict pattern). To reshape `engine` into `stage` + `replica` without forking the entire upstream `__init__`, the wrap uses three coordinated mechanisms:
1. **Class-level metric class slot overrides.** `OmniPrometheusStatLogger` overrides `_gauge_cls`, `_counter_cls`, `_histogram_cls` (which upstream calls via `self._gauge_cls(...)` etc.) with `_RelabelGauge` / `_RelabelCounter` / `_RelabelHistogram` wrapper classes. These intercept the `labelnames` kwarg at metric family creation time and replace `engine` with `("stage", "replica")`.
2. **Property descriptor for `per_engine_labelvalues`.** Upstream builds `self.per_engine_labelvalues = {idx: [model_name, str(idx)]}` inside `__init__` and then captures it into a local variable for `create_metric_per_engine` calls. By making `per_engine_labelvalues` a Python property on the subclass, the setter intercepts upstream's assignment and rewrites each 2-tuple into a 3-tuple `[model_name, stage, replica]` using the `stage_replica_map` supplied at construction time. The captured local then sees the rewritten dict.
3. **Override of `.labels()` on the wrapper classes.** For the five call sites that pass `engine` directly (kwarg or positional, int or str), `_RelabelMixin.labels()` translates the engine value back to `(stage, replica)` via a process-level `_ENGINE_INDEX_MAP` populated by `OmniPrometheusStatLogger.__init__`. This handles `gauge_engine_sleep_state.labels(engine=idx, ...)`, `counter_request_success_base.labels(model_name, str(idx), str(reason))`, `info_gauge.labels(**metrics_info)`, etc.
The three sub-helpers that upstream `PrometheusStatLogger.__init__` constructs (`spec_decoding_prom` / `kv_connector_prom` / `perf_metrics_prom`) use their own `_counter_cls` / `_gauge_cls` / `_histogram_cls` slots and would otherwise build families with the raw 2-element labelnames. `_OmniPerfMetricsProm` / `_OmniSpecDecodingProm` / `_OmniKVConnectorProm` subclass each helper to route the same relabel mixin through their internal family construction.
The `Orchestrator` constructs `stage_replica_map` from the static `stage_pools` configuration at startup:
```python
stage_replica_map = {
flat_idx: (str(stage_id), str(replica_id))
for flat_idx, (stage_id, replica_id) in enumerate(
(s, r)
for s, pool in enumerate(stage_pools)
for r in range(pool.num_replicas)
)
}
```
A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestrator so the per-replica `record(engine_idx=...)` call site can look up the right flat index.
> Dynamic add/remove of replicas at runtime is intentionally out of scope — the upstream `PrometheusStatLogger` materializes per-engine_idx child metrics at init time, and supporting hot-add would require non-trivial intervention into upstream's per-family child dictionaries.
## Gating: `--log-stats`
All metrics — both the 15 `vllm:omni_*` families and the ~65 upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
Behavior with `--log-stats=off` (default):
- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm:omni_*` stay materialized but never have data written to them.
- `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all.
- The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate.
Behavior with `--log-stats=on`: all metric paths fire normally; the orchestrator's per-replica recording is bounded only by `OmniSchedulerMixin.make_stats()`'s per-scheduler 1 Hz throttle (see next section).
The overhead with the flag on is small enough that an A/B benchmark on Qwen3-Omni-30B single replica (30 sequential audio requests) showed a mean latency delta of +0.6% (Welch's t = 0.318, n=30, not statistically significant at α=0.05); the `/metrics` line count drops from 1358 to 124 lines when the flag is off.
## Throttling: `make_stats()` Override
Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, returning a SchedulerStats object for the orchestrator. Under vLLM's architecture, this is fine. But since vLLM-Omni requires that the object be serialized and transferred over ZMQ, receiving a SchedulerStats object on every step can introduce unacceptable overhead to the system.
`OmniSchedulerMixin.make_stats()` (in `vllm_omni/core/sched/omni_scheduler_mixin.py`) throttles stats emission to at most once per second **per scheduler** — i.e. per `(stage, replica)` since each replica owns its own scheduler instance. Between intervals it returns `None`, which the engine core skips serializing. This keeps gauges fresh enough for Prometheus scrapes (typically 15-30s intervals) while eliminating the per-step overhead.
The orchestrator side does not add its own throttle on top: the per-replica recording loop gates only on `raw_outputs.scheduler_stats is not None` (i.e. this replica's scheduler passed its own 1 Hz gate). A previous global `_last_stats_ts` on the orchestrator was removed because it starved every replica other than the first to emit within each second.
## Metric Definitions
### Pipeline (4)
| Metric | Type | Labels | Description |
|--------|------|--------|-------------|
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
### Audio (7)
Labels: `{model_name, stage, replica}` plus the listed extra label.
| Metric | Type | Extra label | Description |
|--------|------|-------------|-------------|
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frames generated |
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
### Cross-stage transfer (4)
Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`.
| Metric | Type | Description |
|--------|------|-------------|
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
### LLM stage-level (wrapped `vllm:*`)
After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — carries `{model_name, stage, replica}` labels. For the full upstream catalog see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md); note that metrics depending on features unsupported in vLLM-Omni (e.g. speculative decoding, LoRA) will not be available.
## Naming Convention
- All time-bearing metrics use the `_s` suffix (values in seconds). Two bucket families are used:
- `SECONDS_BUCKETS` (0.05 s – 300 s) for e2e / generation / TTFP style values.
- `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution.
- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
- Sizes use the `_bytes` suffix.
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
## Logging vs. Prometheus
`OrchestratorAggregator` (in `vllm_omni/metrics/stats.py`) is the logging-oriented metrics path. It collects detailed per-request, per-stage, and per-transfer statistics and prints formatted tables to the `INFO` log. This is designed for development and debugging — individual request traces, transfer bandwidth, inter-stage timing.
`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics` form the Prometheus-oriented path. They record aggregate counters, gauges, and histograms suitable for time-series monitoring and alerting. Both paths share the same source data (`StageRequestStats`, `TransferEdgeStats`) — `OrchestratorAggregator.record_transfer_tx/rx` in particular calls both the existing accumulator code and the Prometheus emit hook in the same method body. The two consumption models can run simultaneously without coupling.
The separation follows upstream vLLM's pattern of `LoggingStatLogger` vs. `PrometheusStatLogger` — same underlying data, different consumption models.
@@ -411,6 +411,8 @@ Notes:
- `runtime.max_batch_size` controls stage-level batching.
- Thinker/Talker commonly use `enforce_eager: false` for CUDA Graph paths.
- Code2Wav often remains eager (`enforce_eager: true`) depending on runtime behavior.
- Qwen3-Omni defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. The Triton has been more stable & faster
than the FlashInfer CUTLASS unquantized MoE backend on recent vLLM rebases.
#### 2) Enable async chunk
+1 -1
View File
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
### Pre-built wheels
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, and 0.21.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, and 0.22.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
=== "NVIDIA CUDA"
@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
uv pip install vllm==0.21.0 --torch-backend=auto
uv pip install vllm==0.22.0 --torch-backend=auto
```
#### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
If you do not need to modify source code of vLLM, you can directly install the stable 0.22.0 release version of the library
```bash
uv pip install vllm==0.21.0 --torch-backend=auto
uv pip install vllm==0.22.0 --torch-backend=auto
```
The 0.21.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 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.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.21.0
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.21.0; the wheel filename has no `+cu130` suffix)
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.21.0/vllm-0.21.0-cp38-abi3-manylinux_2_35_x86_64.whl
#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
```
Install vllm with command below (If you have no existing PyTorch).
```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
--ipc=host \
vllm/vllm-omni:v0.21.0 \
vllm/vllm-omni:v0.22.0 \
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
```
#### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
If you do not need to modify source code of vLLM, you can directly install the stable 0.22.0 release version of the library
```bash
uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
```
The pre-built 0.21.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
The pre-built 0.22.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack 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
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.21.0
git checkout v0.22.0
python3 -m pip install -r requirements/rocm.txt
python3 setup.py develop
```
@@ -130,7 +130,7 @@ docker run --rm \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
vllm/vllm-omni-rocm:v0.21.0 \
vllm/vllm-omni-rocm:v0.22.0 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
@@ -149,7 +149,7 @@ docker run --rm -it \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--entrypoint bash \
vllm/vllm-omni-rocm:v0.21.0
vllm/vllm-omni-rocm:v0.22.0
```
# --8<-- [end:pre-built-images]
+3 -3
View File
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
source .venv/bin/activate
# On CUDA
uv pip install vllm==0.21.0 --torch-backend=auto
uv pip install vllm==0.22.0 --torch-backend=auto
# On ROCm
uv pip install vllm==0.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721
uv pip install vllm==0.22.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm721
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
@@ -35,7 +35,7 @@ For additional installation methods — please see the [installation guide](inst
!!! note
It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.21.0` and vLLM Omni `0.21.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.22.0` and vLLM Omni `0.22.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
## Offline Inference
+2 -1
View File
@@ -119,9 +119,10 @@ def extract_omni_serve_subparser_init():
exec_globals = {
"_FlexibleArgumentParser": _FlexibleArgumentParser,
"FlexibleArgumentParser": _FlexibleArgumentParser,
"TrackingArgumentParser": _FlexibleArgumentParser, # only needs add_argument, not tracking
"TrackingNamespace": None, # type annotation only
"make_arg_parser": lambda parser: parser, # no-op for doc
"_ensure_vllm_platform": lambda: None, # no-op for doc
"nullify_stage_engine_defaults": lambda parser: None, # no-op for doc
"VLLM_SUBCMD_PARSER_EPILOG": "",
"logger": logger,
"DummySubparsers": DummySubparsers,
+3 -1
View File
@@ -33,6 +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` | ✅︎ | | | |
| `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` | ✅︎ | ✅︎ | | |
@@ -41,13 +42,14 @@ th {
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | | | |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `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` | ✅︎ | ✅︎ | | |
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
| `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` | ✅︎ | ✅︎ | | |
+38
View File
@@ -381,6 +381,44 @@ are cached in-process with a shared LRU so repeated requests with the same
all TTS model types; deleting a voice invalidates every model-type slot at
once.
### Precomputed Custom Voices
Qwen3-TTS Base and VoxCPM2 can load offline-precomputed voices at startup.
Generate a directory containing `custom_voice_manifest.json` plus one
`.safetensors` file per voice, then set the pipeline-wide deploy config field:
```yaml
custom_voice_dir: /path/to/custom_voices
```
Qwen3-TTS profiles are created with:
```bash
python examples/online_serving/text_to_speech/qwen3_tts/precompute_custom_voice.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--ref-text "Original transcript of the reference audio" \
--mode icl \
--output-dir /path/to/custom_voices
```
VoxCPM2 profiles are created with:
```bash
python examples/online_serving/text_to_speech/voxcpm2/precompute_custom_voice.py \
--model openbmb/VoxCPM2 \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--mode ref_continuation \
--prompt-text "Original transcript of the reference audio" \
--output-dir /path/to/custom_voices
```
Only profiles whose safetensors payload can be loaded and validated are exposed
by `GET /v1/audio/voices`. Valid precomputed voices can be used in
`POST /v1/audio/speech` by passing `voice="alice"` without `ref_audio`.
**Configuration (environment variables):**
| Variable | Default | Description |
+98
View File
@@ -0,0 +1,98 @@
# Production Metrics
vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. This page covers the text and audio surface; diffusion / image / video metrics are tracked in a follow-up PR.
```bash
vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats
curl http://localhost:8000/metrics
```
**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm:omni_*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
## Metric Namespaces
| Prefix | Source | Present when |
|--------|--------|--------------|
| `vllm:omni_` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
| `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage |
| `http_` / `process_` | Uvicorn / Python runtime | Always |
## Pipeline-Level Metrics (`vllm:omni_`)
Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline.
### Request counts
| Metric | Type | Labels | Description |
|--------|------|--------|-------------|
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
### Latency
| Metric | Type | Labels | Description |
|--------|------|--------|-------------|
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
## Audio Modality Metrics (`vllm:omni_`)
Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label.
| Metric | Type | Extra label | Description |
|--------|------|-------------|-------------|
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition.
## Cross-Stage Transfer Metrics (`vllm:omni_`)
Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed.
| Metric | Type | Description |
|--------|------|-------------|
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
## vLLM Engine Metrics (`vllm:`)
When the pipeline includes an LLM stage, the upstream vLLM engine exposes its full set of ~37 metric families under the `vllm:` prefix.
vLLM-Omni wraps the upstream `vllm.v1.metrics.loggers.PrometheusStatLogger` with `OmniPrometheusStatLogger` so that the original `engine` single label is reshaped into `stage` + `replica`. Every `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — therefore gains per-`(stage, replica)` visibility automatically. No omni-side duplicate is needed for the text path.
```text
# Before wrap:
vllm:num_requests_running{model_name="...", engine="1"} 3.0
# After wrap:
vllm:num_requests_running{model_name="...", stage="1", replica="0"} 2.0
vllm:num_requests_running{model_name="...", stage="1", replica="1"} 1.0
```
For the full list of upstream metrics, see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md).
## Metric Availability by Pipeline Type
| Metric group | Multi-stage LLM (Qwen3-Omni) |
|---|---|
| `vllm:omni_` request tracking + latency | With `--log-stats` |
| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage |
| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
| `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` |
| `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` |
## Naming Convention
- All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values.
- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
- Sizes use the `_bytes` suffix.
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label.
@@ -19,6 +19,8 @@ The full set of backends and their platform defaults is in the **Backend Options
| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
## Configuration
@@ -185,6 +187,29 @@ Quick check:
python -c "import sageattention; print(sageattention.__file__)"
```
## SageAttention3 Installation
vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni.
Build from source:
```bash
git clone https://github.com/thu-ml/SageAttention.git
cd SageAttention/sageattention3_blackwell
python setup.py install
```
Quick check:
```bash
python -c "import sageattn3; print(sageattn3.__file__)"
```
Notes:
- `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class.
- SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness.
## Usage Examples
### Default (auto-route)
@@ -235,6 +260,20 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
--output outputs/wan22_sage.mp4
```
### Enable SageAttention3
Example:
```bash
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_to_video/text_to_video.py \
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
--prompt "A dog running across a field of golden wheat." \
--height 480 --width 832 --num-frames 33 \
--num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
--tensor-parallel-size 2 \
--output outputs/hv15_sage3.mp4
```
### Mixed backends across roles
Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention:
+6 -2
View File
@@ -129,14 +129,16 @@ The following tables show which models support each feature:
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **SenseNova-U1** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **SenseNova-U1** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
> Notes:
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload is not supported; use layerwise offload.
### VideoGen
@@ -147,8 +149,10 @@ The following tables show which models support each feature:
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
**Frame Interpolation Support**
@@ -17,6 +17,9 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded
automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies.
The bundled Qwen3-Omni setup defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. This keeps the Thinker & Talker on vLLM's
Triton unquantized MoE path and avoids the performance regression observed with the FlashInfer CUTLASS unquantized MoE
backend.
Asynchronous chunk streaming is **enabled by default** within the bundled configuration.
To explicitly utilize a custom deployment YAML, specify the configuration path:
@@ -72,6 +75,12 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
```
To experiment with the FlashInfer FP16 MoE path, set `VLLM_USE_FLASHINFER_MOE_FP16=1` before launching the server:
```bash
VLLM_USE_FLASHINFER_MOE_FP16=1 \
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
For the stage-based CLI, you usually do **not** need `--stage-overrides` for
that kind of change. Since each command launches one stage, just pass the knob
directly on that stage command:
@@ -329,6 +329,16 @@ Single-GPU serves now default to the uniproc executor (lower IPC overhead, the B
To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.
### Tuning stage 1 `max_num_seqs` per task type
The bundled `qwen3_tts.yaml` ships stage 1 (Code2Wav) at `max_num_seqs: 10`, tuned for Base voice cloning: stage-1 lifetimes are long (~3 s/req), so admitting up to 10 concurrent codec sequences lets requests progress in parallel in the scheduler — ~2× TTFA p95 at c=4 / c=8 (1× H100, 1.7B-Base, seed-tts) at an 8–12 % audio-throughput cost.
CustomVoice / VoiceDesign have much shorter stage-1 lifetimes (~50–200 ms) and are TTFA-optimal at `max_num_seqs: 1`. Override the default when serving those task types:
```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \
--stage-overrides '{"1": {"max_num_seqs": 1}}'
```
### Sending requests
```bash
# CustomVoice with a predefined speaker
@@ -370,6 +380,31 @@ curl -X POST http://localhost:8091/v1/audio/voices \
```
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
### Precomputed custom voices
For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
```bash
python qwen3_tts/precompute_custom_voice.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--ref-text "Original transcript of the reference audio" \
--mode icl \
--output-dir /path/to/custom_voices
```
`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
```yaml
custom_voice_dir: /path/to/custom_voices
```
Then start the server with that config and call the Speech API with only the voice name:
```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{"input":"Hello from a precomputed voice.","voice":"alice","task_type":"Base"}' \
--output alice.wav
```
### Streaming PCM
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
@@ -409,8 +444,10 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
### Notes
- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
- Under vocoder-bound overload (single-stream `rtf_p99 ≥ 1` at the target concurrency), set `active_stream_window: 2` at the top of the deploy yaml to cap simultaneously active Stage 1 streams. Off by default; trades TTFP for streaming continuity. See [#3592](https://github.com/vllm-project/vllm-omni/pull/3592) for the mechanism and tradeoff numbers.
---
@@ -438,6 +475,23 @@ python voxcpm2/openai_speech_client.py \
```
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
### Precomputed custom voices
For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
```bash
python voxcpm2/precompute_custom_voice.py \
--model openbmb/VoxCPM2 \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--mode ref_continuation \
--prompt-text "Original transcript of the reference audio" \
--output-dir /path/to/custom_voices
```
Add the output directory to the deploy config:
```yaml
custom_voice_dir: /path/to/custom_voices
```
After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
### Gradio demo (gapless streaming via AudioWorklet)
```bash
python voxcpm2/gradio_demo.py
+1 -1
View File
@@ -74,7 +74,7 @@ Online:
vllm serve Qwen/Qwen-Image \
--omni \
--port 8000 \
--quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
--diffusion-quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
```
## Parameters
@@ -195,6 +195,101 @@ stages:
devices: "0,1"
```
### VAE Patch Parallelism
[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) splits Bagel VAE **decode/encode** tiles across multiple GPUs on the **DiT stage**, reducing **per-GPU peak memory during VAE decode**. Use it when high-resolution `text2img` or `img2img` hits VAE OOM or large decode spikes.
**Bagel-specific notes:**
- Implemented in `BagelPipeline` via `DistributedAutoEncoder` (DiT stage only).
- **Single-stage** is the simplest path: one DiT process with TP + VAE patch parallel.
- **Two-stage**: enable on **stage 1 (DiT)** only; stage 0 (Thinker) keeps encoder-only `VAEEncoder` and does not use VAE patch parallel.
- You need a DiT `world_size` ≥ `vae_patch_parallel_size` (typically `tensor_parallel_size=2` on that stage). VAE PP reuses the DiT process group; it is not a standalone second-GPU VAE worker.
**Single-stage via deploy YAML** (recommended for `end2end.py`):
```yaml
pipeline: bagel_single_stage
async_chunk: false
stages:
- stage_id: 0
max_num_batched_tokens: 32768
max_num_seqs: 1
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
devices: "0,1"
vae_use_tiling: true
parallel_config:
tensor_parallel_size: 2
vae_patch_parallel_size: 2
default_sampling_params:
seed: 52
```
```bash
cd examples/offline_inference/bagel
CUDA_VISIBLE_DEVICES=0,1 python end2end.py \
--model /path/to/BAGEL-7B-MoT \
--deploy-config /path/to/bagel_single_stage_vae_pp.yaml \
--modality text2img \
--prompts "A cute cat" \
--steps 10 \
--output ./out_vae_pp
```
**Single-stage via `Omni` kwargs** (same flags as online serving):
```python
from vllm_omni.entrypoints.omni import Omni
omni = Omni(
model="ByteDance-Seed/BAGEL-7B-MoT",
deploy_config="vllm_omni/deploy/bagel_single_stage.yaml",
tensor_parallel_size=2,
vae_patch_parallel_size=2,
vae_use_tiling=True,
)
# Then call omni.generate(...) as in end2end.py
```
**Two-stage (VAE PP on DiT only):**
```yaml
stages:
- stage_id: 0
devices: "0"
# AR Thinker — no vae_patch_parallel here
- stage_id: 1
devices: "0,1"
vae_use_tiling: true
parallel_config:
tensor_parallel_size: 2
vae_patch_parallel_size: 2
```
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--deploy-config /path/to/bagel_vae_pp.yaml \
--modality text2img \
--prompts "A cute cat"
```
**Startup log checks:**
```text
INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
```
| Setting | Role |
| :------ | :--- |
| `parallel_config.tensor_parallel_size` | DiT world size / TP (must be ≥ `vae_patch_parallel_size`) |
| `parallel_config.vae_patch_parallel_size` | Number of ranks for distributed VAE tiles (`1` = off) |
| `vae_use_tiling` | Enable spatial tiling (auto-enabled when `vae_patch_parallel_size > 1`) |
#### Hybrid Sharded Data Parallel (HSDP)
For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
+2 -6
View File
@@ -121,9 +121,6 @@ def parse_args():
help="Temperature for text generation sampling (default: 0.3).",
)
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
args = parser.parse_args()
return args
@@ -154,6 +151,7 @@ def main():
from vllm_omni.entrypoints.omni import Omni
omni_kwargs = vars(args).copy()
omni_kwargs.pop("model", None)
deploy_config = args.deploy_config
if args.think and deploy_config is None:
deploy_config = "vllm_omni/deploy/bagel_think.yaml"
@@ -164,9 +162,7 @@ def main():
if args.quantization:
omni_kwargs["quantization_config"] = args.quantization
# Override CLI --model with the derived model_name.
omni_kwargs["model"] = model_name
omni = Omni(**omni_kwargs)
omni = Omni(model=model_name, **omni_kwargs)
formatted_prompts = []
for p in prompts:
@@ -0,0 +1,304 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import argparse
import uuid
from pathlib import Path
import cv2
import numpy as np
import torch
from PIL import Image
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
WORKER_EXTENSION = "vllm_omni.diffusion.models.dreamzero.video_export_worker.DreamZeroVideoExportWorkerExtension"
ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
DEFAULT_SESSION_PREFIX = "dreamzero-export"
RELATIVE_OFFSETS = [-23, -16, -8, 0]
ACTION_HORIZON = 24
CAMERA_FILES = {
"observation/exterior_image_0_left": "exterior_image_1_left.mp4",
"observation/exterior_image_1_left": "exterior_image_2_left.mp4",
"observation/wrist_image_left": "wrist_image_left.mp4",
}
def _parse_args() -> argparse.Namespace:
repo_root = Path(__file__).resolve().parents[3]
parser = argparse.ArgumentParser(description="Export DreamZero prediction video from downloaded example inputs.")
parser.add_argument("--model", default="GEAR-Dreams/DreamZero-DROID")
parser.add_argument("--deploy-config", type=Path, required=True)
parser.add_argument(
"--video-dir",
type=Path,
default=repo_root / "outputs" / "dreamzero" / "assets",
help="Directory containing the three camera MP4 files.",
)
parser.add_argument(
"--output-dir", type=Path, default=repo_root / "outputs" / "dreamzero" / "generated_predictions"
)
parser.add_argument("--output-stem", default="dreamzero_prediction")
parser.add_argument(
"--prompt",
default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
)
parser.add_argument("--session-id", default=None)
parser.add_argument("--save-input-video", action="store_true")
parser.add_argument("--save-gif", action="store_true")
parser.add_argument("--save-actions", action="store_true")
parser.add_argument("--fps", type=int, default=5)
return parser.parse_args()
def _load_all_frames(video_path: Path) -> np.ndarray:
cap = cv2.VideoCapture(str(video_path))
frames = []
while True:
ok, frame = cap.read()
if not ok:
break
frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
cap.release()
if not frames:
raise RuntimeError(f"No frames loaded from {video_path}")
return np.stack(frames, axis=0)
def _load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
camera_frames: dict[str, np.ndarray] = {}
for camera_key, file_name in CAMERA_FILES.items():
video_path = video_dir / file_name
if not video_path.exists():
raise FileNotFoundError(
f"Missing DreamZero example asset: {video_path}. "
"Download the example videos with: "
f"`hf download {ASSET_REPO_ID} --repo-type dataset --local-dir {video_dir}`"
)
camera_frames[camera_key] = _load_all_frames(video_path)
return camera_frames
def _build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
chunks: list[list[int]] = []
current_frame = 23
for _ in range(num_chunks):
indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
if indices[-1] >= total_frames:
break
chunks.append(indices)
current_frame += ACTION_HORIZON
return chunks
def _make_obs_from_video(
camera_frames: dict[str, np.ndarray],
frame_indices: list[int],
*,
prompt: str,
session_id: str,
) -> dict:
obs: dict = {}
for camera_key, all_frames in camera_frames.items():
selected = all_frames[frame_indices]
obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
obs["prompt"] = prompt
obs["session_id"] = session_id
return obs
def _build_observations(video_dir: Path, prompt: str, session_id: str) -> tuple[dict[str, np.ndarray], list[dict]]:
camera_frames = _load_camera_frames(video_dir)
total_frames = min(frames.shape[0] for frames in camera_frames.values())
chunks = _build_frame_schedule(total_frames, 1)
observations = [
_make_obs_from_video(camera_frames, [0], prompt=prompt, session_id=session_id),
]
if chunks:
observations.append(
_make_obs_from_video(
camera_frames,
chunks[0],
prompt=prompt,
session_id=session_id,
)
)
if len(observations) < 2:
raise RuntimeError("Need at least two DreamZero example observations to export a prediction video.")
return camera_frames, observations[:2]
def _extract_latents(output: OmniRequestOutput) -> torch.Tensor:
if not isinstance(output, OmniRequestOutput):
raise TypeError(f"Expected OmniRequestOutput, got {type(output)!r}")
if not output.images:
raise RuntimeError("DreamZero output does not contain video latents in `images`.")
latents = output.images[0]
if not isinstance(latents, torch.Tensor):
raise TypeError(f"Expected tensor latents, got {type(latents)!r}")
latents = latents.detach().cpu()
if latents.dim() == 4:
latents = latents.unsqueeze(0)
if latents.dim() != 5:
raise ValueError(f"Unexpected latent shape: {tuple(latents.shape)}")
if latents.shape[1] < latents.shape[2]:
latents = latents.transpose(1, 2).contiguous()
return latents
def _write_mp4(path: Path, frames: np.ndarray, fps: int) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
height, width = frames.shape[1:3]
writer = cv2.VideoWriter(
str(path),
cv2.VideoWriter_fourcc(*"mp4v"),
float(fps),
(width, height),
)
if not writer.isOpened():
raise RuntimeError(f"Failed to open video writer for {path}")
try:
for frame in frames:
writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
finally:
writer.release()
def _write_gif(path: Path, frames: np.ndarray, fps: int) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
images = [Image.fromarray(frame) for frame in frames]
duration_ms = max(int(round(1000 / max(fps, 1))), 1)
images[0].save(
path,
save_all=True,
append_images=images[1:],
duration=duration_ms,
loop=0,
)
def _stitch_input_frames(camera_frames: dict[str, np.ndarray]) -> np.ndarray:
total_frames = min(frames.shape[0] for frames in camera_frames.values())
stitched = []
for frame_index in range(total_frames):
left = camera_frames["observation/exterior_image_0_left"][frame_index]
right = camera_frames["observation/exterior_image_1_left"][frame_index]
wrist = camera_frames["observation/wrist_image_left"][frame_index]
pad = np.zeros((left.shape[0], left.shape[1], 3), dtype=np.uint8)
canvas = np.concatenate([left, right], axis=1)
bottom = np.concatenate([wrist, pad], axis=1)
stitched.append(np.concatenate([canvas, bottom], axis=0))
return np.stack(stitched, axis=0)
def _run_generation(
model: str, deploy_config_path: Path, observations: list[dict]
) -> tuple[Omni, list[OmniRequestOutput]]:
omni = Omni(
model=model,
deploy_config=str(deploy_config_path),
enforce_eager=True,
worker_extension_cls=WORKER_EXTENSION,
)
outputs: list[OmniRequestOutput] = []
for index, obs in enumerate(observations):
sampling_params = OmniDiffusionSamplingParams(
extra_args={
"reset": index == 0,
"session_id": obs["session_id"],
"robot_obs": obs,
}
)
result = omni.generate(obs["prompt"], sampling_params_list=[sampling_params])
if not result:
raise RuntimeError(f"No output returned for DreamZero request {index}")
outputs.append(result[0])
return omni, outputs
def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray:
stage_client = omni.engine.stage_clients[0]
engine = getattr(stage_client, "_engine", None)
if engine is None:
raise RuntimeError("DreamZero export requires inline diffusion stage access.")
decoded = engine.executor.collective_rpc(
"decode_video_latents_to_uint8",
args=(full_latents,),
unique_reply_rank=0,
exec_all_ranks=True,
)
if isinstance(decoded, torch.Tensor):
decoded = decoded.numpy()
if not isinstance(decoded, np.ndarray):
raise TypeError(f"Unexpected decoded output type: {type(decoded)!r}")
return decoded
def main() -> None:
args = _parse_args()
session_id = args.session_id or f"{DEFAULT_SESSION_PREFIX}-{uuid.uuid4()}"
camera_frames, observations = _build_observations(
video_dir=args.video_dir,
prompt=args.prompt,
session_id=session_id,
)
args.output_dir.mkdir(parents=True, exist_ok=True)
if args.save_input_video:
input_frames = _stitch_input_frames(camera_frames)
_write_mp4(args.output_dir / f"{args.output_stem}_input.mp4", input_frames, fps=15)
if args.save_gif:
_write_gif(args.output_dir / f"{args.output_stem}_input.gif", input_frames[::3], fps=5)
omni = None
try:
omni, outputs = _run_generation(
model=args.model,
deploy_config_path=args.deploy_config,
observations=observations,
)
latent_steps = [_extract_latents(output) for output in outputs]
full_latents = torch.cat(latent_steps, dim=2)
frames = _decode_with_worker(omni, full_latents)
finally:
if omni is not None:
omni.close()
mp4_path = args.output_dir / f"{args.output_stem}.mp4"
_write_mp4(mp4_path, frames, fps=args.fps)
print(f"SAVED_MP4={mp4_path}")
if args.save_gif:
gif_path = args.output_dir / f"{args.output_stem}.gif"
_write_gif(gif_path, frames, fps=args.fps)
print(f"SAVED_GIF={gif_path}")
if args.save_actions:
npz_path = args.output_dir / f"{args.output_stem}_actions.npz"
np.savez(
npz_path,
step0=np.asarray(outputs[0].multimodal_output.get("actions")),
step1=np.asarray(outputs[1].multimodal_output.get("actions")),
)
print(f"SAVED_ACTIONS={npz_path}")
if __name__ == "__main__":
main()
@@ -42,7 +42,7 @@ pip install \
## 3. Hardware and VRAM Requirements
This example uses a 3-stage pipeline on one GPU by default
([`dynin_omni.yaml`](../../../vllm_omni/model_executor/stage_configs/dynin_omni.yaml)):
([`dynin_omni.yaml`](../../../vllm_omni/deploy/dynin_omni.yaml)):
- Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
- Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`
@@ -18,8 +18,6 @@ import numpy as np
import torch
from PIL import Image
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
TASK_CHOICES = ("t2t", "t2i", "t2s", "i2i", "i2t", "s2t", "v2t")
TASK_DEFAULT_RUNTIME = {
@@ -880,8 +878,8 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
parser.add_argument(
"--stage-config-path",
type=str,
default=str(repo_root / "vllm_omni/model_executor/stage_configs/dynin_omni.yaml"),
help="Path to stage config yaml.",
default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
help="Path to deploy config yaml.",
)
parser.add_argument(
"--dynin-config-path",
@@ -972,8 +970,6 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
parser.add_argument("--vq-model-audio-local-files-only", action=argparse.BooleanOptionalAction, default=None)
parser.add_argument("--disable-hf-xet", action=argparse.BooleanOptionalAction, default=True)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -61,8 +61,8 @@ def parse_args():
parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
parser.add_argument("--guidance-scale", type=float, default=5.0, help="Classifier-free guidance scale.")
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
parser.add_argument("--height", type=int, default=1024, help="Output image height.")
parser.add_argument("--width", type=int, default=1024, help="Output image width.")
parser.add_argument("--height", type=int, default=None, help="Output image height.")
parser.add_argument("--width", type=int, default=None, help="Output image width.")
parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
parser.add_argument(
"--bot-task",
@@ -107,9 +107,6 @@ def parse_args():
),
)
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -229,7 +226,21 @@ def main():
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
ar_stop_token_ids = resolve_stop_token_ids(task=task, bot_task=bot_task, tokenizer=tokenizer)
if (args.height is None) != (args.width is None):
raise ValueError("--height and --width must both be specified or both omitted.")
user_specified_size = args.height is not None and args.width is not None
if args.modality in ("img2text", "text2text"):
ar_image_size = "auto"
elif user_specified_size:
ar_image_size = f"{args.width}x{args.height}"
else:
ar_image_size = None
ar_stop_token_ids = resolve_stop_token_ids(
task=task, bot_task=bot_task, tokenizer=tokenizer, image_size=ar_image_size
)
print(
f"[AR Config] task={task}, bot_task={bot_task}, image_size={ar_image_size}, stop_token_ids={ar_stop_token_ids}"
)
for sp in params_list:
if isinstance(sp, OmniDiffusionSamplingParams):
sp.num_inference_steps = args.steps
@@ -0,0 +1,52 @@
# Lance (ByteDance) — offline inference
[Lance](https://huggingface.co/bytedance-research/Lance) is a 3B unified
autoregressive + diffusion multimodal model on a Qwen2.5-VL backbone. It is
**BAGEL-lineage** (ByteDance Mixture-of-Transformers): the released `Lance_3B`
checkpoint uses the same `*_moe_gen` MoT weight layout as BAGEL, so vLLM-Omni
implements it by reusing the BAGEL transformer core and specializing only the
ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the checkpoint layout.
This example covers all six Lance modalities from the upstream HF model card:
`t2i`, `t2v`, `image_edit`, `video_edit`, `x2t_image` (image understanding) and
`x2t_video` (video understanding).
## Hardware
Single NVIDIA GPU with 16 GB+ VRAM in BF16 (we test on B300 / A100). CUDA ≥ 12.4.
## Run
```bash
# Text-to-image
python examples/offline_inference/lance/end2end.py \
--model bytedance-research/Lance \
--prompts "a corgi astronaut on the moon, cinematic" \
--steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
--height 1024 --width 1024 \
--output ./out
# Text-to-video (uses the Lance_3B_Video subfolder; see ``--modality``
# choices for all six task variants)
python examples/offline_inference/lance/end2end.py \
--model bytedance-research/Lance/Lance_3B_Video --modality text2video \
--num-frames 25 --video-height 480 --video-width 768 \
--prompts "a cat playing piano, cinematic" \
--steps 30 --fps 8 --output ./out
```
`video_edit` requires `--model bytedance-research/Lance/Lance_3B_Video` so
the 3-D `latent_pos_embed` table is loaded; the other paths can point at
the top-level `bytedance-research/Lance` repo and resolve the right
sub-checkpoint automatically.
The HF repo bundles everything (`Lance_3B/`, `Lance_3B_Video/`,
`Qwen2.5-VL-ViT/`, `Wan2.2_VAE.pth`); no separate downloads are required.
## Defaults
Matches upstream `inference_lance.sh`: 30 denoising steps, timestep-shift 3.5,
text CFG 4.0, seed 42, 1024×1024 (override with `--height` / `--width`). For
the understanding paths (`img2text` / `video2text`), sampling is enabled by
default at `--text-temperature 0.8` because Lance's greedy decoder emits an
immediate EOS for many prompts.
+339
View File
@@ -0,0 +1,339 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Lance-3B offline inference (single-stage).
Lance is BAGEL-lineage; this example covers the single-stage paths wired in
this PR:
text2img — image generation from text
img2text — image understanding (caption / VQA)
Examples:
python examples/offline_inference/lance/end2end.py \
--model bytedance-research/Lance \
--prompts "a corgi astronaut on the moon, cinematic" \
--output ./out
python examples/offline_inference/lance/end2end.py \
--model bytedance-research/Lance --modality img2text \
--image-path /path/to/photo.jpg \
--prompts "Describe this image in detail."
"""
import argparse
import os
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--model", default="bytedance-research/Lance", help="HF repo or local path.")
p.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
p.add_argument("--txt-prompts", type=str, default=None, help="File with one prompt per line.")
p.add_argument(
"--modality",
default="text2img",
choices=["text2img", "img2text", "text2video", "video2text", "img2img", "video2video", "image2video"],
help="Lance single-stage modality.",
)
p.add_argument("--image-path", type=str, default=None, help="Input image for img2text.")
p.add_argument("--video-path", type=str, default=None, help="Input video file for video2text.")
p.add_argument("--num-frames", type=int, default=25, help="Number of RGB frames for text2video (max 121).")
p.add_argument("--video-height", type=int, default=480, help="Video frame height.")
p.add_argument("--video-width", type=int, default=768, help="Video frame width.")
p.add_argument("--height", type=int, default=None, help="Image height (t2i). Default = max_hw (1024).")
p.add_argument("--width", type=int, default=None, help="Image width (t2i). Default = max_hw (1024).")
p.add_argument(
"--fps", type=int, default=12, help="Output video FPS when saving MP4 (matches upstream Lance's save_fps=12)."
)
p.add_argument("--output", type=str, default=".", help="Output directory.")
p.add_argument("--steps", type=int, default=30, help="Denoising steps (Lance default 30).")
p.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (Lance default 4.0).")
p.add_argument("--timestep-shift", type=float, default=3.5, help="Flow-match timestep shift (Lance default 3.5).")
p.add_argument("--negative-prompt", type=str, default=None)
p.add_argument("--seed", type=int, default=42)
p.add_argument("--max-text-tokens", type=int, default=512, help="Max tokens for img2text generation.")
p.add_argument(
"--do-sample",
action="store_true",
default=True,
help="Sample (vs greedy) for x2t generation; greedy frequently emits an immediate EOS for Lance.",
)
p.add_argument("--no-sample", dest="do_sample", action="store_false", help="Disable sampling (greedy decoding).")
p.add_argument(
"--text-temperature",
type=float,
default=0.8,
help="Sampling temperature for x2t generation. Lance frequently emits an immediate EOS below ~0.7; 0.8 is a good default.",
)
p.add_argument(
"--system-prompt",
type=str,
default=None,
help=(
"Override the task system prompt (used by x2t_image / x2t_video for "
'per-example QA instructions like "Look at the image carefully and '
'answer the question."; without it x2t falls back to a '
"caption-style default and the model describes instead of answering)."
),
)
# Lance is single-stage diffusion — no deploy YAML needed. Required
# engine knobs (``pipeline``, ``enforce_eager``, ``trust_remote_code``,
# ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
# and ``create_default_diffusion`` materializes the stage config.
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(p)
return p.parse_args()
from vllm_omni.diffusion.models.lance.prompts import ( # noqa: E402
VIDEO_PAD,
VISION_END,
VISION_START,
render_lance_prompt,
)
# Visual block matching upstream Lance: <|vision_start|><|video_pad|><|vision_end|>
# (Lance's renderer uses ``<|video_pad|>`` even for image inputs by default.)
_VISION_BLOCK = f"{VISION_START}{VIDEO_PAD}{VISION_END}"
def _format_text2img_prompts(prompts):
return [{"prompt": render_lance_prompt("t2i", p), "modalities": ["image"]} for p in prompts]
def _format_text2video_prompts(prompts, num_frames, video_h, video_w):
extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
return [{"prompt": render_lance_prompt("t2v", p), "modalities": ["video"], "extra_args": extra} for p in prompts]
def _format_image2video_prompts(prompts, image_path, num_frames, video_h, video_w):
"""Image-to-Video (no first-frame pin): image + text → long video.
Passes the input image via ``multi_modal_data.first_frame``. The pipeline
treats it as a 1-frame reference (VAE+ViT prefill) and generates a fresh
multi-frame video at the requested ``num_frames`` × ``video_h`` ×
``video_w`` shape.
"""
import os as _os
if not image_path or not _os.path.exists(image_path):
raise ValueError(f"image2video requires --image-path pointing to an existing file, got: {image_path}")
from PIL import Image as _Image
img = _Image.open(image_path).convert("RGB")
extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
return [
{
"prompt": render_lance_prompt("i2v", p, vision_token=_VISION_BLOCK),
"multi_modal_data": {"first_frame": img},
"modalities": ["video"],
"extra_args": extra,
}
for p in prompts
]
def _format_img2img_prompts(prompts, image_path):
import os as _os
if not image_path or not _os.path.exists(image_path):
raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {image_path}")
from PIL import Image as _Image
img = _Image.open(image_path).convert("RGB")
return [
{
"prompt": render_lance_prompt("image_edit", p, vision_token=_VISION_BLOCK),
"multi_modal_data": {"img2img": img},
"modalities": ["image"],
}
for p in prompts
]
def _format_video2video_prompts(prompts, video_path, num_frames, video_h, video_w):
import os as _os
if not video_path or not _os.path.exists(video_path):
raise ValueError(f"video2video requires --video-path pointing to an existing file, got: {video_path}")
extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
return [
{
"prompt": render_lance_prompt("video_edit", p, vision_token=_VISION_BLOCK),
"multi_modal_data": {"video": video_path},
"modalities": ["video"],
"extra_args": extra,
}
for p in prompts
]
def _format_video2text_prompts(prompts, video_path, system_prompt=None):
import os as _os
if not video_path or not _os.path.exists(video_path):
raise ValueError(f"video2text requires --video-path pointing to an existing file, got: {video_path}")
import imageio.v3 as iio
video = iio.imread(video_path)
return [
{
"prompt": render_lance_prompt("x2t_video", p, vision_token=_VISION_BLOCK, system_prompt=system_prompt),
"multi_modal_data": {"video": video},
"modalities": ["text"],
}
for p in prompts
]
def _format_img2text_prompts(prompts, image_path, system_prompt=None):
from PIL import Image
if not image_path or not os.path.exists(image_path):
raise ValueError(f"img2text requires --image-path pointing to an existing file, got: {image_path}")
img = Image.open(image_path).convert("RGB")
return [
{
"prompt": render_lance_prompt("x2t_image", p, vision_token=_VISION_BLOCK, system_prompt=system_prompt),
"multi_modal_data": {"image": img},
"modalities": ["text"],
}
for p in prompts
]
def main():
args = parse_args()
os.makedirs(args.output, exist_ok=True)
# Lance aligns more naturally with upstream when using FLASH_ATTN
# (upstream uses flash_attn_varlen_func; SDPA accumulates ~5x more
# numerical drift through the 36-layer Qwen2 stack on B300). Default
# to flash_attn when the user hasn't explicitly set a backend.
os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
if args.txt_prompts:
with open(args.txt_prompts, encoding="utf-8") as f:
prompts = [ln.strip() for ln in f if ln.strip()]
elif args.modality == "img2text":
prompts = args.prompts or ["Describe this image in detail."]
else:
prompts = args.prompts or ["A cute cat, studio lighting, 4k"]
from vllm_omni.entrypoints.omni import Omni
# Video modalities require the Lance_3B_Video checkpoint (3D
# latent_pos_embed shaped (126976, 2048)). Without it, vllm-omni silently
# loads Lance_3B (4096, 2048) — the image-only table — and any t_lat >= 1
# immediately indexes out of bounds. Rewrite the model path so
# LancePipeline._select_video_variant picks the video subfolder.
resolved_model = args.model
if args.modality in ("text2video", "video2video", "video2text", "image2video"):
if not str(resolved_model).rstrip("/").endswith("Lance_3B_Video"):
candidate = os.path.join(resolved_model, "Lance_3B_Video")
if os.path.isdir(candidate):
resolved_model = candidate
omni_kwargs = vars(args).copy()
omni_kwargs["model"] = resolved_model
# Lance single-stage defaults (formerly in vllm_omni/deploy/lance.yaml):
omni_kwargs.setdefault("pipeline", "lance")
omni_kwargs.setdefault("max_num_batched_tokens", 32768)
omni_kwargs.setdefault("max_num_seqs", 1)
omni_kwargs.setdefault("enforce_eager", True)
omni_kwargs.setdefault("trust_remote_code", True)
omni_kwargs.setdefault("enable_prefix_caching", False)
omni_kwargs.setdefault("async_chunk", False)
omni = Omni(**omni_kwargs)
if args.modality == "img2text":
formatted = _format_img2text_prompts(prompts, args.image_path, system_prompt=args.system_prompt)
elif args.modality == "text2video":
formatted = _format_text2video_prompts(prompts, args.num_frames, args.video_height, args.video_width)
elif args.modality == "video2text":
formatted = _format_video2text_prompts(prompts, args.video_path, system_prompt=args.system_prompt)
elif args.modality == "img2img":
formatted = _format_img2img_prompts(prompts, args.image_path)
elif args.modality == "video2video":
formatted = _format_video2video_prompts(
prompts, args.video_path, args.num_frames, args.video_height, args.video_width
)
elif args.modality == "image2video":
formatted = _format_image2video_prompts(
prompts, args.image_path, args.num_frames, args.video_height, args.video_width
)
else:
formatted = _format_text2img_prompts(prompts)
params_list = omni.default_sampling_params_list
diffusion_params = params_list[0] # single-stage: one param set
diffusion_params.num_inference_steps = args.steps # type: ignore
if args.seed is not None:
diffusion_params.seed = args.seed # type: ignore
if args.height is not None:
diffusion_params.height = args.height # type: ignore
if args.width is not None:
diffusion_params.width = args.width # type: ignore
extra = getattr(diffusion_params, "extra_args", {}) or {}
extra["cfg_text_scale"] = args.cfg_text_scale
extra["timestep_shift"] = args.timestep_shift
if args.modality in ("img2text", "video2text"):
extra["max_think_tokens"] = args.max_text_tokens
extra["do_sample"] = args.do_sample
extra["text_temperature"] = args.text_temperature
if args.negative_prompt is not None:
extra["negative_prompt"] = args.negative_prompt
diffusion_params.extra_args = extra # type: ignore
outputs = list(omni.generate(prompts=formatted, sampling_params_list=params_list))
if args.modality in ("img2text", "video2text"):
for i, req_output in enumerate(outputs):
text = getattr(req_output, "output", None) or getattr(req_output, "text", None)
if text is None and hasattr(req_output, "custom_output"):
text = (req_output.custom_output or {}).get("text_output")
print(f"[Output {i}] {text}")
return
if args.modality in ("text2video", "video2video", "image2video"):
for i, req_output in enumerate(outputs):
custom = getattr(req_output, "custom_output", None) or {}
frames = custom.get("video_frames")
if not frames:
print(f"[Output {i}] no video frames returned")
continue
mp4_path = os.path.join(args.output, f"lance_{i}.mp4")
try:
import imageio.v3 as iio
import numpy as np
arr = np.stack([np.asarray(f) for f in frames], axis=0)
iio.imwrite(mp4_path, arr, fps=args.fps, codec="libx264", quality=8)
print(f"[Output {i}] Saved MP4 ({len(frames)} frames, {arr.shape[1]}x{arr.shape[2]}) to {mp4_path}")
except Exception as exc:
print(f"[Output {i}] MP4 encode failed ({exc!r}); writing per-frame ONGs instead")
frame_dir = os.path.join(args.output, f"lance_{i}_frames")
os.makedirs(frame_dir, exist_ok=True)
for j, frame in enumerate(frames):
frame.save(os.path.join(frame_dir, f"{j:04d}.png"))
print(f" → {frame_dir} ({len(frames)} frames)")
return
idx = 0
for req_output in outputs:
images = getattr(req_output, "images", None)
if not images:
continue
for j, img in enumerate(images):
path = os.path.join(args.output, f"lance_{idx}_{j}.png")
img.save(path)
print(f"[Output] Saved image to {path}")
idx += 1
if __name__ == "__main__":
main()
File diff suppressed because it is too large. Load diff
@@ -19,7 +19,6 @@ from vllm import SamplingParams
from vllm.multimodal.image import convert_image_mode
from vllm_omni import Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
DEFAULT_SYSTEM = "You are a helpful assistant."
DEFAULT_QUESTION = "Please summarize the content of this image."
@@ -49,7 +48,6 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -29,7 +29,6 @@ from PIL import Image
from vllm.sampling_params import SamplingParams
from vllm_omni import Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
@@ -118,7 +117,6 @@ def parse_args() -> argparse.Namespace:
)
p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
nullify_stage_engine_defaults(p)
args = p.parse_args()
if not args.prompt:
args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]
@@ -21,11 +21,10 @@ from message_convert import (
to_prompt,
)
from vllm import SamplingParams
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniTokensPrompt
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
SEED = 42
@@ -343,7 +342,7 @@ def main(args):
def parse_args():
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser.add_argument(
"--model-name",
"-m",
@@ -439,7 +438,6 @@ def parse_args():
"vllm_omni/deploy/mimo_audio.yaml based on the HF model_type.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -8,7 +8,7 @@ vLLM-Omni supports three deployment modes:
|------|--------------|--------|
| Thinker + Talker (omni-speech, default) | `vllm_omni/deploy/ming_flash_omni.yaml` | Text + Audio |
| Thinker only (multimodal understanding) | `vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
| Thinker + Imagegen (text-to-image / img2img) | `vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml` | Image (online-serving only at the moment) |
| Thinker + Imagegen (text-to-image / img2img) | `vllm_omni/deploy/ming_flash_omni_image.yaml` | Image (online-serving only at the moment) |
For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
@@ -53,7 +53,7 @@ python examples/offline_inference/ming_flash_omni/end2end.py \
--image-path ./3_0.png
```
### Omni-Speech (Thinker + Talker)
### Omni-Speech (Thinker + Talker) {#omni-speech-thinker--talker}
The default deploy YAML already runs thinker+talker, so spoken output only requires requesting `audio` (or `text,audio`) modalities.
The thinker processes your multimodal input, generates text, then the talker synthesises the response as speech.
@@ -17,14 +17,13 @@ from vllm.assets.image import ImageAsset
from vllm.assets.video import VideoAsset, video_to_ndarrays
from vllm.multimodal.image import convert_image_mode
from vllm.multimodal.media.audio import load_audio
from vllm.utils.argparse_utils import FlexibleArgumentParser
import vllm_omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
# Imports the processor also registers itself
from vllm_omni.transformers_utils.processors.ming import MingFlashOmniProcessor # noqa: F401
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
SEED = 42
MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0"
@@ -397,7 +396,7 @@ def main(args):
def parse_args():
parser = FlexibleArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
parser = TrackingArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
parser.add_argument(
"--query-type",
"-q",
@@ -496,7 +495,6 @@ def parse_args():
help="Output directory for results.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -19,10 +19,9 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
from vllm.multimodal.image import convert_image_mode
from vllm.multimodal.media.audio import load_audio
from vllm.sampling_params import SamplingParams
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
SEED = 42
@@ -432,7 +431,7 @@ def main(args):
def parse_args():
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser.add_argument(
"--model",
type=str,
@@ -554,7 +553,6 @@ def parse_args():
default=False,
help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -19,10 +19,9 @@ from vllm.assets.image import ImageAsset
from vllm.assets.video import VideoAsset, video_to_ndarrays
from vllm.multimodal.image import convert_image_mode
from vllm.multimodal.media.audio import load_audio
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
SEED = 42
@@ -422,7 +421,7 @@ def main(args):
def parse_args():
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser.add_argument(
"--model",
type=str,
@@ -561,7 +560,6 @@ def parse_args():
help="Model dtype (auto, half, float16, bfloat16, float, float32).",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -39,10 +39,9 @@ from vllm.assets.image import ImageAsset
from vllm.assets.video import VideoAsset, video_to_ndarrays
from vllm.multimodal.image import convert_image_mode
from vllm.multimodal.media.audio import load_audio
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
logger = logging.getLogger(__name__)
@@ -383,7 +382,9 @@ async def run_all(args):
print(f"[Info] Creating AsyncOmni with deploy_config={args.deploy_config}")
async_omni = None
try:
async_omni = AsyncOmni(**vars(args))
# ``from_cli_args`` forwards only explicitly-passed CLI args so
# argparse defaults do not silently override deploy YAML values.
async_omni = AsyncOmni.from_cli_args(args, model=args.model)
# Use default sampling params from stage config (they are pre-configured
# in the YAML for each stage).
@@ -453,7 +454,7 @@ async def run_all(args):
def parse_args():
parser = FlexibleArgumentParser(
parser = TrackingArgumentParser(
description=(
"Offline inference with async_chunk enabled via AsyncOmni. "
"Downstream stages start before upstream stages finish, "
@@ -589,7 +590,6 @@ def parse_args():
default=16000,
help="Sampling rate for audio loading.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -37,6 +37,10 @@ python end2end.py --modality text2text \
# Custom resolution (image generation)
python end2end.py --prompt "A futuristic cityscape at sunset" \
--width 2048 --height 1024 --think
# Text-to-image with Cache-DiT acceleration
python end2end.py --prompt "A cute cat" \
--cache-backend cache_dit \
```
> **Note**: Default configuration works on a single **NVIDIA A100 (80GB)** or **H100** GPU.
@@ -135,6 +139,8 @@ python end2end.py \
| `--tensor-parallel-size` | 1 | Number of GPUs for tensor parallelism |
| `--enforce-eager` | False | Disable torch.compile |
| `--enable-cpu-offload` | False | Enable module-wise (sequential) CPU offload to reduce peak VRAM |
| `--cache-backend` | None | Set to `cache_dit` for Cache-DiT acceleration |
| `--enable-cache-dit-summary` | False | Print Cache-DiT cache statistics after generation |
## Reducing GPU Memory Usage
@@ -29,7 +29,6 @@ import os
from PIL import Image
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -155,8 +154,19 @@ def parse_args():
"of extra CPU<->GPU transfers."
),
)
parser.add_argument(
"--cache-backend",
type=str,
default=None,
choices=["cache_dit"],
help="Cache backend to use for image generation acceleration.",
)
parser.add_argument(
"--enable-cache-dit-summary",
action="store_true",
help="Enable cache-dit summary logging after diffusion forward passes.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -182,6 +192,8 @@ def main():
tensor_parallel_size=args.tensor_parallel_size,
enforce_eager=args.enforce_eager,
enable_cpu_offload=args.enable_cpu_offload,
cache_backend=args.cache_backend,
enable_cache_dit_summary=args.enable_cache_dit_summary,
)
extra_args = {
@@ -225,6 +237,7 @@ def main():
print(f" Seed : {args.seed}")
print(f" Think mode : {args.think}")
print(f" TP size : {args.tensor_parallel_size}")
print(f" Cache backend : {args.cache_backend or 'none'}")
print(f"{'=' * 60}\n")
# Build prompt dict
@@ -307,9 +307,6 @@ def parse_args() -> argparse.Namespace:
help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
)
current_omni_platform.pre_register_and_update(parser)
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -271,6 +271,14 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
--lang zh
```
### Seed for Reproducibility
```bash
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
--model k2-fsa/OmniVoice \
--text "Hello, this is a test." \
--seed 42
```
### Notes
- Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking.
- Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz.
@@ -10,7 +10,6 @@ import soundfile as sf
from vllm import SamplingParams
from vllm.multimodal.media.audio import load_audio
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
@@ -65,7 +64,6 @@ def run_e2e():
required=True,
help="Path to tokenizer directory (e.g., <model_path>/CosyVoice-BlankEN).",
)
nullify_stage_engine_defaults(parser)
args = parser.parse_args()
# Ensure tokenizer directory exists
if not os.path.exists(args.tokenizer):
@@ -26,16 +26,14 @@ import torch
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import AsyncOmni, Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.model_executor.models.fish_speech.dac_utils import DAC_HOP_LENGTH, DAC_SAMPLE_RATE
from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
build_fish_text_only_prompt_ids,
estimate_fish_voice_clone_prompt_len_from_normalized,
normalize_fish_voice_clone_texts,
)
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
logger = logging.getLogger(__name__)
@@ -211,7 +209,7 @@ async def main_streaming(args):
def parse_args():
parser = FlexibleArgumentParser(
parser = TrackingArgumentParser(
description="Fish Speech S2 Pro offline inference with vLLM Omni",
)
parser.add_argument(
@@ -267,7 +265,6 @@ def parse_args():
default=False,
help="Stream audio chunks as they arrive via AsyncOmni.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -30,10 +30,9 @@ from pathlib import Path
import numpy as np
import soundfile as sf
import torch
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
SAMPLE_RATE = 24_000
DEFAULT_TEXTS = (
@@ -43,7 +42,7 @@ DEFAULT_TEXTS = (
def parse_args():
parser = FlexibleArgumentParser(description="Offline higgs-audio v2 inference")
parser = TrackingArgumentParser(description="Offline higgs-audio v2 inference")
parser.add_argument(
"--model",
type=str,
@@ -88,7 +87,6 @@ def parse_args():
default=None,
help="Transcript of the reference clip. Required when --ref-audio is set.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -5,11 +5,11 @@ from typing import Any
import soundfile as sf
import torch
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniTokensPrompt
from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
@@ -68,7 +68,7 @@ def save_audio(mm: dict[str, Any], output_path: str) -> None:
def parse_args():
parser = FlexibleArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
parser = TrackingArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
parser.add_argument("--model", type=str, default=MODEL_NAME, help="Model name or local path.")
parser.add_argument(
"--deploy-config",
@@ -89,7 +89,6 @@ def parse_args():
parser.add_argument("--init-timeout", type=int, default=600, help="Engine init timeout in seconds.")
parser.add_argument("--stage-init-timeout", type=int, default=300, help="Single stage init timeout in seconds.")
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -0,0 +1,377 @@
"""Offline inference example for MOSS-TTS variants.
Covers all five models:
- OpenMOSS-Team/MOSS-TTS (8B, general TTS)
- OpenMOSS-Team/MOSS-TTSD-v1.0 (8B, dialogue TTS)
- OpenMOSS-Team/MOSS-SoundEffect (8B, sound effect synthesis)
- OpenMOSS-Team/MOSS-TTS-Realtime (1.7B, low-latency streaming)
- OpenMOSS-Team/MOSS-VoiceGenerator(1.7B, voice design)
Usage
-----
# Standard TTS with voice cloning:
python end2end.py \\
--model OpenMOSS-Team/MOSS-TTS \\
--text "Hello, this is a MOSS-TTS test." \\
--ref-audio path/to/reference.wav \\
--output output.wav
"""
from __future__ import annotations
import argparse
import gc
import sys
import soundfile as sf
import torch
def _load_ref_audio(path: str, target_sr: int = 24000) -> torch.Tensor:
"""Load and resample reference audio to (1, T) at target_sr."""
data, sr = sf.read(path, always_2d=True) # (T, C)
wav = torch.from_numpy(data.T.astype("float32")) # (C, T)
if wav.shape[0] > 1:
wav = wav.mean(dim=0, keepdim=True)
if sr != target_sr:
import torchaudio
wav = torchaudio.functional.resample(wav, sr, target_sr)
return wav # (1, T)
def _build_unified_codes(
model_id: str,
text: str,
ref_audio: torch.Tensor | None,
*,
instruction: str | None = None,
tokens: int | None = None,
quality: str | None = None,
sound_event: str | None = None,
ambient_sound: str | None = None,
language: str | None = None,
) -> tuple[list[int], torch.Tensor, int]:
"""Build the upstream MossTTSDelay unified-codes prompt.
Returns (text_token_ids, audio_codes_LxNQ, n_vq) where
``text_token_ids`` is the prefill prompt for the talker and
``audio_codes_LxNQ`` is the matching delay-pattern audio code grid
(``audio_pad_code`` everywhere except inside the reference-audio block).
"""
from transformers import AutoProcessor
proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
reference: list[torch.Tensor] | None = None
if ref_audio is not None:
# Pre-encode so the processor doesn't need torchaudio.load (which
# now requires torchcodec on recent torchaudio releases).
# encode_audios_from_wav expects (C, T) per item.
wav_2d = ref_audio if ref_audio.dim() == 2 else ref_audio.unsqueeze(0)
codes_list = proc.encode_audios_from_wav([wav_2d], sampling_rate=24000, n_vq=proc.model_config.n_vq)
reference = [codes_list[0]]
user_msg_kwargs: dict = {"text": text}
if reference is not None:
user_msg_kwargs["reference"] = reference
if instruction is not None:
user_msg_kwargs["instruction"] = instruction
if tokens is not None:
user_msg_kwargs["tokens"] = tokens
if quality is not None:
user_msg_kwargs["quality"] = quality
if sound_event is not None:
user_msg_kwargs["sound_event"] = sound_event
if ambient_sound is not None:
user_msg_kwargs["ambient_sound"] = ambient_sound
if language is not None:
user_msg_kwargs["language"] = language
user_msg = proc.build_user_message(**user_msg_kwargs)
batch = proc(conversations=[[user_msg]], mode="generation")
unified = batch["input_ids"][0] # (L, 1+n_vq)
n_vq = int(unified.shape[-1] - 1)
text_ids = unified[:, 0].tolist()
audio_codes = unified[:, 1:].contiguous().to(torch.int64) # (L, n_vq)
# Free the processor (incl. ~1.6B-param audio_tokenizer on CPU) before
# we instantiate Omni — we only needed it for the one-shot encode.
del proc
gc.collect()
return text_ids, audio_codes, n_vq
def _build_realtime_prompt(
model_id: str,
text: str,
ref_audio: torch.Tensor | None,
) -> tuple[list[int], torch.Tensor, int]:
"""Build the prefill prompt for MOSS-TTS-Realtime.
Realtime uses a different prompt format than the delay variants: a system
prompt + (optional) voice-clone context with the encoded reference audio
inlined, followed by the user text and an assistant header. The format is
a ``(L, channels+1=17)`` int grid where col 0 is text/special tokens and
cols 1..16 are RVQ codebook entries (``audio_channel_pad`` outside the
ref-audio block).
The upstream ``MossTTSRealtimeProcessor`` is shipped via
``trust_remote_code`` but is not auto-discovered by ``AutoProcessor``
(no ``processor_config.json``) — we import its module directly from the
snapshot directory and use its ``make_ensemble``/``make_user_prompt``.
"""
import importlib.util
import sys as _sys
from pathlib import Path
import numpy as np
import torch
from huggingface_hub import snapshot_download
from transformers import AutoModel, AutoTokenizer
snap_dir = Path(snapshot_download(repo_id=model_id))
proc_module_path = snap_dir / "processing_mossttsrealtime.py"
if not proc_module_path.exists():
raise FileNotFoundError(f"realtime processor module missing at {proc_module_path}")
# Load the processor module without polluting sys.path globally.
spec = importlib.util.spec_from_file_location("_moss_tts_realtime_proc", proc_module_path)
mod = importlib.util.module_from_spec(spec) # type: ignore[arg-type]
_sys.modules[spec.name] = mod # type: ignore[union-attr]
spec.loader.exec_module(mod) # type: ignore[union-attr]
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
realtime_processor = mod.MossTTSRealtimeProcessor(tokenizer=tokenizer)
# Encode reference audio via the standalone MOSS Audio Tokenizer (CPU is
# fine — it runs once per request).
audio_tokens = None
if ref_audio is not None:
codec = AutoModel.from_pretrained("OpenMOSS-Team/MOSS-Audio-Tokenizer", trust_remote_code=True).eval()
wav_2d = ref_audio if ref_audio.dim() == 2 else ref_audio.unsqueeze(0)
with torch.no_grad():
enc = codec.batch_encode([wav_2d.squeeze(0)], num_quantizers=16)
# enc.audio_codes: (NQ, B, T); we want (T, NQ) numpy.
codes = enc.audio_codes[:, 0, : int(enc.audio_codes_lengths[0].item())]
audio_tokens = codes.transpose(0, 1).contiguous().cpu().numpy() # (T, 16)
del codec
gc.collect()
# Build the inference-style prefill (mirrors upstream
# ``MossTTSRealtimeInference._build_prefill_batch``):
# system_prompt [+ voice_clone_block]
# <|im_start|>assistant\n
# <text_tokens...> (col 0 = text, cols 1..16 = audio_channel_pad)
# <last text token> (col 1 = audio_bos at the last text position)
# Generation continues from here; the local transformer emits real audio
# codes for codebook 0 directly (no BOS frame at gen_step=0).
system_grid = realtime_processor.make_ensemble(prompt_audio_tokens=audio_tokens) # (L_sys, 17)
# Append "<|im_start|>assistant\n" header (the snapshot's make_ensemble
# omits it; inferencer.py's version includes it. We add it here.)
assistant_header_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
assistant_grid = np.full((len(assistant_header_ids), 17), 1024, dtype=np.int64)
assistant_grid[:, 0] = assistant_header_ids
# Append the FIRST 12 text tokens with audio_bos (1025) at the last
# included text position (audio column 0 == grid col 1). The remaining
# text tokens are fed one-per-step during generation via the talker's
# compute_logits → vLLM sampler pipeline (see ``ids.all`` in the
# additional_information). This matches upstream
# ``MossTTSRealtimeInference._build_prefill_batch`` (prefill_max_text=12).
text_ids_only = tokenizer.encode(text, add_special_tokens=False)
if not text_ids_only:
raise ValueError("Empty text after tokenisation")
PREFILL_MAX_TEXT = 12
cur_len = min(len(text_ids_only), PREFILL_MAX_TEXT)
text_grid = np.full((cur_len, 17), 1024, dtype=np.int64)
text_grid[:, 0] = text_ids_only[:cur_len]
text_grid[-1, 1] = 1025 # audio_bos_token at last *prefilled* text token
grid = np.concatenate([system_grid, assistant_grid, text_grid], axis=0)
text_ids = grid[:, 0].tolist()
audio_codes = torch.from_numpy(grid[:, 1:].astype(np.int64).copy()) # (L, 16)
# ``remaining_text_ids`` will be fed one-per-step by the talker.
remaining_text_ids = list(text_ids_only[cur_len:])
return text_ids, audio_codes, 16, remaining_text_ids
def run_tts(args: argparse.Namespace) -> None:
from vllm import SamplingParams
from vllm_omni import Omni
deploy_map = {
"OpenMOSS-Team/MOSS-TTS": "moss_tts",
"OpenMOSS-Team/MOSS-TTSD-v1.0": "moss_ttsd",
"OpenMOSS-Team/MOSS-SoundEffect": "moss_sound_effect",
"OpenMOSS-Team/MOSS-TTS-Realtime": "moss_tts_realtime",
"OpenMOSS-Team/MOSS-VoiceGenerator": "moss_voice_generator",
}
deploy_config = deploy_map.get(args.model, "moss_tts")
is_sound_effect = "SoundEffect" in args.model
is_voice_gen = "VoiceGenerator" in args.model
is_realtime = "Realtime" in args.model
ref_audio = None
builder_kwargs: dict = {}
remaining_text_ids: list[int] | None = None
if is_realtime:
# MOSS-TTS-Realtime uses a different prompt format and has its own
# per-step depth transformer; the delay-style ``MossTTSDelayProcessor``
# path doesn't apply.
if not args.ref_audio:
sys.exit("MOSS-TTS-Realtime currently requires --ref-audio (voice clone).")
ref_audio = _load_ref_audio(args.ref_audio)
text = args.text or ""
text_ids, audio_codes, n_vq, remaining_text_ids = _build_realtime_prompt(args.model, text, ref_audio)
print(
f"Prefill prompt: {len(text_ids)} tokens, {n_vq}-quantizer audio block (realtime); "
f"{len(remaining_text_ids)} text tokens to stream"
)
elif is_sound_effect:
if not args.ambient_sound:
sys.exit("MOSS-SoundEffect requires --ambient-sound")
builder_kwargs["ambient_sound"] = args.ambient_sound
# MOSS-SoundEffect emits ~12.5 codec frames per second; ``tokens``
# tells the AR loop the target duration.
if args.duration_tokens:
builder_kwargs["tokens"] = args.duration_tokens
elif args.duration_seconds:
builder_kwargs["tokens"] = max(1, int(float(args.duration_seconds) * 12.5))
text = args.text or ""
elif is_voice_gen:
# MOSS-VoiceGenerator needs an instruction (voice description) plus
# the text to synthesise; no reference audio.
if not args.instruction:
sys.exit("MOSS-VoiceGenerator requires --instruction (voice description)")
builder_kwargs["instruction"] = args.instruction
text = args.text or ""
else:
if not args.ref_audio:
sys.exit("Voice cloning requires --ref-audio")
ref_audio = _load_ref_audio(args.ref_audio)
text = args.text or ""
if not is_realtime:
text_ids, audio_codes, n_vq = _build_unified_codes(args.model, text, ref_audio, **builder_kwargs)
print(f"Prefill prompt: {len(text_ids)} tokens, {n_vq}-quantizer audio block")
omni = Omni(model=args.model, deploy_config=deploy_config, stage_init_timeout=600)
talker_sp = SamplingParams(
temperature=1.7,
top_p=0.8,
top_k=25,
max_tokens=2048,
seed=42,
detokenize=False,
)
codec_sp = SamplingParams(
temperature=0.0,
top_p=1.0,
top_k=-1,
max_tokens=65536,
seed=42,
detokenize=True,
)
additional_info: dict = {
# Reference-audio codes consumed by the talker prefill (additive
# embedding); shape is (L, n_vq) where pad_code marks non-audio
# positions. The OmniPayload schema reserves ``codes.ref`` for
# exactly this kind of conditioning input.
"codes": {"ref": audio_codes},
}
if remaining_text_ids:
# Realtime variant: feed the remaining text tokens (everything past
# the first ``PREFILL_MAX_TEXT``) one-per-step during decode via the
# talker's compute_logits → vLLM sampler one-hot trick.
additional_info["ids"] = {"all": remaining_text_ids}
request = {
"prompt_token_ids": text_ids,
"additional_information": additional_info,
}
import time as _time
chunks: list[torch.Tensor] = []
t_gen_start = _time.monotonic()
t_first_audio: float | None = None
for stage_outputs in omni.generate(request, [talker_sp, codec_sp]):
# Each yielded ``stage_outputs`` is an OmniRequestOutput; the codec
# stage attaches its waveform under ``multimodal_output["model_outputs"]``.
# Avoid ``a or b`` when ``a`` may be a tensor (truthiness raises).
mm = stage_outputs.multimodal_output or {}
wav = mm.get("audio")
if wav is None:
wav = mm.get("model_outputs")
if wav is None:
continue
if isinstance(wav, list):
non_empty = [w.reshape(-1) for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0]
if non_empty and t_first_audio is None:
t_first_audio = _time.monotonic()
chunks.extend(non_empty)
elif isinstance(wav, torch.Tensor) and wav.numel() > 0:
if t_first_audio is None:
t_first_audio = _time.monotonic()
chunks.append(wav.reshape(-1))
t_gen_end = _time.monotonic()
if not chunks:
print("No audio in output — check model loading logs.")
return
wav_t = torch.cat(chunks).float().cpu()
sf.write(args.output, wav_t.numpy(), 24000)
print(f"Saved {len(wav_t) / 24000:.2f}s of audio to {args.output}")
ttfa_ms = (t_first_audio - t_gen_start) * 1000.0 if t_first_audio is not None else float("nan")
gen_total_ms = (t_gen_end - t_gen_start) * 1000.0
print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={len(wav_t) / 24000:.2f}")
def main() -> None:
parser = argparse.ArgumentParser(description="MOSS-TTS offline inference example")
parser.add_argument(
"--model",
default="OpenMOSS-Team/MOSS-TTS",
help="HuggingFace model ID",
)
parser.add_argument("--text", default="", help="Text to synthesise (TTS variants)")
parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning")
parser.add_argument(
"--ambient-sound",
default=None,
help="Sound description for MOSS-SoundEffect",
)
parser.add_argument(
"--duration-tokens",
type=int,
default=None,
help="Target duration in tokens (1 s ≈ 12.5 tokens) for MOSS-SoundEffect",
)
parser.add_argument(
"--duration-seconds",
type=float,
default=None,
help="Target duration in seconds (converted via ~12.5 tokens/s) for MOSS-SoundEffect",
)
parser.add_argument(
"--instruction",
default=None,
help="Voice description for MOSS-VoiceGenerator (e.g. 'a young woman with a clear voice')",
)
parser.add_argument("--output", default="output.wav", help="Output WAV path")
args = parser.parse_args()
run_tts(args)
if __name__ == "__main__":
main()
@@ -38,7 +38,8 @@ from pathlib import Path
import soundfile as sf
import torch
from vllm import SamplingParams
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
# Prevent multiprocessing from re-importing CUDA in the wrong context.
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
@@ -155,7 +156,7 @@ def main(args) -> None:
def parse_args():
parser = FlexibleArgumentParser(description="MOSS-TTS-Nano offline inference")
parser = TrackingArgumentParser(description="MOSS-TTS-Nano offline inference")
parser.add_argument("--text", default="Hello, this is MOSS-TTS-Nano speaking.", help="Text to synthesize.")
parser.add_argument(
"--ref-audio",
@@ -21,7 +21,6 @@ import os
import numpy as np
import soundfile as sf
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -80,7 +79,12 @@ def run_e2e():
default=600,
help="Stage initialization timeout in seconds",
)
nullify_stage_engine_defaults(parser)
parser.add_argument(
"--seed",
type=int,
default=None,
help="Random seed for generation",
)
args = parser.parse_args()
if not os.path.exists(args.stage_config):
@@ -122,7 +126,7 @@ def run_e2e():
if mm_processor_kwargs:
prompts["mm_processor_kwargs"] = mm_processor_kwargs
sampling_params_list = [OmniDiffusionSamplingParams()]
sampling_params_list = [OmniDiffusionSamplingParams(extra_args={"seed": args.seed})]
print(f"Generating speech for: {args.text}")
@@ -15,10 +15,8 @@ import torch
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import AsyncOmni, Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
logger = logging.getLogger(__name__)
@@ -43,8 +41,8 @@ def _estimate_prompt_len(
"""
try:
from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
Qwen3TTSTalkerForConditionalGeneration,
from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
Qwen3TTSPromptEmbedsBuilder,
)
if model_name not in _cache:
@@ -124,7 +122,7 @@ def _estimate_prompt_len(
except Exception:
return None
return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
additional_information=additional_information,
task_type=task_type,
tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
@@ -435,7 +433,7 @@ async def main_streaming(args):
def parse_args():
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser.add_argument(
"--query-type",
"-q",
@@ -548,7 +546,6 @@ def parse_args():
help="Number of prompts per batch (default: 1, sequential).",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -12,17 +12,16 @@ from pathlib import Path
import soundfile as sf
import torch
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
REPO_ROOT = Path(__file__).resolve().parents[4]
SAMPLE_RATE = 48_000
def parse_args():
parser = FlexibleArgumentParser(description="Offline VoxCPM2 native AR inference")
parser = TrackingArgumentParser(description="Offline VoxCPM2 native AR inference")
parser.add_argument(
"--model",
type=str,
@@ -60,7 +59,6 @@ def parse_args():
default=None,
help="Optional transcript of --ref-audio (enables continuation mode).",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -27,11 +27,10 @@ except ImportError:
"pip install -e ."
)
from vllm import SamplingParams
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import AsyncOmni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
logger = logging.getLogger(__name__)
@@ -230,7 +229,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
def parse_args() -> Namespace:
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
parser.add_argument(
"--model",
type=str,
@@ -313,7 +312,6 @@ def parse_args() -> Namespace:
default=None,
help="Quantization method (e.g. 'fp8'). Applied to the language model only.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
+63
View File
@@ -68,6 +68,67 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size
Or set `tensor_parallel_size` per stage in a custom deploy YAML.
### VAE Patch Parallelism
[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) distributes Bagel VAE **decode/encode** across multiple GPUs by splitting latent tiles. It lowers **per-GPU peak memory during VAE decode**, which helps high-resolution `text2img` / `img2img` when VAE becomes a bottleneck.
**Scope for Bagel:**
| Topology | VAE patch parallel |
| :------- | :----------------- |
| **Single-stage** (DiT only) | Supported on stage 0 (`BagelPipeline` + `DistributedAutoEncoder`) |
| **Two-stage** | Supported on **stage 1 (DiT)** only; stage 0 (Thinker) uses encoder-only VAE and is unrelated |
**Requirements:**
- `vae_patch_parallel_size > 1` and a distributed VAE (`DistributedAutoEncoder` on the DiT pipeline).
- The DiT process group must have at least `vae_patch_parallel_size` ranks. In practice this means the diffusion stage `world_size` must be ≥ 2 (commonly `tensor_parallel_size=2` on that stage).
- `vae_use_tiling` must be enabled. If you set `vae_patch_parallel_size > 1` and omit tiling, the registry auto-enables `vae_use_tiling` at startup.
VAE patch parallel **reuses the DiT process group** (`dit_group`); it does not create a separate VAE-only worker pool. It is not a substitute for single-GPU VAE tiling (`vae_pp=1`).
**Online serving (single-stage, 2 GPUs):**
```bash
CUDA_VISIBLE_DEVICES=0,1 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml \
--tensor-parallel-size 2 \
--vae-patch-parallel-size 2 \
--vae-use-tiling
```
**Online serving (two-stage, VAE PP on DiT stage 1):** use a custom deploy YAML, for example:
```yaml
stages:
- stage_id: 0
devices: "0"
# Thinker (AR) — no VAE patch parallel here
- stage_id: 1
devices: "0,1"
vae_use_tiling: true
parallel_config:
tensor_parallel_size: 2
vae_patch_parallel_size: 2
```
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
--deploy-config /path/to/bagel_vae_pp.yaml
```
**Verify it is active** (check server logs at startup):
```text
INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
```
| CLI flag | Default | Description |
| :------- | :------ | :---------- |
| `--vae-patch-parallel-size` | `1` | Number of DiT ranks used for VAE tile parallelism. Set to `2` or higher to enable. Should be ≤ DiT process group size (typically match `--tensor-parallel-size` on the diffusion stage). |
| `--vae-use-tiling` | off | Enable VAE spatial tiling. Required for VAE patch parallel (auto-enabled when `vae_patch_parallel_size > 1`). |
#### Hybrid Sharded Data Parallel (HSDP)
For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
@@ -323,6 +384,8 @@ python openai_chat_client.py \
| `stages[].gpu_memory_utilization` | per-stage | Fraction of GPU memory to use |
| `stages[].enforce_eager` | per-stage | Disable CUDA graphs |
| `stages[].tensor_parallel_size` | per-stage | TP degree for this stage |
| `stages[].parallel_config.vae_patch_parallel_size` | per-stage (DiT) | VAE tile parallelism degree (DiT stage only) |
| `stages[].vae_use_tiling` | per-stage (DiT) | Enable VAE tiling (required for VAE patch parallel) |
| `connectors` | top-level | Define available connector instances (SHM, Mooncake) |
| `platforms` | top-level | Platform-specific overrides (e.g. `xpu`) |
+132
View File
@@ -0,0 +1,132 @@
# DreamZero
DreamZero is the robot-policy serving example for the OpenPI-compatible `/v1/realtime/robot/openpi` endpoint.
## Supported checkpoints
| Model | HuggingFace repo |
|---|---|
| DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` |
## Quick start
### Start the server
```bash
bash examples/online_serving/dreamzero/run_server.sh
```
By default this launches:
```bash
vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
--served-model-name dreamzero-droid \
--deploy-config vllm_omni/deploy/dreamzero_tp1_cfg2.yaml \
--enforce-eager --disable-log-stats
```
Override `MODEL`, `PORT`, `HOST`, `DEPLOY_CONFIG`, or `SERVED_MODEL_NAME` through the script environment if needed.
### Install example dependencies
The core `pip install -e .` setup does not include the extra packages used by the DreamZero example scripts.
- `openpi_client.py`:
`openpi-client`, `websockets`, `opencv-python`
- `droid_sim_eval_client.py`:
`mediapy`, `websockets`, `openpi-client`
The DROID sim-eval script also needs an Isaac Lab environment that provides `isaaclab`, `isaaclab_tasks`, `sim_evals`, and `gymnasium`.
If you run the DROID client on Python < 3.12, also install `typing-extensions`.
### Configure TP and CFG parallelism
The bundled DreamZero configs intentionally keep only:
| Config | Purpose |
|---|---|
| `vllm_omni/deploy/dreamzero.yaml` | Default TP=1, CFG parallel disabled |
| `vllm_omni/deploy/dreamzero_tp1_cfg2.yaml` | TP=1, CFG parallel size=2 |
For other topologies, use CLI parallelism flags and update stage 0 `devices` with `--stage-overrides`. The number of listed devices must match `tensor_parallel_size * cfg_parallel_size`.
TP=2 with CFG parallel disabled:
```bash
vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
--served-model-name dreamzero-droid \
--deploy-config vllm_omni/deploy/dreamzero.yaml \
--tensor-parallel-size 2 \
--stage-overrides '{"0": {"devices": "0,1"}}' \
--enforce-eager --disable-log-stats
```
TP=2 with CFG parallel size=2:
```bash
vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
--served-model-name dreamzero-droid \
--deploy-config vllm_omni/deploy/dreamzero.yaml \
--tensor-parallel-size 2 \
--cfg-parallel-size 2 \
--stage-overrides '{"0": {"devices": "0,1,2,3"}}' \
--enforce-eager --disable-log-stats
```
### Download example assets
The OpenPI client and DROID sim-eval example expect the three camera MP4 files in `outputs/dreamzero/assets`.
```bash
hf download YangshenDeng/vllm-omni-dreamzero-assets --repo-type dataset --local-dir outputs/dreamzero/assets
```
### Run the OpenPI client
```bash
python examples/online_serving/dreamzero/openpi_client.py \
--host 127.0.0.1 \
--port 8000 \
--video-dir outputs/dreamzero/assets
```
This client uses downloaded example videos and talks to the OpenPI websocket server.
### Run DROID sim eval
```bash
${ISAACLAB_LAUNCHER} -p examples/online_serving/dreamzero/droid_sim_eval_client.py \
--host 127.0.0.1 \
--port 8000 \
--scene 1 \
--episodes 1 \
--headless \
--device cuda:0
```
Set `ISAACLAB_LAUNCHER=path/to/isaaclab.sh` from the vLLM-Omni repository root before running the command.
This launches Isaac Lab / sim-evals and runs the DROID benchmark loop against the same websocket endpoint.
### Export comparison videos offline
```bash
python examples/offline_inference/dreamzero/export_prediction_video.py \
--deploy-config vllm_omni/deploy/dreamzero_tp1_cfg2.yaml \
--save-input-video \
--save-gif
```
The export script uses local `Omni` inference, not the websocket server. It writes the input rollout video and the predicted output video artifacts for side-by-side inspection.
### MolmoSpace demo
```bash
python examples/online_serving/dreamzero/molmospace_dreamzero_eval_demo.py \
--host 127.0.0.1 \
--port 8000 \
--benchmark_dir /path/to/benchmark \
--output_dir /path/to/output
```
This demo adapts DreamZero to the MolmoSpace-style remote policy eval loop.
@@ -0,0 +1,819 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Run one or more DROID sim-eval rollouts against the vLLM DreamZero server.
This script is the vLLM/OpenPI adaptation of the upstream DreamZero
`eval_utils/run_sim_eval.py` sim-eval client.
Behavior intentionally kept close to upstream:
- same DROID observation extraction (`external_cam`, `external_cam_2`,
`wrist_cam`, joint position, gripper position)
- same resize-with-pad preprocessing to `(180, 320)`
- same `open_loop_horizon=8`
- same gripper binarization rule (`> 0.5 -> 1`, else `0`)
- same per-scene language prompts
Unlike upstream DreamZero, vLLM serves the compatible websocket policy endpoint
at `/v1/realtime/robot/openpi`, so this script includes the path suffix in the
client URI.
Run this script through Isaac Lab's launcher from the vLLM-Omni repository
root, for example:
"${ISAACLAB_LAUNCHER}" -p \
examples/online_serving/dreamzero/droid_sim_eval_client.py \
--host 127.0.0.1 \
--port 8000 \
--scene 1 \
--episodes 1 \
--headless \
--device cuda:1
"""
from __future__ import annotations
import argparse
import json
import logging
import sys
import time
import uuid
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
import numpy as np
import torch
try:
import mediapy
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("Optional DROID sim-eval client requires `mediapy`.") from exc
try:
from typing import override
except ImportError:
try:
from typing_extensions import override
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("Optional DROID sim-eval client requires `typing-extensions` on Python < 3.12.") from exc
try:
import websockets.sync.client
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("Optional DROID sim-eval client requires `websockets`.") from exc
# NOTE:
# This directory already contains a local file named `openpi_client.py`.
# However, what we want here is the *installed* `openpi_client` package from
# upstream OpenPI, not the sibling example file. When a script is executed
# directly, Python often puts the script directory at `sys.path[0]`, which
# would cause `import openpi_client` to resolve to the local example file and
# create a circular import.
#
# To avoid that ambiguity, temporarily remove the current example directory
# from the front of `sys.path`, import the real package, and then restore the
# path afterwards.
example_dir = str(Path(__file__).resolve().parent)
removed_path = False
if sys.path and sys.path[0] == example_dir:
sys.path.pop(0)
removed_path = True
try:
from openpi_client import image_tools, msgpack_numpy
from openpi_client.base_policy import BasePolicy
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("Optional DROID sim-eval client requires the `openpi-client` package.") from exc
finally:
if removed_path:
sys.path.insert(0, example_dir)
# -----------------------------------------------------------------------------
# Constants
# -----------------------------------------------------------------------------
#
# These values intentionally mirror the upstream DreamZero sim-eval client
# where possible. The important distinction is:
#
# - ACTION_HORIZON = 24
# The model returns 24 future actions per inference call.
# - DEFAULT_OPEN_LOOP_HORIZON = 8
# The sim client only executes the first 8 actions locally before asking
# the server to replan from a fresh observation.
#
# So a single server call predicts 24x8 actions, but the rollout consumes only
# 8 of them before replanning.
PING_INTERVAL_SECS = 60
PING_TIMEOUT_SECS = 600
DEFAULT_PATH = "/v1/realtime/robot/openpi"
DEFAULT_OPEN_LOOP_HORIZON = 8
ACTION_HORIZON = 24
ACTION_DIM = 8
DEFAULT_OUTPUT_ROOT = Path("runs") / "dreamzero_sim_eval"
SCENE_PROMPTS = {
1: "put the cube in the bowl",
2: "pick up the can and put it in the mug",
3: "put the banana in the bin",
}
def _decode_action_response(response: bytes | str) -> np.ndarray:
if isinstance(response, str):
raise RuntimeError(f"Error in inference server:\n{response}")
decoded = msgpack_numpy.unpackb(response)
if isinstance(decoded, dict) and decoded.get("type") == "error":
message = decoded.get("message", decoded)
raise RuntimeError(f"Error in inference server:\n{message}")
return np.asarray(decoded, dtype=np.float32)
@dataclass(frozen=True)
class StepRecord:
"""One fully materialized rollout step for later JSON export.
The `episode_00.json` artifact is intended to be human-readable and
post-process-friendly. Instead of keeping raw tensors around, each step is
flattened into plain Python types so it can be serialized directly.
"""
# Index of this control step within the episode.
step_index: int
# Whether this step triggered a fresh model call (as opposed to reusing the
# cached open-loop chunk from the previous server response).
used_server_call: bool
# End-to-end latency of the server call that produced the current chunk.
# This is `None` on steps that only reuse cached actions.
chunk_latency_s: float | None
# The concrete 8-D action sent into the simulator at this step.
action: list[float]
# Observed 7-DoF arm joint positions before the next environment step.
joint_position: list[float]
# Observed gripper scalar before the next environment step.
gripper_position: list[float]
# Reward and termination signals directly returned by the simulator.
reward: float
terminated: bool
truncated: bool
# Optional scene object positions for downstream debugging / success
# heuristics. This may be empty if the environment does not expose them.
scene_objects: dict[str, list[float]]
class OpenPIWebsocketClientPolicy(BasePolicy):
"""Minimal websocket client for the DreamZero/OpenPI policy protocol.
Protocol shape:
- connect -> server immediately sends a metadata dict
- infer -> send msgpack observation, receive action chunk
- reset -> send msgpack reset command, receive confirmation string
This class intentionally stays very small because the more interesting
DreamZero-specific behavior lives one layer above, in
`DreamZeroJointPosClient`.
"""
def __init__(
self,
host: str = "127.0.0.1",
port: int = 8000,
path: str = DEFAULT_PATH,
) -> None:
# vLLM serves the robot endpoint under `/v1/realtime/robot/openpi`.
self._uri = f"ws://{host}:{port}{path}"
# Upstream protocol uses msgpack with numpy support, not JSON.
self._packer = msgpack_numpy.Packer()
# Connect immediately and cache the server handshake metadata.
self._ws, self._server_metadata = self._wait_for_server()
def get_server_metadata(self) -> dict[str, Any]:
"""Return a copy of the server handshake metadata."""
return dict(self._server_metadata)
def _wait_for_server(self):
"""Connect to the websocket server and read the initial metadata frame."""
logging.info("Connecting to %s", self._uri)
conn = websockets.sync.client.connect(
self._uri,
compression=None,
max_size=None,
ping_interval=PING_INTERVAL_SECS,
ping_timeout=PING_TIMEOUT_SECS,
)
metadata = msgpack_numpy.unpackb(conn.recv())
if not isinstance(metadata, dict):
raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
return conn, metadata
@override
def infer(self, obs: dict[str, Any]) -> np.ndarray:
"""Send an inference request and return the decoded action chunk."""
# Keep the upstream DreamZero/OpenPI convention that the request itself
# tells the server which logical endpoint is being called.
payload = dict(obs)
payload["endpoint"] = "infer"
self._ws.send(self._packer.pack(payload))
response = self._ws.recv()
return _decode_action_response(response)
@override
def reset(self, reset_info: dict[str, Any] | None = None) -> str:
"""Tell the server to reset its session-side state."""
payload = dict(reset_info or {})
payload["endpoint"] = "reset"
self._ws.send(self._packer.pack(payload))
response = self._ws.recv()
if isinstance(response, str):
return response
decoded = msgpack_numpy.unpackb(response)
if not isinstance(decoded, dict) or decoded.get("status") != "reset successful":
raise RuntimeError(f"Unexpected reset response: {decoded!r}")
return str(decoded["status"])
def close(self) -> None:
"""Close the websocket connection explicitly."""
self._ws.close()
class DreamZeroJointPosClient:
"""DROID sim-eval client that talks to the vLLM OpenPI websocket server.
This is the main compatibility layer between:
- Isaac Lab DROID observations (`obs["policy"][...]`)
- DreamZero/OpenPI websocket payloads
- local open-loop action reuse across several simulator steps
In other words:
simulator obs -> websocket request -> action chunk -> one action per step
"""
def __init__(
self,
remote_host: str = "127.0.0.1",
remote_port: int = 8000,
path: str = DEFAULT_PATH,
open_loop_horizon: int = DEFAULT_OPEN_LOOP_HORIZON,
) -> None:
# Low-level transport client.
self.client = OpenPIWebsocketClientPolicy(remote_host, remote_port, path=path)
# Number of actions to execute locally before replanning.
self.open_loop_horizon = open_loop_horizon
# Cursor into the currently cached action chunk.
self.actions_from_chunk_completed = 0
# Most recent `(ACTION_HORIZON, ACTION_DIM)` server response.
self.pred_action_chunk: np.ndarray | None = None
# Session id is part of the DreamZero serving contract. Changing it
# causes the server side to treat the rollout as a fresh episode.
self.session_id = str(uuid.uuid4())
# Simple runtime stats for reporting.
self.server_calls = 0
self.last_chunk_latency_s: float | None = None
self.last_used_server_call = False
def metadata(self) -> dict[str, Any]:
"""Expose the server metadata to callers / logs."""
return self.client.get_server_metadata()
def reset(self) -> str:
"""Reset local chunk state and remote session state.
Local reset:
- drop cached action chunk
- rewind chunk cursor
- allocate a fresh session id
Remote reset:
- send a websocket `reset` message so the server can clear any
request/session-side state it associates with this client
"""
self.actions_from_chunk_completed = 0
self.pred_action_chunk = None
self.session_id = str(uuid.uuid4())
self.last_chunk_latency_s = None
self.last_used_server_call = False
return self.client.reset({})
def infer(self, obs: dict[str, Any], instruction: str) -> dict[str, Any]:
"""Turn one simulator observation into one executable 8-D action.
Key behavior:
- call the server only when the local chunk cache is empty/exhausted
- otherwise, keep consuming the cached chunk open-loop
- always return exactly one 8-D action for the current simulator step
"""
# Convert Isaac Lab observation structure into a plain numpy-friendly
# record that is easier to serialize and visualize.
curr_obs = self._extract_observation(obs)
self.last_used_server_call = False
# Replan if:
# 1. this is the first step of a rollout / chunk
# 2. we already consumed `open_loop_horizon` actions from the current chunk
# 3. no cached chunk is currently available
if (
self.actions_from_chunk_completed == 0
or self.actions_from_chunk_completed >= self.open_loop_horizon
or self.pred_action_chunk is None
):
self.actions_from_chunk_completed = 0
# Build the exact DreamZero/OpenPI payload expected by the server.
#
# Notes:
# - images are resized/padded to the serving contract's 180x320
# - proprio is cast to float64 to match upstream client behavior
# - cartesian_position is currently unused by DreamZero DROID, so
# a dummy zero vector is sent for protocol completeness
request_data = {
"observation/exterior_image_0_left": image_tools.resize_with_pad(curr_obs["right_image"], 180, 320),
"observation/exterior_image_1_left": image_tools.resize_with_pad(curr_obs["left_image"], 180, 320),
"observation/wrist_image_left": image_tools.resize_with_pad(curr_obs["wrist_image"], 180, 320),
"observation/joint_position": curr_obs["joint_position"].astype(np.float64),
"observation/cartesian_position": np.zeros((6,), dtype=np.float64),
"observation/gripper_position": curr_obs["gripper_position"].astype(np.float64),
"prompt": instruction,
"session_id": self.session_id,
}
# Measure end-to-end server latency for this chunk request.
start = time.perf_counter()
actions = self.client.infer(request_data)
self.last_chunk_latency_s = time.perf_counter() - start
self.last_used_server_call = True
self.server_calls += 1
# DreamZero DROID serving is expected to return an action chunk with
# 24 future actions, each action being 8-D.
if actions.ndim != 2:
raise AssertionError(f"Expected 2D action array, got shape {actions.shape}")
if actions.shape != (ACTION_HORIZON, ACTION_DIM):
raise AssertionError(f"Expected action shape {(ACTION_HORIZON, ACTION_DIM)}, got {actions.shape}")
self.pred_action_chunk = actions
# Consume exactly one action row from the cached chunk for this
# simulator step.
action = np.array(self.pred_action_chunk[self.actions_from_chunk_completed], copy=True)
self.actions_from_chunk_completed += 1
# Upstream DreamZero sim-eval binarizes the gripper command.
action[-1] = 1.0 if action[-1].item() > 0.5 else 0.0
# Produce a human-friendly visualization strip for videos:
# right external | wrist | left external
img1 = image_tools.resize_with_pad(curr_obs["right_image"], 224, 224)
img2 = image_tools.resize_with_pad(curr_obs["wrist_image"], 224, 224)
img3 = image_tools.resize_with_pad(curr_obs["left_image"], 224, 224)
viz = np.concatenate([img1, img2, img3], axis=1)
# Return both the executable action and auxiliary debug info.
return {
"action": action,
"viz": viz,
"joint_position": curr_obs["joint_position"],
"gripper_position": curr_obs["gripper_position"],
"used_server_call": self.last_used_server_call,
"chunk_latency_s": self.last_chunk_latency_s if self.last_used_server_call else None,
}
@staticmethod
def _extract_observation(obs_dict: dict[str, Any]) -> dict[str, np.ndarray]:
"""Extract the pieces DreamZero cares about from Isaac Lab observations.
`sim-evals` exposes camera frames and robot state inside the
`obs["policy"]` group. This helper converts those tensors into numpy
arrays so they can be fed into image preprocessing / websocket packing.
"""
policy = obs_dict["policy"]
# Isaac Lab stores camera observations as batched tensors; use env 0.
right_image = policy["external_cam"][0].clone().detach().cpu().numpy()
left_image = policy["external_cam_2"][0].clone().detach().cpu().numpy()
wrist_image = policy["wrist_cam"][0].clone().detach().cpu().numpy()
# Robot proprioception.
joint_position = policy["arm_joint_pos"].clone().detach().cpu().numpy()
gripper_position = policy["gripper_pos"].clone().detach().cpu().numpy()
return {
"right_image": right_image,
"left_image": left_image,
"wrist_image": wrist_image,
"joint_position": joint_position,
"gripper_position": gripper_position,
}
def _scene_instruction(scene: int) -> str:
"""Map a numeric scene id onto the fixed language prompt used for rollout."""
try:
return SCENE_PROMPTS[scene]
except KeyError as exc:
raise ValueError(f"Unsupported scene {scene}. Available scenes: {sorted(SCENE_PROMPTS)}") from exc
def _capture_scene_objects(env: Any) -> dict[str, list[float]]:
"""Best-effort extraction of scene object root positions.
This is only for debugging / reporting. The rollout logic itself does not
depend on these positions.
"""
objects: dict[str, list[float]] = {}
scene = getattr(env, "scene", None)
if scene is None:
return objects
# Skip non-task entities such as cameras, the robot, and lighting.
for name in scene.keys():
if name in {"robot", "external_cam", "external_cam_2", "wrist_cam", "sphere_light", "scene"}:
continue
entity = scene[name]
data = getattr(entity, "data", None)
root_pos_w = getattr(data, "root_pos_w", None)
if root_pos_w is None:
continue
# Convert tensors to plain lists for JSON serialization.
value = root_pos_w[0].detach().cpu().to(torch.float32).tolist()
objects[str(name)] = [float(x) for x in value]
return objects
def _maybe_infer_success(scene: int, final_objects: dict[str, list[float]]) -> dict[str, Any]:
"""Best-effort geometric heuristic.
The simulator itself does not expose a built-in success term; this function
provides a transparent fallback for human-readable reporting only.
"""
task_pairs = {
1: ("cube", "bowl"),
2: ("can", "mug"),
3: ("banana", "bin"),
}
source_name, target_name = task_pairs.get(scene, (None, None))
if source_name not in final_objects or target_name not in final_objects:
return {
"has_builtin_success": False,
"heuristic_success": None,
"reason": "scene object names unavailable for heuristic",
}
source = np.asarray(final_objects[source_name], dtype=np.float32)
target = np.asarray(final_objects[target_name], dtype=np.float32)
xy_distance = float(np.linalg.norm(source[:2] - target[:2]))
z_delta = float(source[2] - target[2])
heuristic_success = bool(xy_distance < 0.12 and z_delta > -0.08)
return {
"has_builtin_success": False,
"heuristic_success": heuristic_success,
"xy_distance": xy_distance,
"z_delta": z_delta,
"source_object": source_name,
"target_object": target_name,
"reason": "no built-in env success flag; using final object pose heuristic",
}
def _dump_json(path: Path, payload: dict[str, Any]) -> None:
"""Write a JSON file with stable UTF-8 formatting."""
path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
def _scalar_from_env_value(value: Any) -> float:
"""Normalize simulator scalar outputs into a plain Python float.
Isaac Lab / Gym values may come back as tensors, numpy arrays, tuples, or
direct Python scalars depending on the wrapper stack. Centralizing the
conversion here makes the rollout loop cleaner and more robust.
"""
if isinstance(value, torch.Tensor):
return float(value.reshape(-1)[0].detach().cpu().item())
if isinstance(value, np.ndarray):
return float(value.reshape(-1)[0])
if isinstance(value, (list, tuple)):
return float(np.asarray(value).reshape(-1)[0])
return float(value)
def _bool_from_env_value(value: Any) -> bool:
"""Normalize simulator boolean-like outputs into a plain Python bool."""
if isinstance(value, torch.Tensor):
return bool(value.reshape(-1)[0].detach().cpu().item())
if isinstance(value, np.ndarray):
return bool(value.reshape(-1)[0])
if isinstance(value, (list, tuple)):
return bool(np.asarray(value).reshape(-1)[0])
return bool(value)
def _make_output_dir(output_root: Path, scene: int) -> Path:
"""Create a timestamped output directory for one sim-eval run."""
timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
output_dir = output_root / f"scene{scene}_{timestamp}"
output_dir.mkdir(parents=True, exist_ok=True)
return output_dir
def main() -> None:
"""Entry point for one or more DROID simulation rollouts.
High-level flow:
1. parse command-line flags
2. bootstrap Isaac Lab / sim-evals imports
3. create the DROID environment
4. connect the DreamZero websocket client
5. run `episodes` rollouts
6. export videos + JSON summaries
"""
# Script-level arguments.
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--episodes", type=int, default=1, help="Number of episodes to run.")
parser.add_argument("--scene", type=int, default=1, help="DROID scene id (1/2/3).")
parser.add_argument("--host", type=str, default="127.0.0.1", help="vLLM DreamZero server host.")
parser.add_argument("--port", type=int, default=8000, help="vLLM DreamZero server port.")
parser.add_argument("--path", type=str, default=DEFAULT_PATH, help="Websocket path suffix.")
parser.add_argument(
"--open-loop-horizon",
type=int,
default=DEFAULT_OPEN_LOOP_HORIZON,
help="How many actions to consume locally before requesting the next chunk.",
)
parser.add_argument(
"--output-root",
type=Path,
default=DEFAULT_OUTPUT_ROOT,
help="Directory where videos and trajectory logs are stored.",
)
try:
from isaaclab.app import AppLauncher
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError(
"Optional DROID sim-eval client requires Isaac Lab (`isaaclab`). "
"Launch it from an Isaac Lab environment, e.g. via `isaaclab.sh -p`."
) from exc
# Let Isaac Lab inject its own runtime flags (e.g. headless, device).
AppLauncher.add_app_launcher_args(parser)
args = parser.parse_args()
# DreamZero sim-eval always needs camera observations enabled.
args.enable_cameras = True
# Boot Isaac Sim / Isaac Lab.
app_launcher = AppLauncher(args)
simulation_app = app_launcher.app
# Set defaults so the `finally` block can clean up safely even if an
# earlier step fails.
env = None
client = None
# Import simulator modules only *after* the app is launched. This matches
# Isaac Lab's required import ordering.
try:
import gymnasium as gym
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("Optional DROID sim-eval client requires `gymnasium`.") from exc
try:
import sim_evals.environments # noqa: F401
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError(
"Optional DROID sim-eval client requires the external `sim-evals` package or checkout to be importable."
) from exc
try:
from isaaclab_tasks.utils import parse_env_cfg
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("Optional DROID sim-eval client requires `isaaclab_tasks`.") from exc
# Resolve output location and scene prompt.
output_dir = _make_output_dir(args.output_root.expanduser().resolve(), args.scene)
instruction = _scene_instruction(args.scene)
# Build the DROID environment configuration from `sim-evals`.
env_cfg = parse_env_cfg(
"DROID",
device=args.device,
num_envs=1,
use_fabric=True,
)
# Select one of the pre-authored scenes/tasks.
env_cfg.set_scene(args.scene)
env = gym.make("DROID", cfg=env_cfg)
# Upstream sim-evals resets twice so materials / cameras are fully ready.
obs, _ = env.reset()
obs, _ = env.reset()
# Connect the websocket policy client.
client = DreamZeroJointPosClient(
remote_host=args.host,
remote_port=args.port,
path=args.path,
open_loop_horizon=args.open_loop_horizon,
)
# Aggregated per-run results.
all_episode_summaries: list[dict[str, Any]] = []
max_steps = int(env.env.max_episode_length)
logging.info("DreamZero metadata: %s", client.metadata())
logging.info("Scene %s prompt: %s", args.scene, instruction)
logging.info("Writing outputs to %s", output_dir)
try:
# No gradients are needed in inference-only rollout mode.
with torch.no_grad():
for episode_index in range(args.episodes):
# Per-episode collectors.
frames: list[np.ndarray] = []
step_records: list[StepRecord] = []
episode_start = time.perf_counter()
server_time_s = 0.0
final_reward = 0.0
terminated = False
truncated = False
for step_index in range(max_steps):
# Ask the policy for the next action. Internally this may or
# may not trigger a real server request depending on whether
# the local chunk cache has been exhausted.
logging.debug("Episode %d step %d: requesting action", episode_index, step_index)
result = client.infer(obs, instruction)
logging.debug(
"Episode %d step %d: got action (server_call=%s latency=%s)",
episode_index,
step_index,
result["used_server_call"],
result["chunk_latency_s"],
)
# Save one visualization frame per simulator step.
frames.append(result["viz"])
# Isaac Lab expects batched actions, hence `[None]`.
action_tensor = torch.tensor(result["action"], dtype=torch.float32)[None]
logging.debug("Episode %d step %d: stepping env", episode_index, step_index)
obs, reward, term, trunc, info = env.step(action_tensor)
logging.debug("Episode %d step %d: env.step returned", episode_index, step_index)
logging.debug("Episode %d step %d: parsing reward/flags", episode_index, step_index)
logging.debug(
"Episode %d step %d: raw types reward=%s term=%s trunc=%s",
episode_index,
step_index,
type(reward).__name__,
type(term).__name__,
type(trunc).__name__,
)
# Normalize environment outputs into plain Python scalars so
# the rest of the code does not depend on wrapper-specific types.
reward_value = _scalar_from_env_value(reward)
term_value = _bool_from_env_value(term)
trunc_value = _bool_from_env_value(trunc)
logging.debug(
"Episode %d step %d: parsed reward=%s term=%s trunc=%s",
episode_index,
step_index,
reward_value,
term_value,
trunc_value,
)
# Keep scene-object capture optional. It is useful for
# debugging / success heuristics, but the rollout should not
# fail if the environment does not expose object roots.
scene_objects = _capture_scene_objects(env)
# Accumulate total server-side time only on steps that
# triggered a fresh chunk inference.
if result["chunk_latency_s"] is not None:
server_time_s += float(result["chunk_latency_s"])
# Materialize one JSON-serializable trajectory record.
logging.debug("Episode %d step %d: appending trajectory", episode_index, step_index)
step_records.append(
StepRecord(
step_index=step_index,
used_server_call=bool(result["used_server_call"]),
chunk_latency_s=(
float(result["chunk_latency_s"]) if result["chunk_latency_s"] is not None else None
),
action=[float(x) for x in np.asarray(result["action"], dtype=np.float32).tolist()],
joint_position=[
float(x) for x in np.asarray(result["joint_position"], dtype=np.float32).tolist()
],
gripper_position=[
float(x) for x in np.asarray(result["gripper_position"], dtype=np.float32).tolist()
],
reward=reward_value,
terminated=term_value,
truncated=trunc_value,
scene_objects=scene_objects,
)
)
logging.debug("Episode %d step %d: trajectory appended", episode_index, step_index)
# Track final status for summary export.
final_reward = reward_value
terminated = term_value
truncated = trunc_value
if term_value or trunc_value:
# End the rollout early if the environment terminates.
break
# Episode-level timing and video export.
episode_wall_time_s = time.perf_counter() - episode_start
video_path = output_dir / f"episode_{episode_index:02d}.mp4"
logging.info("Episode %d: writing video to %s", episode_index, video_path)
mediapy.write_video(video_path, frames, fps=15)
# Reset the policy server between episodes.
logging.info("Episode %d: sending reset", episode_index)
reset_response = client.reset()
final_objects = step_records[-1].scene_objects if step_records else {}
success_report = _maybe_infer_success(args.scene, final_objects)
# Assemble the per-episode summary that is written to
# `episode_XX.json`.
episode_summary = {
"episode_index": episode_index,
"prompt": instruction,
"video_path": str(video_path),
"steps_executed": len(step_records),
"max_steps": max_steps,
"terminated": terminated,
"truncated": truncated,
"final_reward": final_reward,
"server_calls": client.server_calls,
"server_time_s": server_time_s,
"episode_wall_time_s": episode_wall_time_s,
"avg_server_time_per_call_s": (
server_time_s / client.server_calls if client.server_calls else None
),
"reset_response": reset_response,
"success_report": success_report,
"server_metadata": client.metadata(),
"trajectory": [record.__dict__ for record in step_records],
}
_dump_json(output_dir / f"episode_{episode_index:02d}.json", episode_summary)
all_episode_summaries.append(episode_summary)
logging.info(
"Episode %d done: steps=%d wall=%.2fs server_calls=%d heuristic_success=%s",
episode_index,
len(step_records),
episode_wall_time_s,
client.server_calls,
success_report.get("heuristic_success"),
)
# Reset per-episode counters while keeping the client alive.
client.server_calls = 0
# Top-level run summary across all episodes.
summary = {
"scene": args.scene,
"prompt": instruction,
"episodes": args.episodes,
"host": args.host,
"port": args.port,
"path": args.path,
"device": args.device,
"output_dir": str(output_dir),
"server_metadata": client.metadata(),
"episodes_summary": all_episode_summaries,
}
_dump_json(output_dir / "summary.json", summary)
# Also print the summary to stdout so the caller can capture it in logs.
print(json.dumps(summary, ensure_ascii=False, indent=2))
finally:
# Best-effort cleanup. Avoid masking the main error if cleanup fails.
try:
if client is not None:
client.client.close()
except Exception:
pass
if env is not None:
env.close()
simulation_app.close()
if __name__ == "__main__":
# Keep the script-level logs readable. Per-step rollout details are still
# available via `DEBUG` if needed, but websocket / asyncio internals are
# usually too noisy for normal usage.
logging.basicConfig(level=logging.INFO)
logging.getLogger("websockets").setLevel(logging.WARNING)
logging.getLogger("asyncio").setLevel(logging.WARNING)
main()
@@ -0,0 +1,86 @@
from __future__ import annotations
import argparse
import os
import sys
from pathlib import Path
os.environ.setdefault("MUJOCO_GL", "egl")
os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
# Import base configs at module top level so the subclasses below are pickle-
# resolvable (worker processes import this module fresh via __main__).
from molmo_spaces.configs.policy_configs_baselines import ( # noqa: E402
DreamZeroPolicyConfig,
)
from molmo_spaces.evaluation.configs.evaluation_configs import ( # noqa: E402
DreamZeroPolicyEvalConfig,
)
# We only need to change the backend host and port to the vllm-host!
class DreamZeroVllmOmniPolicyConfig(DreamZeroPolicyConfig):
remote_config: dict = dict(host="127.0.0.1", port=8000)
class DreamZeroVllmOmniEvalConfig(DreamZeroPolicyEvalConfig):
policy_config: DreamZeroVllmOmniPolicyConfig = DreamZeroVllmOmniPolicyConfig()
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument(
"--benchmark_dir",
required=True,
help=(
"Path to a MolmoSpaces benchmark directory, for example "
"$MOLMOSPACES_BENCHMARK_DIR/20260327/ithor/FrankaCloseHardBench/"
"FrankaCloseHardBench_20260206_json_benchmark"
),
)
parser.add_argument("--max_episodes", type=int, default=1)
parser.add_argument("--task_horizon_steps", type=int, default=80)
parser.add_argument(
"--output_dir",
required=True,
help="Directory to write evaluation outputs (created if missing).",
)
parser.add_argument("--episode_idx", type=int, default=None)
args = parser.parse_args()
policy_config = DreamZeroVllmOmniPolicyConfig(remote_config=dict(host=args.host, port=args.port))
DreamZeroVllmOmniPolicyConfig.model_fields["remote_config"].default = policy_config.remote_config
DreamZeroVllmOmniEvalConfig.model_fields["policy_config"].default = policy_config
# Import after env vars are set so MuJoCo picks EGL.
from molmo_spaces.evaluation import run_evaluation
cfg_cls = DreamZeroVllmOmniEvalConfig
output_dir = args.output_dir
Path(output_dir).mkdir(parents=True, exist_ok=True)
print(f"[eval] benchmark_dir={args.benchmark_dir}")
print(f"[eval] max_episodes={args.max_episodes} task_horizon_steps={args.task_horizon_steps}")
print(f"[eval] remote policy: ws://{args.host}:{args.port}/v1/realtime/robot/openpi")
results = run_evaluation(
eval_config_cls=cfg_cls,
benchmark_dir=Path(args.benchmark_dir),
max_episodes=args.max_episodes,
task_horizon_steps=args.task_horizon_steps,
num_workers=1,
use_wandb=False,
output_dir=output_dir,
episode_idx=args.episode_idx,
)
print(f"[eval] success={results.success_count}/{results.total_count} ({results.success_rate:.1%})")
print(f"[eval] output_dir={results.output_dir}")
return 0
if __name__ == "__main__":
sys.exit(main())
+311
View File
@@ -0,0 +1,311 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import argparse
import json
import logging
import sys
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import numpy as np
try:
import cv2
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("DreamZero OpenPI example requires `opencv-python`.") from exc
try:
example_dir = str(Path(__file__).resolve().parent)
removed_path = False
if sys.path and sys.path[0] == example_dir:
sys.path.pop(0)
removed_path = True
try:
from openpi_client.websocket_client_policy import WebsocketClientPolicy
finally:
if removed_path:
sys.path.insert(0, example_dir)
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("DreamZero OpenPI example requires `openpi-client`.") from exc
DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 8000
DEFAULT_PATH = "/v1/realtime/robot/openpi"
DEFAULT_PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
ACTION_HORIZON = 24
DEFAULT_ACTION_DIM = 8
RELATIVE_OFFSETS = [-23, -16, -8, 0]
REPO_ROOT = Path(__file__).resolve().parents[3]
ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
DEFAULT_VIDEO_DIR = REPO_ROOT / "outputs" / "dreamzero" / "assets"
CAMERA_FILES = {
"observation/exterior_image_0_left": "exterior_image_1_left.mp4",
"observation/exterior_image_1_left": "exterior_image_2_left.mp4",
"observation/wrist_image_left": "wrist_image_left.mp4",
}
def _as_action_array(response: Any) -> np.ndarray:
if isinstance(response, dict) and response.get("type") == "error":
message = response.get("message", response)
raise RuntimeError(f"Inference failed: {message}")
return np.asarray(response, dtype=np.float32)
@dataclass(frozen=True)
class DreamZeroServerMetadata:
image_resolution: tuple[int, int]
n_external_cameras: int
needs_wrist_camera: bool
needs_stereo_camera: bool
needs_session_id: bool
action_space: str
@classmethod
def from_dict(cls, payload: dict[str, Any]) -> DreamZeroServerMetadata:
required_keys = (
"image_resolution",
"n_external_cameras",
"needs_wrist_camera",
"needs_stereo_camera",
"needs_session_id",
"action_space",
)
missing_keys = [key for key in required_keys if key not in payload]
if missing_keys:
raise ValueError(f"Missing DreamZero metadata keys: {missing_keys}")
image_resolution = payload["image_resolution"]
if not isinstance(image_resolution, (list, tuple)) or len(image_resolution) != 2:
raise ValueError(f"Invalid image_resolution: {image_resolution!r}")
return cls(
image_resolution=(int(image_resolution[0]), int(image_resolution[1])),
n_external_cameras=int(payload["n_external_cameras"]),
needs_wrist_camera=bool(payload["needs_wrist_camera"]),
needs_stereo_camera=bool(payload["needs_stereo_camera"]),
needs_session_id=bool(payload["needs_session_id"]),
action_space=str(payload["action_space"]),
)
def _openpi_uri(host: str, port: int, path: str) -> str:
if host.startswith(("ws://", "wss://")):
return host
return f"ws://{host}:{port}{path}"
def _close_policy_client(client: WebsocketClientPolicy) -> None:
ws = getattr(client, "_ws", None)
if ws is not None:
ws.close()
def load_all_frames(video_path: Path) -> np.ndarray:
cap = cv2.VideoCapture(str(video_path))
frames = []
while True:
ok, frame = cap.read()
if not ok:
break
frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
cap.release()
if not frames:
raise RuntimeError(f"No frames loaded from {video_path}")
return np.stack(frames, axis=0)
def load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
camera_frames: dict[str, np.ndarray] = {}
for camera_key, file_name in CAMERA_FILES.items():
video_path = video_dir / file_name
if not video_path.exists():
raise FileNotFoundError(
f"Missing DreamZero example asset: {video_path}. "
"Download the example videos with: "
f"`hf download {ASSET_REPO_ID} --repo-type dataset --local-dir {video_dir}`"
)
camera_frames[camera_key] = load_all_frames(video_path)
return camera_frames
def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
chunks: list[list[int]] = []
current_frame = 23
for _ in range(num_chunks):
indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
if indices[-1] >= total_frames:
break
chunks.append(indices)
current_frame += ACTION_HORIZON
return chunks
def make_obs_from_video(
camera_frames: dict[str, np.ndarray],
frame_indices: list[int],
*,
prompt: str,
session_id: str,
) -> dict[str, Any]:
obs: dict[str, Any] = {}
for camera_key, all_frames in camera_frames.items():
selected = all_frames[frame_indices]
obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
obs["prompt"] = prompt
obs["session_id"] = session_id
return obs
def build_demo_observations(
camera_frames: dict[str, np.ndarray],
*,
prompt: str,
session_id: str,
num_chunks: int = 2,
) -> list[dict[str, Any]]:
if num_chunks < 1:
raise ValueError("num_chunks must be at least 1")
total_frames = min(frames.shape[0] for frames in camera_frames.values())
observations = [
make_obs_from_video(
camera_frames,
[0],
prompt=prompt,
session_id=session_id,
)
]
for indices in build_frame_schedule(total_frames, num_chunks - 1):
observations.append(
make_obs_from_video(
camera_frames,
indices,
prompt=prompt,
session_id=session_id,
)
)
return observations
def validate_session_result(
result: dict[str, Any],
*,
expected_action_horizon: int = ACTION_HORIZON,
expected_action_dim: int = DEFAULT_ACTION_DIM,
expected_action_count: int | None = None,
) -> None:
metadata = DreamZeroServerMetadata.from_dict(result["metadata"])
if metadata.image_resolution != (180, 320):
raise AssertionError(f"Unexpected image_resolution: {metadata.image_resolution}")
if metadata.n_external_cameras != 2:
raise AssertionError(f"Unexpected n_external_cameras: {metadata.n_external_cameras}")
if not metadata.needs_wrist_camera:
raise AssertionError("DreamZero example expects wrist camera metadata")
if metadata.action_space != "joint_position":
raise AssertionError(f"Unexpected action_space: {metadata.action_space}")
actions = result["actions"]
if expected_action_count is not None and len(actions) != expected_action_count:
raise AssertionError(f"Expected {expected_action_count} action tensors, got {len(actions)}")
if not actions:
raise AssertionError("Expected at least one action tensor")
for index, action in enumerate(actions):
if action.shape != (expected_action_horizon, expected_action_dim):
raise AssertionError(
f"Action {index} shape mismatch: expected "
f"{(expected_action_horizon, expected_action_dim)}, got {action.shape}"
)
if not np.isfinite(action).all():
raise AssertionError(f"Action {index} contains non-finite values")
def run_policy_session(
*,
host: str = DEFAULT_HOST,
port: int = DEFAULT_PORT,
path: str = DEFAULT_PATH,
video_dir: Path = DEFAULT_VIDEO_DIR,
prompt: str = DEFAULT_PROMPT,
session_id: str | None = None,
num_chunks: int = 2,
) -> dict[str, Any]:
session_id = session_id or str(uuid.uuid4())
camera_frames = load_camera_frames(video_dir)
observations = build_demo_observations(
camera_frames,
prompt=prompt,
session_id=session_id,
num_chunks=num_chunks,
)
uri = _openpi_uri(host, port, path)
logging.info("Connecting to %s", uri)
client = WebsocketClientPolicy(host=uri)
try:
metadata = dict(client.get_server_metadata())
actions = [_as_action_array(client.infer(obs)) for obs in observations]
return {
"metadata": metadata,
"actions": actions,
"session_id": session_id,
}
finally:
_close_policy_client(client)
def format_action_summary(index: int, action: np.ndarray) -> str:
return (
f"Action {index}: shape={tuple(action.shape)} dtype={action.dtype} "
f"min={action.min():.6f} max={action.max():.6f}"
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="DreamZero OpenPI client example with downloaded real videos.")
parser.add_argument("--host", default=DEFAULT_HOST)
parser.add_argument("--port", type=int, default=DEFAULT_PORT)
parser.add_argument("--path", default=DEFAULT_PATH)
parser.add_argument(
"--video-dir", type=Path, default=DEFAULT_VIDEO_DIR, help="Directory containing the three camera MP4 files."
)
parser.add_argument("--prompt", default=DEFAULT_PROMPT)
parser.add_argument("--session-id", default=None)
parser.add_argument("--num-chunks", type=int, default=2)
return parser.parse_args()
def main() -> int:
args = parse_args()
logging.basicConfig(level=logging.INFO)
result = run_policy_session(
host=args.host,
port=args.port,
path=args.path,
video_dir=args.video_dir,
prompt=args.prompt,
session_id=args.session_id,
num_chunks=args.num_chunks,
)
validate_session_result(result, expected_action_count=args.num_chunks)
print("Server metadata:", json.dumps(result["metadata"], sort_keys=True))
for index, action in enumerate(result["actions"]):
print(format_action_summary(index, action))
print("Session ID:", result["session_id"])
return 0
if __name__ == "__main__":
raise SystemExit(main())
+24
View File
@@ -0,0 +1,24 @@
#!/usr/bin/env bash
set -euo pipefail
MODEL="${MODEL:-GEAR-Dreams/DreamZero-DROID}"
HOST="${HOST:-127.0.0.1}"
PORT="${PORT:-8000}"
DEPLOY_CONFIG="${DEPLOY_CONFIG:-vllm_omni/deploy/dreamzero_tp1_cfg2.yaml}"
SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-dreamzero-droid}"
args=(
serve
"$MODEL"
--omni
--host "$HOST"
--port "$PORT"
--served-model-name "$SERVED_MODEL_NAME"
--deploy-config "$DEPLOY_CONFIG"
--enforce-eager
--disable-log-stats
)
ATTENTION_BACKEND="${ATTENTION_BACKEND:-torch}" \
DIFFUSION_ATTENTION_BACKEND="${DIFFUSION_ATTENTION_BACKEND:-TORCH_SDPA}" \
vllm "${args[@]}"
+5 -5
View File
@@ -22,13 +22,14 @@ export PYTHONPATH=<transformers_modules_path>:$PYTHONPATH
export HF_MODULES_CACHE=<transformers_modules_path>
```
Run from repository root:
Run from repository root. The 3-stage pipeline is resolved automatically from
the model type (`vllm_omni/deploy/dynin_omni.yaml`), so no config path is
required. Pass `--deploy-config` only to override the defaults:
```bash
vllm-omni serve snu-aidas/Dynin-Omni \
--omni \
--port 8091 \
--stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
--port 8091
```
If `vllm-omni` is not in PATH, run:
@@ -36,8 +37,7 @@ If `vllm-omni` is not in PATH, run:
```bash
PYTHONPATH="$(pwd)" python -m vllm_omni.entrypoints.cli.main serve snu-aidas/Dynin-Omni \
--omni \
--port 8091 \
--stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
--port 8091
```
Wait until the server logs show both `All stages initialized successfully` and
+40
View File
@@ -0,0 +1,40 @@
# Lance — online serving
OpenAI-compatible chat completions API for Lance, served by `vllm-omni`.
## Start the server
```bash
bash examples/online_serving/lance/run_server.sh
# or, with overrides
MODEL=bytedance-research/Lance \
DEPLOY_CONFIG=vllm_omni/deploy/lance.yaml \
PORT=8091 \
bash examples/online_serving/lance/run_server.sh
```
The default deploy config is the single-stage `lance.yaml`; for `video_edit`
or `text2video` workloads, point ``DEPLOY_CONFIG`` at a video-checkpoint
config (or pass ``--model bytedance-research/Lance/Lance_3B_Video`` to
`run_server.sh`).
## Send a request
```bash
# text-to-image
python examples/online_serving/lance/openai_chat_client.py \
--prompt "A cute corgi astronaut on the moon, cinematic" \
--modality text2img \
--output corgi.png
# image edit
python examples/online_serving/lance/openai_chat_client.py \
--prompt "Convert this into a vibrant cartoon-style illustration" \
--modality img2img \
--image-url path/to/photo.png \
--output edited.png
```
The client is shared with BAGEL — same OpenAI message format, same
``modalities`` and ``num_inference_steps`` / ``seed`` / ``height`` /
``width`` knobs.
+184
View File
@@ -0,0 +1,184 @@
#!/usr/bin/env python3
"""
Lance OpenAI-compatible chat client for image generation and multimodal tasks.
Usage:
python examples/online_serving/lance/openai_chat_client.py --prompt "A cute cat" --output output.png
python examples/online_serving/lance/openai_chat_client.py --prompt "Describe this image" --image-url https://example.com/image.png
"""
import argparse
import base64
from pathlib import Path
import requests
def generate_image(
prompt: str,
server_url: str = "http://localhost:8091",
image_url: str | None = None,
height: int | None = None,
width: int | None = None,
steps: int | None = None,
seed: int | None = None,
negative_prompt: str | None = None,
modality: str = "text2img", # "text2img" (default), "img2img", "img2text", "text2text"
) -> bytes | str | None:
"""Generate an image or text using the chat completions API.
Args:
prompt: Text description or prompt
server_url: Server URL
image_url: URL or path to input image (for img2img/img2text)
height: Image height in pixels
width: Image width in pixels
steps: Number of inference steps
seed: Random seed
negative_prompt: Negative prompt
modality: Task modality hint
Returns:
Image bytes (for image outputs) or Text string (for text outputs) or None if failed
"""
# Construct Message Content
content = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}]
if image_url:
# Check if local file
if Path(image_url).exists():
with open(image_url, "rb") as f:
b64_data = base64.b64encode(f.read()).decode("utf-8")
final_image_url = f"data:image/jpeg;base64,{b64_data}"
else:
final_image_url = image_url
content.append({"type": "image_url", "image_url": {"url": final_image_url}})
messages = [{"role": "user", "content": content}]
# Build request payload with all parameters at top level
# Note: vLLM ignores "extra_body", so we put parameters directly in the payload
payload = {"messages": messages}
# Set output modalities at top level
if modality == "text2img" or modality == "img2img":
payload["modalities"] = ["image"]
elif modality == "img2text" or modality == "text2text":
payload["modalities"] = ["text"]
# Add generation parameters directly to payload
if height is not None:
payload["height"] = height
if width is not None:
payload["width"] = width
if steps is not None:
payload["num_inference_steps"] = steps
if seed is not None:
payload["seed"] = seed
if negative_prompt:
payload["negative_prompt"] = negative_prompt
# Send request
try:
print(f"Sending request to {server_url} with modality {modality}...")
response = requests.post(
f"{server_url}/v1/chat/completions",
headers={"Content-Type": "application/json"},
json=payload,
timeout=300,
)
response.raise_for_status()
data = response.json()
# Extract content - check ALL choices since server may return multiple
# (e.g., text in choices[0], image in choices[1])
choices = data.get("choices", [])
# First pass: look for image output in any choice
for choice in choices:
choice_content = choice.get("message", {}).get("content")
# Handle Image Output
if isinstance(choice_content, list) and len(choice_content) > 0:
first_item = choice_content[0]
if isinstance(first_item, dict) and "image_url" in first_item:
img_url_str = first_item["image_url"].get("url", "")
if img_url_str.startswith("data:image"):
_, b64_data = img_url_str.split(",", 1)
return base64.b64decode(b64_data)
# Second pass: look for text output if no image found
for choice in choices:
choice_content = choice.get("message", {}).get("content")
if isinstance(choice_content, str) and choice_content:
return choice_content
print(f"Unexpected response format: {choices}")
return None
except Exception as e:
print(f"Error: {e}")
return None
def main():
parser = argparse.ArgumentParser(description="Bagel multimodal chat client")
parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt")
parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)")
parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
# Modality Control
parser.add_argument("--image-url", "-i", type=str, help="Input image URL or local path")
parser.add_argument(
"--modality",
"-m",
default="text2img",
choices=["text2img", "img2img", "img2text", "text2text"],
help="Task modality",
)
# Generation Params
parser.add_argument("--height", type=int, default=512, help="Image height")
parser.add_argument("--width", type=int, default=512, help="Image width")
parser.add_argument("--steps", type=int, default=25, help="Inference steps")
parser.add_argument("--seed", type=int, default=42, help="Random seed")
parser.add_argument("--negative", help="Negative prompt")
args = parser.parse_args()
print(f"Mode: {args.modality}")
if args.image_url:
print(f"Input Image: {args.image_url}")
result = generate_image(
prompt=args.prompt,
server_url=args.server,
image_url=args.image_url,
height=args.height,
width=args.width,
steps=args.steps,
seed=args.seed,
negative_prompt=args.negative,
modality=args.modality,
)
if result:
if isinstance(result, bytes):
# It's an image
output_path = Path(args.output)
output_path.write_bytes(result)
print(f"Image saved to: {output_path}")
print(f"Size: {len(result) / 1024:.1f} KB")
elif isinstance(result, str):
# It's text
print("Response:")
print(result)
else:
print("Failed to generate response")
exit(1)
if __name__ == "__main__":
main()
+15
View File
@@ -0,0 +1,15 @@
#!/bin/bash
# Lance online serving startup script
MODEL="${MODEL:-bytedance-research/Lance}"
PORT="${PORT:-8091}"
DEPLOY_CONFIG="${DEPLOY_CONFIG:-vllm_omni/deploy/lance.yaml}"
echo "Starting Lance server..."
echo "Model: $MODEL"
echo "Deploy config: $DEPLOY_CONFIG"
echo "Port: $PORT"
vllm serve "$MODEL" --omni \
--deploy-config "$DEPLOY_CONFIG" \
--port "$PORT"
@@ -12,7 +12,8 @@ from typing import Any
import requests
from openai import OpenAI
from vllm.assets.audio import AudioAsset
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
# Modify OpenAI's API key and API base to use vLLM's API server.
openai_api_key = "EMPTY"
@@ -493,7 +494,7 @@ def run_multimodal_generation(args) -> None:
def parse_args():
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser.add_argument(
"--query-type",
"-q",
@@ -10,7 +10,7 @@ Please refer to [README.md](../../../README.md)
|------|---------------|--------|
| Thinker + Talker (omni-speech, default) | `vllm serve ... --omni` | Text + Audio |
| Thinker only (multimodal understanding) | `vllm serve ... --omni --deploy-config vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
| Thinker + Imagegen (text-to-image / img2img) | `vllm serve ... --omni --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml` | Image |
| Thinker + Imagegen (text-to-image / img2img) | `vllm serve ... --omni --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml` | Image |
For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
@@ -63,34 +63,27 @@ bash run_curl_multimodal_generation.sh use_image_gen
## Image generation (text-to-image)
Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage.
Launch with the dual-stage YAML, which adds an image-gen stage behind the
thinker:
Launch with the image deploy YAML, which adds an image-gen stage behind
the thinker:
```bash
vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
--stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml \
--deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
--stage-init-timeout 1800 \
--port 8188
--init-timeout 1800 \
--port 8091
```
Then request image output by passing `"modalities": ["image"]`:
```bash
curl http://127.0.0.1:8188/v1/chat/completions \
curl -s http://127.0.0.1:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
"messages": [{"role": "user", "content": "Draw a beautiful girl with short black hair and red dress."}],
"messages": [{"role": "user", "content": "Please draw a cute cat."}],
"modalities": ["image"]
}' -o /tmp/ming_response.json
python - <<'PY'
import base64, json
r = json.load(open('/tmp/ming_response.json'))
url = r['choices'][0]['message']['content'][0]['image_url']['url']
open('/tmp/ming_gen.png', 'wb').write(base64.b64decode(url.split(',', 1)[1]))
print('saved /tmp/ming_gen.png')
PY
}' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen.png
```
### Optional knobs
@@ -109,7 +102,7 @@ Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_
Example with all the knobs:
```bash
curl http://127.0.0.1:8188/v1/chat/completions \
curl http://127.0.0.1:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
@@ -122,7 +115,7 @@ curl http://127.0.0.1:8188/v1/chat/completions \
"byte5_text":["理解与生成统一"]}}}
],
"messages": [{"role": "user", "content": "Draw a poster."}]
}' -o /tmp/ming_response.json
}' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen_knobs.png
```
### img2img (reference image)
@@ -142,9 +135,9 @@ into the diffusion stage as `extra[reference_image]`:
### GPU layout
The shipped `ming_flash_omni_dual.yaml` allocates the thinker on GPUs 0–3
The shipped `ming_flash_omni_image.yaml` allocates the thinker on GPUs 0–3
(TP=4) and the diffusion stage on GPU 4 (TP=1). Copy the YAML and edit
`runtime.devices` per stage to relocate; with fewer GPUs available, drop
`devices` per stage to relocate; with fewer GPUs available, drop
the thinker TP to 2 and run the diffusion stage on a free card. Image-gen
warmup takes roughly an extra 30–60 s on top of the thinker — set
`--stage-init-timeout 1800` if the default 300 s is too tight.
@@ -156,7 +149,7 @@ warmup takes roughly an extra 30–60 s on top of the thinker — set
| `["text"]` or omitted | Thinker only | Text |
| `["audio"]` | Thinker + Talker | Audio (speech) |
| `["text", "audio"]` | Thinker + Talker | Text + Audio |
| `["image"]` | Thinker + Imagegen (dual YAML) | Image (PNG, base64 in `choices[0].message.content`) |
| `["image"]` | Thinker + Imagegen (image deploy YAML) | Image (PNG, base64 in `choices[0].message.content`) |
For ready-to-copy curl examples (text / audio / multimodal input, SSE
streaming, reasoning mode), see the recipe at
@@ -103,7 +103,7 @@ echo ""
if [[ "$QUERY_TYPE" == "use_image_gen" ]]; then
# Image-gen branch: dual-stage diffusion endpoint, modalities=["image"].
# The shipped ming_flash_omni_dual.yaml runs the thinker on cards 0-3
# The shipped ming_flash_omni_image.yaml runs the thinker on cards 0-3
# (TP=4) and the diffusion stage on card 4 (TP=1). See the image-gen
# section in README.md for knobs and img2img.
out_path="${OUT_PATH:-/tmp/ming_imagegen.png}"
@@ -0,0 +1,65 @@
# vLLM-Omni · MiniCPM-o 4.5 Online Demo
Gradio-based web UI for **MiniCPM-o 4.5** served via `vllm-omni`'s
OpenAI-compatible endpoints.
The UI supports:
- **Inputs**: text prompt + optional image, audio (file or mic), video.
- **Outputs**: text + speech (WAV player).
## 1. Start the backend server
The deploy config auto-loads via `--omni`; the default
`vllm_omni/deploy/minicpmo_4_5.yaml` targets a 2-GPU layout (thinker on
GPU 0, talker + t2w sharing GPU 1). For other hardware layouts pick
one of the deploy variants below.
| deploy config | GPUs | Notes |
|---|---|---|
| `minicpmo_4_5.yaml` (default) | 2 | Thinker on GPU0, talker+t2w on GPU1. |
| `minicpmo_4_5_3gpu.yaml` | 3 | Thinker 2-way TP on GPU0/1, talker+t2w share GPU2. |
| `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. |
Default (2-GPU):
```bash
vllm serve openbmb/MiniCPM-o-4_5 --omni \
--trust-remote-code \
--host 0.0.0.0 --port 8099
```
Other layouts via `--deploy-config`:
```bash
vllm serve openbmb/MiniCPM-o-4_5 --omni \
--deploy-config vllm_omni/deploy/minicpmo_4_5_8x4090.yaml \
--trust-remote-code \
--host 0.0.0.0 --port 8099
```
## 2. Launch the Gradio demo
```bash
bash examples/online_serving/minicpmo/run_gradio_demo.sh
# Or run the python entry point directly:
python examples/online_serving/minicpmo/gradio_demo.py \
--minicpmo45-api-base http://localhost:8099/v1 \
--minicpmo45-model openbmb/MiniCPM-o-4_5 \
--port 7862
```
Open `http://<host>:7862` in a browser.
## Notes
- **TTS trigger**: the demo sets
`extra_body.chat_template_kwargs.use_tts_template=True`, which appends
`<|tts_bos|>` to the assistant prefix.
- Uncheck **"Generate speech output (TTS)"** to get text-only responses
(faster).
- The audio output is the raw WAV returned by the stage-1 talker +
Token2Wav; sample rate is 24 kHz.
- Video input is forwarded as a base64 `video_url` entry; the server needs
decord/torchvision to decode it.
@@ -0,0 +1,580 @@
"""Gradio demo for MiniCPM-o 4.5 online serving through vllm-omni.
Supports:
- Text / Image / Audio / Video inputs (OpenAI chat-completions multimodal format).
- Text + audio outputs, returned side-by-side in the UI.
Usage:
# run a 4.5 server on :8099
python gradio_demo.py \
--minicpmo45-api-base http://localhost:8099/v1 \
--minicpmo45-model openbmb/MiniCPM-o-4_5 \
--port 7862
"""
from __future__ import annotations
import argparse
import base64
import io
import os
from collections.abc import Iterable
from pathlib import Path
from typing import Any
import gradio as gr
import numpy as np
import soundfile as sf
from openai import OpenAI
from PIL import Image
MINICPMO45 = "MiniCPM-o-4.5"
# Default reference audio per model, relative to the model path. Used to build
# the audio_assistant system prompt so that the LLM adopts the "speak in this
# voice" persona (matches the official MiniCPM-o demo).
_DEFAULT_REF_AUDIO = {
MINICPMO45: "assets/HT_ref_audio.wav",
}
# audio_assistant prompt text, mirrors MiniCPMO.get_sys_prompt("audio_assistant")
# in the MiniCPM-o-Demo reference implementation.
_AUDIO_ASSISTANT_PROMPT = {
"zh": {
"prefix": "模仿音频样本的音色并生成新的内容。",
"suffix": (
"你的任务是用这种声音模式来当一个助手。请认真、高质量地回复用户的问题。"
"请用高自然度的方式和用户聊天。你是由面壁智能开发的人工智能助手:面壁小钢炮。"
),
},
"en": {
"prefix": "Use the voice in the audio prompt to synthesize new content.",
"suffix": "You are a helpful assistant with the above voice style.",
},
}
# ---------------------------------------------------------------------------
# Media helpers
# ---------------------------------------------------------------------------
def image_to_base64_data_url(image: Image.Image) -> str:
buf = io.BytesIO()
if image.mode != "RGB":
image = image.convert("RGB")
image.save(buf, format="JPEG", quality=92)
b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
return f"data:image/jpeg;base64,{b64}"
def audio_to_base64_data_url(audio_np: np.ndarray, sample_rate: int) -> str:
if audio_np.dtype != np.int16:
if audio_np.dtype in (np.float32, np.float64):
audio_np = np.clip(audio_np, -1.0, 1.0)
audio_np = (audio_np * 32767).astype(np.int16)
else:
audio_np = audio_np.astype(np.int16)
buf = io.BytesIO()
sf.write(buf, audio_np, sample_rate, format="WAV")
b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
return f"data:audio/wav;base64,{b64}"
def ref_audio_to_data_url(path: str, target_sr: int = 16000, max_s: float | None = 8.0) -> str:
"""Load a file, downmix to mono, resample to 16kHz, optionally truncate, and
return a data-URL usable as audio_url content in the chat API.
MiniCPM-o expects 16kHz mono audio input; long reference audios can waste
tokens so we cap at max_s seconds by default (3-10s is the recommended range).
"""
data, sr = sf.read(path)
if data.ndim > 1:
data = data[:, 0]
data = data.astype(np.float32)
if sr != target_sr:
try:
# Demo client only; librosa here resamples audio before sending to
# the server. The fallback below covers environments without it.
import librosa # noqa: TID251
data = librosa.resample(data, orig_sr=sr, target_sr=target_sr)
except ImportError:
from math import gcd
g = gcd(sr, target_sr)
up, down = target_sr // g, sr // g
idx = (np.arange(len(data) * up) / up).astype(np.int64)
data = data[idx][::down]
sr = target_sr
if max_s is not None and len(data) > int(max_s * sr):
data = data[: int(max_s * sr)]
return audio_to_base64_data_url(data, sr)
def video_to_base64_data_url(path: str) -> str:
p = Path(path)
if not p.exists():
raise FileNotFoundError(f"Video not found: {path}")
ext = p.suffix.lower()
mime = {
".mp4": "video/mp4",
".webm": "video/webm",
".mov": "video/quicktime",
".avi": "video/x-msvideo",
".mkv": "video/x-matroska",
}.get(ext, "video/mp4")
b64 = base64.b64encode(p.read_bytes()).decode("utf-8")
return f"data:{mime};base64,{b64}"
def process_audio_input(audio_input: Any | None) -> tuple[np.ndarray, int] | None:
"""Normalize Gradio audio input to (np.ndarray mono float32, sample_rate).
Gradio's `gr.Audio(type="numpy")` yields (sr, ndarray); we also accept
filesystem paths and a few nested tuple variants that appear in different
Gradio versions.
"""
if audio_input is None:
return None
def _from_path(p: str) -> tuple[np.ndarray, int] | None:
if not p:
return None
fp = Path(p)
if not fp.exists():
return None
data, sr = sf.read(fp)
if data.ndim > 1:
data = data[:, 0]
return data.astype(np.float32), int(sr)
audio_np: np.ndarray | None = None
sr: int | None = None
if isinstance(audio_input, tuple) and len(audio_input) == 2:
a, b = audio_input
if isinstance(a, int | float) and isinstance(b, np.ndarray):
sr, audio_np = int(a), b
elif isinstance(a, str):
loaded = _from_path(a)
if loaded is not None:
audio_np, sr = loaded
elif isinstance(audio_input, str):
loaded = _from_path(audio_input)
if loaded is not None:
audio_np, sr = loaded
if audio_np is None or sr is None:
return None
if audio_np.ndim > 1:
audio_np = audio_np[:, 0]
return audio_np.astype(np.float32), sr
# ---------------------------------------------------------------------------
# Per-model inference config
# ---------------------------------------------------------------------------
class ModelEndpoint:
"""Endpoint + TTS-trigger convention for one MiniCPM-o variant."""
def __init__(self, name: str, api_base: str, model_path: str):
self.name = name
self.api_base = api_base
self.model_path = model_path
self.client = OpenAI(api_key="EMPTY", base_url=api_base)
self._ref_audio_data_url: str | None = None
@property
def ref_audio_path(self) -> str | None:
rel = _DEFAULT_REF_AUDIO.get(self.name)
if not rel:
return None
full = os.path.join(self.model_path, rel)
return full if os.path.exists(full) else None
def get_ref_audio_data_url(self) -> str | None:
"""Lazily load and cache the default reference audio as a 16kHz mono
data URL. Returns None if the reference audio file is missing."""
if self._ref_audio_data_url is not None:
return self._ref_audio_data_url
p = self.ref_audio_path
if not p:
print(f"[{self.name}] no default reference audio found; falling back to text-only system prompt")
return None
try:
self._ref_audio_data_url = ref_audio_to_data_url(p)
print(f"[{self.name}] loaded reference audio: {p}")
except Exception as e:
print(f"[{self.name}] failed to load reference audio {p}: {e}")
return None
return self._ref_audio_data_url
def build_extras(self) -> dict[str, Any]:
"""Extra kwargs passed to chat.completions.create for this model.
4.5: use_tts_template=True chat-template flag appends <|tts_bos|>.
"""
if self.name == MINICPMO45:
return {
"extra_body": {
"chat_template_kwargs": {"use_tts_template": True},
"modalities": ["text", "audio"],
},
}
return {}
# ---------------------------------------------------------------------------
# System prompt (matches MiniCPMO.get_sys_prompt("audio_assistant") in the
# official MiniCPM-o-Demo). When TTS is disabled we fall back to a plain
# text-only system prompt.
# ---------------------------------------------------------------------------
def build_audio_assistant_system(
ref_audio_data_url: str | None,
language: str = "zh",
) -> dict[str, Any]:
"""Return an OpenAI-style system message carrying the reference audio."""
txt = _AUDIO_ASSISTANT_PROMPT.get(language, _AUDIO_ASSISTANT_PROMPT["zh"])
if ref_audio_data_url is None:
return {
"role": "system",
"content": (
"You are MiniCPM-o, a helpful multimodal assistant that can "
"understand images, audio and video, and respond in text and speech."
),
}
return {
"role": "system",
"content": [
{"type": "text", "text": txt["prefix"]},
{"type": "audio_url", "audio_url": {"url": ref_audio_data_url}},
{"type": "text", "text": txt["suffix"]},
],
}
# ---------------------------------------------------------------------------
# Core inference function
# ---------------------------------------------------------------------------
def run_inference(
endpoint: ModelEndpoint,
user_prompt: str,
audio_tuple: tuple[np.ndarray, int] | None,
image_pil: Image.Image | None,
video_path: str | None,
enable_tts: bool,
max_tokens: int,
temperature: float,
) -> Iterable[tuple[str, tuple[int, np.ndarray] | None]]:
user_prompt = user_prompt or ""
if not user_prompt.strip() and not audio_tuple and not image_pil and not video_path:
yield "Please provide a text prompt or at least one multimodal input.", None
return
try:
content: list[dict[str, Any]] = []
if audio_tuple is not None:
audio_np, sr = audio_tuple
content.append(
{"type": "audio_url", "audio_url": {"url": audio_to_base64_data_url(audio_np, sr)}},
)
if image_pil is not None:
img = image_pil if image_pil.mode == "RGB" else image_pil.convert("RGB")
content.append(
{"type": "image_url", "image_url": {"url": image_to_base64_data_url(img)}},
)
if video_path:
content.append(
{"type": "video_url", "video_url": {"url": video_to_base64_data_url(video_path)}},
)
if user_prompt.strip():
content.append({"type": "text", "text": user_prompt})
if enable_tts:
ref_url = endpoint.get_ref_audio_data_url()
system_msg = build_audio_assistant_system(ref_url, language="zh")
else:
system_msg = {
"role": "system",
"content": (
"You are MiniCPM-o, a helpful multimodal assistant that can "
"understand images, audio and video, and respond in text and speech."
),
}
messages = [
system_msg,
{"role": "user", "content": content},
]
kwargs: dict[str, Any] = {
"model": endpoint.model_path,
"messages": messages,
"max_tokens": int(max_tokens),
"temperature": float(temperature),
}
if enable_tts:
extras = endpoint.build_extras()
# Merge extra_body dicts cleanly.
extra_body = dict(extras.get("extra_body", {}))
kwargs["extra_body"] = extra_body
else:
# Text-only: suppress audio output.
kwargs["extra_body"] = {"modalities": ["text"]}
completion = endpoint.client.chat.completions.create(**kwargs)
text_parts: list[str] = []
audio_out: tuple[int, np.ndarray] | None = None
for choice in completion.choices:
msg = choice.message
if getattr(msg, "content", None):
text_parts.append(msg.content)
audio_obj = getattr(msg, "audio", None)
if audio_obj and getattr(audio_obj, "data", None):
raw = base64.b64decode(audio_obj.data)
if len(raw) > 80:
data, sr = sf.read(io.BytesIO(raw))
if data.ndim > 1:
data = data[:, 0]
audio_out = (int(sr), data.astype(np.float32))
text_response = "\n\n".join(t for t in text_parts if t) or "(no text returned)"
yield text_response, audio_out
except Exception as exc: # noqa: BLE001
yield f"Inference failed: {type(exc).__name__}: {exc}", None
# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
DEFAULT_PROMPTS = [
"你好,请用一句话介绍北京。",
"Describe what you see in the image.",
"Transcribe the audio and translate it into English.",
"What is happening in this video?",
]
def build_interface(endpoints: dict[str, ModelEndpoint], default_model: str) -> gr.Blocks:
model_choices = list(endpoints.keys())
css = """
#generate-btn button { width: 100%; font-size: 1.05rem; }
.media-row > div { min-height: 260px; }
.model-banner { color: #555; font-size: 0.9rem; }
"""
with gr.Blocks(css=css, title="vLLM-Omni · MiniCPM-o Demo") as demo:
gr.Markdown(
"# vLLM-Omni · MiniCPM-o Online Demo\n"
"Multimodal chat (text / image / audio / video in; text + speech out) for "
"MiniCPM-o 4.5 served by vllm-omni."
)
with gr.Row():
model_dd = gr.Dropdown(
label="Model",
choices=model_choices,
value=default_model,
scale=2,
)
endpoint_info = gr.Markdown(
elem_classes="model-banner",
value=f"`{endpoints[default_model].api_base}` → `{endpoints[default_model].model_path}`",
)
with gr.Row():
prompt_box = gr.Textbox(
label="Text Prompt",
placeholder="Ask the model anything. Combine with image / audio / video below.",
lines=3,
value=DEFAULT_PROMPTS[0],
scale=3,
)
with gr.Column(scale=1):
tts_checkbox = gr.Checkbox(label="Generate speech output (TTS)", value=True)
max_tokens_slider = gr.Slider(label="Max tokens", minimum=32, maximum=4096, value=1024, step=32)
temperature_slider = gr.Slider(label="Temperature", minimum=0.0, maximum=1.5, value=0.7, step=0.05)
with gr.Row(elem_classes="media-row"):
image_input = gr.Image(label="Image (optional)", type="pil", sources=["upload"])
audio_input = gr.Audio(
label="Audio (optional)",
type="numpy",
sources=["upload", "microphone"],
)
video_input = gr.Video(label="Video (optional)", sources=["upload"])
with gr.Row():
generate_btn = gr.Button("Generate", variant="primary", elem_id="generate-btn")
clear_btn = gr.Button("Clear")
with gr.Row():
text_output = gr.Textbox(label="Text Output", lines=10, scale=2)
audio_output = gr.Audio(label="Audio Output", interactive=False, scale=1)
gr.Examples(
examples=[
["你好,请用一句话介绍北京。"],
["Describe what you see in detail."],
["请先识别这段语音的内容,再用中文回复说话人。"],
["What is happening in this video?"],
],
inputs=[prompt_box],
label="Quick prompts",
)
def _on_model_change(name: str) -> str:
ep = endpoints[name]
return f"`{ep.api_base}` → `{ep.model_path}`"
model_dd.change(_on_model_change, inputs=[model_dd], outputs=[endpoint_info])
def _run(
model_name: str,
prompt: str,
image_pil: Image.Image | None,
audio_in: Any,
video_path: str | None,
enable_tts: bool,
max_tokens: int,
temperature: float,
):
ep = endpoints[model_name]
audio_tuple = process_audio_input(audio_in)
yield from run_inference(
ep,
prompt,
audio_tuple,
image_pil,
video_path,
enable_tts,
max_tokens,
temperature,
)
generate_btn.click(
_run,
inputs=[
model_dd,
prompt_box,
image_input,
audio_input,
video_input,
tts_checkbox,
max_tokens_slider,
temperature_slider,
],
outputs=[text_output, audio_output],
)
def _clear():
return "", None, None, None, "", None
clear_btn.click(
_clear,
outputs=[
prompt_box,
image_input,
audio_input,
video_input,
text_output,
audio_output,
],
)
demo.queue()
return demo
# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--minicpmo45-api-base", default=os.environ.get("MINICPMO45_API_BASE", "http://localhost:8099/v1"))
p.add_argument(
"--minicpmo45-model",
default=os.environ.get("MINICPMO45_MODEL", "openbmb/MiniCPM-o-4_5"),
)
p.add_argument("--host", default="0.0.0.0")
p.add_argument("--port", type=int, default=7862)
p.add_argument("--share", action="store_true")
p.add_argument(
"--ssl-certfile",
default=os.environ.get("GRADIO_SSL_CERTFILE", ""),
help="Path to TLS certificate (PEM). Browsers require HTTPS to access the microphone.",
)
p.add_argument(
"--ssl-keyfile",
default=os.environ.get("GRADIO_SSL_KEYFILE", ""),
help="Path to TLS private key (PEM).",
)
p.add_argument(
"--ssl-verify",
action="store_true",
default=os.environ.get("GRADIO_SSL_VERIFY", "0") == "1",
help="Require browser to verify the certificate (off by default; self-signed certs should stay off).",
)
return p.parse_args()
def _ping(api_base: str, timeout: float = 3.0) -> bool:
import urllib.request
try:
with urllib.request.urlopen(api_base.rstrip("/").replace("/v1", "/health"), timeout=timeout):
return True
except Exception:
return False
def main() -> None:
args = parse_args()
endpoints: dict[str, ModelEndpoint] = {}
if args.minicpmo45_api_base and args.minicpmo45_model:
ok = _ping(args.minicpmo45_api_base)
status = "reachable" if ok else "not reachable (will error on generate)"
print(f"[4.5] {args.minicpmo45_api_base} ({status})")
endpoints[MINICPMO45] = ModelEndpoint(MINICPMO45, args.minicpmo45_api_base, args.minicpmo45_model)
if not endpoints:
raise SystemExit("No endpoints configured. Pass --minicpmo45-api-base/--minicpmo45-model.")
default_model = next(iter(endpoints.keys()))
demo = build_interface(endpoints, default_model)
launch_kwargs: dict = {
"server_name": args.host,
"server_port": args.port,
"share": args.share,
}
if args.ssl_certfile and args.ssl_keyfile:
launch_kwargs.update(
ssl_certfile=args.ssl_certfile,
ssl_keyfile=args.ssl_keyfile,
ssl_verify=args.ssl_verify,
)
print(f"[tls] cert={args.ssl_certfile} key={args.ssl_keyfile} verify={args.ssl_verify}")
demo.launch(**launch_kwargs)
if __name__ == "__main__":
main()
+36
View File
@@ -0,0 +1,36 @@
#!/bin/bash
# Launch the MiniCPM-o 4.5 gradio demo.
#
# Prereq:
# Start a vllm-omni OpenAI server for MiniCPM-o 4.5 on :8099 (see the
# 8x4090 stage config under vllm_omni/model_executor/stage_configs).
set -e
HERE="$(cd "$(dirname "$0")" && pwd)"
: "${MINICPMO45_API_BASE:=http://localhost:8099/v1}"
: "${MINICPMO45_MODEL:=openbmb/MiniCPM-o-4_5}"
: "${GRADIO_HOST:=0.0.0.0}"
: "${GRADIO_PORT:=7862}"
# HTTPS (browsers require a secure context for microphone access).
# Set GRADIO_SSL_CERTFILE / GRADIO_SSL_KEYFILE to enable TLS.
: "${GRADIO_SSL_CERTFILE:=}"
: "${GRADIO_SSL_KEYFILE:=}"
export MINICPMO45_API_BASE MINICPMO45_MODEL
SSL_ARGS=()
if [ -n "$GRADIO_SSL_CERTFILE" ] && [ -n "$GRADIO_SSL_KEYFILE" ] \
&& [ -f "$GRADIO_SSL_CERTFILE" ] && [ -f "$GRADIO_SSL_KEYFILE" ]; then
SSL_ARGS=(--ssl-certfile "$GRADIO_SSL_CERTFILE" --ssl-keyfile "$GRADIO_SSL_KEYFILE")
echo "HTTPS enabled: cert=$GRADIO_SSL_CERTFILE key=$GRADIO_SSL_KEYFILE"
else
echo "HTTPS disabled (cert/key not found). Microphone won't work on remote browsers."
fi
exec python "$HERE/gradio_demo.py" \
--minicpmo45-api-base "$MINICPMO45_API_BASE" \
--minicpmo45-model "$MINICPMO45_MODEL" \
--host "$GRADIO_HOST" \
--port "$GRADIO_PORT" \
"${SSL_ARGS[@]}"
@@ -6,7 +6,8 @@ from typing import NamedTuple
import requests
from openai import OpenAI
from vllm.assets.audio import AudioAsset
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
SEED = 42
@@ -197,7 +198,14 @@ def _build_prompt_for_query_type(
return query_func(custom_prompt=custom_prompt)
if query_type == "use_audio_in_video":
return query_func(video_path=video_path, custom_prompt=custom_prompt)
# use_mixed_modalities / use_multi_audios
if query_type == "use_mixed_modalities":
return query_func(
video_path=video_path,
image_path=image_path,
audio_path=audio_path,
custom_prompt=custom_prompt,
)
# use_multi_audios
return query_func(custom_prompt=custom_prompt)
@@ -478,7 +486,7 @@ def run_multimodal_generation(args, client: OpenAI) -> None:
def parse_args():
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser.add_argument(
"--query-type",
"-q",
+1 -5
View File
@@ -18,10 +18,6 @@ Asynchronous chunk streaming operates as **enabled by default** within this bund
Additionally, NPU, ROCm, and XPU per-platform configuration deltas are deterministically merged from the
`platforms`: section of the corresponding YAML.
**Note:** The OpenAI-style **`/v1/realtime`** WebSocket interface (facilitating streaming PCM audio input alongside audio and transcription output)
is currently **unsupported** while the `async_chunk` configuration attribute is enabled.
It is requisite to instantiate the default omni architecture or utilize a deployment configuration specifying `async_chunk: false` to facilitate real-time streaming sessions.
To explicitly utilize a custom deployment YAML, mandate the configuration path accordingly:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
@@ -289,7 +285,7 @@ python openai_realtime_client.py \
| `--num-requests` | `1` | Number of sequential sessions (see `--concurrency`) |
| `--concurrency` | `1` | Max concurrent WebSocket sessions when `--num-requests` > 1 |
Ensure the server is running **without** `async_chunk` if you use `/v1/realtime`, for example:
Ensure the server is running, for example:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
@@ -6,11 +6,21 @@
vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091
```
With Cache-DiT acceleration:
```bash
vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091 \
--cache-backend cache_dit
```
Or use the convenience script:
```bash
cd examples/online_serving/sensenova_u1
bash run_server.sh
# Cache-DiT
CACHE_BACKEND=cache_dit bash run_server.sh
```
### Tensor Parallelism (TP)
@@ -3,10 +3,18 @@
MODEL="${MODEL:-SenseNova/SenseNova-U1-8B-MoT}"
PORT="${PORT:-8091}"
CACHE_BACKEND="${CACHE_BACKEND:-none}"
echo "Starting SenseNova-U1 server..."
echo "Model: $MODEL"
echo "Port: $PORT"
echo "Cache backend: $CACHE_BACKEND"
EXTRA_ARGS=()
if [ "$CACHE_BACKEND" != "none" ]; then
EXTRA_ARGS+=(--cache-backend "$CACHE_BACKEND")
fi
vllm serve "$MODEL" --omni \
--port "$PORT"
--port "$PORT" \
"${EXTRA_ARGS[@]}"
@@ -18,7 +18,7 @@ For the full list of supported architectures across all modalities, see
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
| 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` | (offline only) | — | — | — |
| OmniVoice | `k2-fsa/OmniVoice` | ✓ | — | — | — |
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
| VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
@@ -321,16 +321,29 @@ vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
### CLI client
```bash
cd examples/online_serving/text_to_speech/omnivoice
# Text-only (auto voice)
python speech_client.py --text "Hello, how are you?"
# Language hint
python speech_client.py --text "Bonjour, comment allez-vous?" --language French
# Voice cloning (reference audio + optional ref_text)
python speech_client.py \
--text "Bonjour, comment allez-vous?" \
--ref-audio /path/to/ref_audio.wav \
--ref-text "Bonjour, comment allez-vous?"
# Style instruction (voice design-style control)
python speech_client.py \
--text "Bonjour, comment allez-vous?" \
--language French \
--instructions "loud voice"
# Deterministic output with seed parameter
python speech_client.py --text "Hello, how are you?" --seed 42
```
The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--output`.
The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--voice`, `--ref-audio`, `--ref-text`, `--instructions`, `--seed`, and `--output`.
### Notes
- Voice cloning and voice design require offline inference; see the [offline OmniVoice section](../../offline_inference/text_to_speech/README.md#omnivoice).
---
## Qwen3-TTS
@@ -399,6 +412,31 @@ curl -X POST http://localhost:8091/v1/audio/voices \
```
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
### Precomputed custom voices
For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
```bash
python qwen3_tts/precompute_custom_voice.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--ref-text "Original transcript of the reference audio" \
--mode icl \
--output-dir /path/to/custom_voices
```
`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
```yaml
custom_voice_dir: /path/to/custom_voices
```
Then start the server with that config and call the Speech API with only the voice name:
```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{"input":"Hello from a precomputed voice.","voice":"alice","task_type":"Base"}' \
--output alice.wav
```
### Streaming PCM
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
@@ -434,7 +472,8 @@ python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-de
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
### Notes
- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
---
@@ -461,6 +500,23 @@ python voxcpm2/openai_speech_client.py \
```
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
### Precomputed custom voices
For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
```bash
python voxcpm2/precompute_custom_voice.py \
--model openbmb/VoxCPM2 \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--mode ref_continuation \
--prompt-text "Original transcript of the reference audio" \
--output-dir /path/to/custom_voices
```
Add the output directory to the deploy config:
```yaml
custom_voice_dir: /path/to/custom_voices
```
After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
### Gradio demo (gapless streaming via AudioWorklet)
```bash
python voxcpm2/gradio_demo.py
@@ -12,6 +12,8 @@ Examples:
"""
import argparse
import base64
import os
import httpx
@@ -19,6 +21,24 @@ DEFAULT_API_BASE = "http://localhost:8091"
DEFAULT_API_KEY = "EMPTY"
def encode_audio_to_base64(audio_path: str) -> str:
"""Encode a local audio file to 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 = {
"wav": "audio/wav",
"mp3": "audio/mpeg",
"flac": "audio/flac",
"ogg": "audio/ogg",
}.get(ext, "audio/wav")
with open(audio_path, "rb") as f:
b64 = base64.b64encode(f.read()).decode("utf-8")
return f"data:{mime};base64,{b64}"
def run_tts(args) -> None:
"""Generate speech via /v1/audio/speech API."""
payload = {
@@ -26,16 +46,36 @@ def run_tts(args) -> None:
"input": args.text,
"response_format": args.response_format,
}
if args.seed is not None:
payload["extra_params"] = {}
payload["extra_params"]["seed"] = args.seed
if args.voice:
payload["voice"] = args.voice
if args.language:
payload["language"] = args.language
if args.ref_audio:
ref = args.ref_audio
if ref.startswith(("http://", "https://", "data:")):
payload["ref_audio"] = ref
else:
payload["ref_audio"] = encode_audio_to_base64(ref)
if args.ref_text:
payload["ref_text"] = args.ref_text
if args.instructions:
payload["instructions"] = args.instructions
print(f"Model: {args.model}")
print(f"Text: {args.text}")
if args.seed:
print(f"Seed: {args.seed}")
if args.voice:
print(f"Voice: {args.voice}")
if args.language:
print(f"Language: {args.language}")
print("Generating audio...")
@@ -45,7 +85,6 @@ def run_tts(args) -> None:
"Content-Type": "application/json",
"Authorization": f"Bearer {args.api_key}",
}
with httpx.Client(timeout=300.0) as client:
response = client.post(api_url, json=payload, headers=headers)
@@ -86,6 +125,30 @@ def main():
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
help="Audio format (default: wav)",
)
parser.add_argument(
"--ref-audio",
type=str,
default=None,
help="Reference audio for voice cloning (local path, URL, or data: URI)",
)
parser.add_argument(
"--ref-text",
type=str,
default=None,
help="Reference text for voice cloning",
)
parser.add_argument(
"--instructions",
type=str,
default=None,
help="Voice style/emotion instructions",
)
parser.add_argument(
"--seed",
type=int,
default=None,
help="Random seed for generation, default: None for stochastic output)",
)
parser.add_argument("--output", "-o", default=None, help="Output file path")
args = parser.parse_args()
run_tts(args)
@@ -0,0 +1,243 @@
"""Pre-compute Qwen3-TTS custom voice profiles.
The generated directory can be passed to the server via
``custom_voice_dir`` in ``vllm_omni/deploy/qwen3_tts.yaml``. Requests can then
use ``/v1/audio/speech`` with ``voice="<name>"`` and no per-request ref_audio.
"""
from __future__ import annotations
import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any
import numpy as np
import torch
from safetensors import safe_open
from safetensors.torch import save_file
REPO_ROOT = Path(__file__).resolve().parents[4]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from vllm_omni.utils.custom_voice_io import safe_voice_stem # noqa: E402
MANIFEST_NAME = "custom_voice_manifest.json"
_SPEAKER_ENCODER_SAMPLE_RATE = 24000
_SPEAKER_ENCODER_NUM_MELS = 128
_SPEAKER_ENCODER_N_FFT = 1024
_SPEAKER_ENCODER_HOP_SIZE = 256
_SPEAKER_ENCODER_WIN_SIZE = 1024
_SPEAKER_ENCODER_FMIN = 0
_SPEAKER_ENCODER_FMAX = 12000
def _resolve_model_dir(model: str) -> str:
if os.path.isdir(model):
return model
from transformers.utils.hub import cached_file
config_path = cached_file(model, "config.json")
if config_path is None:
raise FileNotFoundError(f"Could not resolve config.json for {model}")
return os.path.dirname(config_path)
def _read_audio_mono(path: str) -> tuple[np.ndarray, int]:
import soundfile as sf
wav, sr = sf.read(path, dtype="float32")
if wav.ndim > 1:
wav = wav.mean(axis=-1)
return np.asarray(wav, dtype=np.float32), int(sr)
def _resample(wav: np.ndarray, sr: int, target_sr: int) -> np.ndarray:
if sr == target_sr:
return wav.astype(np.float32)
from vllm.multimodal.audio import AudioResampler
return AudioResampler(target_sr=target_sr).resample(wav.astype(np.float32), orig_sr=int(sr))
def _load_speaker_encoder(model_dir: str, device: torch.device):
from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSSpeakerEncoder
config = Qwen3TTSConfig.from_pretrained(model_dir)
encoder = Qwen3TTSSpeakerEncoder(config.speaker_encoder_config)
index_path = Path(model_dir) / "model.safetensors.index.json"
state: dict[str, torch.Tensor] = {}
if index_path.exists():
index = json.loads(index_path.read_text(encoding="utf-8"))
weight_map = index.get("weight_map", {})
shard_names = sorted({v for k, v in weight_map.items() if k.startswith("speaker_encoder.")})
shard_paths = [Path(model_dir) / name for name in shard_names]
else:
shard_paths = sorted(Path(model_dir).glob("model*.safetensors"))
for shard in shard_paths:
with safe_open(str(shard), framework="pt", device="cpu") as f:
for key in f.keys():
if key.startswith("speaker_encoder."):
state[key.removeprefix("speaker_encoder.")] = f.get_tensor(key)
if not state:
raise RuntimeError(
f"No speaker_encoder.* tensors found under {model_dir}. "
"Use a Qwen3-TTS Base checkpoint for pre-computing custom voice profiles."
)
encoder.load_state_dict(state)
encoder.to(device=device, dtype=torch.bfloat16).eval()
return config, encoder
def _speaker_embedding(
encoder: torch.nn.Module,
speaker_encoder_config: object,
wav: np.ndarray,
sr: int,
device: torch.device,
) -> torch.Tensor:
from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import mel_spectrogram
target_sr = int(getattr(speaker_encoder_config, "sample_rate", _SPEAKER_ENCODER_SAMPLE_RATE))
mel_dim = int(getattr(speaker_encoder_config, "mel_dim", _SPEAKER_ENCODER_NUM_MELS))
if target_sr != _SPEAKER_ENCODER_SAMPLE_RATE or mel_dim != _SPEAKER_ENCODER_NUM_MELS:
raise ValueError(
"Unsupported Qwen3-TTS speaker encoder mel config for precompute: "
f"sample_rate={target_sr}, mel_dim={mel_dim}. "
f"Expected sample_rate={_SPEAKER_ENCODER_SAMPLE_RATE}, mel_dim={_SPEAKER_ENCODER_NUM_MELS} "
"to match the serving-time speaker embedding path."
)
wav = _resample(wav, sr, target_sr)
wav_t = torch.from_numpy(wav).to(device=device, dtype=torch.float32)
mels = mel_spectrogram(
wav_t.unsqueeze(0),
n_fft=_SPEAKER_ENCODER_N_FFT,
num_mels=_SPEAKER_ENCODER_NUM_MELS,
sampling_rate=_SPEAKER_ENCODER_SAMPLE_RATE,
hop_size=_SPEAKER_ENCODER_HOP_SIZE,
win_size=_SPEAKER_ENCODER_WIN_SIZE,
fmin=_SPEAKER_ENCODER_FMIN,
fmax=_SPEAKER_ENCODER_FMAX,
).transpose(1, 2)
with torch.inference_mode():
return encoder(mels.to(device=device, dtype=torch.bfloat16))[0].float().cpu().contiguous()
def _ref_code(model_dir: str, wav: np.ndarray, sr: int, device: torch.device) -> torch.Tensor:
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
tok = Qwen3TTSTokenizer.from_pretrained(
str(Path(model_dir) / "speech_tokenizer"),
torch_dtype=torch.bfloat16,
)
try:
del tok.model.decoder
tok.model.decoder = None
tok.model.encoder.to(device)
tok.device = device
except Exception:
tok.device = device
with torch.inference_mode():
enc = tok.encode(wav, sr=int(sr), return_dict=True)
codes = getattr(enc, "audio_codes", None)
if isinstance(codes, list):
codes = codes[0] if codes else None
if not isinstance(codes, torch.Tensor):
raise RuntimeError("Speech tokenizer did not return audio_codes")
if codes.ndim == 3:
codes = codes[0]
return codes.to(dtype=torch.int32, device="cpu").contiguous()
def _load_manifest(output_dir: Path, model: str, hidden_size: int) -> dict[str, Any]:
path = output_dir / MANIFEST_NAME
if path.exists():
return json.loads(path.read_text(encoding="utf-8"))
return {
"schema_version": 1,
"model_type": "qwen3_tts",
"model": model,
"hidden_size": hidden_size,
"voices": {},
}
def _write_voice(
*,
model: str,
output_dir: Path,
voice_name: str,
ref_audio: str,
ref_text: str | None,
mode: str,
speaker_description: str | None,
device: torch.device,
) -> None:
model_dir = _resolve_model_dir(model)
config, encoder = _load_speaker_encoder(model_dir, device)
wav, sr = _read_audio_mono(ref_audio)
if wav.size < 1024:
raise ValueError(f"Reference audio too short: {wav.size} samples")
tensors: dict[str, torch.Tensor] = {
"speaker_embedding": _speaker_embedding(encoder, config.speaker_encoder_config, wav, sr, device),
}
if mode == "icl":
if not ref_text or not ref_text.strip():
raise ValueError("--ref-text is required for --mode icl")
tensors["ref_code"] = _ref_code(model_dir, wav, sr, device)
output_dir.mkdir(parents=True, exist_ok=True)
filename = f"{safe_voice_stem(voice_name)}.safetensors"
save_file(tensors, str(output_dir / filename))
hidden_size = int(getattr(config.talker_config, "hidden_size", tensors["speaker_embedding"].numel()))
manifest = _load_manifest(output_dir, model, hidden_size)
entry: dict[str, Any] = {
"name": voice_name,
"file": filename,
"mode": mode,
"embedding_dim": int(tensors["speaker_embedding"].numel()),
}
if "ref_code" in tensors:
entry["ref_code_length"] = int(tensors["ref_code"].shape[0])
if ref_text:
entry["ref_text"] = ref_text
if speaker_description:
entry["speaker_description"] = speaker_description
manifest.setdefault("voices", {})[voice_name] = entry
(output_dir / MANIFEST_NAME).write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
print(f"Wrote {output_dir / filename}")
print(f"Updated {output_dir / MANIFEST_NAME}")
def main() -> None:
parser = argparse.ArgumentParser(description="Pre-compute Qwen3-TTS custom voice profile")
parser.add_argument("--model", required=True, help="Qwen3-TTS Base model path or Hugging Face ID")
parser.add_argument("--voice-name", required=True)
parser.add_argument("--ref-audio", required=True)
parser.add_argument("--ref-text", default=None)
parser.add_argument("--mode", choices=["xvec", "icl"], default="xvec")
parser.add_argument("--speaker-description", default=None)
parser.add_argument("--output-dir", required=True)
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
args = parser.parse_args()
_write_voice(
model=args.model,
output_dir=Path(args.output_dir),
voice_name=args.voice_name,
ref_audio=args.ref_audio,
ref_text=args.ref_text,
mode=args.mode,
speaker_description=args.speaker_description,
device=torch.device(args.device),
)
if __name__ == "__main__":
main()
@@ -0,0 +1,128 @@
"""Pre-compute VoxCPM2 custom voice profiles.
The generated directory can be passed to the server via
``custom_voice_dir`` in ``vllm_omni/deploy/voxcpm2.yaml``. Requests can then
use ``/v1/audio/speech`` with ``voice="<name>"`` and no per-request ref_audio.
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
from typing import Any
import torch
from safetensors.torch import save_file
REPO_ROOT = Path(__file__).resolve().parents[4]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from vllm_omni.utils.custom_voice_io import safe_voice_stem # noqa: E402
MANIFEST_NAME = "custom_voice_manifest.json"
def _load_tts(model: str, device: torch.device):
from vllm_omni.model_executor.models.voxcpm2.voxcpm2_import_utils import import_voxcpm2_core
VoxCPM = import_voxcpm2_core()
native = VoxCPM.from_pretrained(model, load_denoiser=False, optimize=False)
return native.tts_model.to(device).eval()
def _load_manifest(output_dir: Path, model: str) -> dict[str, Any]:
path = output_dir / MANIFEST_NAME
if path.exists():
return json.loads(path.read_text(encoding="utf-8"))
return {
"schema_version": 1,
"model_type": "voxcpm2",
"model": model,
"voices": {},
}
def _write_voice(
*,
model: str,
output_dir: Path,
voice_name: str,
ref_audio: str,
prompt_text: str | None,
mode: str,
speaker_description: str | None,
device: torch.device,
) -> None:
if mode in ("continuation", "ref_continuation") and not prompt_text:
raise ValueError("--prompt-text is required for continuation/ref_continuation modes")
tts = _load_tts(model, device)
tensors: dict[str, torch.Tensor] = {}
with torch.inference_mode():
if mode in ("reference", "ref_continuation"):
tensors["ref_audio_feat"] = tts._encode_wav(ref_audio, padding_mode="right").float().cpu().contiguous()
if mode in ("continuation", "ref_continuation"):
tensors["audio_feat"] = tts._encode_wav(ref_audio, padding_mode="left").float().cpu().contiguous()
output_dir.mkdir(parents=True, exist_ok=True)
filename = f"{safe_voice_stem(voice_name)}.safetensors"
save_file(tensors, str(output_dir / filename))
manifest = _load_manifest(output_dir, model)
entry: dict[str, Any] = {
"name": voice_name,
"file": filename,
"mode": mode,
}
if "ref_audio_feat" in tensors:
entry["ref_audio_feat_len"] = int(tensors["ref_audio_feat"].shape[0])
if "audio_feat" in tensors:
entry["audio_feat_len"] = int(tensors["audio_feat"].shape[0])
if prompt_text:
entry["prompt_text"] = prompt_text
if speaker_description:
entry["speaker_description"] = speaker_description
manifest.setdefault("voices", {})[voice_name] = entry
(output_dir / MANIFEST_NAME).write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
print(f"Wrote {output_dir / filename}")
print(f"Updated {output_dir / MANIFEST_NAME}")
def main() -> None:
parser = argparse.ArgumentParser(description="Pre-compute VoxCPM2 custom voice profile")
parser.add_argument("--model", default="openbmb/VoxCPM2", help="VoxCPM2 model path or Hugging Face ID")
parser.add_argument("--voice-name", required=True)
parser.add_argument("--ref-audio", required=True)
parser.add_argument(
"--prompt-text",
default=None,
help="Transcript of ref audio for continuation/ref_continuation modes",
)
parser.add_argument(
"--mode",
choices=["reference", "continuation", "ref_continuation"],
default="reference",
)
parser.add_argument("--speaker-description", default=None)
parser.add_argument("--output-dir", required=True)
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
args = parser.parse_args()
_write_voice(
model=args.model,
output_dir=Path(args.output_dir),
voice_name=args.voice_name,
ref_audio=args.ref_audio,
prompt_text=args.prompt_text,
mode=args.mode,
speaker_description=args.speaker_description,
device=torch.device(args.device),
)
if __name__ == "__main__":
main()
+23
View File
@@ -76,6 +76,18 @@ local = [
"FlagEmbedding",
]
# Optional runtime dependencies for the MiniCPM-o 4.5 talker pipeline.
# The talker loads MiniCPMTTS via `trust_remote_code` from the model repo,
# and that module performs `from stepaudio2 import Token2wav` against the
# MiniCPM-o-flavored Token2wav vocoder (PyPI: `stepaudio2-minicpmo`, NOT
# upstream `stepfun-ai/Step-Audio2` — the upstream signature rejects
# `n_timesteps` and crashes at construction). A clean `vllm-omni` install
# does not pull this in by default to avoid forcing the dep onto users
# who only need other models; opt in with `pip install vllm-omni[minicpmo]`.
minicpmo = [
"stepaudio2-minicpmo",
]
docs = [
"mkdocs>=1.5.0",
"mkdocs-api-autonav",
@@ -101,6 +113,13 @@ Documentation = "https://vllm-omni.readthedocs.io"
[project.scripts]
vllm-omni = "vllm_omni.entrypoints.cli.main:main"
# Register vllm-omni's model registrations on every vllm subprocess. vllm
# auto-loads anything advertised under `vllm.general_plugins` in its own
# `load_general_plugins()`, so worker subprocesses spawned by vllm (which
# only import vllm, not vllm_omni) will still pick up the omni archs.
[project.entry-points."vllm.general_plugins"]
vllm_omni_register_models = "vllm_omni.engine.arg_utils:register_omni_models_to_vllm"
[tool.setuptools.packages.find]
where = ["."]
@@ -149,6 +168,10 @@ ignore = [
[tool.ruff.lint.per-file-ignores]
"examples/**" = ["E501"] # Allow long lines in examples
"tests/**" = ["E501"] # Allow long lines in tests
# MiniCPM-o reference modeling code vendored from upstream HF/transformers/MiniCPM-o.
# These files keep upstream conventions (star imports, long lines, math-symbol
# argument names) so they stay easy to diff against upstream releases.
"vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py" = ["E402", "E501", "F403", "F405", "N803"]
[tool.ruff.lint.flake8-tidy-imports.banned-api]
"librosa".msg = "The librosa module is banned, use vllm.multimodal helpers instead"
+177
View File
@@ -0,0 +1,177 @@
# Lance
> Unified autoregressive + diffusion multimodal (text / image / video)
## Summary
- Vendor: ByteDance
- Model: [`bytedance-research/Lance`](https://huggingface.co/bytedance-research/Lance) (`Lance_3B`, `Lance_3B_Video`)
- Task: text2img, text2video, img2img (image edit), video2video (video edit),
img2text (image understanding), video2text (video understanding)
- Mode: Offline inference, Online serving (OpenAI-compatible API)
- Maintainer: Community
## When to use this recipe
Use this recipe to run Lance — a 3B unified AR + diffusion multimodal model
on a Qwen2.5-VL backbone — via vLLM-Omni. Lance is **BAGEL-lineage**: the
released checkpoint uses the same `*_moe_gen` Mixture-of-Transformers weight
layout as BAGEL, so vLLM-Omni reuses BAGEL's transformer core and specializes
only the ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the checkpoint layout.
All six single-stage modalities are supported on the `Lance_3B` checkpoint;
`text2video` and `video2video` additionally need `Lance_3B_Video` for the 3-D
`latent_pos_embed` table. The two-stage AR + DiT topology is deferred and
needs `LanceConfig` / `LanceProcessor` registered in the `vllm` package (a
separate upstream PR).
## References
- Offline example:
[`examples/offline_inference/lance/end2end.py`](../../examples/offline_inference/lance/end2end.py)
- Online serving:
[`examples/online_serving/lance/`](../../examples/online_serving/lance/)
- Deploy config:
[`vllm_omni/deploy/lance.yaml`](../../vllm_omni/deploy/lance.yaml)
- E2E tests:
[`tests/e2e/online_serving/test_lance.py`](../../tests/e2e/online_serving/test_lance.py)
- HuggingFace model page:
[bytedance-research/Lance](https://huggingface.co/bytedance-research/Lance)
- PR: [vllm-project/vllm-omni#3710](https://github.com/vllm-project/vllm-omni/pull/3710)
(closes [#3697](https://github.com/vllm-project/vllm-omni/issues/3697))
## Hardware Support
## GPU
### 1x NVIDIA B300 / A100 80GB
#### Environment
- OS: Linux
- Python: 3.12
- Driver / runtime: CUDA ≥ 12.4
- vLLM-Omni version: 0.18.x.dev
#### Command — text-to-image (default)
```bash
python examples/offline_inference/lance/end2end.py \
--model bytedance-research/Lance \
--prompts "a corgi astronaut on the moon, cinematic" \
--steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
--height 1024 --width 1024 \
--seed 42 --output ./out
```
Defaults match upstream `inference_lance.sh`: 30 denoising steps,
`timestep-shift 3.5`, text CFG 4.0, seed 42, 1024×1024.
#### Command — image edit (img2img)
```bash
python examples/offline_inference/lance/end2end.py \
--model bytedance-research/Lance --modality img2img \
--image-path /path/to/input.png \
--prompts "Convert this into a vibrant cartoon-style illustration" \
--steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
--output ./out
```
The Lance-native VAE prefill scatters Wan2.2 latents into the LLM query
sequence; no separate image encoder is needed.
#### Command — text-to-video
```bash
python examples/offline_inference/lance/end2end.py \
--model bytedance-research/Lance/Lance_3B_Video --modality text2video \
--num-frames 25 --video-height 480 --video-width 768 \
--prompts "a cat playing piano, cinematic" \
--steps 30 --fps 8 --output ./out
```
Use the `Lance_3B_Video` subfolder for any video path so the 3-D
`latent_pos_embed` table is loaded; image / understanding paths can point at
the top-level repo and resolve the right sub-checkpoint automatically.
#### Command — image / video understanding
```bash
# Image → text (caption / VQA)
python examples/offline_inference/lance/end2end.py \
--model bytedance-research/Lance --modality img2text \
--image-path /path/to/photo.jpg \
--prompts "Describe this image in detail." \
--do-sample --text-temperature 0.8
# Video → text
python examples/offline_inference/lance/end2end.py \
--model bytedance-research/Lance --modality video2text \
--video-path /path/to/clip.mp4 \
--prompts "What is happening in this video?"
```
Sampling is enabled by default at `--text-temperature 0.8` for the
understanding paths because Lance's greedy decoder emits an immediate EOS for
many prompts.
#### Verification
```bash
pytest -s -v tests/e2e/online_serving/test_lance.py
```
#### Notes
- BF16 footprint: ~7 GB LLM + Qwen2.5-VL ViT + Wan2.2 VAE; comfortably fits
on a single 16 GB+ GPU for `Lance_3B`.
- `rope_scaling = {"type": "mrope", "mrope_section": [16, 24, 24]}` is wired
through `BagelRotaryEmbedding`. `text2img` uses scalar position ids
(BAGEL-equivalent); `img2text` / `video` / edit paths thread per-axis 3-D
position ids.
- `video_edit` quality is more abstract than `text2video` at the same
resolution (known position-id offset between VAE-prefill and gen-latent
blocks). Functionally correct end-to-end.
## Online Serving
Lance supports all single-stage modalities via the OpenAI-compatible
`/v1/chat/completions` API.
### Launch
```bash
bash examples/online_serving/lance/run_server.sh
# or, with overrides
MODEL=bytedance-research/Lance \
DEPLOY_CONFIG=vllm_omni/deploy/lance.yaml \
PORT=8091 \
bash examples/online_serving/lance/run_server.sh
```
For `text2video` / `video2video`, set
`MODEL=bytedance-research/Lance/Lance_3B_Video`.
### Send requests
```bash
# Text-to-image
python examples/online_serving/lance/openai_chat_client.py \
--prompt "A cute corgi astronaut on the moon, cinematic" \
--modality text2img --output corgi.png
# Image edit
python examples/online_serving/lance/openai_chat_client.py \
--prompt "Convert this into a vibrant cartoon-style illustration" \
--modality img2img --image-url path/to/photo.png \
--output edited.png
# Image understanding
python examples/online_serving/lance/openai_chat_client.py \
--prompt "Describe this image" \
--modality img2text --image-url photo.jpg
```
The client is shared with BAGEL — same OpenAI message format, same
`modalities` / `num_inference_steps` / `seed` / `height` / `width` knobs.
+112
View File
@@ -0,0 +1,112 @@
# GLM-Image for text-to-image and image editing
## Summary
- Vendor: Z.ai
- Model: `GLM/GLM-Image`
- Task: Text-to-image (T2I) and image-to-image
- Mode: Online serving with the OpenAI-compatible API
- Maintainer: Community
## When to use this recipe
Use this recipe when you want a known-good starting point for serving
`GLM/GLM-Image` with vLLM-Omni on **two 80 GB NVIDIA A800** GPUs (Ampere-class,
same default layout as the upstream **2×A100 80GB** example: Stage 0 AR on GPU 0,
Stage 1 diffusion on GPU 1) and validate the deployment with the existing
`examples/online_serving/glm_image` clients.
## References
- Upstream or canonical docs:
[`docs/user_guide/examples/online_serving/glm_image.md`](../../docs/user_guide/examples/online_serving/glm_image.md)
- Related issue or discussion:
[#2888](https://github.com/vllm-project/vllm-omni/pull/2888)
## Hardware Support
This recipe documents **dual-GPU** CUDA layouts on A800 80 GB
for the same software stack. Add more platforms (for example ROCm / NPU) as
community validation lands.
## GPU
### 2× A800 80GB
#### Environment
These versions were taken from a working **editable** install: activate `vllm-omni/.venv` (or your equivalent), then align `pip` / Git with the rows below when reproducing this recipe.
- OS: Linux
- Python: 3.12
- Driver / runtime: NVIDIA CUDA stack with **two** A800 80 GB GPUs visible (set `CUDA_VISIBLE_DEVICES` on your host if needed)
- vLLM: **0.19.0**
- vLLM-Omni: **0.19.0rc2.dev138+g38d5f2d53** (editable install from this repo; Git **`38d5f2d5`**, `git describe` ≈ **`v0.19.0rc1-138-g38d5f2d5`**)
- Transformers: **5.5.4** (same `.venv` as above; required so `glm_image` configs load for Stage 0)
#### Command
Start the server from the repository root:
```bash
vllm serve zai-org/GLM-Image --omni --port 8091
```
To use the bundled stage config explicitly (same default as above):
```bash
vllm serve zai-org/GLM-Image \
--omni \
--port 8091 \
--deploy-config vllm_omni/deploy/glm_image.yaml
```
#### Verification
Run one of the existing example clients after the server is ready:
```bash
curl -s http://172.18.69.133:8000/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"messages": [
{"role": "user", "content": "A beautiful landscape painting"}
],
"extra_body": {
"height": 1920,
"width": 1920,
"num_inference_steps": 50,
"true_cfg_scale": 1.5,
"seed": 42
}
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > land.png
```
After the command finishes, check for the output files:
```bash
ls output.png
```
#### Sample end-to-end metrics
One representative **offline** GLM-Image E2E run on this recipe’s **2× A800 80GB**.
Overall summary from the run’s metrics. Rough wall-time split: **Stage 0 (AR)** ~**25 s**,
**Stage 1 (diffusion)** ~**34 s** (see `e2e_stage_*_wall_time_ms` below).
| Field | Value |
| --- | ---: |
| e2e_requests | 1 |
| e2e_wall_time_ms | 61,148.679 |
| e2e_total_tokens | 1,300 |
| e2e_avg_time_per_request_ms | 61,148.679 |
| e2e_avg_tokens_per_s | 21.260 |
| e2e_stage_0_wall_time_ms | 24,708.760 |
| e2e_stage_1_wall_time_ms | 33,787.442 |
#### Notes
- Memory usage: Roughly **~38 GiB + KV** on Stage 0 (AR) and **~20 GiB** on Stage 1 (DiT+VAE) per the user guide; two 80 GB cards match the default split.
- Key flags: `--omni` is required; `--stage-configs-path` is optional unless you use a custom YAML (for example single-GPU).
- Keep **Transformers ≥ 5.5.1** (this recipe used **5.5.4**) so `glm_image` configs resolve; otherwise Stage 0 can fail at `ModelConfig` validation.
- Known limitations: This starter recipe follows the dual-GPU online path documented under `examples/online_serving/glm_image`. The first request may be slower due to warmup.
- Generation time: about **61 s** wall time end-to-end for the sample above (50 inference steps, 1024×1024).
+109
View File
@@ -0,0 +1,109 @@
# Helios for video generation
## Summary
- Vendor: Helios
- Model: `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled`
- Task: Text-to-video, image-to-video, and video-to-video generation
- Mode: Offline inference with the bundled Helios example
- Maintainer: Community
## When to use this recipe
Use this recipe when you want a known-good starting point for running Helios
video generation with vLLM-Omni. The concrete command below focuses on
`BestWishYsh/Helios-Base` text-to-video generation on one NVIDIA H20 GPU. The
same example script also supports Helios-Mid, Helios-Distilled, image-to-video,
and video-to-video; see the linked example docs for those model-specific flags.
## References
- Upstream repository: <https://github.com/PKU-YuanGroup/Helios>
- Model weights:
<https://huggingface.co/BestWishYsh/Helios-Base>
- Upstream or canonical docs:
[`docs/user_guide/examples/offline_inference/helios.md`](../../docs/user_guide/examples/offline_inference/helios.md)
- Related example under `examples/`:
[`examples/offline_inference/helios/README.md`](../../examples/offline_inference/helios/README.md)
## Hardware Support
## GPU
### 1x NVIDIA H20
#### Environment
- OS: Ubuntu Linux x86_64
- Python: 3.12.12
- Driver / runtime: NVIDIA driver `580.126.20`, CUDA `13.0`
- Hardware: 1x NVIDIA H20 GPU from an 8x H20 host
- vLLM version: `0.19.0`
- vLLM-Omni version or commit: `a3903810`
#### Command
Run the baseline Helios-Base text-to-video example from the repository root:
```bash
cd examples/offline_inference/helios
python end2end.py \
--model BestWishYsh/Helios-Base \
--sample-type t2v \
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
--guidance-scale 5.0 \
--output helios_t2v_base.mp4
```
To use cache-dit acceleration, run on a vLLM-Omni checkout that includes Helios cache-dit support and enable the cache backend:
```bash
cd examples/offline_inference/helios
python end2end.py \
--cache-backend cache_dit \
--enable-cache-dit-summary \
--model BestWishYsh/Helios-Base \
--sample-type t2v \
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
--guidance-scale 5.0 \
--output helios_t2v_base.mp4
```
#### Verification
The script should print the resolved generation configuration, total generation
time, and output path. A successful run ends with output similar to:
```text
Total generation time: <seconds> seconds (<milliseconds> ms)
Saved generated video to helios_t2v_base.mp4
```
#### Important flags and stage configs
- `--model BestWishYsh/Helios-Base` selects the base Helios checkpoint used by
this recipe.
- `--sample-type t2v` selects text-to-video generation. Use `i2v` with
`--image-path` for image-to-video, or `v2v` with `--video-path` for
video-to-video.
- `--guidance-scale 5.0` matches the Helios-Base recommendation. Use
`--guidance-scale 1.0` for Helios-Distilled.
- `--cache-backend cache_dit` enables the cache-dit acceleration path.
- `--enable-cache-dit-summary` prints cache-dit summary information after
diffusion forward passes.
- No separate stage config is required for this offline recipe; the bundled
Helios example configures the pipeline through `end2end.py` arguments.
- For Helios-Mid, add `--is-enable-stage2`,
`--pyramid-num-inference-steps-list 20 20 20`, `--use-cfg-zero-star`,
`--use-zero-init`, and `--zero-steps 1`.
- For Helios-Distilled, add `--is-enable-stage2`,
`--pyramid-num-inference-steps-list 2 2 2`, and
`--is-amplify-first-chunk`.
#### Known limitations
- Helios generates video in 33-frame chunks. For best performance, set
`--num-frames` to a multiple of `33`; non-multiple values are rounded up to
the nearest multiple of `33`.
+229
View File
@@ -0,0 +1,229 @@
# MiniCPM-o 4.5
> Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech)
## Summary
- Vendor: OpenBMB
- Model: [`openbmb/MiniCPM-o-4_5`](https://huggingface.co/openbmb/MiniCPM-o-4_5)
- Task: Omni multimodal chat — accepts text / image / audio / video input;
emits text and 24 kHz mono speech in the same response
- Mode: Online serving via the OpenAI-compatible `/v1/chat/completions`
API, plus a bundled Gradio demo (text + speech UI)
- Maintainer: [`@tc-mb`](https://github.com/tc-mb) (MiniCPM-V / MiniCPM-o team)
## When to use this recipe
Use this recipe as a known-good starting point for serving
`openbmb/MiniCPM-o-4_5` on vLLM-Omni. MiniCPM-o 4.5 is the omni member
of the MiniCPM-o family — it pairs a multimodal-understanding thinker
LLM with a streaming `MiniCPMTTS + Token2Wav` talker so a single
`/v1/chat/completions` call can return text and 24 kHz speech in one
shot. The recipe covers three shipped GPU layouts (2 / 3 / 8 GPUs)
selected via `--deploy-config`.
## References
- Default deploy configs (auto-loaded by HF `model_type=minicpmo` +
`hf_config.version="4.5"`):
- 2-GPU layout (default):
[`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
- 3-GPU layout (thinker TP=2):
[`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml)
- 8x RTX 4090 layout:
[`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml)
- Online example + Gradio demo:
[`examples/online_serving/minicpmo/`](../../examples/online_serving/minicpmo/)
- Pipeline / talker source:
[`vllm_omni/model_executor/models/minicpmo_4_5/`](../../vllm_omni/model_executor/models/minicpmo_4_5/)
- Stage-input processor (thinker → talker bridge):
[`vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py`](../../vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py)
- Upstream model card:
[`openbmb/MiniCPM-o-4_5`](https://huggingface.co/openbmb/MiniCPM-o-4_5)
- Integration PR:
[vllm-project/vllm-omni#3642](https://github.com/vllm-project/vllm-omni/pull/3642)
## Hardware Support
Three GPU layouts ship with default deploy configs. Pick the layout that
matches your hardware and pass it via `--deploy-config`; the talker
(`MiniCPMTTS + Token2Wav`) always lives on its own GPU because of the
in-process vocoder, and the thinker is the part that scales out via TP.
| Layout | Thinker | Talker + Token2Wav | Typical hardware |
| --- | --- | --- | --- |
| 2-GPU (default) | GPU 0 | GPU 1 | 2x A100/H100/H200 80GB |
| 3-GPU (thinker TP=2) | GPU 0,1 (TP=2) | GPU 2 | 3x mid-tier GPUs |
| 8x RTX 4090 24GB | GPU 0–3 (TP=4) | GPU 4 | 8x RTX 4090 consumer |
## GPU
### 2 x GPU (default — single command)
The default
[`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
puts the thinker on GPU 0 (`~70 %` memory, `enforce_eager: true`,
`max_num_seqs: 1`) and the talker + Token2Wav vocoder on GPU 1
(`~75 %` memory). This is the recommended starting layout — works on
any pair of 80GB-class GPUs (A100, H100, H200) and on most 40GB+
pairs as long as the thinker model weights fit.
#### Environment
- OS: Linux
- Python: 3.10+
- vLLM / vLLM-Omni: >= 0.21.0 (or current `main`)
- Optional Talker dep: `stepaudio2-minicpmo` (see Notes for why this is
required and how to install it)
#### Command
```bash
vllm serve openbmb/MiniCPM-o-4_5 --omni \
--trust-remote-code \
--host 0.0.0.0 --port 8099
```
The deploy config is auto-loaded by the model registry — no
`--deploy-config` flag needed for this default 2-GPU layout.
#### Verification
**Quick smoke test (text-only output)**:
```bash
curl http://localhost:8099/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "openbmb/MiniCPM-o-4_5",
"messages": [{"role": "user", "content": "Briefly introduce yourself."}],
"modalities": ["text"]
}'
```
**Text + speech in one response** (the headline 4.5 feature). The TTS
path is gated by a Jinja flag on the chat template, passed through
`extra_body.chat_template_kwargs.use_tts_template=true`:
```bash
curl http://localhost:8099/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "openbmb/MiniCPM-o-4_5",
"messages": [{"role": "user", "content": "Say hello, then introduce vLLM in one sentence."}],
"modalities": ["text", "audio"],
"extra_body": {"chat_template_kwargs": {"use_tts_template": true}}
}'
```
Response carries text in `choices[0].message.content` and base64 WAV
in `choices[0].message.audio.data` (24 kHz mono, see Notes).
**Gradio demo (text + image + audio + video UI)**:
```bash
bash examples/online_serving/minicpmo/run_gradio_demo.sh
# or run the python entry point directly:
python examples/online_serving/minicpmo/gradio_demo.py \
--minicpmo45-api-base http://localhost:8099/v1 \
--minicpmo45-model openbmb/MiniCPM-o-4_5 \
--port 7862
```
Open `http://<host>:7862` and try a text prompt with the **"Generate
speech output (TTS)"** checkbox on / off.
#### Notes
- Memory budget: thinker weights occupy GPU 0 at `gpu_memory_utilization:
0.7`; talker + Token2Wav vocoder share GPU 1 at `0.75`.
- `--trust-remote-code` is required — the HF repo ships a custom
`MiniCPMO` config / model class.
- Pin: `enforce_eager: true` on both stages (CUDA graph capture is off
by design for the talker's Token2Wav path).
- Stage 1 (talker) is hard-capped to `max_num_seqs: 1`: the talker
only consumes `runtime_additional_information[0]`, so any value > 1
makes concurrent requests share request-0's audio. This is the same
cap baked into the deploy config.
### 3 x GPU (thinker TP=2)
Use
[`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml)
when you have a third GPU available and want the thinker on 2-way
tensor parallel for higher throughput; the talker stays on its own
GPU (talker has its own in-process Token2Wav vocoder, so co-locating
it with the thinker risks OOM under load).
#### Command
```bash
vllm serve openbmb/MiniCPM-o-4_5 --omni \
--deploy-config vllm_omni/deploy/minicpmo_4_5_3gpu.yaml \
--trust-remote-code \
--host 0.0.0.0 --port 8099
```
Verification and Notes mirror the 2-GPU section; thinker latency
roughly halves under load thanks to TP=2.
### 8 x RTX 4090 24GB (consumer-GPU layout)
Use
[`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml)
on an 8x RTX 4090 host. Thinker uses 4-way TP across GPUs 0–3
(`~85 %` mem each ≈ 20.4 GiB/card), talker + Token2Wav lives on GPU 4
(`~90 %` mem). GPUs 5–7 are left free.
#### Command
```bash
vllm serve openbmb/MiniCPM-o-4_5 --omni \
--deploy-config vllm_omni/deploy/minicpmo_4_5_8x4090.yaml \
--trust-remote-code \
--host 0.0.0.0 --port 8099
```
#### Notes
- `max_model_len` is capped at 4096 in this layout — 8192 still OOMs on
4090s. Raise it if your cards have more headroom (e.g. 4090 D /
custom 32 GB SKUs), but verify with a long-prompt run before
promoting.
- All other knobs match the 2-GPU section; the only difference is the
per-card memory pressure on the thinker shards.
## Notes (applies to all layouts)
- **Talker dependency**: the `MiniCPM-o 4.5` talker calls
`from stepaudio2 import Token2wav` against the MiniCPM-o-flavored
vocoder (PyPI package `stepaudio2-minicpmo` — NOT the upstream
`stepfun-ai/Step-Audio2`, whose `Token2wav.__init__` signature
rejects `n_timesteps`). Install via the published extra:
```bash
pip install 'vllm-omni[minicpmo]'
```
Equivalent direct install: `pip install stepaudio2-minicpmo`. A
missing dep raises `ImportError` at first request with the same
install hint instead of silently emitting empty audio.
- **TTS trigger**: speech output is only emitted when the client passes
`extra_body.chat_template_kwargs.use_tts_template=true`. Without it,
the response is text-only (which is also faster).
- **Output audio**: 24 kHz mono WAV inside the OpenAI-style
`message.audio.data` (base64). The Gradio demo's WAV player decodes
this automatically.
- **Routing**: MiniCPM-o 4.5 and 2.6 both ship `architectures=
["MiniCPMO"]` in HF config; routing is disambiguated by
`hf_config.version == "4.5"` via the
`hf_config_predicate` on the 4.5 pipeline. A 2.6 checkpoint loaded
with this recipe's `--deploy-config` will be rejected at startup
rather than silently misrouted.
- **Async chunking**: disabled in all three deploy configs
(`async_chunk: false`) — the talker batches a single full thinker
output, not chunks.
+249
View File
@@ -0,0 +1,249 @@
# VoxCPM2
> Online + offline TTS with native AR pipeline (48 kHz, 30+ languages)
## Summary
- Vendor: OpenBMB
- Model: `openbmb/VoxCPM2`
- Task: Text-to-speech (zero-shot synthesis and reference-audio voice cloning)
- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API, plus
an offline end-to-end example
- Maintainer: Community
## When to use this recipe
Use this recipe as a known-good starting point for serving `openbmb/VoxCPM2`
on vLLM-Omni on consumer-class GPUs. VoxCPM2 is a 2B-parameter native AR
TTS model that emits 48 kHz mono audio and supports 30+ languages out of
the box. It runs as a single-stage pipeline (`MiniCPM4 base LM → FSQ →
MiniCPM4 residual LM → LocDiT CFM solver → AudioVAE`) and supports both
zero-shot synthesis and reference-audio voice cloning, with optional
gapless streaming via the bundled Gradio AudioWorklet demo.
Because the model is only 2B params, a single 24 GB consumer GPU is
sufficient — no tensor parallelism is required.
## References
- Example guide (online):
[`examples/online_serving/text_to_speech/README.md`](../../examples/online_serving/text_to_speech/README.md#voxcpm2)
- Online client / Gradio demo:
[`examples/online_serving/text_to_speech/voxcpm2/`](../../examples/online_serving/text_to_speech/voxcpm2/)
- Offline end-to-end script:
[`examples/offline_inference/text_to_speech/voxcpm2/end2end.py`](../../examples/offline_inference/text_to_speech/voxcpm2/end2end.py)
- Default deploy config:
[`vllm_omni/deploy/voxcpm2.yaml`](../../vllm_omni/deploy/voxcpm2.yaml)
- Talker / pipeline source:
[`vllm_omni/model_executor/models/voxcpm2/`](../../vllm_omni/model_executor/models/voxcpm2/)
- Speech API reference:
[`docs/serving/speech_api.md`](../../docs/serving/speech_api.md#voxcpm2)
- Related issue or discussion:
[RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
## Hardware Support
This recipe documents one tested 24 GB consumer-GPU configuration.
Larger-VRAM (H20 / H100 / A100) and other vendor sections (ROCm, NPU) are
welcome as community validation lands.
## GPU
### 1 x RTX 4090 24GB (Single GPU, Minimum Recommended)
VoxCPM2 (2B params in bfloat16) fits comfortably on a single 24 GB GPU.
The bundled default config at
[`vllm_omni/deploy/voxcpm2.yaml`](../../vllm_omni/deploy/voxcpm2.yaml)
(`gpu_memory_utilization: 0.9`, `max_num_seqs: 4`, `enforce_eager: true`)
loads cleanly on a 4090 with ~4.9 GiB for model weights, ~15.2 GiB for
KV cache, and ~2 GiB for the talker's CUDA-Graph / CFM / VAE buffers,
for a total resident footprint of about **22 GiB / 24 GiB** in steady
state.
#### Environment
- OS: Linux (Ubuntu 22.04)
- Python: 3.12
- Driver / runtime: NVIDIA driver 575.57.08, CUDA 12.9 runtime
- torch: 2.11.0+cu128
- vLLM: 0.21.0
- vLLM-Omni: 0.20.1.dev (current `main`)
- Native VoxCPM2 source: `voxcpm` >= 2.0
#### Command
Start the server from the repository root. On a clean 4090 with no other
tenant on the GPU, the defaults from `voxcpm2.yaml` work as-is:
```bash
vllm serve openbmb/VoxCPM2 --omni \
--host 0.0.0.0 --port 8000
```
The deploy config at
[`vllm_omni/deploy/voxcpm2.yaml`](../../vllm_omni/deploy/voxcpm2.yaml)
is loaded automatically by the model registry (HF `model_type=voxcpm2`).
Pass `--deploy-config <path>` to override.
#### Verification
**Server cold-start**: ~60 s from `vllm serve` to `Application startup
complete` (subprocess fork + vLLM 0.21 init + model load + flashinfer JIT
+ torch.compile of LocDiT / feat_encoder / AudioVAE + CUDA-Graph warmup).
The first request after startup pays a small additional cost; steady-state
requests are much faster.
**T1 — zero-shot synthesis (curl, WAV)**:
```bash
curl -X POST http://127.0.0.1:8000/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"model": "voxcpm2",
"input": "Hello, this is VoxCPM2 served by vLLM-Omni on a single RTX 4090.",
"voice": "default",
"response_format": "wav"
}' --output output.wav
```
Observed: `output.wav` 491 KB, 5.12 s @ 48 kHz mono, peak amplitude 0.996.
Server-to-client wall time: 2.6 s (server-side RTF ~0.5 after warmup).
**T2 — bundled Python client (zero-shot)**:
```bash
python examples/online_serving/text_to_speech/voxcpm2/openai_speech_client.py \
--text "This audio was generated by the bundled openai_speech_client.py script against a vLLM-Omni server." \
--api-base http://127.0.0.1:8000 \
--output output.wav
```
Observed: `output.wav` 968 KB, 10.08 s @ 48 kHz mono, peak amplitude 0.949.
**T3 — voice cloning from a reference WAV**:
```bash
python examples/online_serving/text_to_speech/voxcpm2/openai_speech_client.py \
--text "This should sound like the reference speaker, but saying a different sentence on a single RTX 4090." \
--ref-audio /path/to/reference.wav \
--api-base http://127.0.0.1:8000 \
--output cloned.wav
```
Observed (using the T2 output `output.wav` as `--ref-audio`):
`cloned.wav` 568 KB, 5.92 s @ 48 kHz mono, peak amplitude 0.990;
end-to-end wall time 3.16 s.
`--ref-audio` accepts local file paths (auto-base64), HTTP/HTTPS URLs, or
`data:audio/wav;base64,...` data URIs. The `voice` field is required by
the OpenAI schema but ignored by VoxCPM2 unless it refers to an
uploaded voice — voice cloning is driven entirely by `ref_audio`.
**T4 — streaming PCM (raw, 48 kHz s16le)**:
```bash
curl -X POST http://127.0.0.1:8000/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"model": "voxcpm2",
"input": "Streaming PCM test on a single RTX 4090 with VoxCPM2.",
"voice": "default",
"stream": true,
"response_format": "pcm"
}' --no-buffer | play -t raw -r 48000 -e signed -b 16 -c 1 -
```
Observed: 230400 int16 samples = 4.80 s of PCM at 48 kHz, peak 32525 /
32767 (~99 % full scale). The player sample rate is **48 kHz**, not
24 kHz.
**T5 — gapless streaming (Gradio AudioWorklet demo)**:
```bash
python examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py \
--api-base http://127.0.0.1:8000
```
Observed: the demo binds to `http://0.0.0.0:7860` in ~4 s after launch.
The UI ("VoxCPM2 Streaming Demo") exposes three modes — **Voice Design**
(control instruction only), **Controllable Cloning** (ref audio +
optional style control), and **Ultimate Cloning** (ref audio +
transcript for audio continuation) — plus a `Stream (gapless)`
checkbox (default on) that drives the AudioWorklet streaming path. The
right pane includes runnable example prompts and a `Ready` status
indicator. End-to-end browser playback through `Generate Speech` was
verified working on the same vLLM-Omni server backing T1–T4 above.
**T6 — offline inference (no server, single process per request)**:
```bash
python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
--text "Hello, this is VoxCPM2 native AR synthesis running on vLLM-Omni and a single RTX 4090." \
--output-dir output_audio
```
`end2end.py` is a one-shot script (init → one `generate()` → exit), so
its reported RTF includes per-process compile + CUDA-Graph capture +
first-request runtime warmup. Note that the default `voxcpm2.yaml` sets
`enforce_eager: true`, so vLLM's engine-level CUDA graphs are off by
design; the "CUDA-Graph capture" referenced throughout this recipe is the
talker's own `_CapturedGraph` over the CFM solver / AudioVAE, not the LM
decode loop. Flipping `enforce_eager=false` does **not** remove this cost.
To separate the warmup cost from the
talker's true per-step throughput, the table below reports **5
consecutive `engine.generate([prompt])` calls in the same process**
after a one-off ~28 s engine init:
| call | duration | inference | RTF | notes |
| ---- | -------- | --------- | ----- | ------------------------------------------------ |
| #1 | 6.72 s | 11.97 s | 1.782 | cold: torch.compile + CUDA-Graph capture |
| #2 | 6.24 s | 11.43 s | 1.831 | still runtime warmup |
| #3 | 6.88 s | 0.82 s | 0.120 | ⚡ steady-state |
| #4 | 6.56 s | 0.78 s | 0.119 | steady-state |
| #5 | 5.76 s | 0.70 s | 0.121 | steady-state |
**Steady-state RTF (mean of calls #3–#5): 0.120**, i.e. ~8× real-time on
a single 4090. The high call-#1 / call-#2 numbers are one-time warmup
costs that **amortise to zero after the second request** — they are not
the model's per-step throughput.
A literal one-shot `python end2end.py` invocation in a fresh process
only sees call #1 (~RTF 1.8). For repeated workloads, either reuse the
same engine across calls (as in the table above) or use the online
serving path (T1–T5); a long-lived server amortises the warmup across
all subsequent requests (T1 wall-time RTF was 0.51 with a warm server).
Voice cloning + continuation mode (`ref_audio` + transcript):
```bash
python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
--text "Hello, this is the offline voice cloning test on a single RTX 4090." \
--ref-audio /path/to/reference.wav \
--ref-text "Original transcript of the reference audio."
```
Same 5-call methodology as the zero-shot table, this time with a
10.08 s reference WAV supplied via `--ref-audio` (and matching
`--ref-text`):
| call | duration | inference | RTF | notes |
| ---- | -------- | --------- | ----- | ------------------------------------------------ |
| #1 | 5.44 s | 12.38 s | 2.276 | cold: compile + CUDA-Graph capture + ref encode |
| #2 | 5.12 s | 2.47 s | 0.482 | most warmup done |
| #3 | 5.44 s | 0.76 s | 0.139 | ⚡ steady-state |
| #4 | 4.96 s | 0.68 s | 0.137 | steady-state |
| #5 | 5.28 s | 0.71 s | 0.134 | steady-state |
**Steady-state cloning RTF (mean of calls #3–#5): 0.137**, i.e. only
**+0.017 RTF** above the zero-shot steady-state of 0.120. That delta is
the per-request `AudioVAE.encode` over the reference clip (~85 ms for a
10 s ref WAV) plus the ref-continuation prompt build; the AR decode
steps themselves run at the same per-step rate as zero-shot.
#### Notes
- Prerequisites: `pip install voxcpm soundfile httpx ninja`. `ninja` must be on `PATH` because `flashinfer` JIT-compiles a top-k sampling kernel during the first profile pass (without it, startup dies with `FileNotFoundError: 'ninja'`).
- Memory usage: ~4.9 GiB model weights + ~15.2 GiB KV cache ≈ 22 GiB / 24 GiB resident with the default `voxcpm2.yaml`. On a shared GPU, pass `--gpu-memory-utilization 0.75` (or lower) if startup fails the free-memory check.
- Key flags: `--omni` is required. `--trust-remote-code` is not needed — the HF config registers under `model_type=voxcpm2`.
- Output: 48 kHz mono. When streaming PCM to a player, use `-r 48000`.
- Cold start: ~60 s server boot + ~25 s first-inference overhead (torch.compile + flashinfer JIT + the talker's own `_CapturedGraph` capture over CFM / AudioVAE — not vLLM engine CUDA graphs, which `enforce_eager: true` disables). Steady-state inference RTF is **~0.12** (~8× real-time on a single 4090; see T6 table for the per-call breakdown); online server requests run at ~RTF 0.5 wall-time (includes HTTP round-trip). Keep the server warm for interactive use.
+145
View File
@@ -0,0 +1,145 @@
# MOSS-TTS
## Summary
- Vendor: OpenMOSS
- Models: `OpenMOSS-Team/MOSS-TTS` (8B), `OpenMOSS-Team/MOSS-TTS-v1.5` (8B),
`OpenMOSS-Team/MOSS-TTS-Realtime` (1.7B), `OpenMOSS-Team/MOSS-TTSD-v1.0` (8B),
`OpenMOSS-Team/MOSS-SoundEffect` (8B), `OpenMOSS-Team/MOSS-VoiceGenerator` (1.7B)
- Task: Text-to-speech synthesis, sound effect generation, zero-shot voice design
- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API; offline inference
- Maintainer: Community
## When to use this recipe
Use this recipe for 24 kHz multilingual TTS with voice cloning (20 languages including
Chinese and English). Choose a variant based on your latency and quality requirements:
| Model | Params | Use case |
|---|---|---|
| MOSS-TTS | 8B | General TTS, highest quality |
| MOSS-TTS-v1.5 | 8B | General TTS upgrade of 1.0: 31 languages, steadier cloning, `[pause Xs]` markers (set `language` for best results); same `MossTTSDelay` API |
| MOSS-TTS-Realtime | 1.7B | Lowest latency (TTFB ~180 ms), streaming-first |
| MOSS-TTSD-v1.0 | 8B | Multi-turn dialogue TTS |
| MOSS-SoundEffect | 8B | Sound effect synthesis from text description |
| MOSS-VoiceGenerator | 1.7B | Zero-shot voice design |
All variants share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and
output 24 kHz mono audio.
## References
- Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)
- Deploy configs: [`vllm_omni/deploy/moss_tts.yaml`](../../vllm_omni/deploy/moss_tts.yaml) and variants
- HuggingFace org: <https://huggingface.co/OpenMOSS-Team>
## Hardware Support
### GPU
#### 1x H100 80GB — MOSS-TTS (8B)
##### Environment
- OS: Linux
- Python: 3.11+
- CUDA 12.8
- vLLM-Omni version: see `vllm_omni/__version__.py`
##### Command
```bash
# The codec is loaded automatically from OpenMOSS-Team/MOSS-Audio-Tokenizer.
# Override the path with MOSS_TTS_CODEC_PATH if you have a local copy.
vllm serve OpenMOSS-Team/MOSS-TTS --omni --port 8091
```
##### Verification
Voice cloning (provide a reference audio clip):
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"model": "OpenMOSS-Team/MOSS-TTS",
"input": "Hello, this is a voice cloning test.",
"voice": "default",
"ref_audio": "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/zh_1.wav",
"response_format": "wav"
}' --output output.wav
```
##### Notes
- Peak GPU memory: ~18 GB for the talker (8B) + ~8 GB for the codec decoder on the same device.
Use `gpu_memory_utilization: 0.85` in `moss_tts.yaml` (default).
- Output: 24 kHz mono WAV.
- The `MOSS_TTS_CODEC_PATH` environment variable overrides the codec checkpoint location.
---
#### 1x A10G 24GB — MOSS-TTS-Realtime (1.7B)
##### Environment
- OS: Linux
- Python: 3.11+
- CUDA 12.8
##### Command
```bash
vllm serve OpenMOSS-Team/MOSS-TTS-Realtime --omni --port 8091
```
##### Verification
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"model": "OpenMOSS-Team/MOSS-TTS-Realtime",
"input": "This is a low-latency streaming TTS test.",
"voice": "default",
"ref_audio": "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/zh_1.wav",
"response_format": "wav",
"stream": true
}' --output output.wav
```
##### Notes
- Peak GPU memory: ~6 GB for the talker (1.7B) + ~8 GB for the codec decoder.
- First-audio latency (TTFB): ~180 ms on A10G.
- `codec_chunk_frames: 15` in `moss_tts_realtime.yaml` for lower TTFA than the 8B variant.
---
#### 1x A10G 24GB — MOSS-SoundEffect (8B, sound effect synthesis)
##### Command
```bash
vllm serve OpenMOSS-Team/MOSS-SoundEffect --omni --port 8091
```
##### Verification
Sound effect synthesis takes a text description instead of reference audio:
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"model": "OpenMOSS-Team/MOSS-SoundEffect",
"input": "Thunder rumbling, rain pattering on a tin roof.",
"response_format": "wav"
}' --output thunder.wav
```
##### Notes
- No `ref_audio` required or accepted for MOSS-SoundEffect.
- Input field maps to the `ambient_sound` parameter in the upstream processor.
- Rate: ~12.5 tokens per second; longer descriptions produce longer audio.
+9
View File
@@ -29,10 +29,17 @@ recipes/
|--------|------|----------|
| [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB |
| [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
| [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
| [`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 |
| [`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 |
| [`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 |
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay | 1x A100 80GB |
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
@@ -44,6 +51,8 @@ recipes/
| [`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 |
| [`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 |
|
Within a single recipe file, include different hardware support sections such
as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like
+107
View File
@@ -0,0 +1,107 @@
# Voxtral TTS for text-to-speech
## Summary
- Vendor: Mistral AI
- Model: `mistralai/Voxtral-4B-TTS-2603`
- Task: Text-to-speech with model-provided voice presets or voice cloning
- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API
- Maintainer: Community
## When to use this recipe
Use this recipe when you want a known-good starting point for serving Voxtral TTS
models with vLLM-Omni and validate the deployment with the existing TTS client
examples in this repository.
## References
- Upstream model card: <https://huggingface.co/mistralai/Voxtral-4B-TTS-2603>
- Related examples under `examples/`:
[`examples/offline_inference/text_to_speech/voxtral_tts/end2end.py`](../../examples/offline_inference/text_to_speech/voxtral_tts/end2end.py)
[`examples/online_serving/text_to_speech/voxtral_tts/gradio_demo.py`](../../examples/online_serving/text_to_speech/voxtral_tts/gradio_demo.py)
- Related issue or discussion:
[RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
## Hardware Support
## GPU
### 1x RTX 4090 24GB
#### Environment
- OS: Ubuntu 22.04.5 LTS
- Python: 3.12.11
- PyTorch: 2.11.0+cu130
- Driver / runtime: CUDA 13.0
- GPU: NVIDIA GeForce RTX 4090, 24 GB
- vLLM version: 0.20.0
- vLLM-Omni version or commit: 0.20.0rc2.dev99+g857356d5b
- Required package: mistral_common >= 1.10.0
#### Command
Online:
```bash
vllm serve mistralai/Voxtral-4B-TTS-2603 \
--omni \
--port 8091
```
The deploy config is auto-loaded from `vllm_omni/deploy/voxtral_tts.yaml`
#### Verification
**Quick smoke test with curl:**
```bash
curl -X POST http://127.0.0.1:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"model": "mistralai/Voxtral-4B-TTS-2603",
"input": "Hello, this is Voxtral TTS running with vLLM-Omni.",
"voice": "casual_female",
"language": "English",
"response_format": "wav"
}' \
--output voxtral.wav
```
**Streaming audio:**
```bash
curl -X POST http://127.0.0.1:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"model": "mistralai/Voxtral-4B-TTS-2603",
"input": "Hello, this is Voxtral TTS streaming PCM.",
"voice": "casual_female",
"language": "English",
"stream": true,
"response_format": "pcm"
}' \
--output voxtral_stream.pcm
```
Convert the raw PCM output to WAV
```bash
ffmpeg -f s16le -ar 24000 -ac 1 -i voxtral_stream.pcm voxtral_stream.wav -y
ffprobe -hide_banner voxtral_stream.wav
```
#### Notes
- Memory usage: Stage 0 (`audio_generation`) ~18.95 GiB, Stage 1
(`audio_tokenizer`) ~1.55 GiB. The observed server-startup peak was
`21006 MiB / 24564 MiB` (`20.51 GiB / 23.99 GiB`, about 85.5%).
Both stages share GPU 0 via the deploy config (`gpu_memory_utilization: 0.8`
for Stage 0, `0.1` for Stage 1).
- Key flags: `--omni` is required. `--port 8091` is optional and only selects
the HTTP port.
- Known limitation: Voxtral TTS outputs mono audio at 24 kHz. And with the public
mistralai/Voxtral-4B-TTS-2603 checkpoint. running the
voice-cloning e2e test ([`Voice cloning (capability gated upstream)`](../../examples/offline_inference/text_to_speech/README.md#voice-cloning-capability-gated-upstream)
) or adding a `ref_audio` request parameter will fail
because the public checkpoint does not provide the encoder weights needed to
convert reference audio into conditioning features. The expected failure is: `RuntimeError: encode_waveforms requires encoder weights which are not available in the open-source checkpoint.`
+245
View File
@@ -0,0 +1,245 @@
# Cosmos3-Nano
> Text-to-image, text-to-video, and image-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)
- 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:
- **T2I** — `POST /v1/images/generations` (or a prompt carrying `modalities=["image"]`).
- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image.
- **I2V** — `POST /v1/videos/sync` with a reference image (`input_reference` file
upload, or `image_reference` JSON).
- **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.
## References
- Model card (authoritative usage + example assets): <https://huggingface.co/nvidia/Cosmos3-Nano>
- Example inputs/outputs live in the repo's `assets/` (`example_t2v_prompt.json`,
`example_i2v_prompt.json`, `example_i2v_input.jpg`, `negative_prompt.json`;
audio examples: `example_t2vs_prompt.json`, `example_t2vs_output.mp4`,
`example_i2vs_output.mp4`).
- Prompt upsampling (recommended for quality): the model expects JSON-upsampled
structured prompts; see NVIDIA's `cosmos-framework` prompt-upsampling docs.
- Pipeline: [`vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py`](../../vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py)
- Smoke tests (canonical request formats): [`tests/e2e/accuracy/test_cosmos3_similarity.py`](../../tests/e2e/accuracy/test_cosmos3_similarity.py)
## Hardware Support
## GPU
### 1x H200 141GB / B300 (Online serving)
#### Environment
- OS: Ubuntu 22.04+
- Python: 3.12+
- Driver / runtime: NVIDIA CUDA environment
- vLLM version: match the repository requirements from your current checkout
- vLLM-Omni version or commit: use the commit you are deploying from
#### Command
Requires the `vllm-omni` package (or the `vllm/vllm-omni:cosmos3` container),
which provides the `vllm serve … --omni` entrypoint used below.
Safety guardrails are **on by default** (NVIDIA Open Model License). They load
the **gated** `nvidia/Cosmos-1.0-Guardrail` model, so to keep them on you must:
1. `pip install cosmos-guardrail`
2. Accept the license at <https://huggingface.co/nvidia/Cosmos-1.0-Guardrail>
3. Export a token with access: `export HF_TOKEN=hf_...`
Then launch the recommended server:
```bash
vllm serve nvidia/Cosmos3-Nano \
--omni \
--host 0.0.0.0 --port 8000 \
--init-timeout 1800
```
To run **without** guardrails (you are responsible for license compliance),
add `--no-guardrails` (no token/`cosmos-guardrail` needed). For extra GPUs use
`--ulysses-degree N` (context parallel) or `--tensor-parallel-size N`;
`--enable-layerwise-offload` reduces VRAM on smaller GPUs. The pipeline
auto-resolves from `model_index.json`; pass
`--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly.
#### Verification
Best quality uses the JSON-upsampled prompts from `assets/` (download with
`hf download nvidia/Cosmos3-Nano assets/ --local-dir Cosmos3-Nano`). Minimal
self-contained examples:
```bash
curl http://localhost:8000/v1/models
# Text-to-image -> /v1/images/generations (1024x1024, 50 steps; base64 PNG)
curl -sS -X POST http://localhost:8000/v1/images/generations \
-H "Content-Type: application/json" \
-d '{
"model": "nvidia/Cosmos3-Nano",
"prompt": "A photorealistic red sports car on a city street at golden hour, cinematic lighting.",
"negative_prompt": "blurry, distorted, low quality",
"size": "1024x1024", "n": 1, "response_format": "b64_json",
"num_inference_steps": 50, "guidance_scale": 7.0, "seed": 42
}' | python -c "import sys,json,base64; open('cosmos3_t2i.png','wb').write(base64.b64decode(json.load(sys.stdin)['data'][0]['b64_json']))"
# Text-to-video -> /v1/videos/sync (720p, 189 frames @ 24fps; official params)
curl -sS -X POST http://localhost:8000/v1/videos/sync \
-H "Accept: video/mp4" \
-F "model=nvidia/Cosmos3-Nano" \
-F "prompt=A robot arm is cleaning a plate in the kitchen" \
-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}' \
-F "seed=123" \
-o cosmos3_t2v.mp4
# Image-to-video -> /v1/videos/sync with an uploaded reference image
curl -sS -X POST http://localhost:8000/v1/videos/sync \
-H "Accept: video/mp4" \
-F "model=nvidia/Cosmos3-Nano" \
-F "prompt=The scene comes to life with smooth, natural motion." \
-F "negative_prompt=blurry, distorted, low quality" \
-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}' \
-F "seed=1111" \
-F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
-o cosmos3_i2v.mp4
# Text-to-video-with-sound
curl -sS -X POST http://localhost:8000/v1/videos/sync \
-H "Accept: video/mp4" \
-F "model=nvidia/Cosmos3-Nano" \
-F "prompt=The video opens with a view of a well-lit indoor fruit display. A robotic arm picks up a pear, an orange, and a carambola one by one, placing each into a plastic bag in a shopping cart with red handles. The video is 7.875 seconds long, 24 FPS, and 1280x720. Audio description: soft servo whirs, gentle fruit thuds, plastic bag rustling, and a faint refrigeration hum." \
-F "negative_prompt=blurry, distorted, low quality" \
-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 "seed=0" \
-F "generate_sound=true" \
-F "sound_duration=7.875" \
-F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
-o cosmos3_t2v_with_sound.mp4
```
#### Notes
- **Measured latency (1x B300, bf16, guardrails off):**
- 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**
- 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)
~33 GB on disk.
- **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`.
- **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.
- **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.
### 1x GPU (Offline generation)
#### Environment
- OS: Ubuntu 22.04+
- Python: 3.12+
- Driver / runtime: NVIDIA CUDA environment
- vLLM-Omni version or commit: use the commit you are deploying from
#### Command
```python
# cosmos3_offline.py — run with: python cosmos3_offline.py
import torch
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
def main():
omni = Omni(
model="nvidia/Cosmos3-Nano",
model_class_name="Cosmos3OmniDiffusersPipeline",
trust_remote_code=True,
enforce_eager=True,
# Guardrails are disabled here for a quick local run; install
# cosmos-guardrail + gated-repo access and drop this to enable them.
model_config={"guardrails": False},
)
gen = torch.Generator(device="cpu").manual_seed(42)
# Text-to-image (modalities=["image"]). For T2V use modalities=["video"]
# plus num_frames/fps; for I2V add multi_modal_data={"image": <PIL.Image>}.
outputs = omni.generate(
{
"prompt": "A photorealistic red sports car at golden hour, cinematic lighting.",
"negative_prompt": "blurry, distorted, low quality",
"modalities": ["image"],
},
OmniDiffusionSamplingParams(
height=1024, width=1024, generator=gen,
guidance_scale=7.0, num_inference_steps=50, num_outputs_per_prompt=1,
),
)
outputs[0].request_output.images[0].save("cosmos3_t2i.png")
omni.close()
if __name__ == "__main__":
main()
```
#### Verification
```bash
python cosmos3_offline.py
python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image', im.size, im.mode)"
```
#### Notes
- 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
`(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.
Loaded 100 of 431 files, more files were not shown because too many files have changed in this diff. Show more