diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 985b50fc7..1151da467 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -320,6 +320,29 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" + - label: "OmniVoice E2E Test" + timeout_in_minutes: 20 + depends_on: upload-ready-pipeline + commands: + - | + timeout 20m bash -c ' + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_WORKER_MULTIPROC_METHOD=spawn + pytest -s -v tests/e2e/online_serving/test_omnivoice.py -m "core_model" --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" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + - label: "Voxtral-TTS E2E Test" timeout_in_minutes: 20 depends_on: upload-ready-pipeline diff --git a/examples/offline_inference/omnivoice/README.md b/examples/offline_inference/omnivoice/README.md new file mode 100644 index 000000000..d804b61b5 --- /dev/null +++ b/examples/offline_inference/omnivoice/README.md @@ -0,0 +1,73 @@ +# OmniVoice + +This directory contains an offline demo for running OmniVoice TTS models with vLLM Omni. It generates speech from text and saves WAV files locally. + +## Model Overview + +[OmniVoice](https://huggingface.co/k2-fsa/OmniVoice) is a zero-shot multilingual TTS model supporting 600+ languages. It uses a diffusion language model (Qwen3-0.6B backbone) with iterative masked unmasking to generate speech. + +Three inference modes are supported: + +- **Auto Voice**: Generate speech without any reference — the model picks a voice automatically. +- **Voice Clone**: Clone a voice from a reference audio + transcription. +- **Voice Design**: Control voice style via natural language instruction (e.g., "female, low pitch, british accent"). + +## Setup + +Ensure the model is downloaded: + +```bash +huggingface-cli download k2-fsa/OmniVoice +``` + +> **Note:** Voice cloning requires `transformers>=5.3.0` for `HiggsAudioV2TokenizerModel`. Auto voice and voice design modes work with `transformers>=4.57.0`. + +## Quick Start + +Auto voice (text only): + +```bash +python end2end.py --model k2-fsa/OmniVoice --text "Hello, this is a test." +``` + +Voice design (with style instruction): + +```bash +python end2end.py --model k2-fsa/OmniVoice \ + --text "Hello, this is a test." \ + --instruct "female, low pitch, british accent" +``` + +Voice clone (with reference audio): + +```bash +python end2end.py --model k2-fsa/OmniVoice \ + --text "Hello, this is a test." \ + --ref-audio ref.wav \ + --ref-text "This is the reference transcription." +``` + +## Language Support + +Specify a language for improved quality: + +```bash +python end2end.py --model k2-fsa/OmniVoice \ + --text "你好，这是一个测试。" \ + --lang zh +``` + +## Architecture + +OmniVoice uses a two-stage pipeline: + +- **Stage 0 (Generator)**: Qwen3-0.6B transformer with 32-step iterative unmasking and classifier-free guidance. Generates 8-codebook audio tokens from text. +- **Stage 1 (Decoder)**: HiggsAudioV2 RVQ quantizer + DAC acoustic decoder. Converts tokens to 24kHz waveform. + +Both stages use `GPUGenerationWorker` with `OmniGenerationScheduler`. + +## Notes + +- Output audio is saved to `output.wav` by default. Use `--output` to change the path. +- The model estimates duration from text automatically via `RuleDurationEstimator`. +- Use `--stage-init-timeout` to increase the stage initialization timeout for first-time model downloads. diff --git a/examples/offline_inference/omnivoice/end2end.py b/examples/offline_inference/omnivoice/end2end.py new file mode 100644 index 000000000..b41379b01 --- /dev/null +++ b/examples/offline_inference/omnivoice/end2end.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end OmniVoice TTS inference via vLLM-Omni. + +Supports: +- Auto voice mode: text only → generated speech +- Voice cloning mode: text + reference audio → cloned voice speech + +Usage: + # Auto voice + python end2end.py --model k2-fsa/OmniVoice --text "Hello world" + + # Voice cloning + python end2end.py --model k2-fsa/OmniVoice --text "Hello" \ + --ref-audio ref.wav --ref-text "reference transcription" +""" + +import argparse +import os + +import numpy as np +import soundfile as sf + +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + +def run_e2e(): + parser = argparse.ArgumentParser(description="OmniVoice E2E TTS inference") + parser.add_argument( + "--model", + type=str, + default="k2-fsa/OmniVoice", + help="Model name or path (HuggingFace or local)", + ) + parser.add_argument( + "--stage-config", + type=str, + default="vllm_omni/model_executor/stage_configs/omnivoice.yaml", + ) + parser.add_argument( + "--text", + type=str, + default="Hello, this is a test of the OmniVoice text to speech system.", + ) + parser.add_argument( + "--ref-audio", + type=str, + default=None, + help="Reference audio for voice cloning (WAV file)", + ) + parser.add_argument( + "--ref-text", + type=str, + default=None, + help="Transcription of reference audio", + ) + parser.add_argument( + "--lang", + type=str, + default=None, + help="Language code (e.g., 'en', 'zh')", + ) + parser.add_argument( + "--instruct", + type=str, + default=None, + help="Voice design instruction (e.g., 'female, low pitch, british accent')", + ) + parser.add_argument( + "--output", + type=str, + default="output.wav", + help="Output audio file path", + ) + parser.add_argument( + "--stage-init-timeout", + type=int, + default=600, + help="Stage initialization timeout in seconds", + ) + args = parser.parse_args() + + if not os.path.exists(args.stage_config): + raise FileNotFoundError(f"Stage config not found: {args.stage_config}") + + print(f"Initializing OmniVoice with model={args.model}") + + omni = Omni( + model=args.model, + stage_configs_path=args.stage_config, + trust_remote_code=True, + log_stats=True, + ) + + print("Model initialized. Preparing inputs...") + + # Build prompt + mm_processor_kwargs = {} + multi_modal_data = {} + + if args.ref_audio: + if not os.path.exists(args.ref_audio): + raise FileNotFoundError(f"Reference audio not found: {args.ref_audio}") + + import librosa + + audio_signal, sr = librosa.load(args.ref_audio, sr=None) + multi_modal_data["audio"] = (audio_signal.astype(np.float32), sr) + mm_processor_kwargs["ref_text"] = args.ref_text or "" + mm_processor_kwargs["sample_rate"] = sr + + if args.lang: + mm_processor_kwargs["lang"] = args.lang + if args.instruct: + mm_processor_kwargs["instruct"] = args.instruct + + prompts = {"prompt": args.text} + if multi_modal_data: + prompts["multi_modal_data"] = multi_modal_data + if mm_processor_kwargs: + prompts["mm_processor_kwargs"] = mm_processor_kwargs + + sampling_params_list = [OmniDiffusionSamplingParams()] + + print(f"Generating speech for: {args.text}") + + outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list)) + + print(f"Received {len(outputs)} outputs.") + for i, output in enumerate(outputs): + try: + ro = output.request_output + if ro is None: + print("No request_output found.") + continue + + mm = getattr(ro, "multimodal_output", None) + if not mm and ro.outputs: + mm = getattr(ro.outputs[0], "multimodal_output", None) + + if mm: + print(f"Multimodal output keys: {mm.keys()}") + if "audio" in mm: + audio_out = mm["audio"] + sr = mm.get("sr", 24000) + if isinstance(audio_out, np.ndarray): + audio_np = audio_out + else: + audio_np = audio_out.cpu().numpy().squeeze() + out_path = args.output if i == 0 else f"output_{i}.wav" + sf.write(out_path, audio_np, sr) + print(f"Saved audio to {out_path} ({sr}Hz, {len(audio_np) / sr:.2f}s)") + else: + print("No multimodal output found.") + except Exception as e: + print(f"Error inspecting output: {e}") + + omni.close() + print("Done.") + + +if __name__ == "__main__": + run_e2e() diff --git a/examples/online_serving/omnivoice/README.md b/examples/online_serving/omnivoice/README.md new file mode 100644 index 000000000..1d8f00421 --- /dev/null +++ b/examples/online_serving/omnivoice/README.md @@ -0,0 +1,131 @@ +# OmniVoice + +## Model Overview + +| Model | Description | +|-------|-------------| +| `k2-fsa/OmniVoice` | Zero-shot multilingual TTS (600+ languages) with diffusion language model (Qwen3-0.6B backbone) | + +> **Note:** Requires `transformers>=5.3.0` for voice cloning (HiggsAudioV2 tokenizer). Auto voice and voice design work with `transformers>=4.57.0`. + +## Launch the Server + +```bash +vllm serve k2-fsa/OmniVoice \ + --omni \ + --port 8091 \ + --trust-remote-code +``` + +Or use the convenience script: + +```bash +./run_server.sh +``` + +## Send TTS Request + +### Using curl + +```bash +# Basic TTS (auto voice) +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, how are you?", + "voice": "default", + "response_format": "wav" + }' --output output.wav +``` + +### Using Python + +```python +import httpx + +response = httpx.post( + "http://localhost:8091/v1/audio/speech", + json={ + "input": "Hello, how are you?", + "voice": "default", + "response_format": "wav", + }, + timeout=300.0, +) + +with open("output.wav", "wb") as f: + f.write(response.content) +``` + +### Using OpenAI SDK + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8091/v1", api_key="none") + +response = client.audio.speech.create( + model="k2-fsa/OmniVoice", + voice="default", + input="Hello, how are you?", +) + +response.stream_to_file("output.wav") +``` + +### Using the CLI Client + +```bash +cd examples/online_serving/omnivoice + +# Basic TTS +python speech_client.py --text "Hello, how are you?" + +# Specify language for improved quality +python speech_client.py --text "Bonjour, comment allez-vous?" --language French +``` + +The CLI client supports: + +- `--api-base`: API base URL (default: `http://localhost:8091`) +- `--model` (or `-m`): Model name (default: `k2-fsa/OmniVoice`) +- `--text`: Text to synthesize (required) +- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav) +- `--language`: Language hint (default: Auto) +- `--output` (or `-o`): Output file path (default: `omnivoice_output.wav`) + +## Inference Modes + +OmniVoice supports three inference modes. Currently, **auto voice** is supported through the online Speech API. Voice cloning and voice design are available via offline inference (see `examples/offline_inference/omnivoice/`). + +| Mode | Description | Online API | Offline | +|------|-------------|:----------:|:-------:| +| Auto Voice | Generate speech without reference | Yes | Yes | +| Voice Clone | Clone from reference audio + transcript | - | Yes | +| Voice Design | Control style via natural language instruction | - | Yes | + +## Architecture + +OmniVoice uses a single-stage diffusion pipeline: + +- **Stage 0 (Generator)**: Qwen3-0.6B transformer with 32-step iterative masked unmasking and classifier-free guidance. Generates 8-codebook audio tokens from text, then decodes to 24kHz waveform via HiggsAudioV2 RVQ quantizer + DAC acoustic decoder. + +## API Parameters + +OmniVoice uses the standard `/v1/audio/speech` endpoint. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full documentation. + +Key parameters: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `input` | string | **required** | Text to synthesize | +| `voice` | string | "default" | Voice name | +| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus | +| `speed` | float | 1.0 | Playback speed (0.25-4.0) | + +## Troubleshooting + +1. **TTS model did not produce audio output**: Ensure the model is fully downloaded (`huggingface-cli download k2-fsa/OmniVoice`) +2. **Connection refused**: Make sure the server is running on the correct port +3. **Out of memory**: Reduce `--gpu-memory-utilization` (default stage config uses 0.5) +4. **Slow first request**: The model performs warmup on first inference; subsequent requests are faster diff --git a/examples/online_serving/omnivoice/run_server.sh b/examples/online_serving/omnivoice/run_server.sh new file mode 100755 index 000000000..abe9bb798 --- /dev/null +++ b/examples/online_serving/omnivoice/run_server.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Launch vLLM-Omni server for OmniVoice TTS +# +# Usage: +# ./run_server.sh +# CUDA_VISIBLE_DEVICES=0 ./run_server.sh + +set -e + +MODEL="${MODEL:-k2-fsa/OmniVoice}" +PORT="${PORT:-8091}" + +echo "Starting OmniVoice server with model: $MODEL" + +vllm serve "$MODEL" \ + --host 0.0.0.0 \ + --port "$PORT" \ + --trust-remote-code \ + --omni diff --git a/examples/online_serving/omnivoice/speech_client.py b/examples/online_serving/omnivoice/speech_client.py new file mode 100644 index 000000000..b8e6f3889 --- /dev/null +++ b/examples/online_serving/omnivoice/speech_client.py @@ -0,0 +1,84 @@ +"""Client for OmniVoice TTS via /v1/audio/speech endpoint. + +Examples: + # Basic TTS (auto voice) + python speech_client.py --text "Hello, how are you?" + + # Specify language + python speech_client.py --text "Bonjour, comment allez-vous?" --language French +""" + +import argparse + +import httpx + +DEFAULT_API_BASE = "http://localhost:8091" +DEFAULT_API_KEY = "EMPTY" + + +def run_tts(args) -> None: + """Generate speech via /v1/audio/speech API.""" + payload = { + "model": args.model, + "input": args.text, + "voice": "default", + "response_format": args.response_format, + } + + if args.language: + payload["language"] = args.language + + print(f"Model: {args.model}") + print(f"Text: {args.text}") + if args.language: + print(f"Language: {args.language}") + print("Generating audio...") + + api_url = f"{args.api_base}/v1/audio/speech" + headers = { + "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) + + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.text) + return + + try: + text = response.content.decode("utf-8") + if text.startswith('{"error"'): + print(f"Error: {text}") + return + except UnicodeDecodeError: + pass + + output_path = args.output or "omnivoice_output.wav" + with open(output_path, "wb") as f: + f.write(response.content) + print(f"Audio saved to: {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description="OmniVoice TTS client") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--model", "-m", default="k2-fsa/OmniVoice", help="Model name") + parser.add_argument("--text", required=True, help="Text to synthesize") + parser.add_argument("--language", default=None, help="Language hint (e.g., English, Chinese, French)") + parser.add_argument( + "--response-format", + default="wav", + choices=["wav", "mp3", "flac", "pcm", "aac", "opus"], + help="Audio format (default: wav)", + ) + parser.add_argument("--output", "-o", default=None, help="Output file path") + args = parser.parse_args() + run_tts(args) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 15e7c6305..e49aa6e32 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -223,3 +223,4 @@ extend-ignore-identifiers-re = [ ue = "ue" semantics = "semantics" fullset = "fullset" +Vai = "Vai" diff --git a/tests/e2e/offline_inference/test_omnivoice.py b/tests/e2e/offline_inference/test_omnivoice.py new file mode 100644 index 000000000..4b093e357 --- /dev/null +++ b/tests/e2e/offline_inference/test_omnivoice.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E offline tests for OmniVoice TTS model with text input and audio output. + +Uses GPUGenerationWorker for both stages (iterative unmasking + DAC decoder). +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +from pathlib import Path + +import numpy as np +import pytest + +from tests.utils import hardware_test + +MODEL = "k2-fsa/OmniVoice" + + +def get_stage_config(): + return str( + Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "omnivoice.yaml" + ) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_omnivoice_text_to_audio() -> None: + """ + Test OmniVoice text-to-audio generation via offline Omni runner. + Deploy Setting: omnivoice.yaml (enforce_eager=true) + Input Modal: text + Output Modal: audio + """ + from vllm_omni.entrypoints.omni import Omni + + omni = Omni( + model=MODEL, + stage_configs_path=get_stage_config(), + trust_remote_code=True, + log_stats=True, + ) + + try: + prompts = {"prompt": "Hello, this is a test for text to audio."} + + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + sampling_params_list = [OmniDiffusionSamplingParams()] + + outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list)) + + assert len(outputs) > 0, "No outputs generated" + + # Check final output has audio + final_output = outputs[-1] + ro = final_output.request_output + assert ro is not None, "No request_output" + + mm = getattr(ro, "multimodal_output", None) + if not mm and ro.outputs: + mm = getattr(ro.outputs[0], "multimodal_output", None) + + assert mm is not None, "No multimodal_output" + assert "audio" in mm, f"No 'audio' key in multimodal_output: {mm.keys()}" + + audio = mm["audio"] + if isinstance(audio, np.ndarray): + audio_np = audio + else: + audio_np = audio.cpu().numpy().squeeze() + + assert audio_np.size > 0, "Audio output is empty" + rms = np.sqrt(np.mean(audio_np**2)) + assert rms > 0.01, f"Audio RMS too low ({rms:.4f}), likely silence" + + print(f"Generated audio: {len(audio_np) / 24000:.2f}s, rms={rms:.4f}") + finally: + omni.close() diff --git a/tests/e2e/online_serving/test_omnivoice.py b/tests/e2e/online_serving/test_omnivoice.py new file mode 100644 index 000000000..ec1981aab --- /dev/null +++ b/tests/e2e/online_serving/test_omnivoice.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E Online tests for OmniVoice TTS model via /v1/audio/speech endpoint. + +Tests verify that the OmniVoice model generates valid audio when +accessed through the standard OpenAI-compatible speech API. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +from pathlib import Path + +import httpx +import pytest + +from tests.conftest import OmniServerParams +from tests.utils import hardware_test + +MODEL = "k2-fsa/OmniVoice" + +STAGE_CONFIG = str( + Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "omnivoice.yaml" +) +EXTRA_ARGS = [ + "--trust-remote-code", + "--disable-log-stats", +] +TEST_PARAMS = [ + OmniServerParams( + model=MODEL, + stage_config_path=STAGE_CONFIG, + server_args=EXTRA_ARGS, + ) +] + +MIN_AUDIO_BYTES = 5000 + + +def make_speech_request( + host: str, + port: int, + text: str, + timeout: float = 180.0, +) -> httpx.Response: + """Make a request to the /v1/audio/speech endpoint for OmniVoice.""" + url = f"http://{host}:{port}/v1/audio/speech" + payload = {"input": text} + + with httpx.Client(timeout=timeout) as client: + return client.post(url, json=payload) + + +def verify_wav_audio(content: bytes) -> bool: + """Verify that content is valid WAV audio data.""" + if len(content) < 44: + return False + return content[:4] == b"RIFF" and content[8:12] == b"WAVE" + + +@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) +class TestOmniVoiceTTS: + """E2E tests for OmniVoice TTS model.""" + + @pytest.mark.core_model + @pytest.mark.omni + @hardware_test(res={"cuda": "L4"}, num_cards=1) + def test_speech_auto_voice(self, omni_server) -> None: + """Test auto voice TTS generation (text only, no reference audio).""" + response = make_speech_request( + host=omni_server.host, + port=omni_server.port, + text="Hello, this is a test of the OmniVoice text to speech system.", + ) + + assert response.status_code == 200, f"Request failed: {response.text}" + assert response.headers.get("content-type") == "audio/wav" + assert verify_wav_audio(response.content), "Response is not valid WAV audio" + assert len(response.content) > MIN_AUDIO_BYTES, ( + f"Audio too small ({len(response.content)} bytes), expected > {MIN_AUDIO_BYTES}" + ) diff --git a/vllm_omni/diffusion/models/omnivoice/__init__.py b/vllm_omni/diffusion/models/omnivoice/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/diffusion/models/omnivoice/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py new file mode 100644 index 000000000..568e2f516 --- /dev/null +++ b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +OmniVoice TTS Pipeline for vLLM-Omni diffusion engine. + +Single-stage pipeline that runs the full text-to-speech flow: + text → tokenize → 32-step iterative unmasking → 8-codebook tokens → DAC decode → 24kHz audio + +Uses request-mode execution (all steps in one forward() call). +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Iterable +from typing import ClassVar + +import torch +from tokenizers import Tokenizer as HFTokenizer +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.models.interface import SupportAudioOutput +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig +from vllm_omni.model_executor.models.omnivoice.duration import RuleDurationEstimator +from vllm_omni.model_executor.models.omnivoice.omnivoice_decoder import OmniVoiceDecoder +from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import OmniVoiceGenerator + +logger = init_logger(__name__) + + +def get_omnivoice_post_process_func(od_config: OmniDiffusionConfig): + """Post-processing: convert audio tensor to numpy for WAV encoding.""" + + def post_process_func(audio: torch.Tensor, output_type: str = "np"): + if output_type == "pt": + return audio + return audio.cpu().float().numpy() + + return post_process_func + + +class OmniVoicePipeline(nn.Module, SupportAudioOutput): + """OmniVoice text-to-speech pipeline for the diffusion engine. + + Wraps OmniVoiceGenerator (32-step iterative unmasking) and + OmniVoiceDecoder (HiggsAudioV2 RVQ + DAC) into a single forward() call. + """ + + support_audio_output: ClassVar[bool] = True + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__() + self.od_config = od_config + self.device = get_local_device() + self.model_path = od_config.model + + # Resolve model path (HF hub ID → local cache) + if not os.path.isdir(self.model_path): + from huggingface_hub import snapshot_download + + self.model_path = snapshot_download(self.model_path) + + # Load OmniVoice config + config_path = os.path.join(self.model_path, "config.json") + with open(config_path) as f: + hf_config = json.load(f) + self.config = OmniVoiceConfig(**hf_config) + + # Build generator and decoder + self.generator = OmniVoiceGenerator(self.config) + self.decoder = OmniVoiceDecoder(self.config) + + # Tokenizer (low-level, avoids HF tokenizer extra_special_tokens issue) + tokenizer_path = os.path.join(self.model_path, "tokenizer.json") + self.tokenizer = HFTokenizer.from_file(tokenizer_path) + + # Duration estimator + self.duration_estimator = RuleDurationEstimator() + + # Generation parameters + self.num_step = self.config.num_step + self.guidance_scale = self.config.guidance_scale + self.t_shift = self.config.t_shift + self.layer_penalty_factor = self.config.layer_penalty_factor + self.position_temperature = self.config.position_temperature + self.class_temperature = self.config.class_temperature + self.sample_rate = self.config.sample_rate + + @torch.inference_mode() + def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + """Generate speech audio from text. + + Args: + req: Diffusion request containing text prompt(s). + + Returns: + DiffusionOutput with audio tensor in .output + """ + # Extract text from request + prompt = req.prompts[0] if req.prompts else "" + if isinstance(prompt, dict): + text = prompt.get("input", prompt.get("text", str(prompt))) + else: + text = str(prompt) + + if not text: + return DiffusionOutput(error="Empty text prompt") + + device = self.device + num_cb = self.config.num_audio_codebook + mask_id = self.config.audio_mask_id + + # Estimate target duration + target_len = self.duration_estimator.estimate_duration(text, "Nice to meet you.", 25) + target_len = max(1, int(target_len)) + + # Tokenize with control tokens + style = "<|denoise|><|lang_start|>None<|lang_end|><|instruct_start|>None<|instruct_end|>" + full_prompt = f"{style}<|text_start|>{text}<|text_end|>" + encoding = self.tokenizer.encode(full_prompt) + text_tokens = torch.tensor(encoding.ids, dtype=torch.long, device=device) + text_len = text_tokens.shape[0] + + # Build conditional + unconditional batches [2, 8, max_len] + text_ids = text_tokens.unsqueeze(0).repeat(num_cb, 1) + target_ids = torch.full((num_cb, target_len), mask_id, dtype=torch.long, device=device) + cond_ids = torch.cat([text_ids, target_ids], dim=1) + cond_len = cond_ids.shape[1] + + uncond_ids = target_ids.clone() + uncond_len = target_len + max_len = max(cond_len, uncond_len) + if uncond_len < max_len: + pad = torch.full( + (num_cb, max_len - uncond_len), + mask_id, + dtype=torch.long, + device=device, + ) + uncond_ids = torch.cat([uncond_ids, pad], dim=1) + + batch_input_ids = torch.stack([cond_ids, uncond_ids]) + + batch_audio_mask = torch.zeros(2, max_len, dtype=torch.bool, device=device) + batch_audio_mask[0, text_len:cond_len] = True + batch_audio_mask[1, :uncond_len] = True + + batch_attn_mask = torch.zeros(2, 1, max_len, max_len, dtype=torch.bool, device=device) + batch_attn_mask[0, :, :cond_len, :cond_len] = True + batch_attn_mask[1, :, :uncond_len, :uncond_len] = True + + # Run 32-step iterative unmasking + tokens = self.generator( + input_ids=batch_input_ids, + audio_mask=batch_audio_mask, + attention_mask=batch_attn_mask, + target_lens=[target_len], + num_step=self.num_step, + guidance_scale=self.guidance_scale, + t_shift=self.t_shift, + layer_penalty_factor=self.layer_penalty_factor, + position_temperature=self.position_temperature, + class_temperature=self.class_temperature, + ) + + # Decode tokens to audio + audio = self.decoder(tokens) # [1, 1, samples] + + return DiffusionOutput(output=audio) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights from model directory (not from the iterator). + + The diffusion model loader passes HF safetensors weights, but OmniVoice + has custom weight names (llm.* → generator.*, audio_tokenizer.* → decoder.*). + We load from model_path directly and return all param names to satisfy + the loader's "all weights initialized" check. + """ + # Consume the iterator (required by the loader contract) + for _ in weights: + pass + + device = self.device + self.generator.load_weights(self.model_path, device) + self.generator = self.generator.to(device).eval() + self.decoder.load_weights(self.model_path, device) + logger.info("OmniVoice pipeline loaded on %s", device) + + # Return all parameter names to indicate they're initialized + return {name for name, _ in self.named_parameters()} diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index db8805722..c1f48137e 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -173,6 +173,16 @@ _DIFFUSION_MODELS = { "pipeline_hunyuan_video_1_5_i2v", "HunyuanVideo15I2VPipeline", ), + "OmniVoicePipeline": ( + "omnivoice", + "pipeline_omnivoice", + "OmniVoicePipeline", + ), + "OmniVoice": ( + "omnivoice", + "pipeline_omnivoice", + "OmniVoicePipeline", + ), } @@ -358,6 +368,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "Flux2Pipeline": "get_flux2_post_process_func", "HunyuanVideo15Pipeline": "get_hunyuan_video_15_post_process_func", "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_post_process_func", + "OmniVoicePipeline": "get_omnivoice_post_process_func", } _DIFFUSION_PRE_PROCESS_FUNCS = { diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index f4a082cff..a1dc373dd 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -18,6 +18,7 @@ def _register_omni_hf_configs() -> None: from transformers import AutoConfig from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config + from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( Qwen3TTSConfig, ) @@ -31,6 +32,7 @@ def _register_omni_hf_configs() -> None: for model_type, config_cls in [ ("qwen3_tts", Qwen3TTSConfig), ("cosyvoice3", CosyVoice3Config), + ("omnivoice", OmniVoiceConfig), ("voxtral_tts", VoxtralTTSConfig), ]: try: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 0ffe33abd..acf45b4fe 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -510,6 +510,13 @@ async def omni_init_app_state( stage_configs=diffusion_stage_configs, ) + state.openai_serving_speech = OmniOpenAIServingSpeech.for_diffusion( + diffusion_engine=engine_client, + model_name=model_name, + stage_configs=diffusion_stage_configs, + ) + state.openai_streaming_speech = None + state.enable_server_load_tracking = getattr(args, "enable_server_load_tracking", False) state.server_load_metrics = 0 logger.info("Pure diffusion API server initialized for model: %s", model_name) diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index b483181fd..75279f075 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -47,7 +47,10 @@ logger = init_logger(__name__) _VOXTRAL_TTS_MODEL_STAGES = {"audio_generation"} _QWEN3_TTS_MODEL_STAGES = {"qwen3_tts"} _FISH_TTS_MODEL_STAGES = {"fish_speech_slow_ar"} -_TTS_MODEL_STAGES: set[str] = _VOXTRAL_TTS_MODEL_STAGES | _QWEN3_TTS_MODEL_STAGES | _FISH_TTS_MODEL_STAGES +_OMNIVOICE_TTS_MODEL_STAGES = {"omnivoice_generator"} +_TTS_MODEL_STAGES: set[str] = ( + _VOXTRAL_TTS_MODEL_STAGES | _QWEN3_TTS_MODEL_STAGES | _FISH_TTS_MODEL_STAGES | _OMNIVOICE_TTS_MODEL_STAGES +) _TTS_LANGUAGES: set[str] = { "Auto", "Chinese", @@ -145,6 +148,27 @@ def _validate_path_within_directory(file_path: Path, directory: Path) -> bool: class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): + _diffusion_mode: bool = False + + @classmethod + def for_diffusion( + cls, + diffusion_engine: "Any", + model_name: str, + stage_configs: "list[Any] | None" = None, + ) -> "OmniOpenAIServingSpeech": + """Create a speech serving instance for pure diffusion TTS models. + + Bypasses OpenAIServing.__init__ which requires a fully configured + engine client that pure diffusion engines don't provide. + """ + instance = cls.__new__(cls) + instance._diffusion_mode = True + instance._diffusion_engine = diffusion_engine + instance._diffusion_model_name = model_name + instance._diffusion_stage_configs = stage_configs + return instance + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # Initialize uploaded speakers storage @@ -240,6 +264,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return "voxtral_tts" if model_stage in _FISH_TTS_MODEL_STAGES: return "fish_tts" + if model_stage in _OMNIVOICE_TTS_MODEL_STAGES: + return "omnivoice" return None def _compute_max_instructions_length(self) -> int: @@ -1203,6 +1229,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ref_audio_data = (wav_list, sr) prompt = self._build_fish_speech_prompt(request, ref_audio_data=ref_audio_data) tts_params = {} + elif self._tts_model_type == "omnivoice": + tts_params = {} + prompt = request.input # Diffusion engine takes raw text elif self._is_tts: validation_error = self._validate_tts_request(request) if validation_error: @@ -1324,6 +1353,79 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): audio_response: AudioResponse = self.create_audio(audio_obj) return audio_response.audio_data, audio_response.media_type + async def _create_diffusion_speech( + self, + request: OpenAICreateSpeechRequest, + ) -> Response: + """Handle speech generation for pure diffusion TTS models (e.g. OmniVoice).""" + from vllm_omni.outputs import OmniRequestOutput + + try: + request_id = f"speech-{random_uuid()}" + prompt = request.input + + logger.info( + "Diffusion TTS speech request %s: text=%r", + request_id, + prompt[:50] + "..." if len(prompt) > 50 else prompt, + ) + + generator = self._diffusion_engine.generate( + prompt=prompt, + request_id=request_id, + sampling_params_list=self._diffusion_engine.default_sampling_params_list, + output_modalities=["audio"], + ) + + final_output: OmniRequestOutput | None = None + async for res in generator: + final_output = res + + if final_output is None: + raise ValueError("No output generated from the model.") + + audio_output, audio_key = self._extract_audio_output(final_output) + if audio_key is None: + raise ValueError("TTS model did not produce audio output.") + + audio_tensor = audio_output[audio_key] + sr_raw = audio_output.get("sr", 24000) + sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw + sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val) + + if isinstance(audio_tensor, list): + non_empty = [c for c in audio_tensor if c.numel() > 0] + audio_tensor = torch.cat(non_empty, dim=-1) if non_empty else np.zeros((0,), dtype=np.float32) + if hasattr(audio_tensor, "float"): + audio_tensor = audio_tensor.float().detach().cpu().numpy() + if audio_tensor.ndim > 1: + audio_tensor = audio_tensor.squeeze() + + audio_obj = CreateAudio( + audio_tensor=audio_tensor, + sample_rate=sample_rate, + response_format=request.response_format or "wav", + speed=request.speed or 1.0, + stream_format=request.stream_format, + base64_encode=False, + ) + audio_response: AudioResponse = self.create_audio(audio_obj) + return Response(content=audio_response.audio_data, media_type=audio_response.media_type) + + except asyncio.CancelledError: + return self._diffusion_error_response("Client disconnected") + except ValueError as e: + return self._diffusion_error_response(str(e)) + except Exception as e: + logger.exception("Diffusion speech generation failed: %s", e) + return self._diffusion_error_response(f"Speech generation failed: {e}") + + @staticmethod + def _diffusion_error_response(message: str) -> Response: + """Create a JSON error response without depending on OpenAIServing.""" + error_body = json.dumps({"error": {"message": message, "type": "server_error", "param": None, "code": 500}}) + return Response(content=error_body, media_type="application/json", status_code=500) + async def create_speech( self, request: OpenAICreateSpeechRequest, @@ -1349,6 +1451,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): Each Code2Wav chunk is yielded as raw audio bytes as soon as it is decoded. For WAV format, a header with placeholder size values is emitted first. """ + if self._diffusion_mode: + return await self._create_diffusion_speech(request) + error_check_ret = await self._check_model(request) if error_check_ret is not None: logger.error("Error with model %s", error_check_ret) @@ -1426,6 +1531,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): batch_request: BatchSpeechRequest, ) -> BatchSpeechResponse | ErrorResponse: """Generate speech for multiple items concurrently.""" + if self._diffusion_mode: + raise ValueError("Batch speech is not supported in diffusion mode") if len(batch_request.items) > self._batch_max_items: raise ValueError( f"Batch contains {len(batch_request.items)} items, exceeding the maximum of {self._batch_max_items}." diff --git a/vllm_omni/model_executor/models/omnivoice/__init__.py b/vllm_omni/model_executor/models/omnivoice/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/model_executor/models/omnivoice/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/model_executor/models/omnivoice/config.py b/vllm_omni/model_executor/models/omnivoice/config.py new file mode 100644 index 000000000..a24176bcf --- /dev/null +++ b/vllm_omni/model_executor/models/omnivoice/config.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""OmniVoice configuration for vLLM-Omni two-stage pipeline.""" + +from transformers.configuration_utils import PretrainedConfig + + +class OmniVoiceConfig(PretrainedConfig): + """Configuration for OmniVoice model in vLLM-Omni. + + This mirrors the HuggingFace OmniVoiceConfig but adds fields needed + for the two-stage serving pipeline. + """ + + model_type = "omnivoice" + + def get_text_config(self, **kwargs): + """Return self so vLLM uses our top-level config (which has + num_attention_heads etc.) instead of trying to extract a sub-config.""" + return self + + def __init__(self, **kwargs): + # HF repos (e.g. k2-fsa/OmniVoice) may nest generation hyperparameters. + gen_cfg = kwargs.pop("generation_config", None) + if isinstance(gen_cfg, dict): + for k, v in gen_cfg.items(): + kwargs.setdefault(k, v) + + super().__init__(**kwargs) + + # Audio codec params (prefer values set by PretrainedConfig from config.json) + self.audio_vocab_size = getattr(self, "audio_vocab_size", 1025) + self.audio_mask_id = getattr(self, "audio_mask_id", 1024) + self.num_audio_codebook = getattr(self, "num_audio_codebook", 8) + self.audio_codebook_weights = getattr( + self, + "audio_codebook_weights", + [8, 8, 6, 6, 4, 4, 2, 2], + ) + + # LLM backbone params (Qwen3-0.6B defaults from HF config) + llm_config = getattr(self, "llm_config", None) or {} + if isinstance(llm_config, PretrainedConfig): + llm_config = llm_config.to_dict() + elif not isinstance(llm_config, dict): + llm_config = {} + self.llm_hidden_size = llm_config.get("hidden_size", 1024) + self.llm_num_hidden_layers = llm_config.get("num_hidden_layers", 28) + self.llm_num_attention_heads = llm_config.get("num_attention_heads", 16) + self.llm_num_key_value_heads = llm_config.get("num_key_value_heads", 8) + self.llm_intermediate_size = llm_config.get("intermediate_size", 3072) + self.llm_vocab_size = llm_config.get("vocab_size", 151676) + self.llm_max_position_embeddings = llm_config.get("max_position_embeddings", 40960) + self.llm_rope_theta = llm_config.get("rope_theta", 1000000.0) + self.llm_rms_norm_eps = llm_config.get("rms_norm_eps", 1e-6) + self.llm_head_dim = llm_config.get("head_dim", self.llm_hidden_size // self.llm_num_attention_heads) + + # Expose LLM params at top level for vLLM ModelConfig compatibility + # (vLLM expects num_attention_heads, hidden_size, etc. on the config) + self.num_attention_heads = self.llm_num_attention_heads + self.num_key_value_heads = self.llm_num_key_value_heads + self.num_hidden_layers = self.llm_num_hidden_layers + self.hidden_size = self.llm_hidden_size + self.head_dim = self.llm_head_dim + if not hasattr(self, "vocab_size"): + self.vocab_size = self.llm_vocab_size + + # Generation params (defaults from OmniVoiceGenerationConfig) + self.num_step = getattr(self, "num_step", 32) + self.guidance_scale = getattr(self, "guidance_scale", 2.0) + self.t_shift = getattr(self, "t_shift", 0.1) + self.layer_penalty_factor = getattr(self, "layer_penalty_factor", 5.0) + self.position_temperature = getattr(self, "position_temperature", 5.0) + self.class_temperature = getattr(self, "class_temperature", 0.0) + + # Audio output + self.sample_rate = getattr(self, "sample_rate", 24000) + self.frame_rate = getattr(self, "frame_rate", 25) + + # Serving + self.speculative_config = None diff --git a/vllm_omni/model_executor/models/omnivoice/duration.py b/vllm_omni/model_executor/models/omnivoice/duration.py new file mode 100644 index 000000000..8343362a2 --- /dev/null +++ b/vllm_omni/model_executor/models/omnivoice/duration.py @@ -0,0 +1,281 @@ +#!/usr/bin/env python3 +# Copyright 2026 Xiaomi Corp. (authors: Han Zhu) +# +# See ../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Text duration estimation for TTS generation. + +Provides ``RuleDurationEstimator``, which estimates audio duration from text +using character phonetic weights across 600+ languages. Used by +``OmniVoice.generate()`` to determine output length when no duration is specified. +""" + +import bisect +import unicodedata +from functools import lru_cache + + +class RuleDurationEstimator: + def __init__(self): + # ========================================== + # 1. Phonetic Weights Table + # ========================================== + # The weight represents the relative speaking time compared to + # a standard Latin letter. + # Benchmark: 1.0 = One Latin Character (~40-50ms) + self.weights = { + # --- Logographic (1 char = full syllable/word) --- + "cjk": 3.0, # Chinese, Japanese Kanji, etc. + # --- Syllabic / Blocks + "hangul": 2.5, # Korean Hangul + "kana": 2.2, # Japanese Hiragana/Katakana + "ethiopic": 3.0, # Amharic/Ge'ez + "yi": 3.0, # Yi script + # --- Abugida (Consonant-Vowel complexes) --- + "indic": 1.8, # Hindi, Bengali, Tamil, etc. + "thai_lao": 1.5, # Thai, Lao + "khmer_myanmar": 1.8, # Khmer, Myanmar + # --- Abjad (Consonant-heavy) --- + "arabic": 1.5, # Arabic, Persian, Urdu + "hebrew": 1.5, # Hebrew + # --- Alphabet (Segmental) --- + "latin": 1.0, # English, Spanish, French, Vietnamese, etc. (Baseline) + "cyrillic": 1.0, # Russian, Ukrainian + "greek": 1.0, # Greek + "armenian": 1.0, # Armenian + "georgian": 1.0, # Georgian + # --- Symbols & Misc --- + "punctuation": 0.5, # Pause capability + "space": 0.2, # Word boundary/Breath (0.05 / 0.22) + "digit": 3.5, # Numbers + "mark": 0.0, # Diacritics/Accents (Silent modifiers) + "default": 1.0, # Fallback for unknown scripts + } + + # ========================================== + # 2. Unicode Range Mapping + # ========================================== + # Format: (End_Codepoint, Type_Key) + # Used for fast binary search (bisect). + self.ranges = [ + (0x02AF, "latin"), # Latin (Basic, Supplement, Ext, IPA) + (0x03FF, "greek"), # Greek & Coptic + (0x052F, "cyrillic"), # Cyrillic + (0x058F, "armenian"), # Armenian + (0x05FF, "hebrew"), # Hebrew + (0x077F, "arabic"), # Arabic, Syriac, Arabic Supplement + (0x089F, "arabic"), # Arabic Extended-B (+ Syriac Supp) + (0x08FF, "arabic"), # Arabic Extended-A + (0x097F, "indic"), # Devanagari + (0x09FF, "indic"), # Bengali + (0x0A7F, "indic"), # Gurmukhi + (0x0AFF, "indic"), # Gujarati + (0x0B7F, "indic"), # Oriya + (0x0BFF, "indic"), # Tamil + (0x0C7F, "indic"), # Telugu + (0x0CFF, "indic"), # Kannada + (0x0D7F, "indic"), # Malayalam + (0x0DFF, "indic"), # Sinhala + (0x0EFF, "thai_lao"), # Thai & Lao + (0x0FFF, "indic"), # Tibetan (Abugida) + (0x109F, "khmer_myanmar"), # Myanmar + (0x10FF, "georgian"), # Georgian + (0x11FF, "hangul"), # Hangul Jamo + (0x137F, "ethiopic"), # Ethiopic + (0x139F, "ethiopic"), # Ethiopic Supplement + (0x13FF, "default"), # Cherokee + (0x167F, "default"), # Canadian Aboriginal Syllabics + (0x169F, "default"), # Ogham + (0x16FF, "default"), # Runic + (0x171F, "default"), # Tagalog (Baybayin) + (0x173F, "default"), # Hanunoo + (0x175F, "default"), # Buhid + (0x177F, "default"), # Tagbanwa + (0x17FF, "khmer_myanmar"), # Khmer + (0x18AF, "default"), # Mongolian + (0x18FF, "default"), # Canadian Aboriginal Syllabics Ext + (0x194F, "indic"), # Limbu + (0x19DF, "indic"), # Tai Le & New Tai Lue + (0x19FF, "khmer_myanmar"), # Khmer Symbols + (0x1A1F, "indic"), # Buginese + (0x1AAF, "indic"), # Tai Than + (0x1B7F, "indic"), # Balinese + (0x1BBF, "indic"), # Sundanese + (0x1BFF, "indic"), # Batak + (0x1C4F, "indic"), # Lepcha + (0x1C7F, "indic"), # Ol Chiki (Santali) + (0x1C8F, "cyrillic"), # Cyrillic Extended-C + (0x1CBF, "georgian"), # Georgian Extended + (0x1CCF, "indic"), # Sundanese Supplement + (0x1CFF, "indic"), # Vedic Extensions + (0x1D7F, "latin"), # Phonetic Extensions + (0x1DBF, "latin"), # Phonetic Extensions Supplement + (0x1DFF, "default"), # Combining Diacritical Marks Supplement + (0x1EFF, "latin"), # Latin Extended Additional (Vietnamese) + (0x309F, "kana"), # Hiragana + (0x30FF, "kana"), # Katakana + (0x312F, "cjk"), # Bopomofo (Pinyin) + (0x318F, "hangul"), # Hangul Compatibility Jamo + (0x9FFF, "cjk"), # CJK Unified Ideographs (Main) + (0xA4CF, "yi"), # Yi Syllables + (0xA4FF, "default"), # Lisu + (0xA63F, "default"), # Vai + (0xA69F, "cyrillic"), # Cyrillic Extended-B + (0xA6FF, "default"), # Bamum + (0xA7FF, "latin"), # Latin Extended-D + (0xA82F, "indic"), # Syloti Nagri + (0xA87F, "default"), # Phags-pa + (0xA8DF, "indic"), # Saurashtra + (0xA8FF, "indic"), # Devanagari Extended + (0xA92F, "indic"), # Kayah Li + (0xA95F, "indic"), # Rejang + (0xA97F, "hangul"), # Hangul Jamo Extended-A + (0xA9DF, "indic"), # Javanese + (0xA9FF, "khmer_myanmar"), # Myanmar Extended-B + (0xAA5F, "indic"), # Cham + (0xAA7F, "khmer_myanmar"), # Myanmar Extended-A + (0xAADF, "indic"), # Tai Viet + (0xAAFF, "indic"), # Meetei Mayek Extensions + (0xAB2F, "ethiopic"), # Ethiopic Extended-A + (0xAB6F, "latin"), # Latin Extended-E + (0xABBF, "default"), # Cherokee Supplement + (0xABFF, "indic"), # Meetei Mayek + (0xD7AF, "hangul"), # Hangul Syllables + (0xFAFF, "cjk"), # CJK Compatibility + (0xFDFF, "arabic"), # Arabic Presentation Forms-A + (0xFE6F, "default"), # Variation Selectors + (0xFEFF, "arabic"), # Arabic Presentation Forms-B + (0xFFEF, "latin"), # Fullwidth Latin + ] + self.breakpoints = [r[0] for r in self.ranges] + + @lru_cache(maxsize=4096) + def _get_char_weight(self, char): + """Determines the weight of a single character.""" + code = ord(char) + if (65 <= code <= 90) or (97 <= code <= 122): + return self.weights["latin"] + if code == 32: + return self.weights["space"] + + # Ignore arabic Tatweel + if code == 0x0640: + return self.weights["mark"] + + category = unicodedata.category(char) + + if category.startswith("M"): + return self.weights["mark"] + + if category.startswith("P") or category.startswith("S"): + return self.weights["punctuation"] + + if category.startswith("Z"): + return self.weights["space"] + + if category.startswith("N"): + return self.weights["digit"] + + # 3. Binary search for Unicode Block (此时区间里绝不会再混进标点符号) + idx = bisect.bisect_left(self.breakpoints, code) + if idx < len(self.ranges): + script_type = self.ranges[idx][1] + return self.weights.get(script_type, self.weights["default"]) + + # 4. Handle upper planes (CJK Ext B/C/D, Historic scripts) + if code > 0x20000: + return self.weights["cjk"] + + return self.weights["default"] + + def calculate_total_weight(self, text): + """Sums up the normalized weights for a string.""" + return sum(self._get_char_weight(c) for c in text) + + def estimate_duration( + self, + target_text: str, + ref_text: str, + ref_duration: float, + low_threshold: float | None = 50, + boost_strength: float = 3, + ) -> float: + """ + + Args: + target_text (str): The text for which we want to estimate the duration. + ref_text (str): The reference text that was used to measure + the ref_duration. + ref_duration (float): The actual duration it took + to speak the ref_text. + low_threshold (float): The minimum duration threshold below which the + estimation will be considered unreliable. + boost_strength (float): Controls the power-curve boost for short durations. + Higher values boost small durations more aggressively. + 1 = no boost (linear), 2 = sqrt-like + + Returns: + float: The estimated duration for the target_text based + on the ref_text and ref_duration. + """ + if ref_duration <= 0 or not ref_text: + return 0.0 + + ref_weight = self.calculate_total_weight(ref_text) + if ref_weight == 0: + return 0.0 + + speed_factor = ref_weight / ref_duration + target_weight = self.calculate_total_weight(target_text) + + estimated_duration = target_weight / speed_factor + if low_threshold is not None and estimated_duration < low_threshold: + alpha = 1.0 / boost_strength + return low_threshold * (estimated_duration / low_threshold) ** alpha + else: + return estimated_duration + + +# ========================================== +# Example Usage +# ========================================== +if __name__ == "__main__": + estimator = RuleDurationEstimator() + + ref_txt = "Hello, world." + ref_dur = 1.5 + + test_cases = [ + ("Hindi (With complex marks)", "नमस्ते दुनिया"), + ("Arabic (With vowels)", "مَرْحَبًا بِالْعَالَم"), + ("Vietnamese (Lots of diacritics)", "Chào thế giới"), + ("Chinese", "你好，世界！"), + ("Mixed Emoji", "Hello 🌍! This is fun 🎉"), + ] + + print("--- Reference ---") + print(f"Reference Text: '{ref_txt}'") + print(f"Reference Duration: {ref_dur}s") + print("-" * 30) + + for lang, txt in test_cases: + est_time = estimator.estimate_duration(txt, ref_txt, ref_dur) + weight = estimator.calculate_total_weight(txt) + + print(f"[{lang}]") + print(f"Text: {txt}") + print(f"Total Weight: {weight:.2f}") + print(f"Estimated Duration: {est_time:.2f} s") + print("-" * 30) diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py new file mode 100644 index 000000000..a3603a3c3 --- /dev/null +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice.py @@ -0,0 +1,520 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +OmniVoice model for vLLM-Omni two-stage TTS pipeline. + +Stage 0 (Generator): Qwen3 backbone + iterative unmasking → 8-codebook tokens +Stage 1 (Decoder): HiggsAudioV2 decoder → 24kHz waveform +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable, Mapping, Sequence + +import numpy as np +import torch +import torch.nn as nn +from transformers.feature_extraction_utils import BatchFeature +from vllm.config import VllmConfig +from vllm.config.multimodal import BaseDummyOptions +from vllm.logger import init_logger +from vllm.multimodal.inputs import ( + MultiModalDataDict, + MultiModalFieldConfig, + MultiModalKwargsItems, +) +from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + ProcessorInputs, + PromptIndexTargets, + PromptInsertion, + PromptUpdate, +) +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# Multimodal processing +# --------------------------------------------------------------------------- + + +class OmniVoiceMultiModalProcessingInfo(BaseProcessingInfo): + def get_hf_config(self): + return self.ctx.get_hf_config(OmniVoiceConfig) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"audio": None} + + def get_data_parser(self): + return MultiModalDataParser( + target_sr=self.ctx.get_hf_config().sample_rate, + expected_hidden_size=self._get_expected_hidden_size(), + ) + + +class OmniVoiceMultiModalProcessor(BaseMultiModalProcessor[OmniVoiceMultiModalProcessingInfo]): + """Processes text + optional reference audio for OmniVoice. + + For voice cloning: text + reference audio → tokenized reference + For auto voice: text only + """ + + def _ensure_cached_runtime_components(self, model_dir: str, config: OmniVoiceConfig) -> None: + cached_model_dir = getattr(self, "_cached_model_dir", None) + if cached_model_dir == model_dir: + return + + from transformers import AutoTokenizer + + self.text_tokenizer = AutoTokenizer.from_pretrained(model_dir) + + # Audio tokenizer for encoding reference audio + audio_tokenizer_path = os.path.join(model_dir, "audio_tokenizer") + if os.path.isdir(audio_tokenizer_path): + try: + from transformers import ( + AutoFeatureExtractor, + HiggsAudioV2TokenizerModel, + ) + except ImportError as e: + raise ImportError( + "OmniVoice voice cloning requires transformers with " + "HiggsAudioV2TokenizerModel. Upgrade transformers or " + "use text-only mode (no reference audio)." + ) from e + + self.audio_tokenizer = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_path, device_map="cpu") + self.feature_extractor = AutoFeatureExtractor.from_pretrained(audio_tokenizer_path) + self.audio_tokenizer.eval() + else: + self.audio_tokenizer = None + self.feature_extractor = None + logger.warning( + "audio_tokenizer not found at %s, voice cloning disabled", + audio_tokenizer_path, + ) + + self._cached_model_dir = model_dir + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + config = self.info.ctx.get_hf_config() + model_dir = self.info.ctx.model_config.model + self._ensure_cached_runtime_components(model_dir, config) + + audio = mm_data.get("audio", None) + if audio is None: + audio = mm_data.get("audios") + if audio is not None: + audio = audio[0], config.sample_rate + + # Build text prompt with control tokens + lang = mm_kwargs.get("lang", None) + instruct = mm_kwargs.get("instruct", None) + denoise = mm_kwargs.get("denoise", True) + ref_text = mm_kwargs.get("ref_text", None) + + # Construct the style + text portion + style_text = "" + if denoise: + style_text += "<|denoise|>" + lang_str = lang if lang else "None" + instruct_str = instruct if instruct else "None" + style_text += f"<|lang_start|>{lang_str}<|lang_end|>" + style_text += f"<|instruct_start|>{instruct_str}<|instruct_end|>" + + # Combine ref_text and main text + if ref_text: + full_text = f"{ref_text} {prompt}" + else: + full_text = prompt + + text_prompt = f"{style_text}<|text_start|>{full_text}<|text_end|>" + text_tokens = self.text_tokenizer(text_prompt, return_tensors="pt").input_ids.squeeze(0) # [N_text] + + if audio is None: + # Text-only path (auto voice mode) + return BatchFeature( + { + "input_ids": text_tokens, + "input_len": [len(text_tokens)], + } + ) + + # Voice cloning: encode reference audio to tokens + audio_signal, sr = audio + if isinstance(audio_signal, np.ndarray): + audio_signal = torch.from_numpy(audio_signal).float() + if audio_signal.dim() == 1: + audio_signal = audio_signal.unsqueeze(0) + + # Resample to tokenizer sample rate if needed + if self.feature_extractor is not None: + target_sr = self.feature_extractor.sampling_rate + if sr != target_sr: + import torchaudio + + audio_signal = torchaudio.functional.resample(audio_signal, sr, target_sr) + + # Encode reference audio to 8-codebook tokens + if self.audio_tokenizer is not None: + with torch.inference_mode(): + ref_audio_tokens = self.audio_tokenizer.encode(audio_signal) # [8, T_ref] + if ref_audio_tokens.dim() == 3: + ref_audio_tokens = ref_audio_tokens.squeeze(0) # [8, T_ref] + else: + raise RuntimeError( + "Audio tokenizer not available for voice cloning. Ensure audio_tokenizer/ exists in model directory." + ) + + ft = BatchFeature( + { + "input_ids": text_tokens, + "ref_audio_tokens": ref_audio_tokens, # [8, T_ref] + "ref_audio_len": [ref_audio_tokens.shape[1]], + } + ) + return ft + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + return { + "ref_audio_tokens": MultiModalFieldConfig.batched("audio"), + "ref_audio_len": MultiModalFieldConfig.batched("audio"), + } + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return False + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + def insertion_end(item_idx): + if "audio" in out_mm_kwargs and out_mm_kwargs["audio"]: + ref_len = out_mm_kwargs["audio"][0]["ref_audio_len"].data[0].item() + return [1] * ref_len + return [] + + return [ + PromptInsertion( + modality="audio", + target=PromptIndexTargets.start(), + insertion=insertion_end, + ), + ] + + +class OmniVoiceDummyInputsBuilder(BaseDummyInputsBuilder[OmniVoiceMultiModalProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return "Hello, this is a test of the OmniVoice system." + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions] | None = None, + ) -> MultiModalDataDict: + num_audios = mm_counts.get("audio") + max_prompt_seconds = 10 + prompt_sample_rate = 24000 + target_audio_length = max_prompt_seconds * prompt_sample_rate + + audio_overrides = mm_options.get("audio") if mm_options else None + mm_data = { + "audio": ( + self._get_dummy_audios( + length=target_audio_length, + num_audios=num_audios, + overrides=audio_overrides, + )[0], + 24000, + ), + } + return mm_data + + def get_dummy_processor_inputs( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions] | None = None, + ) -> ProcessorInputs: + inputs = super().get_dummy_processor_inputs(seq_len, mm_counts, mm_options) + inputs.hf_processor_mm_kwargs = {"ref_text": "Testing voice cloning."} + return inputs + + +# --------------------------------------------------------------------------- +# Main model class +# --------------------------------------------------------------------------- + + +class OmniVoiceModel( + nn.Module, +): + """OmniVoice model for vLLM-Omni two-stage pipeline. + + Routes to generator (Stage 0) or decoder (Stage 1) based on model_stage. + """ + + requires_raw_input_tokens = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_config + self.have_multimodal_outputs = True + self.model_stage = vllm_config.model_config.model_stage + self.model_dir = vllm_config.model_config.model + + if self.model_stage == "omnivoice_generator": + from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import ( + OmniVoiceGenerator, + ) + + self.generator = OmniVoiceGenerator(self.config) + self.model = self.generator + elif self.model_stage == "omnivoice_decoder": + from vllm_omni.model_executor.models.omnivoice.omnivoice_decoder import ( + OmniVoiceDecoder, + ) + + self.decoder = OmniVoiceDecoder(self.config) + self.model = self.decoder + else: + raise ValueError(f"Unsupported model_stage: {self.model_stage}") + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + is_multimodal=None, + ) -> torch.Tensor: + if self.model_stage == "omnivoice_generator": + # Generator handles its own embedding in forward() + hidden = int(self.config.llm_hidden_size) + return torch.zeros((input_ids.shape[0], hidden), device=input_ids.device) + elif self.model_stage == "omnivoice_decoder": + hidden = int(self.config.llm_hidden_size) + return torch.zeros((input_ids.shape[0], hidden), device=input_ids.device) + else: + raise RuntimeError(f"embed_input_ids not valid for {self.model_stage}") + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + additional_information: dict[str, object] | None = None, + **kwargs: object, + ) -> OmniOutput: + if self.model_stage == "omnivoice_generator": + return self._forward_generator(input_ids, kwargs) + elif self.model_stage == "omnivoice_decoder": + return self._forward_decoder(input_ids, kwargs) + else: + raise ValueError(f"Unsupported model_stage: {self.model_stage}") + + def _forward_generator(self, input_ids: torch.Tensor, kwargs: dict) -> OmniOutput: + """Run generator stage: text → 8-codebook audio tokens.""" + runtime_info = kwargs.get("runtime_additional_information", []) + + if not runtime_info: + # Profiling / dummy run — return a plain tensor (not OmniOutput) + # so the v1 model runner's _dummy_run can index into it. + return torch.zeros( + (input_ids.shape[0], self.config.llm_hidden_size), + device=input_ids.device, + dtype=torch.float32, + ) + + info = runtime_info[0] + device = input_ids.device + num_codebooks = self.config.num_audio_codebook + mask_id = self.config.audio_mask_id + + # Extract text tokens from input_ids + text_tokens = input_ids # [N_text] + text_len = text_tokens.shape[0] + + # Estimate target length using RuleDurationEstimator + # (same formula as reference OmniVoice: weight * 25 / 14.1) + from vllm_omni.model_executor.models.omnivoice.duration import ( + RuleDurationEstimator, + ) + + if not hasattr(self, "_duration_estimator"): + self._duration_estimator = RuleDurationEstimator() + raw_text = info.get("raw_text", "") + if raw_text: + target_len = self._duration_estimator.estimate_duration(raw_text, "Nice to meet you.", 25) + target_len = max(1, int(target_len)) + else: + # Fallback: use character weight formula on text tokens + # approximate ~1.77 frames per text token (25/14.1) + target_len = max(int(text_len * 1.77), 25) + + # Get reference audio tokens if available + ref_audio_tokens = info.get("ref_audio_tokens", None) + + # Build input_ids tensor: [2*B, 8, S] + # B=1, conditional + unconditional + + # Replicate text tokens across 8 codebooks + text_ids = text_tokens.unsqueeze(0).repeat(num_codebooks, 1) # [8, N_text] + + # Target: all MASK + target_ids = torch.full((num_codebooks, target_len), mask_id, dtype=torch.long, device=device) + + # Conditional: [text] [ref_audio?] [target_mask] + if ref_audio_tokens is not None: + ref_tokens = ref_audio_tokens.to(device) # [8, T_ref] + cond_ids = torch.cat([text_ids, ref_tokens, target_ids], dim=1) + cond_audio_start = text_ids.shape[1] + else: + cond_ids = torch.cat([text_ids, target_ids], dim=1) + cond_audio_start = text_ids.shape[1] + + cond_len = cond_ids.shape[1] + + # Unconditional: [target_mask only] + uncond_ids = target_ids.clone() + uncond_len = target_len + + # Pad to same length + max_len = max(cond_len, uncond_len) + if cond_len < max_len: + pad = torch.full( + (num_codebooks, max_len - cond_len), + mask_id, + dtype=torch.long, + device=device, + ) + cond_ids = torch.cat([cond_ids, pad], dim=1) + if uncond_len < max_len: + pad = torch.full( + (num_codebooks, max_len - uncond_len), + mask_id, + dtype=torch.long, + device=device, + ) + uncond_ids = torch.cat([uncond_ids, pad], dim=1) + + batch_input_ids = torch.stack([cond_ids, uncond_ids], dim=0) # [2, 8, max_len] + + # Audio mask: True for audio positions + batch_audio_mask = torch.zeros((2, max_len), dtype=torch.bool, device=device) + batch_audio_mask[0, cond_audio_start:cond_len] = True + batch_audio_mask[1, :uncond_len] = True + + # Attention mask: [2, 1, S, S] + batch_attention_mask = torch.zeros((2, 1, max_len, max_len), dtype=torch.bool, device=device) + batch_attention_mask[0, :, :cond_len, :cond_len] = True + batch_attention_mask[1, :, :uncond_len, :uncond_len] = True + + # Run iterative generation + tokens = self.generator( + input_ids=batch_input_ids, + audio_mask=batch_audio_mask, + attention_mask=batch_attention_mask, + target_lens=[target_len], + num_step=self.config.num_step, + guidance_scale=self.config.guidance_scale, + t_shift=self.config.t_shift, + layer_penalty_factor=self.config.layer_penalty_factor, + position_temperature=self.config.position_temperature, + class_temperature=self.config.class_temperature, + ) # [1, 8, target_len] + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"audio_tokens": tokens}, + ) + + def _forward_decoder(self, input_ids: torch.Tensor, kwargs: dict) -> OmniOutput: + """Run decoder stage: 8-codebook tokens → audio waveform.""" + runtime_info = kwargs.get("runtime_additional_information", []) + + if not runtime_info: + # Profiling / dummy run — return plain tensor for v1 runner compat + return torch.zeros( + (input_ids.shape[0], self.config.llm_hidden_size), + device=input_ids.device, + dtype=torch.float32, + ) + + info = runtime_info[0] + audio_tokens = info.get("audio_tokens", None) + + if audio_tokens is None: + raise RuntimeError("No audio_tokens received from generator stage") + + if isinstance(audio_tokens, np.ndarray): + audio_tokens = torch.from_numpy(audio_tokens) + + # audio_tokens: [B, 8, T]; buffer may be CPU — move to decoder weights + if audio_tokens.dim() == 2: + audio_tokens = audio_tokens.unsqueeze(0) # Add batch dim + + dec_device = next(self.decoder.parameters()).device + audio_tokens = audio_tokens.to(device=dec_device, dtype=torch.long) + + tts_speech = self.decoder(audio_tokens) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={ + "audio": tts_speech, + "sr": self.config.sample_rate, + }, + ) + + def _resolve_model_dir(self) -> str: + """Resolve model directory to local path (handles HF hub IDs).""" + model_dir = self.model_dir + if os.path.isdir(model_dir): + return model_dir + # HF hub model ID — resolve to local cache + from huggingface_hub import snapshot_download + + return snapshot_download(model_dir) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + try: + device = next(self.parameters()).device + except StopIteration: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + model_dir = self._resolve_model_dir() + + if self.model_stage == "omnivoice_generator": + self.generator.load_weights(model_dir, device) + elif self.model_stage == "omnivoice_decoder": + self.decoder.load_weights(model_dir, device) + else: + raise ValueError(f"{self.model_stage} not supported!") diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice_decoder.py b/vllm_omni/model_executor/models/omnivoice/omnivoice_decoder.py new file mode 100644 index 000000000..cf69f2658 --- /dev/null +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice_decoder.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +OmniVoice Decoder (Stage 1) - Audio token to waveform conversion. + +Implements the HiggsAudioV2 decode path using transformers' DacModel decoder +and a custom RVQ quantizer, compatible with transformers 4.x. + +Decode path: + audio_codes [B, 8, T] + → RVQ codebook lookup + project_out → sum → [B, 1024, T] + → fc2 Linear(1024, 256) → [B, 256, T] + → DAC acoustic decoder (conv transpose upsampling) → [B, 1, T*960] + → 24kHz waveform (25fps × 960 samples/frame) +""" + +from __future__ import annotations + +import json +import os + +import torch +import torch.nn as nn +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig + +logger = init_logger(__name__) + + +class HiggsAudioVQLayer(nn.Module): + """Single VQ layer: codebook lookup + project_out.""" + + def __init__(self, codebook_size: int = 1024, codebook_dim: int = 64, hidden_size: int = 1024): + super().__init__() + self.codebook = nn.Embedding(codebook_size, codebook_dim) + self.project_out = nn.Linear(codebook_dim, hidden_size) + + def decode(self, indices: torch.Tensor) -> torch.Tensor: + """indices: [B, T] → [B, hidden_size, T]""" + quantized = self.codebook(indices) # [B, T, codebook_dim] + quantized = self.project_out(quantized) # [B, T, hidden_size] + return quantized.permute(0, 2, 1) # [B, hidden_size, T] + + +class HiggsAudioRVQ(nn.Module): + """Residual Vector Quantizer with 8 codebook layers.""" + + def __init__( + self, num_quantizers: int = 8, codebook_size: int = 1024, codebook_dim: int = 64, hidden_size: int = 1024 + ): + super().__init__() + self.quantizers = nn.ModuleList( + [HiggsAudioVQLayer(codebook_size, codebook_dim, hidden_size) for _ in range(num_quantizers)] + ) + + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """codes: [num_quantizers, B, T] → [B, hidden_size, T]""" + result = torch.zeros( + codes.shape[1], + self.quantizers[0].project_out.out_features, + codes.shape[2], + device=codes.device, + dtype=torch.float32, + ) + for i, quantizer in enumerate(self.quantizers): + result = result + quantizer.decode(codes[i]) + return result + + +class OmniVoiceDecoder(nn.Module): + """OmniVoice Stage 1: Token-to-audio decoder. + + Uses DAC acoustic decoder from transformers + custom HiggsAudio RVQ + quantizer to convert 8-codebook tokens into 24kHz waveform. + """ + + def __init__(self, config: OmniVoiceConfig): + super().__init__() + self.config = config + self.sample_rate = config.sample_rate + self._loaded = False + + # These are populated by load_weights + self.quantizer = None + self.fc2 = None + self.acoustic_decoder = None + + @torch.inference_mode() + def forward(self, audio_codes: torch.Tensor) -> torch.Tensor: + """Decode audio tokens to waveform. + + Args: + audio_codes: [B, 8, T] - 8-codebook audio token IDs + + Returns: + waveform: [B, 1, audio_samples] at 24kHz + """ + if not self._loaded: + raise RuntimeError("Decoder not loaded. Call load_weights() first.") + + device = audio_codes.device + + # Transpose: [B, 8, T] → [8, B, T] + codes = audio_codes.transpose(0, 1).long() + + # RVQ decode: sum codebook embeddings → [B, 1024, T] + quantized = self.quantizer.decode(codes) + + # Project: [B, 1024, T] → fc2 → [B, 256, T] + quantized = self.fc2(quantized.transpose(1, 2)).transpose(1, 2) + + # Acoustic decoder: [B, 256, T] → [B, 1, T*960] + audio = self.acoustic_decoder(quantized) + + # Ensure [B, 1, samples] + if audio.dim() == 2: + audio = audio.unsqueeze(1) + + return audio.to(device) + + def _adjust_output_padding(self, decoder: nn.Module): + """Adjust ConvTranspose1d output_padding (HiggsAudioV2 modification).""" + for module in decoder.modules(): + if isinstance(module, nn.ConvTranspose1d): + stride = module.stride[0] if isinstance(module.stride, tuple) else module.stride + module.output_padding = (stride % 2,) + + def load_weights(self, model_dir: str, device: torch.device) -> None: + """Load decoder components from audio_tokenizer/model.safetensors.""" + from safetensors.torch import load_file + from transformers import DacConfig, DacModel + + audio_tokenizer_path = os.path.join(model_dir, "audio_tokenizer") + config_path = os.path.join(audio_tokenizer_path, "config.json") + weights_path = os.path.join(audio_tokenizer_path, "model.safetensors") + + if not os.path.exists(weights_path): + raise FileNotFoundError(f"Audio tokenizer weights not found at {weights_path}") + + with open(config_path) as f: + tokenizer_config = json.load(f) + + state_dict = load_file(weights_path, device=str(device)) + + # 1. Build RVQ quantizer + codebook_dim = tokenizer_config.get("codebook_dim", 64) + codebook_size = tokenizer_config.get("codebook_size", 1024) + # Hidden size = quantizer project_out output dim + hidden_size = state_dict["quantizer.quantizers.0.project_out.weight"].shape[0] + num_quantizers = sum( + 1 for k in state_dict if k.startswith("quantizer.quantizers.") and k.endswith(".codebook.embed") + ) + + self.quantizer = HiggsAudioRVQ( + num_quantizers=num_quantizers, + codebook_size=codebook_size, + codebook_dim=codebook_dim, + hidden_size=hidden_size, + ).to(device) + + # Load quantizer weights + for i in range(num_quantizers): + prefix = f"quantizer.quantizers.{i}" + embed_key = f"{prefix}.codebook.embed" + if embed_key in state_dict: + self.quantizer.quantizers[i].codebook.weight.data.copy_(state_dict[embed_key]) + proj_out_w = f"{prefix}.project_out.weight" + proj_out_b = f"{prefix}.project_out.bias" + if proj_out_w in state_dict: + self.quantizer.quantizers[i].project_out.weight.data.copy_(state_dict[proj_out_w]) + if proj_out_b in state_dict: + self.quantizer.quantizers[i].project_out.bias.data.copy_(state_dict[proj_out_b]) + + # 2. Build fc2 projection + fc2_w = state_dict["fc2.weight"] + fc2_b = state_dict["fc2.bias"] + self.fc2 = nn.Linear(fc2_w.shape[1], fc2_w.shape[0]).to(device) + self.fc2.weight.data.copy_(fc2_w) + self.fc2.bias.data.copy_(fc2_b) + + # 3. Build DAC acoustic decoder + dac_cfg = DacConfig(**tokenizer_config["acoustic_model_config"]) + dac_model = DacModel(dac_cfg) + self.acoustic_decoder = dac_model.decoder.to(device) + + # Load acoustic decoder weights + loaded = 0 + for name, param in self.acoustic_decoder.named_parameters(): + higgs_name = f"acoustic_decoder.{name}" + if higgs_name in state_dict: + param.data.copy_(state_dict[higgs_name]) + loaded += 1 + + # Apply HiggsAudioV2 output padding adjustment + self._adjust_output_padding(self.acoustic_decoder) + + # Remove tanh if present (HiggsAudioV2 uses Identity instead) + if hasattr(self.acoustic_decoder, "tanh"): + self.acoustic_decoder.tanh = nn.Identity() + + self.acoustic_decoder.eval() + self._loaded = True + + logger.info( + "Loaded OmniVoice decoder: %d quantizers, fc2(%d→%d), acoustic decoder (%d weights)", + num_quantizers, + fc2_w.shape[1], + fc2_w.shape[0], + loaded, + ) diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py new file mode 100644 index 000000000..32fe42272 --- /dev/null +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py @@ -0,0 +1,588 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +OmniVoice Generator (Stage 0) - Iterative unmasking with Qwen3 backbone. + +Generates 8-codebook audio tokens from text via 32-step non-autoregressive +iterative masked prediction with classifier-free guidance. + +Uses vLLM-Omni's DiffusionAttention for optimized full (bidirectional) attention +via FlashAttention/SageAttention/SDPA backends. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# Unmasking schedule helpers +# --------------------------------------------------------------------------- + + +def _get_time_steps( + t_start: float, + t_end: float, + num_step: int, + t_shift: float, +) -> torch.Tensor: + """Compute the unmasking schedule with time shift. + + Returns cumulative proportions [0, ..., 1] of length num_step. + Formula: r_n = t_shift * (n/N) / (1 + (t_shift - 1) * (n/N)) + """ + steps = torch.linspace(t_start, t_end, num_step) + shifted = t_shift * steps / (1.0 + (t_shift - 1.0) * steps) + return shifted + + +def _gumbel_sample(logits: torch.Tensor, temperature: float) -> torch.Tensor: + """Add Gumbel noise for stochastic position selection.""" + noise = -torch.log(-torch.log(torch.rand_like(logits).clamp(min=1e-8))) + return logits / max(temperature, 1e-8) + noise + + +# --------------------------------------------------------------------------- +# Qwen3-style transformer blocks using DiffusionAttention +# --------------------------------------------------------------------------- + + +class OmniVoiceRMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + self.eps) + return self.weight * x.to(self.weight.dtype) + + +class OmniVoiceAttention(nn.Module): + """Qwen3-style GQA attention using DiffusionAttention backend.""" + + def __init__(self, config: OmniVoiceConfig): + super().__init__() + self.hidden_size = config.llm_hidden_size + self.num_heads = config.llm_num_attention_heads + self.num_kv_heads = config.llm_num_key_value_heads + self.head_dim = config.llm_head_dim + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + # Qwen3 uses per-head QK norm + self.q_norm = OmniVoiceRMSNorm(self.head_dim) + self.k_norm = OmniVoiceRMSNorm(self.head_dim) + + self.scale = 1.0 / math.sqrt(self.head_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + cos: torch.Tensor | None = None, + sin: torch.Tensor | None = None, + ) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape + + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + q = q.view(batch_size, seq_len, self.num_heads, self.head_dim) + k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) + v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) + + # Per-head QK norm (Qwen3) + q = self.q_norm(q) + k = self.k_norm(k) + + # Apply RoPE + if cos is not None and sin is not None: + q = _apply_rotary_pos_emb(q, cos, sin) + k = _apply_rotary_pos_emb(k, cos, sin) + + # Expand KV heads for GQA (8 KV heads → 16 Q heads) + if self.num_kv_heads != self.num_heads: + repeat_factor = self.num_heads // self.num_kv_heads + k = k.repeat_interleave(repeat_factor, dim=2) + v = v.repeat_interleave(repeat_factor, dim=2) + + # Full bidirectional attention via SDPA with proper mask support + # Permute to (batch, heads, seq, head_dim) for SDPA + q = q.permute(0, 2, 1, 3) + k = k.permute(0, 2, 1, 3) + v = v.permute(0, 2, 1, 3) + + # Convert [B, 1, S, S] bool mask to float mask for SDPA + sdpa_mask = None + if attention_mask is not None: + sdpa_mask = attention_mask.to(dtype=q.dtype) + sdpa_mask = sdpa_mask.masked_fill(~attention_mask, float("-inf")) + sdpa_mask = sdpa_mask.masked_fill(attention_mask, 0.0) + + out = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=sdpa_mask, + scale=1.0 / math.sqrt(self.head_dim), + ) + + # Back to (batch, seq, heads * head_dim) + out = out.permute(0, 2, 1, 3).contiguous() + out = out.view(batch_size, seq_len, self.num_heads * self.head_dim) + return self.o_proj(out) + + +class OmniVoiceMLP(nn.Module): + """Qwen3-style MLP with SwiGLU.""" + + def __init__(self, config: OmniVoiceConfig): + super().__init__() + self.gate_proj = nn.Linear(config.llm_hidden_size, config.llm_intermediate_size, bias=False) + self.up_proj = nn.Linear(config.llm_hidden_size, config.llm_intermediate_size, bias=False) + self.down_proj = nn.Linear(config.llm_intermediate_size, config.llm_hidden_size, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class OmniVoiceTransformerBlock(nn.Module): + """Single Qwen3 transformer block with DiffusionAttention.""" + + def __init__(self, config: OmniVoiceConfig): + super().__init__() + self.input_layernorm = OmniVoiceRMSNorm(config.llm_hidden_size, eps=config.llm_rms_norm_eps) + self.self_attn = OmniVoiceAttention(config) + self.post_attention_layernorm = OmniVoiceRMSNorm(config.llm_hidden_size, eps=config.llm_rms_norm_eps) + self.mlp = OmniVoiceMLP(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + cos: torch.Tensor | None = None, + sin: torch.Tensor | None = None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, attention_mask=attention_mask, cos=cos, sin=sin) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +# --------------------------------------------------------------------------- +# RoPE +# --------------------------------------------------------------------------- + + +def _precompute_rope( + head_dim: int, + max_seq_len: int, + theta: float = 1000000.0, + device: torch.device | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Precompute RoPE cos/sin tensors.""" + inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim)) + t = torch.arange(max_seq_len, device=device, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + cos = freqs.cos() + sin = freqs.sin() + return cos, sin + + +def _apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """Apply rotary position embedding. x shape: (B, S, H, D).""" + seq_len = x.shape[1] + cos = cos[:seq_len].unsqueeze(0).unsqueeze(2) # (1, S, 1, D/2) + sin = sin[:seq_len].unsqueeze(0).unsqueeze(2) + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + rotated = torch.cat([-x2, x1], dim=-1) + return x * torch.cat([cos, cos], dim=-1) + rotated * torch.cat([sin, sin], dim=-1) + + +# --------------------------------------------------------------------------- +# Generator model +# --------------------------------------------------------------------------- + + +class OmniVoiceGenerator(nn.Module): + """OmniVoice Stage 0: Iterative unmasking generator. + + Architecture: + - Text embedding (from Qwen3 vocab) + Audio embedding (8*1025 entries) + - 28-layer Qwen3 transformer with full bidirectional attention + - 8-codebook prediction head (single linear: hidden → 8*1025) + - 32-step iterative unmasking with classifier-free guidance + + Optimizations: + - DiffusionAttention (FlashAttn/SageAttn/SDPA auto-selected) + - TeaCache / Cache-DiT compatible (hook-based, non-intrusive) + - regionally_compile() compatible for torch.compile on repeated blocks + - Sequence parallelism via SP hooks for multi-GPU + """ + + # For regionally_compile() support + _repeated_blocks = ["layers"] + + def __init__(self, config: OmniVoiceConfig): + super().__init__() + self.config = config + + # Text embedding (shared with LLM) + self.text_embedding = nn.Embedding(config.llm_vocab_size, config.llm_hidden_size) + + # Audio embedding: 8 codebooks * 1025 tokens + self.audio_embeddings = nn.Embedding( + config.num_audio_codebook * config.audio_vocab_size, + config.llm_hidden_size, + ) + self.register_buffer( + "codebook_layer_offsets", + torch.arange(config.num_audio_codebook) * config.audio_vocab_size, + ) + + # Transformer layers + self.layers = nn.ModuleList([OmniVoiceTransformerBlock(config) for _ in range(config.llm_num_hidden_layers)]) + self.norm = OmniVoiceRMSNorm(config.llm_hidden_size, eps=config.llm_rms_norm_eps) + + # Prediction head: hidden → 8 * 1025 + self.audio_heads = nn.Linear( + config.llm_hidden_size, + config.num_audio_codebook * config.audio_vocab_size, + bias=False, + ) + + # Precompute RoPE + self._rope_cos = None + self._rope_sin = None + + def _ensure_rope(self, seq_len: int, device: torch.device) -> None: + """Lazily compute RoPE cos/sin if needed.""" + if self._rope_cos is None or self._rope_cos.shape[0] < seq_len: + max_len = max(seq_len, 4096) + self._rope_cos, self._rope_sin = _precompute_rope( + self.config.llm_head_dim, + max_len, + theta=self.config.llm_rope_theta, + device=device, + ) + + def _prepare_embeddings( + self, + input_ids: torch.Tensor, + audio_mask: torch.Tensor, + ) -> torch.Tensor: + """Prepare mixed text+audio embeddings. + + Args: + input_ids: [B, 8, S] - text tokens replicated across codebooks, + audio positions have per-codebook token IDs + audio_mask: [B, S] - True for audio positions, False for text + + Returns: + embeddings: [B, S, hidden_size] + """ + # Text embeddings from first codebook row (all rows identical for text) + text_embeds = self.text_embedding(input_ids[:, 0, :]) + + # Audio embeddings: offset per codebook, then sum across codebooks + shifted_ids = (input_ids * audio_mask.unsqueeze(1)) + self.codebook_layer_offsets.view(1, -1, 1) + audio_embeds = self.audio_embeddings(shifted_ids).sum(dim=1) + + # Merge: audio where audio_mask=True, text elsewhere + return torch.where(audio_mask.unsqueeze(-1), audio_embeds, text_embeds) + + def _transformer_forward( + self, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run through transformer layers. + + Args: + inputs_embeds: [B, S, hidden_size] + attention_mask: [B, 1, S, S] or None + + Returns: + hidden_states: [B, S, hidden_size] + """ + device = inputs_embeds.device + seq_len = inputs_embeds.shape[1] + self._ensure_rope(seq_len, device) + + hidden_states = inputs_embeds + cos = self._rope_cos.to(device=device, dtype=hidden_states.dtype) + sin = self._rope_sin.to(device=device, dtype=hidden_states.dtype) + + for layer in self.layers: + hidden_states = layer( + hidden_states, + attention_mask=attention_mask, + cos=cos, + sin=sin, + ) + + return self.norm(hidden_states) + + def _get_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Project hidden states to per-codebook logits. + + Args: + hidden_states: [B, S, hidden_size] + + Returns: + logits: [B, 8, S, 1025] + """ + batch_size, seq_len, _ = hidden_states.shape + logits_flat = self.audio_heads(hidden_states) # [B, S, 8*1025] + return logits_flat.view( + batch_size, + seq_len, + self.config.num_audio_codebook, + self.config.audio_vocab_size, + ).permute(0, 2, 1, 3) # [B, 8, S, 1025] + + @torch.inference_mode() + def forward( + self, + input_ids: torch.Tensor, + audio_mask: torch.Tensor, + attention_mask: torch.Tensor, + target_lens: list[int], + num_step: int = 32, + guidance_scale: float = 2.0, + t_shift: float = 0.1, + layer_penalty_factor: float = 5.0, + position_temperature: float = 5.0, + class_temperature: float = 0.0, + ) -> torch.Tensor: + """Run the full 32-step iterative unmasking generation. + + Args: + input_ids: [2*B, 8, S] - conditional (0:B) + unconditional (B:2B) + audio_mask: [2*B, S] - True for audio positions + attention_mask: [2*B, 1, S, S] - attention mask + target_lens: List of target audio lengths per batch item + num_step: Number of unmasking steps + guidance_scale: CFG scale + t_shift: Time shift for schedule + layer_penalty_factor: Penalty for later codebooks + position_temperature: Gumbel temperature for position selection + class_temperature: Temperature for token prediction (0=greedy) + + Returns: + tokens: [B, 8, max_target_len] - generated audio tokens + """ + B = len(target_lens) + device = input_ids.device + max_target_len = max(target_lens) + mask_id = self.config.audio_mask_id + num_codebooks = self.config.num_audio_codebook + + # Initialize all target tokens as [MASK] + tokens = torch.full( + (B, num_codebooks, max_target_len), + mask_id, + dtype=torch.long, + device=device, + ) + + # Compute unmasking schedule + timesteps = _get_time_steps(0.0, 1.0, num_step + 1, t_shift).tolist() + schedules = [] + for t_len in target_lens: + total_mask = t_len * num_codebooks + rem = total_mask + sched = [] + for step in range(num_step): + num = ( + rem + if step == num_step - 1 + else min( + math.ceil(total_mask * (timesteps[step + 1] - timesteps[step])), + rem, + ) + ) + sched.append(int(num)) + rem -= int(num) + schedules.append(sched) + + layer_ids = torch.arange(num_codebooks, device=device).view(1, -1, 1) + + # Compute c_lens for extracting target region from full sequence + c_lens = [] + for i in range(B): + # Conditional sequence length = number of non-padding positions + c_len = attention_mask[i, 0, 0].sum().item() + c_lens.append(int(c_len)) + + # Main iterative loop + for step in range(num_step): + # Prepare embeddings and run transformer + inputs_embeds = self._prepare_embeddings(input_ids, audio_mask) + hidden_states = self._transformer_forward(inputs_embeds, attention_mask) + batch_logits = self._get_logits(hidden_states).to(torch.float32) + # batch_logits: [2*B, 8, S, 1025] + + for i in range(B): + k = schedules[i][step] + if k <= 0: + continue + + c_len = c_lens[i] + t_len = target_lens[i] + + # Extract logits for target region + c_logits = batch_logits[i : i + 1, :, c_len - t_len : c_len, :] # [1, 8, T, 1025] + u_logits = batch_logits[B + i : B + i + 1, :, :t_len, :] # [1, 8, T, 1025] + + # Classifier-free guidance + if guidance_scale != 0: + c_log_probs = F.log_softmax(c_logits, dim=-1) + u_log_probs = F.log_softmax(u_logits, dim=-1) + log_probs = torch.log_softmax( + c_log_probs + guidance_scale * (c_log_probs - u_log_probs), + dim=-1, + ) + else: + log_probs = F.log_softmax(c_logits, dim=-1) + + # Prevent predicting [MASK] + log_probs[..., mask_id] = -float("inf") + + # Token prediction + if class_temperature > 0.0: + pred_tokens = _gumbel_sample(log_probs, class_temperature).argmax(dim=-1) + else: + pred_tokens = log_probs.argmax(dim=-1) # [1, 8, T] + + # Confidence scores + scores = log_probs.max(dim=-1)[0] # [1, 8, T] + + # Layer penalty (earlier codebooks get higher priority) + scores = scores - (layer_ids * layer_penalty_factor) + + # Gumbel noise for position selection + if position_temperature > 0.0: + scores = _gumbel_sample(scores, position_temperature) + + # Mask out already unmasked positions + sample_tokens = tokens[i : i + 1, :, :t_len] + scores.masked_fill_(sample_tokens != mask_id, -float("inf")) + + # Select top-k positions to unmask + _, topk_idx = torch.topk(scores.flatten(), k) + flat_tokens = sample_tokens.flatten().clone() + flat_tokens[topk_idx] = pred_tokens.flatten()[topk_idx] + sample_tokens.copy_(flat_tokens.view_as(sample_tokens)) + + # Update tokens and batch inputs for next iteration + tokens[i : i + 1, :, :t_len] = sample_tokens + input_ids = input_ids.clone() + input_ids[i, :, c_len - t_len : c_len] = sample_tokens.squeeze(0) + input_ids[B + i, :, :t_len] = sample_tokens.squeeze(0) + + return tokens + + def load_weights(self, model_dir: str, device: torch.device) -> None: + """Load weights from HuggingFace OmniVoice model.safetensors. + + The HF checkpoint contains: + - llm.* -> Qwen3 transformer weights + - audio_embeddings.* -> audio embedding table + - audio_heads.* -> prediction head + """ + import os + + from safetensors.torch import load_file + + weights_path = os.path.join(model_dir, "model.safetensors") + if not os.path.exists(weights_path): + raise FileNotFoundError(f"Model weights not found at {weights_path}") + + state_dict = load_file(weights_path, device=str(device)) + + # Map HF weight names to our module names + loaded_keys = set() + + # 1. Text embedding: llm.embed_tokens.weight -> text_embedding.weight + text_emb_key = "llm.embed_tokens.weight" + if text_emb_key in state_dict: + self.text_embedding.weight.data.copy_(state_dict[text_emb_key]) + loaded_keys.add(text_emb_key) + + # 2. Audio embeddings + for key in ["audio_embeddings.weight"]: + if key in state_dict: + self.audio_embeddings.weight.data.copy_(state_dict[key]) + loaded_keys.add(key) + + # 3. Audio heads + for key in ["audio_heads.weight"]: + if key in state_dict: + self.audio_heads.weight.data.copy_(state_dict[key]) + loaded_keys.add(key) + + # 4. Transformer layers: llm.layers.N.* -> layers.N.* + for key, value in state_dict.items(): + if key.startswith("llm.layers."): + # llm.layers.0.self_attn.q_proj.weight -> layers.0.self_attn.q_proj.weight + our_key = key.replace("llm.layers.", "layers.") + parts = our_key.split(".") + module = self + try: + for part in parts[:-1]: + if part.isdigit(): + module = module[int(part)] + else: + module = getattr(module, part) + param_name = parts[-1] + param = getattr(module, param_name) + if isinstance(param, nn.Parameter): + param.data.copy_(value) + elif isinstance(param, torch.Tensor): + param.copy_(value) + loaded_keys.add(key) + except (AttributeError, IndexError, KeyError) as e: + logger.warning("Failed to load weight %s: %s", key, e) + + # 5. Final norm: llm.norm.weight -> norm.weight + norm_key = "llm.norm.weight" + if norm_key in state_dict: + self.norm.weight.data.copy_(state_dict[norm_key]) + loaded_keys.add(norm_key) + + unloaded = set(state_dict.keys()) - loaded_keys + # Filter out audio_tokenizer weights (loaded in decoder stage) + unloaded = {k for k in unloaded if not k.startswith("audio_tokenizer.")} + if unloaded: + logger.info( + "Generator: %d/%d weights loaded, %d skipped (decoder weights)", + len(loaded_keys), + len(state_dict), + len(unloaded), + ) + else: + logger.info("Generator: all %d weights loaded", len(loaded_keys)) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index b5ef92616..139892345 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -57,6 +57,11 @@ _OMNI_MODELS = { "cosyvoice3", "CosyVoice3Model", ), + "OmniVoiceModel": ( + "omnivoice", + "omnivoice", + "OmniVoiceModel", + ), "MammothModa2Qwen2ForCausalLM": ( "mammoth_moda2", "mammoth_moda2", diff --git a/vllm_omni/model_executor/stage_configs/omnivoice.yaml b/vllm_omni/model_executor/stage_configs/omnivoice.yaml new file mode 100644 index 000000000..49f11e967 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/omnivoice.yaml @@ -0,0 +1,20 @@ +# Stage config for OmniVoice TTS via diffusion engine +# Single stage: text → 32-step iterative unmasking → 8-codebook tokens → DAC decode → 24kHz audio + +stage_args: + - stage_id: 0 + stage_type: diffusion + is_comprehension: true + runtime: + devices: 0 + engine_args: + model_stage: dit + model_class_name: "OmniVoicePipeline" + gpu_memory_utilization: 0.5 + enforce_eager: true + trust_remote_code: true + engine_output_type: audio + distributed_executor_backend: "mp" + dtype: "float32" + final_output: true + final_output_type: audio diff --git a/vllm_omni/model_executor/stage_input_processors/omnivoice.py b/vllm_omni/model_executor/stage_input_processors/omnivoice.py new file mode 100644 index 000000000..b7f5c102e --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/omnivoice.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inter-stage processor for OmniVoice: Generator → Decoder.""" + +from typing import Any + +from vllm.inputs import TextPrompt + +from vllm_omni.inputs.data import OmniTokensPrompt + + +def tokens2audio( + stage_list: list[Any], + engine_input_source: list[int], + prompt: OmniTokensPrompt | TextPrompt = None, + requires_multimodal_data: bool = True, +): + """Build stage-1 (decoder) inputs from stage-0 (generator) outputs. + + Takes the 8-codebook audio tokens from the generator and packages + them for the HiggsAudioV2 decoder. + """ + source_stage_id = engine_input_source[0] + source_outputs = stage_list[source_stage_id].engine_outputs + + if not isinstance(prompt, list): + prompt = [prompt] + + source_output = source_outputs[0] + output = source_output.outputs[0] + + multi_modal_data = output.multimodal_output + if multi_modal_data is None: + raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}") + + # Pass audio_tokens from generator to decoder + engine_input = OmniTokensPrompt( + prompt_token_ids=output.token_ids, + additional_information=multi_modal_data, + ) + return [engine_input]