anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Merge branch 'main' into lance-model-support

Browse Source
This commit is contained in:
SYLAR authored and GitHub committed 2026-05-22 04:13:58 +08:00
2 parents 7b885f9d13 e949ccf087
commit f0cec3ecd4
224 files changed
+19202 -7026
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

examples/offline_inference/hunyuan_image3/end2end.py
+24
View File
@@ -77,6 +77,24 @@ def parse_args():
parser.add_argument("--log-stats", action="store_true", default=False)
parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.")
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
parser.add_argument(
"--diffusion-kv-cache-dtype",
type=str,
default=None,
help="Diffusion attention KV cache dtype, for example 'fp8'. Separate from vLLM --kv-cache-dtype.",
)
parser.add_argument(
"--diffusion-kv-cache-skip-steps",
type=str,
default=None,
help="Denoising step selector to keep diffusion KV cache in native dtype, for example '0,1,4-6'.",
)
parser.add_argument(
"--diffusion-kv-cache-skip-layers",
type=str,
default=None,
help="Transformer layer selector to keep diffusion KV cache in native dtype, for example '0-2,10'.",
)
parser.add_argument(
"--additional-config",
type=str,
@@ -140,6 +158,9 @@ def main():
"init_timeout": args.init_timeout,
"enforce_eager": args.enforce_eager,
"mode": _MODALITY_MODE[args.modality],
"diffusion_kv_cache_dtype": args.diffusion_kv_cache_dtype,
"diffusion_kv_cache_skip_steps": args.diffusion_kv_cache_skip_steps,
"diffusion_kv_cache_skip_layers": args.diffusion_kv_cache_skip_layers,
}
if additional_config is not None:
@@ -237,6 +258,9 @@ def main():
print(f" Inference steps: {args.steps}")
print(f" Guidance scale: {args.guidance_scale}")
print(f" Seed: {args.seed}")
print(f" diffusion_kv_cache_dtype: {args.diffusion_kv_cache_dtype}")
print(f" diffusion_kv_cache_skip_steps: {args.diffusion_kv_cache_skip_steps}")
print(f" diffusion_kv_cache_skip_layers: {args.diffusion_kv_cache_skip_layers}")
if args.modality == "text2img":
print(f" Output size: {args.width}x{args.height}")
if args.image_path:
examples/offline_inference/image_to_video/image_to_video.py
+43 -33
View File
@@ -142,6 +142,42 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable layerwise (blockwise) offloading on DiT modules.",
)
parser.add_argument(
"--enforce-eager",
action="store_true",
help="Disable torch.compile and force eager execution.",
)
parser.add_argument(
"--audio-sample-rate",
type=int,
default=24000,
help="Sample rate for audio output when saved (default: 24000).",
)
parser.add_argument(
"--cache-backend",
type=str,
default=None,
choices=["cache_dit", "tea_cache"],
help=(
"Cache backend to use for acceleration. "
"Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
"Default: None (no cache acceleration)."
),
)
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
parser.add_argument(
"--quantization",
type=str,
default=None,
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
)
# Distributed and parallel execution
parser.add_argument(
"--ulysses-degree",
type=int,
@@ -173,28 +209,6 @@ def parse_args() -> argparse.Namespace:
default=1,
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
)
parser.add_argument(
"--enforce-eager",
action="store_true",
help="Disable torch.compile and force eager execution.",
)
parser.add_argument(
"--audio-sample-rate",
type=int,
default=24000,
help="Sample rate for audio output when saved (default: 24000).",
)
parser.add_argument(
"--cache-backend",
type=str,
default=None,
choices=["cache_dit", "tea_cache"],
help=(
"Cache backend to use for acceleration. "
"Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
"Default: None (no cache acceleration)."
),
)
parser.add_argument(
"--use-hsdp",
action="store_true",
@@ -219,16 +233,10 @@ def parse_args() -> argparse.Namespace:
),
)
parser.add_argument(
"--quantization",
type=str,
default=None,
choices=["fp8", "mxfp8", "int8", "gguf"],
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
)
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
"--pipeline-parallel-size",
type=int,
default=1,
help="Number of pipeline parallel stages.",
)
parser.add_argument(
"--profiler-config",
@@ -326,6 +334,7 @@ def main():
use_hsdp=args.use_hsdp,
hsdp_shard_size=args.hsdp_shard_size,
hsdp_replicate_size=args.hsdp_replicate_size,
pipeline_parallel_size=args.pipeline_parallel_size,
)
omni_kwargs = dict(
model=args.model,
@@ -366,7 +375,8 @@ def main():
print(f" diffusion_kv_cache_skip_layers(config): {args.diffusion_kv_cache_skip_layers}")
print(
f" Parallel configuration: cfg_parallel_size={args.cfg_parallel_size},"
f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}"
f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size},"
f" pipeline_parallel_size={args.pipeline_parallel_size}"
)
print(f" Video size: {args.width}x{args.height}")
print(f"{'=' * 60}\n")
examples/offline_inference/text_to_speech/README.md
+2 -53
View File
@@ -14,13 +14,12 @@ list of supported architectures across all modalities, see
| Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
|---|---|---|---|---|---|---|
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 22.05 kHz |
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
| MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | single (AR + codec) | ✓ (required) | ✓ | voice_clone, continuation | 48 kHz |
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
| VoxCPM | `openbmb/VoxCPM-0.5B` | split | ✓ | ✓ | — | 24 kHz |
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
@@ -41,7 +40,7 @@ python examples/offline_inference/text_to_speech/<model>/end2end.py \
## CosyVoice3
2-stage TTS pipeline (`talker` + `code2wav`) at 22.05 kHz.
2-stage TTS pipeline (`talker` + `code2wav`) at 24 kHz.
### Prerequisites
```bash
@@ -308,56 +307,6 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
---
## VoxCPM
Split-stage TTS. The hub example covers single TTS, single voice cloning, and streaming. Use `benchmarks/voxcpm/` for warmup, batch JSONL prompts, profiler injection, and offline TTFP / RTF measurement.
### Prerequisites
```bash
pip install voxcpm soundfile
# or use a local source tree:
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
```
### Quick start
```bash
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
--model openbmb/VoxCPM-0.5B \
--text "This is a split-stage VoxCPM synthesis example running on vLLM Omni."
```
### Voice cloning
```bash
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
--model openbmb/VoxCPM-0.5B \
--text "This sentence is synthesized with a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "The exact transcript spoken in reference.wav."
```
### Streaming
Pass `--streaming` together with the legacy async-chunk stage config:
```bash
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
--model openbmb/VoxCPM-0.5B \
--streaming \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
--text "This is a split-stage VoxCPM streaming example running on vLLM Omni."
```
### Persistent HF config (optional)
The engine auto-patches a missing `model_type` in the checkpoint's `config.json` in-memory at startup (see `vllm_omni/engine/arg_utils.py`), so no manual setup is required.
If you'd rather maintain a persistent patched config dir (e.g. to share across runs or to add fields the auto-patcher doesn't set), export `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH` to its location — both the bundled deploy and async-chunk yamls read it via `${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,}`.
### Notes
- Non-streaming auto-loads `vllm_omni/deploy/voxcpm.yaml` from the model directory name / HF `model_type`. Streaming uses the legacy `voxcpm_async_chunk.yaml` passed via `--stage-configs-path`.
- Streaming is currently single-request oriented.
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades quality.
- For online serving, see the [VoxCPM section in the online hub](../../online_serving/text_to_speech/README.md#voxcpm). For benchmark reporting, see [`benchmarks/voxcpm`](../../../benchmarks/voxcpm/README.md).
---
## VoxCPM2
Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`.
examples/offline_inference/text_to_speech/voxcpm/end2end.py
-220
View File
@@ -1,220 +0,0 @@
"""Minimal offline VoxCPM example for vLLM Omni."""
from __future__ import annotations
import asyncio
import time
from pathlib import Path
from typing import Any
import soundfile as sf
import torch
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import AsyncOmni, Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
def _build_prompt(args) -> dict[str, Any]:
additional_information: dict[str, list[Any]] = {
"text": [args.text],
"cfg_value": [args.cfg_value],
"inference_timesteps": [args.inference_timesteps],
"min_len": [args.min_len],
"max_new_tokens": [args.max_new_tokens],
}
if args.streaming_prefix_len is not None:
additional_information["streaming_prefix_len"] = [args.streaming_prefix_len]
if args.ref_audio is not None:
additional_information["ref_audio"] = [args.ref_audio]
if args.ref_text is not None:
additional_information["ref_text"] = [args.ref_text]
return {
"prompt_token_ids": [1],
"additional_information": additional_information,
}
def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor:
audio = mm.get("audio", mm.get("model_outputs"))
if audio is None:
raise ValueError("No audio output found in multimodal output.")
if isinstance(audio, list):
parts = [torch.as_tensor(item).float().cpu().reshape(-1) for item in audio]
audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0)
if not isinstance(audio, torch.Tensor):
audio = torch.as_tensor(audio)
return audio.float().cpu().reshape(-1)
def _extract_sample_rate(mm: dict[str, Any]) -> int:
sr_raw = mm.get("sr", 24000)
if isinstance(sr_raw, list) and sr_raw:
sr_raw = sr_raw[-1]
if hasattr(sr_raw, "item"):
return int(sr_raw.item())
return int(sr_raw)
def _save_audio(audio: torch.Tensor, sample_rate: int, output_dir: Path, request_id: str) -> Path:
output_dir.mkdir(parents=True, exist_ok=True)
output_path = output_dir / f"output_{request_id}.wav"
sf.write(
output_path,
audio.float().cpu().clamp(-1.0, 1.0).numpy(),
sample_rate,
format="WAV",
subtype="PCM_16",
)
return output_path
async def _run_streaming(args) -> Path:
prompt = _build_prompt(args)
output_dir = Path(args.output_dir) if args.output_dir is not None else Path("output_audio_streaming")
request_id = "streaming_example"
sample_rate = 24000
buffered_samples = 0
chunks: list[torch.Tensor] = []
started = time.perf_counter()
omni = AsyncOmni(
model=args.model,
stage_configs_path=args.stage_configs_path,
deploy_config=args.deploy_config,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
try:
async for stage_output in omni.generate(prompt, request_id=request_id):
mm = getattr(stage_output, "multimodal_output", None)
if not isinstance(mm, dict):
request_output = getattr(stage_output, "request_output", None)
if request_output is None:
continue
mm = getattr(request_output, "multimodal_output", None)
if not isinstance(mm, dict) and getattr(request_output, "outputs", None):
mm = getattr(request_output.outputs[0], "multimodal_output", None)
if not isinstance(mm, dict):
continue
audio = _extract_audio_tensor(mm)
if audio.numel() == 0:
continue
sample_rate = _extract_sample_rate(mm)
if audio.numel() > buffered_samples:
delta = audio[buffered_samples:]
buffered_samples = int(audio.numel())
else:
delta = audio
buffered_samples += int(delta.numel())
if delta.numel() > 0:
chunks.append(delta)
if not chunks:
raise RuntimeError("No streaming audio chunks received from VoxCPM.")
output_audio = torch.cat(chunks, dim=0)
output_path = _save_audio(output_audio, sample_rate, output_dir, request_id)
print(f"Saved streaming audio to: {output_path} ({time.perf_counter() - started:.2f}s)")
return output_path
finally:
omni.shutdown()
def _run_sync(args) -> Path:
prompt = _build_prompt(args)
output_dir = Path(args.output_dir) if args.output_dir is not None else Path("output_audio")
request_id = "sync_example"
started = time.perf_counter()
last_mm: dict[str, Any] | None = None
omni = Omni(
model=args.model,
stage_configs_path=args.stage_configs_path,
deploy_config=args.deploy_config,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
for stage_outputs in omni.generate(prompt):
request_output = getattr(stage_outputs, "request_output", None)
if request_output is None:
continue
outputs = getattr(request_output, "outputs", None)
if outputs:
for output in outputs:
mm = getattr(output, "multimodal_output", None)
if isinstance(mm, dict):
last_mm = mm
mm = getattr(request_output, "multimodal_output", None)
if isinstance(mm, dict):
last_mm = mm
if last_mm is None:
raise RuntimeError("No audio output received from VoxCPM.")
output_path = _save_audio(
_extract_audio_tensor(last_mm),
_extract_sample_rate(last_mm),
output_dir,
request_id,
)
print(f"Saved audio to: {output_path} ({time.perf_counter() - started:.2f}s)")
return output_path
def parse_args():
parser = FlexibleArgumentParser(description="Minimal offline VoxCPM example for vLLM Omni.")
parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.")
parser.add_argument(
"--deploy-config",
type=str,
default=None,
help=("Override the deploy config path."),
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help=(
"Legacy stage_args yaml path. Required for streaming "
"(vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml); "
"leave unset for non-streaming to use the auto-resolved deploy config."
),
)
parser.add_argument(
"--streaming",
action="store_true",
default=False,
help="Stream audio chunks as they arrive via AsyncOmni.",
)
parser.add_argument("--text", type=str, required=True, help="Input text for synthesis.")
parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.")
parser.add_argument("--ref-text", type=str, default=None, help="Transcript of the reference audio.")
parser.add_argument("--output-dir", type=str, default=None, help="Output directory for generated wav files.")
parser.add_argument("--cfg-value", type=float, default=2.0, help="Guidance value passed to VoxCPM.")
parser.add_argument("--inference-timesteps", type=int, default=10, help="Number of diffusion timesteps.")
parser.add_argument("--min-len", type=int, default=2, help="Minimum latent length.")
parser.add_argument("--max-new-tokens", type=int, default=4096, help="Maximum latent length.")
parser.add_argument(
"--streaming-prefix-len",
type=int,
default=3,
help="Streaming prefix length used by voxcpm_async_chunk.yaml.",
)
parser.add_argument("--stage-init-timeout", type=int, default=600, help="Stage initialization timeout in seconds.")
parser.add_argument("--log-stats", action="store_true", help="Enable vLLM Omni stats logging.")
nullify_stage_engine_defaults(parser)
args = parser.parse_args()
if (args.ref_audio is None) != (args.ref_text is None):
raise ValueError("Voice cloning requires --ref-audio and --ref-text together.")
return args
def main(args) -> None:
route = "streaming" if args.streaming else "sync"
print(f"Model: {args.model}")
print(f"Deploy config: {args.deploy_config or '<auto from HF model_type>'}")
print(f"Stage configs path: {args.stage_configs_path or '<unused>'}")
print(f"Route: {route}")
if route == "streaming":
asyncio.run(_run_streaming(args))
else:
_run_sync(args)
if __name__ == "__main__":
main(parse_args())
examples/offline_inference/text_to_video/text_to_video.py
+60 -25
View File
@@ -142,6 +142,32 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable layerwise (blockwise) offloading on DiT modules.",
)
parser.add_argument(
"--audio-sample-rate",
type=int,
default=24000,
help="Sample rate for audio output when saved (default: 24000).",
)
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
parser.add_argument(
"--profiler-config",
type=parse_profiler_config,
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
parser.add_argument(
"--quantization",
type=str,
default=None,
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
)
# Distributed and parallel execution
parser.add_argument(
"--ulysses-mode",
type=str,
@@ -174,41 +200,23 @@ def parse_args() -> argparse.Namespace:
default=1,
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
)
parser.add_argument(
"--audio-sample-rate",
type=int,
default=24000,
help="Sample rate for audio output when saved (default: 24000).",
)
parser.add_argument(
"--vae-patch-parallel-size",
type=int,
default=1,
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
)
parser.add_argument(
"--pipeline-parallel-size",
type=int,
default=1,
help="Number of pipeline parallel stages.",
)
parser.add_argument(
"--enable-expert-parallel",
action="store_true",
help="Enable expert parallelism for MoE layers.",
)
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
parser.add_argument(
"--profiler-config",
type=parse_profiler_config,
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
parser.add_argument(
"--quantization",
type=str,
default=None,
choices=["fp8", "mxfp8", "int8", "gguf"],
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
)
parser.add_argument(
"--use-hsdp",
action="store_true",
@@ -229,6 +237,27 @@ def parse_args() -> argparse.Namespace:
return parser.parse_args()
def _extract_peak_memory_mb(result: Any) -> float:
"""Pull worker-reported peak VRAM (MiB) generation result.
Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb.
"""
if isinstance(result, list):
result = result[0] if result else None
if result is None:
return 0.0
val = getattr(result, "peak_memory_mb", 0.0)
if not val:
inner = getattr(result, "request_output", None)
if isinstance(inner, list):
inner = inner[0] if inner else None
val = getattr(inner, "peak_memory_mb", 0.0)
try:
return float(val or 0.0)
except (TypeError, ValueError):
return 0.0
def main():
args = parse_args()
model_class_name = args.model_class_name
@@ -262,6 +291,7 @@ def main():
cfg_parallel_size=args.cfg_parallel_size,
tensor_parallel_size=args.tensor_parallel_size,
vae_patch_parallel_size=args.vae_patch_parallel_size,
pipeline_parallel_size=args.pipeline_parallel_size,
enable_expert_parallel=args.enable_expert_parallel,
)
@@ -307,7 +337,8 @@ def main():
print(
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree},"
f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size},"
f" vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}"
f" vae_patch_parallel_size={args.vae_patch_parallel_size}, pipeline_parallel_size={args.pipeline_parallel_size},"
f" enable_expert_parallel={args.enable_expert_parallel}"
)
print(f" Video size: {args.width}x{args.height}")
print(f"{'=' * 60}\n")
@@ -338,6 +369,10 @@ def main():
# Print profiling results
print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
peak_mb = _extract_peak_memory_mb(frames)
if peak_mb:
print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
audio = None
if isinstance(frames, list):
frames = frames[0] if frames else None