[Model] Add Voxtral TTS model (#1803)
Signed-off-by: Chen-Yo Sun <chenyo.sun@mistral.ai>
This commit is contained in:
Chen-Yo Sun authored and GitHub committed
1 parent
b766c4702c
commit
e9fe78e28c
29 files changed
+5135
-51
No files matched your search
@@ -340,3 +340,59 @@ steps:
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# - label: "Voxtral TTS E2E Test"
|
||||
# timeout_in_minutes: 20
|
||||
# depends_on: upload-merge-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_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
# '
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 1
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
- label: "CUDA Unit Test"
|
||||
timeout_in_minutes: 10
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- pytest -v -s -m 'core_model and cuda'
|
||||
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
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
@@ -11,6 +11,9 @@ steps:
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
EXIT2=$$?
|
||||
# TODO: uncomment when mistralai/tts-model is public
|
||||
# pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
# EXIT3=$$?
|
||||
exit $$((EXIT1 | EXIT2))
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
|
||||
@@ -234,6 +234,45 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
# TODO: uncomment when mistralai/tts-model is public
|
||||
# - label: "Voxtral-TTS E2E Test"
|
||||
# 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_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
# '
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 1
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
# - label: "Diffusion Image Edit Test with H100 (1 GPU)"
|
||||
# depends_on: upload-ready-pipeline
|
||||
# commands:
|
||||
|
||||
@@ -141,6 +141,7 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]:
|
||||
"vllm_omni.diffusion.models.qwen_image",
|
||||
"vllm_omni.entrypoints.async_diffusion",
|
||||
"vllm_omni.entrypoints.openai",
|
||||
"vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts",
|
||||
]
|
||||
if any(module_path.startswith(prefix) for prefix in excluded_prefixes):
|
||||
continue
|
||||
|
||||
@@ -52,6 +52,7 @@ th {
|
||||
|`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
|
||||
|`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` |
|
||||
|`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` |
|
||||
|`VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/tts-model` |
|
||||
|
||||
|
||||
## List of Supported Models for NPU
|
||||
|
||||
@@ -4,6 +4,7 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
|
||||
|
||||
- **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
|
||||
- **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
|
||||
- **Voxtral TTS** (`mistralai/tts-model`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
|
||||
|
||||
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
|
||||
|
||||
@@ -28,6 +29,14 @@ vllm-omni serve fishaudio/s2-pro \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.9
|
||||
|
||||
# Voxtral TTS
|
||||
vllm serve mistralai/tts-model \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
### Generate Speech
|
||||
@@ -347,6 +356,12 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
|
||||
Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech online serving example](../user_guide/examples/online_serving/fish_speech.md) for details.
|
||||
|
||||
### Voxtral TTS
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| `mistralai/tts-model` | 3B AR + FlowMatching TTS. Supports text-to-speech with preset voices. |
|
||||
|
||||
## Error Responses
|
||||
|
||||
### 400 Bad Request
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
# Voxtral TTS Offline Inference
|
||||
|
||||
`end2end.py` runs Voxtral TTS end-to-end offline inference using vLLM. It supports both blocking (`Omni`) and streaming (`AsyncOmni`) generation, batched prompts with configurable concurrency, and voice selection via preset name or reference audio file.
|
||||
|
||||
When `mistral_common` has `SpeechRequest` support, prompt token IDs are built via `encode_speech_request`. Otherwise, the script falls back to manual token construction.
|
||||
|
||||
## Usage Examples
|
||||
|
||||
```bash
|
||||
# Basic single-prompt with reference audio
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--write-audio \
|
||||
--model mistralai/tts-model \
|
||||
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?" \
|
||||
--audio-path path/to/reference_audio.wav
|
||||
|
||||
# 32 replicate prompts with reference audio
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--num-prompts 32 --write-audio \
|
||||
--model mistralai/tts-model \
|
||||
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?" \
|
||||
--audio-path path/to/reference_audio.wav
|
||||
|
||||
# Short debug prompt with reference audio
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--write-audio \
|
||||
--model mistralai/tts-model \
|
||||
--text "This is a test message." \
|
||||
--audio-path path/to/reference_audio.wav
|
||||
|
||||
# Streaming with neutral_female voice preset
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--streaming --write-audio --voice neutral_female \
|
||||
--model mistralai/tts-model
|
||||
|
||||
# 32 prompts, 8 concurrent requests per wave, streaming with casual_male voice
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--num-prompts 32 --concurrency 8 --streaming --write-audio --voice casual_male \
|
||||
--model mistralai/tts-model
|
||||
```
|
||||
|
||||
## Arguments
|
||||
|
||||
| Argument | Description |
|
||||
|---|---|
|
||||
| `--model PATH` | HuggingFace repo ID or local directory path (default: `mistralai/tts-model`) |
|
||||
| `--text TEXT` | Text to synthesize (default: `"This is a test message."`) |
|
||||
| `--audio-path PATH` | Path to reference audio file for voice cloning |
|
||||
| `--output-dir DIR` | Directory to write output WAV files (default: `output_audio`) |
|
||||
| `--stage-configs-path PATH` | Path to stage configs YAML (auto-resolved from model if not set) |
|
||||
| `--num-prompts N` | Number of replicate prompts to run for measuring performance (default: 1) |
|
||||
| `--streaming` | Use streaming generation via `AsyncOmni` (default: blocking `Omni`) |
|
||||
| `--concurrency N` | Max concurrent requests per wave (must be used with `--streaming`, must evenly divide `--num-prompts`) |
|
||||
| `--voice NAME` | Voice preset to use instead of reference audio file (e.g., casual_female, casual_male, cheerful_female, neutral_female, neutral_male) |
|
||||
| `--write-audio` | Write generated audio to WAV files |
|
||||
| `--profiling-mode` | Enable profiling mode (reduces max tokens to 50) |
|
||||
| `--log-stats` | Enable detailed statistics logging |
|
||||
@@ -0,0 +1,377 @@
|
||||
"""
|
||||
This example shows how to use vLLM for running Voxtral TTS
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import gc
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from argparse import Namespace
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from mistral_common.protocol.instruct.chunk import TextChunk
|
||||
|
||||
try:
|
||||
from mistral_common.protocol.speech.request import SpeechRequest
|
||||
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Could not import SpeechRequest or MistralTokenizer from mistral_common. "
|
||||
"Please pull the latest mistral-common code and install it with: "
|
||||
"pip install -e ."
|
||||
)
|
||||
from vllm import SamplingParams
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import AsyncOmni
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---- Streaming version ----
|
||||
async def run_streaming(inputs, sampling_params_list, model_name, args, output_dir):
|
||||
async_omni = AsyncOmni(
|
||||
model=model_name,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
)
|
||||
|
||||
# Normalize to a list so batch and single-input share the same code path
|
||||
if isinstance(inputs, list):
|
||||
inputs_list = inputs
|
||||
else:
|
||||
inputs_list = [inputs]
|
||||
|
||||
total_audio_dur = 0.0
|
||||
total_gen_time = 0.0
|
||||
total_ttfa = 0.0
|
||||
ttfa_count = 0
|
||||
total_waits = 0
|
||||
total_chunks = 0
|
||||
results_lock = asyncio.Lock()
|
||||
|
||||
async def _generate_one(batch_idx, single_input):
|
||||
nonlocal total_audio_dur, total_gen_time, total_ttfa, ttfa_count
|
||||
nonlocal total_waits, total_chunks
|
||||
request_id = str(uuid.uuid4())
|
||||
all_audio_chunks = []
|
||||
chunk_arrival_times = []
|
||||
chunk_durations = []
|
||||
chunk_idx = 0
|
||||
gen_start = time.time()
|
||||
ttfa = None
|
||||
accumulated_sample = 0
|
||||
|
||||
async for stage_output in async_omni.generate(single_input, request_id, sampling_params_list):
|
||||
mm_output = stage_output.multimodal_output
|
||||
finished = stage_output.finished
|
||||
if not mm_output or "audio" not in mm_output:
|
||||
continue
|
||||
|
||||
now = time.time()
|
||||
if ttfa is None:
|
||||
ttfa = now - gen_start
|
||||
|
||||
audio_chunk = mm_output["audio"]
|
||||
# audio_chunk is a 1-D tensor from decode_helper_async
|
||||
if isinstance(audio_chunk, torch.Tensor):
|
||||
if finished:
|
||||
# Last chunk may return whole audio instead of chunk delta.
|
||||
# Cut accumulated samples from previous chunks.
|
||||
audio_numpy = audio_chunk[accumulated_sample:].float().detach().cpu().numpy()
|
||||
else:
|
||||
audio_numpy = audio_chunk.float().detach().cpu().numpy()
|
||||
elif isinstance(audio_chunk, list):
|
||||
# Audio_chunk list contain all previous chunks. Use chunk_idx to index
|
||||
audio_numpy = audio_chunk[chunk_idx].float().detach().cpu().numpy()
|
||||
else:
|
||||
audio_numpy = audio_chunk
|
||||
|
||||
accumulated_sample += len(audio_numpy)
|
||||
all_audio_chunks.append(audio_numpy)
|
||||
chunk_arrival_times.append(now)
|
||||
chunk_durations.append(len(audio_numpy) / 24000)
|
||||
chunk_idx += 1
|
||||
|
||||
gen_elapsed = time.time() - gen_start
|
||||
|
||||
# Analyze wait/no-wait per chunk:
|
||||
# A "wait" means the client has finished playing all previous audio
|
||||
# before this chunk arrived, so there is a playback gap.
|
||||
# buffer_time > 0: client still has audio to play (no wait)
|
||||
# buffer_time <= 0: client ran out of audio and had to wait
|
||||
chunk_labels = []
|
||||
accumulated_audio_dur = 0.0
|
||||
if chunk_arrival_times:
|
||||
first_arrival = chunk_arrival_times[0]
|
||||
for i in range(len(chunk_arrival_times)):
|
||||
if i == 0:
|
||||
chunk_labels.append("no_wait")
|
||||
else:
|
||||
playback_elapsed = chunk_arrival_times[i] - first_arrival
|
||||
buffer_time = accumulated_audio_dur - playback_elapsed
|
||||
if buffer_time <= 0:
|
||||
chunk_labels.append("wait")
|
||||
else:
|
||||
chunk_labels.append("no_wait")
|
||||
accumulated_audio_dur += chunk_durations[i]
|
||||
|
||||
req_wait_count = sum(1 for chunk_label in chunk_labels if chunk_label == "wait")
|
||||
req_wait_rate = req_wait_count / len(chunk_labels) if chunk_labels else 0.0
|
||||
|
||||
# Concatenate all chunks for this request's audio
|
||||
if all_audio_chunks:
|
||||
full_audio = np.concatenate(all_audio_chunks)
|
||||
output_audio_dur = len(full_audio) / 24000
|
||||
output_path = os.path.join(output_dir, f"tts_output_{batch_idx}.wav")
|
||||
if args.write_audio:
|
||||
sf.write(output_path, full_audio, 24000)
|
||||
print(
|
||||
f"Request {batch_idx}: saved {len(full_audio)} samples ({output_audio_dur:.2f}s) to {output_path}"
|
||||
)
|
||||
# Per-chunk wait details
|
||||
for i, label in enumerate(chunk_labels):
|
||||
dur_ms = chunk_durations[i] * 1000
|
||||
arrival_ms = (chunk_arrival_times[i] - first_arrival) * 1000 if i > 0 else 0.0
|
||||
print(
|
||||
f" Request {batch_idx} chunk {i}: {label} | arrived={arrival_ms:.1f}ms | chunk_dur={dur_ms:.1f}ms"
|
||||
)
|
||||
print(
|
||||
f"Request {batch_idx}: TTFA={ttfa:.4f}s | "
|
||||
f"Generation={gen_elapsed:.4f}s | "
|
||||
f"Audio={output_audio_dur:.2f}s | "
|
||||
f"RTF={output_audio_dur / gen_elapsed:.4f} | "
|
||||
f"WaitRate={req_wait_rate:.2%} ({req_wait_count}/{len(chunk_labels)})"
|
||||
)
|
||||
async with results_lock:
|
||||
total_audio_dur += output_audio_dur
|
||||
total_gen_time += gen_elapsed
|
||||
total_waits += req_wait_count
|
||||
total_chunks += len(chunk_labels)
|
||||
if ttfa is not None:
|
||||
total_ttfa += ttfa
|
||||
ttfa_count += 1
|
||||
|
||||
# Launch requests in waves of `concurrency` size
|
||||
concurrency = args.concurrency or len(inputs_list)
|
||||
gen_start_all = time.time()
|
||||
for wave_start in range(0, len(inputs_list), concurrency):
|
||||
wave = inputs_list[wave_start : wave_start + concurrency]
|
||||
wave_idx_offset = wave_start
|
||||
print(f"\n--- Wave {wave_start // concurrency + 1} (requests {wave_start}-{wave_start + len(wave) - 1}) ---")
|
||||
await asyncio.gather(*[_generate_one(wave_idx_offset + i, inp) for i, inp in enumerate(wave)])
|
||||
generation_time = time.time() - gen_start_all
|
||||
avg_ttfa = total_ttfa / ttfa_count if ttfa_count else float("nan")
|
||||
overall_wait_rate = total_waits / total_chunks if total_chunks else float("nan")
|
||||
print(
|
||||
f"\nAll requests: Generation={generation_time:.4f}s | "
|
||||
f"TotalAudio={total_audio_dur:.2f}s | "
|
||||
f"Concurrency={concurrency} | "
|
||||
f"AvgTTFA={avg_ttfa:.4f}s | "
|
||||
f"RTF(total)={total_audio_dur / generation_time:.4f} | "
|
||||
f"RTF(per-request)={total_audio_dur / total_gen_time:.4f} | "
|
||||
f"WaitRate={overall_wait_rate:.2%} ({total_waits}/{total_chunks})"
|
||||
)
|
||||
|
||||
async_omni.shutdown()
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
# ---- Non-streaming version ----
|
||||
def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir):
|
||||
llm = Omni(
|
||||
model=model_name,
|
||||
log_stats=args.log_stats,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
)
|
||||
|
||||
if args.profiling_mode:
|
||||
llm.start_profile()
|
||||
|
||||
vllm_start = time.time()
|
||||
outputs = llm.generate(inputs, sampling_params_list)
|
||||
vllm_elapsed = time.time() - vllm_start
|
||||
|
||||
if args.profiling_mode:
|
||||
llm.stop_profile()
|
||||
time.sleep(10)
|
||||
|
||||
print(f"vLLM run time: {vllm_elapsed:.4f}s")
|
||||
output_audio_dur = 0.0
|
||||
|
||||
for batch_idx, o in enumerate(outputs):
|
||||
audio_tensor = torch.cat(o.request_output[0].outputs[0].multimodal_output["audio"])
|
||||
audio_array = audio_tensor.tolist()
|
||||
output_audio_dur += float(len(audio_array)) / 24000
|
||||
if args.write_audio:
|
||||
output_path = os.path.join(output_dir, f"tts_output_{batch_idx}.wav")
|
||||
sf.write(output_path, audio_array, 24000)
|
||||
print(f"Audio saved to {output_path}")
|
||||
|
||||
print(f"Total audio duration: {output_audio_dur:.2f}s")
|
||||
print(f"RTF: {output_audio_dur / vllm_elapsed:.4f}")
|
||||
|
||||
del llm
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
def parse_args() -> Namespace:
|
||||
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="mistralai/tts-model",
|
||||
help="Model name or path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
type=str,
|
||||
default="This is a test message.",
|
||||
help="Text to synthesize.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to reference audio file for voice cloning.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default="output_audio",
|
||||
help="Directory to write output wav files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to stage configs YAML. Auto-resolved from model if not set.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-prompts", type=int, default=1, help="Number of replicate prompts to run for measuring performance"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiling-mode",
|
||||
default=False,
|
||||
action="store_true",
|
||||
help="Set max_num_tokens=2 to reduce profiling time",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable writing detailed statistics (default: disabled)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--write-audio",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Write audio output to files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--streaming",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Use streaming generation via AsyncOmni instead of blocking Omni.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--concurrency",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Max concurrent requests per wave (default: all at once). Must evenly divide --num-prompts.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--voice",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Voice to use instead of audio file.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def compose_request(
|
||||
model_name: str,
|
||||
text_chunk: TextChunk,
|
||||
audio_prompt_file: str,
|
||||
args: Any,
|
||||
) -> dict:
|
||||
"""Build the full TTS input dict (prompt_token_ids, multi_modal_data or additional_information)."""
|
||||
inputs: dict[str, Any] = {}
|
||||
if Path(model_name).is_dir():
|
||||
mistral_tokenizer = MistralTokenizer.from_file(str(Path(model_name) / "tekken.json"))
|
||||
else:
|
||||
mistral_tokenizer = MistralTokenizer.from_hf_hub(model_name)
|
||||
instruct_tokenizer = mistral_tokenizer.instruct_tokenizer
|
||||
|
||||
if args.voice is not None:
|
||||
tokenized = instruct_tokenizer.encode_speech_request(SpeechRequest(input=text_chunk.text, voice=args.voice))
|
||||
inputs["additional_information"] = {"voice": [args.voice]}
|
||||
inputs["prompt_token_ids"] = tokenized.tokens
|
||||
else:
|
||||
with open(audio_prompt_file, "rb") as f:
|
||||
ref_audio_bytes = f.read()
|
||||
tokenized = instruct_tokenizer.encode_speech_request(
|
||||
SpeechRequest(input=text_chunk.text, ref_audio=ref_audio_bytes)
|
||||
)
|
||||
audio = tokenized.audios[0]
|
||||
inputs["multi_modal_data"] = {"audio": [(audio.audio_array, audio.sampling_rate)]}
|
||||
inputs["prompt_token_ids"] = tokenized.tokens
|
||||
|
||||
return inputs
|
||||
|
||||
|
||||
def main(args: Any) -> None:
|
||||
max_num_tokens = 50 if args.profiling_mode else 2500
|
||||
|
||||
model_name = args.model
|
||||
output_dir = args.output_dir
|
||||
|
||||
if args.voice is None and args.audio_path is None:
|
||||
raise ValueError("Either --voice or --audio-path must be provided.")
|
||||
|
||||
audio_prompt_file = args.audio_path
|
||||
text_chunk = TextChunk(text=args.text)
|
||||
|
||||
if args.write_audio:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
inputs = compose_request(model_name, text_chunk, audio_prompt_file, args)
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
max_tokens=max_num_tokens,
|
||||
)
|
||||
sampling_params_list = [
|
||||
sampling_params,
|
||||
sampling_params,
|
||||
]
|
||||
|
||||
if args.num_prompts > 1:
|
||||
inputs = [inputs] * args.num_prompts
|
||||
|
||||
if args.concurrency is not None:
|
||||
assert args.streaming, "--concurrency must be used with --streaming on AsyncOmni"
|
||||
assert args.num_prompts % args.concurrency == 0, (
|
||||
f"--num-prompts ({args.num_prompts}) must be divisible by --concurrency ({args.concurrency})"
|
||||
)
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
if args.streaming:
|
||||
asyncio.run(run_streaming(inputs, sampling_params_list, model_name, args, output_dir))
|
||||
else:
|
||||
run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,255 @@
|
||||
import argparse
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import socket
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import gradio as gr
|
||||
import httpx
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
logger = logging.getLogger()
|
||||
|
||||
LOGFORMAT = "%(asctime)s - %(levelname)s - %(message)s"
|
||||
TIMEFORMAT = "%Y-%m-%d %H:%M:%S"
|
||||
|
||||
logging.basicConfig(level=logging.INFO, force=True, format=LOGFORMAT, datefmt=TIMEFORMAT)
|
||||
logger.setLevel(logging.INFO)
|
||||
|
||||
|
||||
_BASE_URL = f"{socket.gethostname()}:7860"
|
||||
|
||||
|
||||
def run_inference(
|
||||
voice_name: str,
|
||||
text_prompt: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
) -> tuple[int, np.ndarray]:
|
||||
"""Call /v1/audio/speech and return (sample_rate, audio_array)."""
|
||||
text_prompt = text_prompt.strip()
|
||||
if not text_prompt:
|
||||
raise gr.Error("Please enter a text prompt.")
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"input": text_prompt,
|
||||
"model": model,
|
||||
"response_format": "wav",
|
||||
"voice": voice_name,
|
||||
}
|
||||
|
||||
response = httpx.post(
|
||||
f"{base_url}/audio/speech",
|
||||
json=payload,
|
||||
timeout=120.0,
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
audio_array, sr = sf.read(io.BytesIO(response.content), dtype="float32")
|
||||
return sr, audio_array
|
||||
|
||||
|
||||
def _save_example(
|
||||
outputs_dir: Path,
|
||||
voice_name: str,
|
||||
text_prompt: str,
|
||||
sr: int,
|
||||
audio_array: np.ndarray,
|
||||
) -> tuple[str, str]:
|
||||
"""
|
||||
Save inputs/outputs for sharing.
|
||||
Returns (share_id, saved_audio_path)
|
||||
"""
|
||||
share_id = uuid.uuid4().hex
|
||||
|
||||
# Save generated audio
|
||||
saved_audio_path = outputs_dir / f"{share_id}_gen.wav"
|
||||
sf.write(str(saved_audio_path), audio_array, sr)
|
||||
|
||||
meta = {
|
||||
"id": share_id,
|
||||
"created_at": datetime.utcnow().isoformat(),
|
||||
"voice_name": voice_name,
|
||||
"text_prompt": text_prompt,
|
||||
"generated_audio_path": str(saved_audio_path),
|
||||
}
|
||||
|
||||
with open(outputs_dir / f"{share_id}.json", "w") as f:
|
||||
json.dump(meta, f, ensure_ascii=False)
|
||||
|
||||
return share_id, str(saved_audio_path)
|
||||
|
||||
|
||||
def _load_from_share(
|
||||
outputs_dir: Path | None,
|
||||
request: gr.Request,
|
||||
) -> tuple[str | None, str, str | None, dict[str, Any], str]:
|
||||
"""
|
||||
Called on page load. If ?share_id=... is present, load stored example.
|
||||
Returns: (voice_name, text_prompt, output_audio, submit_btn_update, share_link_text)
|
||||
"""
|
||||
if outputs_dir is None:
|
||||
return "Neutral_Male", "", None, gr.update(interactive=False), ""
|
||||
|
||||
share_id = None
|
||||
if request and request.query_params:
|
||||
share_id = request.query_params.get("share_id")
|
||||
|
||||
if not share_id:
|
||||
return "Neutral_Male", "", None, gr.update(interactive=False), ""
|
||||
|
||||
meta_path = outputs_dir / f"{share_id}.json"
|
||||
if not meta_path.exists():
|
||||
logger.warning("No stored example for share_id=%s", share_id)
|
||||
return "Neutral_Male", "", None, gr.update(interactive=False), ""
|
||||
|
||||
with open(meta_path) as f:
|
||||
meta = json.load(f)
|
||||
|
||||
voice = meta.get("voice_name", "Neutral_Male")
|
||||
text_prompt = meta.get("text_prompt", "")
|
||||
gen_path = meta.get("generated_audio_path")
|
||||
|
||||
if _BASE_URL:
|
||||
share_link = f"http://{_BASE_URL}?share_id={share_id}"
|
||||
else:
|
||||
share_link = f"Use this query on your current URL: ?share_id={share_id}"
|
||||
|
||||
return voice, text_prompt, gen_path, gr.update(interactive=True), share_link
|
||||
|
||||
|
||||
def main(
|
||||
model: str,
|
||||
output_dir: str | None = None,
|
||||
debug: bool = False,
|
||||
) -> None:
|
||||
base_url = "http://localhost:8091/v1"
|
||||
print(f"Using speech API at: {base_url}/audio/speech")
|
||||
|
||||
outputs_dir: Path | None = None
|
||||
if output_dir is not None:
|
||||
outputs_dir = Path(output_dir)
|
||||
outputs_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with gr.Blocks(title="Voxtral TTS", fill_height=True) as demo:
|
||||
gr.Markdown("## Voxtral TTS")
|
||||
|
||||
with gr.Row():
|
||||
with gr.Column():
|
||||
voice_name = gr.Dropdown(
|
||||
["Neutral_Male", "Neutral_Female", "Cheerful_Female", "Casual_Male", "Casual_Female"],
|
||||
label="Voice",
|
||||
value="Neutral_Male",
|
||||
)
|
||||
text_prompt = gr.Textbox(
|
||||
label="Text prompt",
|
||||
placeholder="Enter the text you want to synthesize...",
|
||||
lines=4,
|
||||
)
|
||||
with gr.Row():
|
||||
reset_btn = gr.Button("Clear")
|
||||
submit_btn = gr.Button("Generate audio", interactive=False)
|
||||
|
||||
with gr.Column():
|
||||
output_audio = gr.Audio(
|
||||
label="Generated audio",
|
||||
show_download_button=True,
|
||||
interactive=False,
|
||||
autoplay=True,
|
||||
type="filepath",
|
||||
)
|
||||
share_link_box = gr.Textbox(
|
||||
label="Shareable link",
|
||||
interactive=False,
|
||||
show_copy_button=True,
|
||||
visible=outputs_dir is not None,
|
||||
)
|
||||
|
||||
# --- UI logic: disable submit until text is non-empty ---
|
||||
def _toggle_submit(text: str):
|
||||
enabled = bool(text.strip())
|
||||
return gr.update(interactive=enabled)
|
||||
|
||||
text_prompt.change(
|
||||
fn=_toggle_submit,
|
||||
inputs=[text_prompt],
|
||||
outputs=submit_btn,
|
||||
)
|
||||
|
||||
# --- Wiring inference + persistence to the button ---
|
||||
def _on_submit(voice: str, text: str):
|
||||
assert text.strip() != ""
|
||||
sr, audio_array = run_inference(voice, text, base_url, model)
|
||||
if outputs_dir is not None:
|
||||
share_id, saved_audio_path = _save_example(
|
||||
outputs_dir,
|
||||
voice_name=voice,
|
||||
text_prompt=text,
|
||||
sr=sr,
|
||||
audio_array=audio_array,
|
||||
)
|
||||
share_link = f"{_BASE_URL}?share_id={share_id}"
|
||||
return saved_audio_path, share_link
|
||||
return (sr, audio_array), ""
|
||||
|
||||
submit_btn.click(
|
||||
fn=_on_submit,
|
||||
inputs=[voice_name, text_prompt],
|
||||
outputs=[output_audio, share_link_box],
|
||||
)
|
||||
|
||||
# --- Clear everything and disable submit again ---
|
||||
def _on_reset():
|
||||
return (
|
||||
"Neutral_Male", # voice_name
|
||||
"", # text_prompt
|
||||
None, # output_audio
|
||||
gr.update(interactive=False), # submit_btn
|
||||
"", # share_link_box
|
||||
)
|
||||
|
||||
reset_btn.click(
|
||||
fn=_on_reset,
|
||||
inputs=[],
|
||||
outputs=[voice_name, text_prompt, output_audio, submit_btn, share_link_box],
|
||||
)
|
||||
|
||||
demo.load(
|
||||
fn=lambda request: _load_from_share(outputs_dir, request),
|
||||
inputs=[],
|
||||
outputs=[voice_name, text_prompt, output_audio, submit_btn, share_link_box],
|
||||
)
|
||||
|
||||
launch_kwargs: dict[str, Any] = {
|
||||
"server_name": "0.0.0.0",
|
||||
"share": True,
|
||||
}
|
||||
if outputs_dir is not None:
|
||||
launch_kwargs["allowed_paths"] = [str(outputs_dir)]
|
||||
demo.launch(**launch_kwargs)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Voxtral TTS Gradio Demo")
|
||||
parser.add_argument("--model", type=str, required=True, help="Path to the model checkpoint")
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory to save generated audio and share links. "
|
||||
"If not provided, save/share functionality is disabled.",
|
||||
)
|
||||
parser.add_argument("--debug", action="store_true", default=False, help="Enable debug mode")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
main(
|
||||
model=args.model,
|
||||
output_dir=args.output_dir,
|
||||
debug=args.debug,
|
||||
)
|
||||
@@ -1236,8 +1236,6 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st
|
||||
port = params.port
|
||||
stage_config_path = params.stage_config_path
|
||||
if run_level == "advanced_model" and stage_config_path is not None:
|
||||
# Dynamically detect stages from config to avoid KeyError
|
||||
# for models with fewer stages (e.g., BAGEL has 2, Qwen3-Omni has 3)
|
||||
with open(stage_config_path, encoding="utf-8") as f:
|
||||
_cfg = yaml.safe_load(f) or {}
|
||||
_stage_ids = [s["stage_id"] for s in _cfg.get("stage_args", []) if "stage_id" in s]
|
||||
|
||||
@@ -0,0 +1,192 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
E2E Online tests for Voxtral TTS model with text input and audio output.
|
||||
|
||||
These tests verify the /v1/audio/speech endpoint works correctly with
|
||||
actual model inference, not mocks.
|
||||
"""
|
||||
|
||||
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 = "mistralai/tts-model"
|
||||
|
||||
STAGE_CONFIG = str(
|
||||
Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "voxtral_tts.yaml"
|
||||
)
|
||||
EXTRA_ARGS = ["--trust-remote-code", "--enforce-eager", "--disable-log-stats"]
|
||||
TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG, server_args=EXTRA_ARGS)]
|
||||
|
||||
|
||||
def make_speech_request(
|
||||
host: str,
|
||||
port: int,
|
||||
text: str,
|
||||
voice: str = "casual_female",
|
||||
language: str = "English",
|
||||
timeout: float = 120.0,
|
||||
) -> httpx.Response:
|
||||
"""Make a request to the /v1/audio/speech endpoint."""
|
||||
url = f"http://{host}:{port}/v1/audio/speech"
|
||||
payload = {
|
||||
"input": text,
|
||||
"voice": voice,
|
||||
"language": language,
|
||||
}
|
||||
|
||||
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."""
|
||||
# WAV files start with "RIFF" header
|
||||
if len(content) < 44: # Minimum WAV header size
|
||||
return False
|
||||
return content[:4] == b"RIFF" and content[8:12] == b"WAVE"
|
||||
|
||||
|
||||
# Minimum expected audio size for a short sentence (~1 second of 24kHz 16-bit mono WAV)
|
||||
MIN_AUDIO_BYTES = 10000
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
class TestVoxtralTTSFixedVoice:
|
||||
"""E2E tests for Voxtral TTS model."""
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_speech_english_basic(self, omni_server) -> None:
|
||||
"""Test basic English TTS generation."""
|
||||
response = make_speech_request(
|
||||
host=omni_server.host,
|
||||
port=omni_server.port,
|
||||
text="Hello, how are you?",
|
||||
voice="casual_female",
|
||||
language="English",
|
||||
)
|
||||
|
||||
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 content too small ({len(response.content)} bytes), expected at least {MIN_AUDIO_BYTES} bytes"
|
||||
)
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_speech_different_voices(self, omni_server) -> None:
|
||||
"""Test TTS with different voice presets."""
|
||||
voices = ["casual_female", "neutral_male"]
|
||||
for voice in voices:
|
||||
response = make_speech_request(
|
||||
host=omni_server.host,
|
||||
port=omni_server.port,
|
||||
text="Testing voice selection.",
|
||||
voice=voice,
|
||||
language="English",
|
||||
)
|
||||
|
||||
assert response.status_code == 200, f"Request failed for voice {voice}: {response.text}"
|
||||
assert verify_wav_audio(response.content), f"Invalid WAV for voice {voice}"
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_speech_invalid_voice_rejected(self, omni_server) -> None:
|
||||
"""Request with a non-existent voice should return an error response."""
|
||||
response = make_speech_request(
|
||||
host=omni_server.host,
|
||||
port=omni_server.port,
|
||||
text="This voice does not exist.",
|
||||
voice="nonexistent_voice_xyz",
|
||||
language="English",
|
||||
)
|
||||
|
||||
# vLLM returns ErrorResponse as JSON with HTTP 200;
|
||||
# verify the body is a JSON error, not valid WAV audio.
|
||||
data = response.json()
|
||||
assert "error" in data or "message" in data, f"Expected error response for invalid voice, got: {data}"
|
||||
assert not verify_wav_audio(response.content), "Should not return valid audio for invalid voice"
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_speech_binary_response_not_utf8_error(self, omni_server) -> None:
|
||||
"""
|
||||
Regression test: Verify binary audio is returned, not UTF-8 error.
|
||||
|
||||
This test ensures the multimodal_output property correctly retrieves
|
||||
audio from completion outputs, preventing the "TTS model did not
|
||||
produce audio output" error.
|
||||
"""
|
||||
response = make_speech_request(
|
||||
host=omni_server.host,
|
||||
port=omni_server.port,
|
||||
text="This should return binary audio, not a JSON error.",
|
||||
voice="casual_female",
|
||||
language="English",
|
||||
)
|
||||
|
||||
# Should NOT be a JSON error response
|
||||
assert response.status_code == 200, f"Request failed: {response.text}"
|
||||
|
||||
# Verify it's binary audio, not JSON
|
||||
try:
|
||||
# If this succeeds and starts with {"error", it's a bug
|
||||
text = response.content.decode("utf-8")
|
||||
assert not text.startswith('{"error"'), f"Got error response instead of audio: {text}"
|
||||
except UnicodeDecodeError:
|
||||
# This is expected - binary audio can't be decoded as UTF-8
|
||||
pass
|
||||
|
||||
assert verify_wav_audio(response.content), "Response is not valid WAV audio"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
class TestVoxtralTTSAPIEndpoints:
|
||||
"""Test API endpoint functionality."""
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_list_voices_endpoint(self, omni_server) -> None:
|
||||
"""Test the /v1/audio/voices endpoint returns available voices."""
|
||||
url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/voices"
|
||||
|
||||
with httpx.Client(timeout=30.0) as client:
|
||||
response = client.get(url)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "voices" in data
|
||||
assert isinstance(data["voices"], list)
|
||||
assert len(data["voices"]) > 0
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_models_endpoint(self, omni_server) -> None:
|
||||
"""Test the /v1/models endpoint returns loaded model."""
|
||||
url = f"http://{omni_server.host}:{omni_server.port}/v1/models"
|
||||
|
||||
with httpx.Client(timeout=30.0) as client:
|
||||
response = client.get(url)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "data" in data
|
||||
assert len(data["data"]) > 0
|
||||
Whitespace-only changes.
@@ -0,0 +1,161 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for parse_batched_audio_input and apply_ctx_frames_cutting.
|
||||
|
||||
These functions parse the [ctx_frames, context_length, ...tokens] wire format
|
||||
produced by generator2tokenizer_async_chunk and cut leading context samples
|
||||
from decoded audio arrays.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts import (
|
||||
apply_ctx_frames_cutting,
|
||||
parse_batched_audio_input,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
NUM_CODEBOOKS = 37
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
# parse_batched_audio_input tests
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _make_input_ids(*requests: tuple[int, int, torch.Tensor]) -> torch.Tensor:
|
||||
"""Build a flat input_ids tensor from (ctx_frames, context_length, tokens) tuples."""
|
||||
parts = []
|
||||
for ctx_frames, context_length, tokens in requests:
|
||||
parts.append(torch.tensor([ctx_frames, context_length]))
|
||||
parts.append(tokens.flatten())
|
||||
return torch.cat(parts)
|
||||
|
||||
|
||||
def test_parse_single_request():
|
||||
"""Single request is parsed correctly."""
|
||||
ctx_frames = 3
|
||||
context_length = 5
|
||||
total_frames = ctx_frames + context_length
|
||||
tokens = torch.arange(total_frames * NUM_CODEBOOKS).reshape(total_frames, NUM_CODEBOOKS)
|
||||
input_ids = _make_input_ids((ctx_frames, context_length, tokens))
|
||||
|
||||
all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS)
|
||||
|
||||
assert len(all_audio_tokens) == 1
|
||||
assert len(all_ctx_frames) == 1
|
||||
assert all_ctx_frames[0] == ctx_frames
|
||||
assert all_audio_tokens[0].shape == (total_frames, NUM_CODEBOOKS)
|
||||
torch.testing.assert_close(all_audio_tokens[0], tokens)
|
||||
|
||||
|
||||
def test_parse_multiple_requests():
|
||||
"""Multiple requests in a batch are parsed correctly."""
|
||||
req1 = (2, 4, torch.ones((6, NUM_CODEBOOKS), dtype=torch.long))
|
||||
req2 = (0, 3, torch.full((3, NUM_CODEBOOKS), 7, dtype=torch.long))
|
||||
input_ids = _make_input_ids(req1, req2)
|
||||
|
||||
all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS)
|
||||
|
||||
assert len(all_audio_tokens) == 2
|
||||
assert all_ctx_frames == [2, 0]
|
||||
assert all_audio_tokens[0].shape == (6, NUM_CODEBOOKS)
|
||||
assert all_audio_tokens[1].shape == (3, NUM_CODEBOOKS)
|
||||
assert torch.all(all_audio_tokens[0] == 1)
|
||||
assert torch.all(all_audio_tokens[1] == 7)
|
||||
|
||||
|
||||
def test_parse_zero_ctx_frames():
|
||||
"""ctx_frames=0 means no context, just chunk data."""
|
||||
ctx_frames = 0
|
||||
context_length = 10
|
||||
tokens = torch.randint(0, 100, (context_length, NUM_CODEBOOKS))
|
||||
input_ids = _make_input_ids((ctx_frames, context_length, tokens))
|
||||
|
||||
all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS)
|
||||
|
||||
assert all_ctx_frames[0] == 0
|
||||
assert all_audio_tokens[0].shape == (context_length, NUM_CODEBOOKS)
|
||||
|
||||
|
||||
def test_parse_fails_on_misaligned_input():
|
||||
"""Input not divisible by num_codebooks should raise AssertionError."""
|
||||
# 2 header elements + 5 data elements (not divisible by 37)
|
||||
input_ids = torch.tensor([0, 1] + [0] * 5)
|
||||
with pytest.raises(AssertionError, match="divisible by"):
|
||||
parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS)
|
||||
|
||||
|
||||
def test_parse_custom_num_codebooks():
|
||||
"""Works with a non-default num_codebooks value."""
|
||||
cb = 4
|
||||
ctx_frames = 1
|
||||
context_length = 2
|
||||
total_frames = ctx_frames + context_length
|
||||
tokens = torch.arange(total_frames * cb).reshape(total_frames, cb)
|
||||
input_ids = _make_input_ids((ctx_frames, context_length, tokens))
|
||||
|
||||
all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=cb)
|
||||
|
||||
assert all_audio_tokens[0].shape == (total_frames, cb)
|
||||
torch.testing.assert_close(all_audio_tokens[0], tokens)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
# apply_ctx_frames_cutting tests
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_cut_removes_leading_samples():
|
||||
"""Context frames are removed from the front of the audio array."""
|
||||
downsample_factor = 240
|
||||
ctx_frames = 5
|
||||
total_samples = 2400 # 10 frames * 240
|
||||
audio = torch.arange(total_samples, dtype=torch.float32)
|
||||
|
||||
result = apply_ctx_frames_cutting([audio], [ctx_frames], downsample_factor)
|
||||
|
||||
expected_cut = ctx_frames * downsample_factor # 1200
|
||||
assert len(result) == 1
|
||||
assert result[0].shape[0] == total_samples - expected_cut
|
||||
torch.testing.assert_close(result[0], audio[expected_cut:])
|
||||
|
||||
|
||||
def test_cut_zero_ctx_frames_unchanged():
|
||||
"""ctx_frames=0 leaves the audio unchanged."""
|
||||
downsample_factor = 240
|
||||
audio = torch.randn(2400)
|
||||
|
||||
result = apply_ctx_frames_cutting([audio], [0], downsample_factor)
|
||||
|
||||
assert result[0].shape == audio.shape
|
||||
torch.testing.assert_close(result[0], audio)
|
||||
|
||||
|
||||
def test_cut_multiple_requests():
|
||||
"""Each request in the batch gets its own ctx_frames cut."""
|
||||
downsample_factor = 100
|
||||
audio1 = torch.arange(1000, dtype=torch.float32) # 10 frames
|
||||
audio2 = torch.arange(500, dtype=torch.float32) # 5 frames
|
||||
|
||||
result = apply_ctx_frames_cutting([audio1, audio2], [3, 0], downsample_factor)
|
||||
|
||||
assert len(result) == 2
|
||||
assert result[0].shape[0] == 1000 - 3 * 100 # 700
|
||||
assert result[1].shape[0] == 500 # unchanged
|
||||
torch.testing.assert_close(result[0], audio1[300:])
|
||||
torch.testing.assert_close(result[1], audio2)
|
||||
|
||||
|
||||
def test_cut_all_frames_returns_empty():
|
||||
"""Cutting all frames returns an empty tensor."""
|
||||
downsample_factor = 100
|
||||
ctx_frames = 5
|
||||
audio = torch.randn(500) # exactly 5 frames
|
||||
|
||||
result = apply_ctx_frames_cutting([audio], [ctx_frames], downsample_factor)
|
||||
|
||||
assert result[0].numel() == 0
|
||||
@@ -0,0 +1,329 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for CUDA Graph acoustic transformer wrapper numerical equivalence.
|
||||
|
||||
Verifies that CUDA Graph-accelerated decoding produces results equivalent
|
||||
to eager mode, with special attention to padding cases where zero-padding
|
||||
may introduce small numerical differences due to the flow matching ODE.
|
||||
|
||||
Architecture note: the real FlowMatchingAudioTransformer uses a multi-step
|
||||
Euler ODE with CFG. The synthetic model here uses simple linear layers
|
||||
to exercise the wrapper mechanism while keeping the test lightweight.
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
import os
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.cuda,
|
||||
]
|
||||
|
||||
DEVICE = torch.device("cuda:0")
|
||||
HIDDEN_DIM = 64
|
||||
N_ACOUSTIC_CODEBOOK = 7
|
||||
SEMANTIC_CODEBOOK_SIZE = 128
|
||||
ACOUSTIC_EMBEDDINGS_LEVELS = 1024
|
||||
|
||||
# Load CUDAGraphAcousticTransformerWrapper: try package import first, fall back to direct file load
|
||||
try:
|
||||
from vllm_omni.model_executor.models.voxtral_tts.cuda_graph_acoustic_transformer_wrapper import (
|
||||
CUDAGraphAcousticTransformerWrapper,
|
||||
)
|
||||
from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation import (
|
||||
AudioSpecialTokens,
|
||||
)
|
||||
except Exception:
|
||||
_WRAPPER_PATH = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
os.pardir,
|
||||
os.pardir,
|
||||
os.pardir,
|
||||
os.pardir,
|
||||
"vllm_omni",
|
||||
"model_executor",
|
||||
"models",
|
||||
"voxtral_tts",
|
||||
"cuda_graph_acoustic_transformer_wrapper.py",
|
||||
)
|
||||
_spec = importlib.util.spec_from_file_location(
|
||||
"cuda_graph_acoustic_transformer_wrapper", os.path.abspath(_WRAPPER_PATH)
|
||||
)
|
||||
_mod = importlib.util.module_from_spec(_spec)
|
||||
_spec.loader.exec_module(_mod)
|
||||
CUDAGraphAcousticTransformerWrapper = _mod.CUDAGraphAcousticTransformerWrapper
|
||||
|
||||
_AUDIO_GEN_PATH = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
os.pardir,
|
||||
os.pardir,
|
||||
os.pardir,
|
||||
os.pardir,
|
||||
"vllm_omni",
|
||||
"model_executor",
|
||||
"models",
|
||||
"voxtral_tts",
|
||||
"voxtral_tts_audio_generation.py",
|
||||
)
|
||||
_spec2 = importlib.util.spec_from_file_location("voxtral_tts_audio_generation", os.path.abspath(_AUDIO_GEN_PATH))
|
||||
_mod2 = importlib.util.module_from_spec(_spec2)
|
||||
_spec2.loader.exec_module(_mod2)
|
||||
AudioSpecialTokens = _mod2.AudioSpecialTokens
|
||||
|
||||
|
||||
class SyntheticModelArgs:
|
||||
"""Mimics MultimodalAudioModelArgs interface."""
|
||||
|
||||
def __init__(self):
|
||||
self.semantic_codebook_size = SEMANTIC_CODEBOOK_SIZE
|
||||
self.n_acoustic_codebook = N_ACOUSTIC_CODEBOOK
|
||||
|
||||
|
||||
class SyntheticAcousticTransformer(nn.Module):
|
||||
"""A small acoustic transformer mimicking FlowMatchingAudioTransformer's interface.
|
||||
|
||||
Uses simple linear layers so the test stays lightweight while exercising
|
||||
the full wrapper mechanism (semantic logit, flow matching ODE, quantization).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.model_args = SyntheticModelArgs()
|
||||
self.acoustic_embeddings_levels = ACOUSTIC_EMBEDDINGS_LEVELS
|
||||
|
||||
# semantic_codebook_output: hidden_dim -> padded_codebook_size
|
||||
padded_semantic_size = len(AudioSpecialTokens) + SEMANTIC_CODEBOOK_SIZE
|
||||
# Round up to multiple of 128 like the real model
|
||||
padded_semantic_size = 128 * ((padded_semantic_size + 127) // 128)
|
||||
self.semantic_codebook_output = nn.Linear(HIDDEN_DIM, padded_semantic_size, bias=False)
|
||||
|
||||
# time_embedding: (B, 1) -> (B, dim)
|
||||
self.time_embedding = nn.Linear(1, HIDDEN_DIM, bias=False)
|
||||
|
||||
# _predict_velocity: takes x_t (B, C), llm_output (B, D), t_emb (B, D) -> (B, C)
|
||||
self._velocity_proj = nn.Linear(N_ACOUSTIC_CODEBOOK + HIDDEN_DIM + HIDDEN_DIM, N_ACOUSTIC_CODEBOOK, bias=False)
|
||||
|
||||
def _predict_velocity(self, x_t: torch.Tensor, llm_output: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
|
||||
combined = torch.cat([x_t, llm_output, t_emb], dim=-1)
|
||||
return self._velocity_proj(combined)
|
||||
|
||||
|
||||
class SyntheticModel(nn.Module):
|
||||
"""Mimics VoxtralTTSAudioGenerationForConditionalGeneration interface."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.acoustic_transformer = SyntheticAcousticTransformer()
|
||||
end_audio_id = AudioSpecialTokens.id(AudioSpecialTokens.end_audio)
|
||||
empty_audio_id = AudioSpecialTokens.id(AudioSpecialTokens.empty_audio)
|
||||
self.end_audio_id = end_audio_id
|
||||
self.empty_audio_id = empty_audio_id
|
||||
|
||||
def compute_mm_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
mm_sampling_tensors=None,
|
||||
):
|
||||
"""Eager fallback path: replicate what the wrapper does."""
|
||||
at = self.acoustic_transformer
|
||||
B = hidden_states.shape[0]
|
||||
|
||||
empty_audio_id = self.empty_audio_id
|
||||
end_audio_id = self.end_audio_id
|
||||
semantic_mask_start = len(AudioSpecialTokens) + at.model_args.semantic_codebook_size
|
||||
|
||||
# Semantic logits
|
||||
semantic_logit = at.semantic_codebook_output(hidden_states).float()
|
||||
semantic_logit[:, empty_audio_id] = -float("inf")
|
||||
semantic_logit[:, semantic_mask_start:] = -float("inf")
|
||||
semantic_code = semantic_logit.argmax(dim=-1, keepdim=True)
|
||||
|
||||
# Flow matching Euler ODE
|
||||
should_decode = semantic_code.squeeze(1) != end_audio_id
|
||||
x = torch.randn(B, at.model_args.n_acoustic_codebook, device=hidden_states.device, dtype=hidden_states.dtype)
|
||||
hidden_zero = torch.zeros_like(hidden_states)
|
||||
timesteps = torch.linspace(0, 1, 16, device=hidden_states.device, dtype=hidden_states.dtype)
|
||||
|
||||
for i in range(len(timesteps) - 1):
|
||||
t = timesteps[i]
|
||||
dt = timesteps[i + 1] - timesteps[i]
|
||||
t_emb = at.time_embedding(t.view(-1, 1).repeat(B, 1)).to(hidden_states.dtype)
|
||||
x_batched = torch.cat([x, x], dim=0)
|
||||
llm_batched = torch.cat([hidden_states, hidden_zero], dim=0)
|
||||
t_emb_batched = t_emb.repeat(2, 1)
|
||||
v_all = at._predict_velocity(x_t=x_batched, llm_output=llm_batched, t_emb=t_emb_batched)
|
||||
v_t, uncond_v_t = v_all[:B], v_all[B:]
|
||||
cfg_alpha = 1.2
|
||||
v_t = cfg_alpha * v_t + (1 - cfg_alpha) * uncond_v_t
|
||||
x = x + v_t * dt
|
||||
|
||||
sampled = torch.clamp(x, -1, 1)
|
||||
scaled_x = ((sampled + 1) / 2) * (at.acoustic_embeddings_levels - 1)
|
||||
output_codes = scaled_x.round().long()
|
||||
output_codes[~should_decode] = empty_audio_id
|
||||
acoustic_codes = output_codes + len(AudioSpecialTokens)
|
||||
|
||||
audio_codes = torch.cat([semantic_code, acoustic_codes], dim=1)
|
||||
fake_eos = torch.where(
|
||||
audio_codes[:, 0] == end_audio_id,
|
||||
torch.tensor(1.0, dtype=hidden_states.dtype, device=hidden_states.device),
|
||||
torch.tensor(0.0, dtype=hidden_states.dtype, device=hidden_states.device),
|
||||
)
|
||||
|
||||
audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0))
|
||||
mm_tokens = {"audio": audio_list}
|
||||
return fake_eos, mm_tokens
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def model():
|
||||
"""Create a synthetic model on CUDA with fixed weights in bfloat16."""
|
||||
torch.manual_seed(42)
|
||||
return SyntheticModel().to(device=DEVICE, dtype=torch.bfloat16).eval()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def wrapper(model):
|
||||
"""Create a warmed-up CUDAGraphAcousticTransformerWrapper."""
|
||||
w = CUDAGraphAcousticTransformerWrapper(
|
||||
model=model,
|
||||
capture_sizes=[1, 2, 4, 8, 16, 32],
|
||||
)
|
||||
w._warmup_and_capture(
|
||||
device=DEVICE,
|
||||
dtype=torch.bfloat16,
|
||||
hidden_dim=HIDDEN_DIM,
|
||||
)
|
||||
return w
|
||||
|
||||
|
||||
def _random_hidden(batch_size, device=DEVICE, dtype=torch.bfloat16):
|
||||
return torch.randn(batch_size, HIDDEN_DIM, device=device, dtype=dtype)
|
||||
|
||||
|
||||
def _unpack_audio_codes(result):
|
||||
"""Unpack (fake_eos, {"audio": [list of tensors]}) into (fake_eos, audio_codes)."""
|
||||
fake_eos, mm_tokens = result
|
||||
audio_list = mm_tokens["audio"]
|
||||
# Each element is (1, 1, 1+n_acoustic_codebook), concat along batch dim
|
||||
audio_codes = torch.cat(audio_list, dim=0).squeeze(1) # (B, 1+n_acoustic)
|
||||
return fake_eos, audio_codes
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
# 1. Exact-size inputs (no padding) → correct format and bounds
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16, 32])
|
||||
def test_exact_size_output_format(model, wrapper, batch_size):
|
||||
"""Graph path returns correctly shaped and bounded outputs."""
|
||||
hidden = _random_hidden(batch_size)
|
||||
with torch.no_grad():
|
||||
graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
|
||||
assert graph_eos.shape == (batch_size,)
|
||||
assert graph_codes.shape == (batch_size, 1 + N_ACOUSTIC_CODEBOOK)
|
||||
# fake_eos should be 0.0 or 1.0
|
||||
assert torch.all((graph_eos == 0.0) | (graph_eos == 1.0))
|
||||
# Audio codes should be non-negative
|
||||
assert torch.all(graph_codes >= 0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16, 32])
|
||||
def test_exact_size_deterministic(model, wrapper, batch_size):
|
||||
"""Same input + same RNG state produces identical CUDA graph output."""
|
||||
hidden = _random_hidden(batch_size)
|
||||
with torch.no_grad():
|
||||
torch.manual_seed(42)
|
||||
eos1, codes1 = _unpack_audio_codes(wrapper(hidden))
|
||||
torch.manual_seed(42)
|
||||
eos2, codes2 = _unpack_audio_codes(wrapper(hidden))
|
||||
torch.testing.assert_close(eos1, eos2, atol=0, rtol=0)
|
||||
torch.testing.assert_close(codes1, codes2, atol=0, rtol=0)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
# 2. Padded inputs (zero-padding to nearest capture size)
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [3, 5, 7, 10, 24])
|
||||
def test_padded_output_shape(model, wrapper, batch_size):
|
||||
"""Padded decode must return output trimmed to actual batch size."""
|
||||
hidden = _random_hidden(batch_size)
|
||||
with torch.no_grad():
|
||||
graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
|
||||
assert graph_eos.shape == (batch_size,)
|
||||
assert graph_codes.shape == (batch_size, 1 + N_ACOUSTIC_CODEBOOK)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [3, 5, 7, 10, 24])
|
||||
def test_padded_output_bounded(model, wrapper, batch_size):
|
||||
"""Padded output audio codes should be non-negative integers."""
|
||||
hidden = _random_hidden(batch_size)
|
||||
with torch.no_grad():
|
||||
graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
|
||||
# fake_eos should be 0.0 or 1.0
|
||||
assert torch.all((graph_eos == 0.0) | (graph_eos == 1.0))
|
||||
# Audio codes should be non-negative
|
||||
assert torch.all(graph_codes >= 0)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
# 3. Batch size exceeds all capture sizes so fallback to eager
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [33, 48, 64])
|
||||
def test_fallback_eager_exact_match(model, wrapper, batch_size):
|
||||
"""Cudagraph fallback to eager. Two eager runs should produce identical results."""
|
||||
hidden = _random_hidden(batch_size)
|
||||
with torch.no_grad():
|
||||
torch.manual_seed(100)
|
||||
eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden))
|
||||
torch.manual_seed(100)
|
||||
graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
|
||||
torch.testing.assert_close(graph_eos, eager_eos, atol=0, rtol=0)
|
||||
torch.testing.assert_close(graph_codes, eager_codes, atol=0, rtol=0)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
# 4. Disabled wrapper so fallback to eager
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_disabled_wrapper_matches_eager(model, wrapper):
|
||||
"""Cudagraph fallback to eager. Two eager runs should produce identical results."""
|
||||
hidden = _random_hidden(4)
|
||||
wrapper.enabled = False
|
||||
with torch.no_grad():
|
||||
torch.manual_seed(200)
|
||||
eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden))
|
||||
torch.manual_seed(200)
|
||||
graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
|
||||
wrapper.enabled = True
|
||||
torch.testing.assert_close(graph_eos, eager_eos, atol=0, rtol=0)
|
||||
torch.testing.assert_close(graph_codes, eager_codes, atol=0, rtol=0)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
# 5. Determinism - same input + same RNG state on cudagraph runs
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_deterministic_across_calls(model, wrapper):
|
||||
"""Same input + same RNG state. Two cudagraph runs should produce identical results."""
|
||||
hidden = _random_hidden(4)
|
||||
with torch.no_grad():
|
||||
torch.manual_seed(300)
|
||||
eos1, codes1 = _unpack_audio_codes(wrapper(hidden))
|
||||
torch.manual_seed(300)
|
||||
eos2, codes2 = _unpack_audio_codes(wrapper(hidden))
|
||||
torch.testing.assert_close(eos1, eos2, atol=0, rtol=0)
|
||||
torch.testing.assert_close(codes1, codes2, atol=0, rtol=0)
|
||||
@@ -0,0 +1,281 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections import defaultdict
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.stage_input_processors.voxtral_tts import (
|
||||
generator2tokenizer,
|
||||
generator2tokenizer_async_chunk,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
# ---- Helpers for generator2tokenizer (non-streaming) ----
|
||||
|
||||
|
||||
def _make_stage(engine_outputs):
|
||||
return SimpleNamespace(engine_outputs=engine_outputs)
|
||||
|
||||
|
||||
def _make_output(audio_tensors):
|
||||
"""Build a single generator output with multimodal_output["audio"] = audio_tensors."""
|
||||
return SimpleNamespace(
|
||||
outputs=[SimpleNamespace(multimodal_output={"audio": audio_tensors})],
|
||||
)
|
||||
|
||||
|
||||
# ---- Tests for generator2tokenizer ----
|
||||
|
||||
|
||||
class TestGenerator2Tokenizer:
|
||||
def test_single_output_single_chunk(self):
|
||||
"""Single output with one audio tensor produces correct prompt_token_ids."""
|
||||
audio = torch.tensor([10, 20, 30, 40])
|
||||
stage = _make_stage([_make_output([audio])])
|
||||
|
||||
result = generator2tokenizer([stage], engine_input_source=[0])
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["prompt_token_ids"] == [10, 20, 30, 40]
|
||||
assert result[0]["multi_modal_data"] is None
|
||||
|
||||
def test_2d_audio_tensors_flattened(self):
|
||||
"""2D audio tensors (e.g., multi-codebook frames) are flattened."""
|
||||
audio = torch.tensor([[1, 2, 3], [4, 5, 6]]) # (2, 3)
|
||||
stage = _make_stage([_make_output([audio])])
|
||||
|
||||
result = generator2tokenizer([stage], engine_input_source=[0])
|
||||
|
||||
assert result[0]["prompt_token_ids"] == [1, 2, 3, 4, 5, 6]
|
||||
|
||||
def test_empty_engine_input_source_raises(self):
|
||||
"""Empty engine_input_source raises ValueError."""
|
||||
stage = _make_stage([_make_output([torch.tensor([1])])])
|
||||
with pytest.raises(ValueError, match="cannot be empty"):
|
||||
generator2tokenizer([stage], engine_input_source=[])
|
||||
|
||||
def test_invalid_stage_id_raises(self):
|
||||
"""Stage id beyond stage_list length raises IndexError."""
|
||||
stage = _make_stage([_make_output([torch.tensor([1])])])
|
||||
with pytest.raises(IndexError, match="Invalid stage_id"):
|
||||
generator2tokenizer([stage], engine_input_source=[5])
|
||||
|
||||
def test_no_outputs_yet_raises(self):
|
||||
"""Stage with engine_outputs=None raises RuntimeError."""
|
||||
stage = _make_stage(engine_outputs=None)
|
||||
with pytest.raises(RuntimeError, match="no outputs yet"):
|
||||
generator2tokenizer([stage], engine_input_source=[0])
|
||||
|
||||
|
||||
# ---- Helpers for generator2tokenizer_async_chunk (streaming) ----
|
||||
|
||||
|
||||
def _req(external_req_id: str, *, finished: bool):
|
||||
return SimpleNamespace(
|
||||
external_req_id=external_req_id,
|
||||
is_finished=lambda: finished,
|
||||
)
|
||||
|
||||
|
||||
def _make_transfer_manager(codec_chunk_frames=25, codec_left_context_frames=25, codec_chunk_frames_at_begin=5):
|
||||
return SimpleNamespace(
|
||||
code_prompt_token_ids=defaultdict(list),
|
||||
connector=SimpleNamespace(
|
||||
config={
|
||||
"extra": {
|
||||
"codec_chunk_frames": codec_chunk_frames,
|
||||
"codec_left_context_frames": codec_left_context_frames,
|
||||
"codec_chunk_frames_at_begin": codec_chunk_frames_at_begin,
|
||||
}
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_empty_chunk_when_not_finished():
|
||||
"""Returns None when no audio data and not finished."""
|
||||
transfer_manager = _make_transfer_manager()
|
||||
request = _req("rid-empty", finished=False)
|
||||
|
||||
payload = generator2tokenizer_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output={"audio": torch.zeros((0,))},
|
||||
request=request,
|
||||
)
|
||||
|
||||
assert payload is None
|
||||
|
||||
|
||||
def test_flush_tail_when_finished():
|
||||
"""Emits remaining frames on finish with finished=True."""
|
||||
transfer_manager = _make_transfer_manager()
|
||||
request_id = "rid-tail"
|
||||
# Pre-populate with 24 frames of 8-element codes
|
||||
transfer_manager.code_prompt_token_ids[request_id] = [[1, 2, 3, 4, 5, 6, 7, 8] for _ in range(24)]
|
||||
request = _req(request_id, finished=True)
|
||||
|
||||
payload = generator2tokenizer_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None, # e.g. EOS step with no audio
|
||||
request=request,
|
||||
)
|
||||
|
||||
assert payload is not None
|
||||
assert payload["finished"].item() is True
|
||||
codes = payload["code_predictor_codes"]
|
||||
# Format: [ctx_frames, context_length, ...flat_codes]
|
||||
assert len(codes) >= 2 # At least ctx_frames + context_length header
|
||||
ctx_frames = codes[0]
|
||||
context_length = codes[1]
|
||||
flat_codes = codes[2:]
|
||||
# Total frames in window = ctx_frames + context_length
|
||||
total_window_frames = ctx_frames + context_length
|
||||
assert len(flat_codes) == total_window_frames * 8
|
||||
|
||||
|
||||
def test_eof_marker_when_finished_with_no_frames():
|
||||
"""Emits EOF marker with empty codes when finished with no accumulated frames."""
|
||||
transfer_manager = _make_transfer_manager()
|
||||
request = _req("rid-eof", finished=True)
|
||||
|
||||
payload = generator2tokenizer_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
)
|
||||
|
||||
assert payload == {
|
||||
"code_predictor_codes": [],
|
||||
"finished": torch.tensor(True, dtype=torch.bool),
|
||||
}
|
||||
|
||||
|
||||
def test_normal_chunk_emission():
|
||||
"""Emits chunks at codec_chunk_frames boundary."""
|
||||
transfer_manager = _make_transfer_manager(
|
||||
codec_chunk_frames=25,
|
||||
codec_left_context_frames=25,
|
||||
codec_chunk_frames_at_begin=5,
|
||||
)
|
||||
request_id = "rid-chunk"
|
||||
request = _req(request_id, finished=False)
|
||||
|
||||
# Feed 25 frames one by one
|
||||
for i in range(25):
|
||||
pooling_output = {"audio": torch.tensor([float(i)] * 4)}
|
||||
payload = generator2tokenizer_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=pooling_output,
|
||||
request=request,
|
||||
)
|
||||
|
||||
# A chunk should be emitted
|
||||
assert payload is not None
|
||||
codes = payload["code_predictor_codes"]
|
||||
ctx_frames = codes[0]
|
||||
context_length = codes[1]
|
||||
assert ctx_frames == 20 # 25 - 5(chunk_size_at_begin)
|
||||
assert context_length == 5 # 25 - 20(chunk_size_at_begin * 4)
|
||||
|
||||
|
||||
def test_small_initial_chunks():
|
||||
"""Uses codec_chunk_frames_at_begin for initial frames when length <= chunk_size."""
|
||||
transfer_manager = _make_transfer_manager(
|
||||
codec_chunk_frames=25,
|
||||
codec_left_context_frames=25,
|
||||
codec_chunk_frames_at_begin=5,
|
||||
)
|
||||
request_id = "rid-begin"
|
||||
request = _req(request_id, finished=False)
|
||||
|
||||
# Feed 5 frames (should trigger emission because codec_chunk_frames_at_begin=5)
|
||||
for i in range(5):
|
||||
pooling_output = {"audio": torch.tensor([float(i + 1)] * 3)}
|
||||
payload = generator2tokenizer_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=pooling_output,
|
||||
request=request,
|
||||
)
|
||||
|
||||
assert payload is not None
|
||||
codes = payload["code_predictor_codes"]
|
||||
ctx_frames = codes[0]
|
||||
context_length = codes[1]
|
||||
assert ctx_frames == 0
|
||||
assert context_length == 5 # chunk_size_at_begin
|
||||
|
||||
|
||||
def test_no_emission_between_boundaries():
|
||||
"""No chunk emitted when not at chunk boundary and not finished."""
|
||||
transfer_manager = _make_transfer_manager(
|
||||
codec_chunk_frames=25,
|
||||
codec_left_context_frames=25,
|
||||
codec_chunk_frames_at_begin=5,
|
||||
)
|
||||
request_id = "rid-mid"
|
||||
request = _req(request_id, finished=False)
|
||||
|
||||
# Feed 3 frames (3 % 5 != 0, should not emit)
|
||||
for i in range(3):
|
||||
pooling_output = {"audio": torch.tensor([float(i)] * 4)}
|
||||
payload = generator2tokenizer_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=pooling_output,
|
||||
request=request,
|
||||
)
|
||||
|
||||
assert payload is None
|
||||
|
||||
|
||||
def test_context_handling_format():
|
||||
"""Verifies [ctx_frames, context_length, ...codes] format."""
|
||||
transfer_manager = _make_transfer_manager(
|
||||
codec_chunk_frames=10,
|
||||
codec_left_context_frames=5,
|
||||
codec_chunk_frames_at_begin=5,
|
||||
)
|
||||
request_id = "rid-ctx"
|
||||
request = _req(request_id, finished=False)
|
||||
|
||||
# Feed 5 frames to trigger a chunk (chunk_size_at_begin=5)
|
||||
for i in range(5):
|
||||
pooling_output = {"audio": torch.tensor([float(i + 10)] * 2)} # codebook_dim=2
|
||||
payload = generator2tokenizer_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=pooling_output,
|
||||
request=request,
|
||||
)
|
||||
|
||||
assert payload is not None
|
||||
codes = payload["code_predictor_codes"]
|
||||
# First two elements are ctx_frames and context_length
|
||||
ctx_frames = codes[0]
|
||||
context_length = codes[1]
|
||||
flat_codes = codes[2:]
|
||||
# The window has ctx_frames + context_length frames
|
||||
assert isinstance(ctx_frames, int)
|
||||
assert isinstance(context_length, int)
|
||||
assert ctx_frames >= 0
|
||||
assert context_length > 0
|
||||
# flat_codes = total_window_frames * codebook_dim
|
||||
total_window_frames = ctx_frames + context_length
|
||||
assert len(flat_codes) == total_window_frames * 2 # codebook_dim=2
|
||||
|
||||
|
||||
def test_none_pooling_output_not_finished_returns_none():
|
||||
"""None pooling_output when not finished returns None."""
|
||||
transfer_manager = _make_transfer_manager()
|
||||
request = _req("rid-none", finished=False)
|
||||
|
||||
payload = generator2tokenizer_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
)
|
||||
|
||||
assert payload is None
|
||||
@@ -222,6 +222,7 @@ class StageConfigFactory:
|
||||
"qwen2_5_omni": "qwen2_5_omni",
|
||||
"bagel": "bagel",
|
||||
"qwen3_tts": "qwen3_tts",
|
||||
"voxtral_tts": "voxtral_tts",
|
||||
"mimo_audio": "mimo_audio",
|
||||
"glm-image": "glm_image",
|
||||
"cosyvoice3": "cosyvoice3",
|
||||
|
||||
@@ -22,16 +22,23 @@ def _register_omni_hf_configs() -> None:
|
||||
from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
|
||||
Qwen3TTSConfig,
|
||||
)
|
||||
from vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts import (
|
||||
VoxtralTTSConfig,
|
||||
)
|
||||
except Exception as exc: # pragma: no cover - best-effort optional registration
|
||||
logger.warning("Skipping omni HF config registration due to import error: %s", exc)
|
||||
return
|
||||
|
||||
try:
|
||||
AutoConfig.register("qwen3_tts", Qwen3TTSConfig)
|
||||
AutoConfig.register("cosyvoice3", CosyVoice3Config)
|
||||
except ValueError:
|
||||
# Already registered elsewhere; ignore.
|
||||
return
|
||||
for model_type, config_cls in [
|
||||
("qwen3_tts", Qwen3TTSConfig),
|
||||
("cosyvoice3", CosyVoice3Config),
|
||||
("mistral", VoxtralTTSConfig),
|
||||
]:
|
||||
try:
|
||||
AutoConfig.register(model_type, config_cls)
|
||||
except ValueError:
|
||||
# Already registered elsewhere; ignore.
|
||||
pass
|
||||
|
||||
|
||||
def register_omni_models_to_vllm():
|
||||
|
||||
@@ -30,7 +30,10 @@ from vllm_omni.outputs import OmniRequestOutput
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# TTS Configuration
|
||||
_TTS_MODEL_STAGES: set[str] = {"qwen3_tts", "fish_speech_slow_ar"}
|
||||
_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
|
||||
_TTS_LANGUAGES: set[str] = {
|
||||
"Auto",
|
||||
"Chinese",
|
||||
@@ -145,6 +148,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
)
|
||||
self._fish_speech_tokenizer = None
|
||||
|
||||
# Determine TTS model type or None
|
||||
self._tts_model_type = self._detect_tts_model_type()
|
||||
|
||||
# Cache TTS configuration values (computed once, reused per request)
|
||||
self._max_instructions_length = self._compute_max_instructions_length()
|
||||
|
||||
@@ -203,6 +209,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
return stage
|
||||
return None
|
||||
|
||||
def _detect_tts_model_type(self) -> str | None:
|
||||
"""Detect TTS model type from the stage's model_stage attribute."""
|
||||
if self._tts_stage is None:
|
||||
return None
|
||||
model_stage = getattr(self._tts_stage.engine_args, "model_stage", None)
|
||||
if model_stage in _QWEN3_TTS_MODEL_STAGES:
|
||||
return "qwen3_tts"
|
||||
if model_stage in _VOXTRAL_TTS_MODEL_STAGES:
|
||||
return "voxtral_tts"
|
||||
if model_stage in _FISH_TTS_MODEL_STAGES:
|
||||
return "fish_tts"
|
||||
return None
|
||||
|
||||
def _compute_max_instructions_length(self) -> int:
|
||||
"""Compute max instructions length with precedence: CLI > stage config > default.
|
||||
|
||||
@@ -225,16 +244,22 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
def _load_supported_speakers(self) -> set[str]:
|
||||
"""Load supported speakers (case-insensitive) from the model configuration."""
|
||||
try:
|
||||
talker_config = self.engine_client.model_config.hf_config.talker_config
|
||||
if self._tts_model_type == "voxtral_tts":
|
||||
config = self.engine_client.model_config.hf_config.audio_config
|
||||
else:
|
||||
# Default is qwen3_tts path
|
||||
config = self.engine_client.model_config.hf_config.talker_config
|
||||
|
||||
# Check for speakers in either spk_id or speaker_id
|
||||
for attr_name in ["spk_id", "speaker_id"]:
|
||||
speakers_dict = getattr(talker_config, attr_name, None)
|
||||
if isinstance(config, dict):
|
||||
speakers_dict = config.get(attr_name)
|
||||
else:
|
||||
speakers_dict = getattr(config, attr_name, None)
|
||||
if speakers_dict and isinstance(speakers_dict, dict):
|
||||
# Normalize to lowercase for case-insensitive matching
|
||||
return {speaker.lower() for speaker in speakers_dict.keys()}
|
||||
|
||||
logger.warning("No speakers found in talker_config (checked spk_id and speaker_id)")
|
||||
logger.warning("No speakers found in config (checked spk_id and speaker_id)")
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not load speakers from model config: {e}")
|
||||
|
||||
@@ -497,6 +522,49 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
|
||||
def _validate_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
|
||||
"""Validate TTS request parameters. Returns error message or None."""
|
||||
if self._tts_model_type == "voxtral_tts":
|
||||
return self._validate_voxtral_tts_request(request)
|
||||
return self._validate_qwen_tts_request(request)
|
||||
|
||||
def _validate_ref_audio_format(self, ref_audio: str) -> str | None:
|
||||
"""Validate ref_audio is a supported URI format. Returns error or None."""
|
||||
if not (
|
||||
ref_audio.startswith(("http://", "https://"))
|
||||
or ref_audio.startswith("data:")
|
||||
or ref_audio.startswith("file://")
|
||||
):
|
||||
return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)"
|
||||
return None
|
||||
|
||||
def _validate_voxtral_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
|
||||
"""Validate Voxtral TTS request parameters. Returns error message or None."""
|
||||
if not request.input or not request.input.strip():
|
||||
return "Input text cannot be empty"
|
||||
|
||||
# Voxtral TTS requires either a preset voice or ref_audio for voice cloning.
|
||||
if request.voice is None and request.ref_audio is None:
|
||||
return "Either 'voice' (preset speaker) or 'ref_audio' (voice cloning) must be provided"
|
||||
|
||||
if request.ref_audio is not None:
|
||||
fmt_err = self._validate_ref_audio_format(request.ref_audio)
|
||||
if fmt_err:
|
||||
return fmt_err
|
||||
|
||||
if request.voice is not None:
|
||||
request.voice = request.voice.lower()
|
||||
if self.supported_speakers and request.voice not in self.supported_speakers:
|
||||
return f"Invalid speaker '{request.voice}'. Supported: {', '.join(sorted(self.supported_speakers))}"
|
||||
|
||||
if request.max_new_tokens is not None:
|
||||
if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN:
|
||||
return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}"
|
||||
if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX:
|
||||
return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}"
|
||||
|
||||
return None
|
||||
|
||||
def _validate_qwen_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
|
||||
"""Validate Qwen TTS request parameters. Returns error message or None."""
|
||||
# Infer Base task when ref_audio or ref_text is provided without explicit task_type.
|
||||
if request.task_type is None and (request.ref_audio is not None or request.ref_text is not None):
|
||||
request.task_type = "Base"
|
||||
@@ -530,13 +598,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
if request.voice is None:
|
||||
if request.ref_audio is None:
|
||||
return "Base task requires 'ref_audio' for voice cloning"
|
||||
# Validate ref_audio format (include file:// from upstream)
|
||||
if not (
|
||||
request.ref_audio.startswith(("http://", "https://"))
|
||||
or request.ref_audio.startswith("data:")
|
||||
or request.ref_audio.startswith("file://")
|
||||
):
|
||||
return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)"
|
||||
fmt_err = self._validate_ref_audio_format(request.ref_audio)
|
||||
if fmt_err:
|
||||
return fmt_err
|
||||
# In-context voice cloning (default) requires non-empty ref_text.
|
||||
# x_vector_only_mode skips in-context and only uses speaker embedding.
|
||||
if not request.x_vector_only_mode:
|
||||
@@ -560,15 +624,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
return (
|
||||
f"Base task with built-in speaker '{request.voice}' requires 'ref_audio' for voice cloning"
|
||||
)
|
||||
# Validate ref_audio format for built-in speaker
|
||||
if not (
|
||||
request.ref_audio.startswith(("http://", "https://"))
|
||||
or request.ref_audio.startswith("data:")
|
||||
or request.ref_audio.startswith("file://")
|
||||
):
|
||||
return (
|
||||
"ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)"
|
||||
)
|
||||
fmt_err = self._validate_ref_audio_format(request.ref_audio)
|
||||
if fmt_err:
|
||||
return fmt_err
|
||||
|
||||
# Validate cross-parameter dependencies
|
||||
if task_type != "Base":
|
||||
@@ -775,6 +833,40 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
|
||||
return params
|
||||
|
||||
# ---- Voxtral TTS helpers ----
|
||||
|
||||
async def _build_voxtral_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
|
||||
"""Build Voxtral TTS engine prompt from shared TTS parameters."""
|
||||
from mistral_common.protocol.speech.request import SpeechRequest
|
||||
|
||||
text = request.input
|
||||
voice = request.voice
|
||||
ref_audio = request.ref_audio
|
||||
assert voice or ref_audio, "Either voice or ref_audio must be provided"
|
||||
# Strip data URI prefix — mistral_common expects raw base64
|
||||
if ref_audio is not None and isinstance(ref_audio, str) and ref_audio.startswith("data:"):
|
||||
_, _, ref_audio = ref_audio.partition(",")
|
||||
if self._tts_tokenizer is None:
|
||||
from vllm.tokenizers import cached_tokenizer_from_config
|
||||
|
||||
mistral_tokenizer = cached_tokenizer_from_config(self.engine_client.model_config)
|
||||
self._tts_tokenizer = mistral_tokenizer.instruct
|
||||
if voice is not None:
|
||||
tokens = self._tts_tokenizer.encode_speech_request(SpeechRequest(input=text, voice=voice)).tokens
|
||||
return {
|
||||
"prompt_token_ids": tokens,
|
||||
"additional_information": {"voice": [voice]},
|
||||
}
|
||||
else:
|
||||
tokenized = self._tts_tokenizer.encode_speech_request(SpeechRequest(input=text, ref_audio=ref_audio))
|
||||
audio = tokenized.audios[0]
|
||||
return {
|
||||
"prompt_token_ids": tokenized.tokens,
|
||||
"multi_modal_data": {"audio": [(audio.audio_array, audio.sampling_rate)]},
|
||||
}
|
||||
|
||||
# ---- Fish Speech helpers ----
|
||||
|
||||
def _build_fish_speech_prompt(
|
||||
self,
|
||||
request: OpenAICreateSpeechRequest,
|
||||
@@ -848,6 +940,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
|
||||
# ---- Common speech generation helpers ----
|
||||
|
||||
async def _prepare_speech_generation(
|
||||
self,
|
||||
request: OpenAICreateSpeechRequest,
|
||||
@@ -871,25 +965,30 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
if validation_error:
|
||||
raise ValueError(validation_error)
|
||||
|
||||
tts_params = self._build_tts_params(request)
|
||||
if request.ref_audio is not None:
|
||||
wav_list, sr = await self._resolve_ref_audio(request.ref_audio)
|
||||
tts_params["ref_audio"] = [[wav_list, sr]]
|
||||
if self._tts_model_type == "voxtral_tts":
|
||||
prompt = await self._build_voxtral_prompt(request)
|
||||
tts_params = {}
|
||||
else:
|
||||
tts_params = self._build_tts_params(request)
|
||||
if request.ref_audio is not None:
|
||||
wav_list, sr = await self._resolve_ref_audio(request.ref_audio)
|
||||
tts_params["ref_audio"] = [[wav_list, sr]]
|
||||
|
||||
ph_len = self._estimate_prompt_len(tts_params)
|
||||
prompt = {"prompt_token_ids": [1] * ph_len, "additional_information": tts_params}
|
||||
ph_len = self._estimate_prompt_len(tts_params)
|
||||
prompt = {"prompt_token_ids": [1] * ph_len, "additional_information": tts_params}
|
||||
else:
|
||||
tts_params = {}
|
||||
prompt = {"prompt": request.input}
|
||||
|
||||
request_id = f"speech-{random_uuid()}"
|
||||
model_type = (
|
||||
"fish_speech"
|
||||
if self._is_fish_speech
|
||||
else tts_params.get("task_type", ["unknown"])[0]
|
||||
if self._is_tts
|
||||
else "generic"
|
||||
)
|
||||
if self._is_fish_speech:
|
||||
model_type = "fish_speech"
|
||||
elif self._tts_model_type == "voxtral_tts":
|
||||
model_type = "voxtral_tts"
|
||||
elif self._is_tts:
|
||||
model_type = tts_params.get("task_type", ["unknown"])[0]
|
||||
else:
|
||||
model_type = "generic"
|
||||
logger.info(
|
||||
"TTS speech request %s: text=%r, model=%s",
|
||||
request_id,
|
||||
|
||||
@@ -140,6 +140,18 @@ _OMNI_MODELS = {
|
||||
"fish_speech_dac_decoder",
|
||||
"FishSpeechDACDecoder",
|
||||
),
|
||||
## Voxtral TTS
|
||||
"VoxtralTTSForConditionalGeneration": (
|
||||
"voxtral_tts",
|
||||
"voxtral_tts",
|
||||
"VoxtralTTSForConditionalGeneration",
|
||||
),
|
||||
"VoxtralTTSAudioGeneration": (
|
||||
"voxtral_tts",
|
||||
"voxtral_tts_audio_generation",
|
||||
"VoxtralTTSAudioGenerationForConditionalGeneration",
|
||||
),
|
||||
"VoxtralTTSAudioTokenizer": ("voxtral_tts", "voxtral_tts_audio_tokenizer", "VoxtralTTSAudioTokenizer"),
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from transformers import PretrainedConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.transformers_utils.config import MistralConfigParser, register_config_parser
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class VoxtralTTSConfig(PretrainedConfig):
|
||||
"""HuggingFace-style config for Voxtral TTS models."""
|
||||
|
||||
model_type = "voxtral_tts"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
text_config: PretrainedConfig | dict | None = None,
|
||||
audio_config: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
super().__init__(**kwargs)
|
||||
|
||||
if isinstance(text_config, PretrainedConfig):
|
||||
self.text_config = text_config
|
||||
elif isinstance(text_config, dict):
|
||||
self.text_config = PretrainedConfig.from_dict(text_config)
|
||||
else:
|
||||
self.text_config = PretrainedConfig()
|
||||
|
||||
self.audio_config = audio_config or {}
|
||||
|
||||
def get_text_config(self, **kwargs: Any) -> PretrainedConfig:
|
||||
return self.text_config
|
||||
|
||||
|
||||
@register_config_parser("mistral")
|
||||
class VoxtralTTSConfigParser(MistralConfigParser):
|
||||
"""Config parser that extends the base Mistral parser with TTS support.
|
||||
|
||||
This only support voxtral_tts for now.
|
||||
"""
|
||||
|
||||
def _remap_mistral_audio_args(self, config_dict: dict) -> dict:
|
||||
encoder_args = config_dict["multimodal"].pop("audio_model_args")
|
||||
audio_tokenizer_args = config_dict["multimodal"].pop("audio_tokenizer_args", None)
|
||||
audio_config = {}
|
||||
if encoder_args is not None:
|
||||
audio_config = {
|
||||
"sampling_rate": encoder_args["audio_encoding_args"]["sampling_rate"],
|
||||
"codec_args": audio_tokenizer_args,
|
||||
"audio_model_args": encoder_args,
|
||||
"speaker_id": audio_tokenizer_args.get("voice", {}),
|
||||
}
|
||||
return audio_config
|
||||
|
||||
def parse(
|
||||
self,
|
||||
model: str | Path,
|
||||
trust_remote_code: bool,
|
||||
revision: str | None = None,
|
||||
code_revision: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> tuple[dict, PretrainedConfig]:
|
||||
from vllm.transformers_utils.config import (
|
||||
_download_mistral_config_file,
|
||||
)
|
||||
|
||||
config_dict = _download_mistral_config_file(model, revision)
|
||||
|
||||
from vllm.transformers_utils.configs.mistral import (
|
||||
_remap_general_mistral_args,
|
||||
_remap_mistral_quantization_args,
|
||||
)
|
||||
|
||||
# Extract audio config before building text config
|
||||
audio_config = {}
|
||||
if (config_dict.get("multimodal") or {}).get("audio_model_args"):
|
||||
audio_config = self._remap_mistral_audio_args(config_dict)
|
||||
|
||||
# Build text_config from the top-level keys
|
||||
non_text_keys = {"multimodal"}
|
||||
text_config = {k: v for k, v in config_dict.items() if k not in non_text_keys}
|
||||
text_config = _remap_general_mistral_args(text_config)
|
||||
if text_config.get("quantization"):
|
||||
text_config = _remap_mistral_quantization_args(text_config)
|
||||
|
||||
# The text sub-model is a plain MistralForCausalLM
|
||||
text_config.setdefault("architectures", ["MistralForCausalLM"])
|
||||
|
||||
config = VoxtralTTSConfig(
|
||||
text_config=PretrainedConfig.from_dict(text_config),
|
||||
audio_config=audio_config,
|
||||
architectures=config_dict.get("architectures", ["VoxtralTTSForConditionalGeneration"]),
|
||||
)
|
||||
|
||||
return config_dict, config
|
||||
+256
@@ -0,0 +1,256 @@
|
||||
"""
|
||||
CUDA Graph wrapper for AcousticTransformer in VoxtralTTS.
|
||||
|
||||
Currently it support FlowMatching implementation only.
|
||||
|
||||
Captures the AcousticTransformer forward pass (semantic logit +
|
||||
n-step Euler ODE with CFG) into CUDA graphs for fixed batch sizes,
|
||||
eliminating kernel launch overhead on every decode step.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch.cuda import CUDAGraph
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation import (
|
||||
AudioSpecialTokens,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class CUDAGraphAcousticTransformerWrapper:
|
||||
"""
|
||||
CUDA Graph wrapper for the acoustic transformer.
|
||||
|
||||
Replaces the vLLM Sampler path (which has Python-level branching and dynamic
|
||||
tensor allocation) with a CUDA-graph-compatible path using torch.argmax
|
||||
(equivalent to top_k=1 greedy sampling) and pre-allocated buffers.
|
||||
"""
|
||||
|
||||
DEFAULT_CAPTURE_SIZES = [1, 2, 4, 8, 16, 32]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
capture_sizes: list[int] | None = None,
|
||||
):
|
||||
self.model = model
|
||||
self.acoustic_transformer = model.acoustic_transformer
|
||||
|
||||
self.capture_sizes = sorted(capture_sizes or self.DEFAULT_CAPTURE_SIZES)
|
||||
|
||||
# Pre-compute constants from the acoustic transformer
|
||||
self.empty_audio_token_id = AudioSpecialTokens.id(AudioSpecialTokens.empty_audio)
|
||||
self.end_audio_token_id = AudioSpecialTokens.id(AudioSpecialTokens.end_audio)
|
||||
self.semantic_mask_start = len(AudioSpecialTokens) + self.acoustic_transformer.model_args.semantic_codebook_size
|
||||
self.n_acoustic_codebook = self.acoustic_transformer.model_args.n_acoustic_codebook
|
||||
self.acoustic_embeddings_levels = self.acoustic_transformer.acoustic_embeddings_levels
|
||||
|
||||
self.cfg_alpha = 1.2
|
||||
self.n_steps = 16
|
||||
|
||||
# Graph storage
|
||||
self.graphs: dict[int, CUDAGraph] = {}
|
||||
self.static_inputs: dict[int, torch.Tensor] = {}
|
||||
self.static_noise: dict[int, torch.Tensor] = {}
|
||||
self.static_fake_eos: dict[int, torch.Tensor] = {}
|
||||
self.static_audio_codes: dict[int, torch.Tensor] = {}
|
||||
|
||||
self.enabled = False
|
||||
self._warmed_up = False
|
||||
|
||||
def _warmup_and_capture(self, device: torch.device, dtype: torch.dtype, hidden_dim: int):
|
||||
"""Perform eager warmup and CUDA graph capture for all bucket sizes."""
|
||||
if self._warmed_up:
|
||||
logger.warning("CUDAGraphAcousticTransformerWrapper already warmed up, skipping")
|
||||
return
|
||||
|
||||
logger.info(
|
||||
"CUDAGraphAcousticTransformerWrapper: starting warmup and capture for sizes %s",
|
||||
self.capture_sizes,
|
||||
)
|
||||
|
||||
# Pre-create persistent buffers
|
||||
self.timesteps = torch.linspace(0, 1, self.n_steps, device=device, dtype=dtype)
|
||||
self.fake_eos_one = torch.tensor(1.0, dtype=dtype, device=device)
|
||||
self.fake_eos_zero = torch.tensor(0.0, dtype=dtype, device=device)
|
||||
|
||||
# Phase 1: Eager warmup for ALL capture sizes
|
||||
for size in self.capture_sizes:
|
||||
dummy = torch.zeros(size, hidden_dim, device=device, dtype=dtype)
|
||||
with torch.no_grad():
|
||||
self._forward_cudagraph_compatible(dummy)
|
||||
|
||||
torch.cuda.synchronize(device)
|
||||
|
||||
# Phase 2: Capture graphs
|
||||
for size in self.capture_sizes:
|
||||
try:
|
||||
self._capture_graph_for_size(size, device, dtype, hidden_dim)
|
||||
logger.info(" Captured CUDA Graph for batch_size=%d", size)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
" Failed to capture CUDA Graph for batch_size=%d",
|
||||
size,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
self.enabled = True
|
||||
self._warmed_up = True
|
||||
logger.info(
|
||||
"CUDAGraphAcousticTransformerWrapper warmup complete. Captured %d/%d graphs.",
|
||||
len(self.graphs),
|
||||
len(self.capture_sizes),
|
||||
)
|
||||
|
||||
def _forward_cudagraph_compatible(self, hidden_states: torch.Tensor, noise: torch.Tensor | None = None):
|
||||
"""
|
||||
The actual computation captured by the CUDA graph.
|
||||
|
||||
This replaces the full compute_mm_logits -> acoustic_transformer.forward() path
|
||||
with a graph-compatible version:
|
||||
- Uses argmax instead of vLLM Sampler (equivalent for top_k=1)
|
||||
- Uses pre-created timesteps buffer instead of torch.linspace
|
||||
- Uses pre-created scalar tensors for torch.where
|
||||
- Calls _predict_velocity directly
|
||||
- Uses a pre-allocated noise buffer to avoid baking random state
|
||||
into the CUDA graph
|
||||
"""
|
||||
at = self.acoustic_transformer
|
||||
B = hidden_states.shape[0]
|
||||
|
||||
# --- Semantic logits via linear projection ---
|
||||
semantic_logit = at.semantic_codebook_output(hidden_states).float()
|
||||
semantic_logit[:, self.empty_audio_token_id] = -float("inf")
|
||||
semantic_logit[:, self.semantic_mask_start :] = -float("inf")
|
||||
|
||||
# argmax == top_k=1 greedy sampling
|
||||
semantic_code = semantic_logit.argmax(dim=-1, keepdim=True) # (B, 1)
|
||||
|
||||
# --- Flow matching: Euler ODE ---
|
||||
should_decode = semantic_code.squeeze(1) != self.end_audio_token_id
|
||||
|
||||
if noise is not None:
|
||||
x = noise
|
||||
else:
|
||||
x = torch.randn(B, self.n_acoustic_codebook, device=hidden_states.device, dtype=hidden_states.dtype)
|
||||
|
||||
# Pre-compute zero hidden states for unconditional CFG branch
|
||||
hidden_states_zero = torch.zeros_like(hidden_states)
|
||||
|
||||
timesteps = self.timesteps
|
||||
for i in range(len(timesteps) - 1):
|
||||
t = timesteps[i]
|
||||
dt = timesteps[i + 1] - timesteps[i]
|
||||
|
||||
# Batch conditional + unconditional velocity in a single forward pass
|
||||
t_emb = at.time_embedding(t.view(-1, 1).repeat(B, 1)).to(hidden_states.dtype)
|
||||
x_batched = torch.cat([x, x], dim=0) # (2B, C)
|
||||
llm_batched = torch.cat([hidden_states, hidden_states_zero], dim=0) # (2B, D)
|
||||
t_emb_batched = t_emb.repeat(2, 1) # (2B, D)
|
||||
|
||||
v_all = at._predict_velocity(x_t=x_batched, llm_output=llm_batched, t_emb=t_emb_batched)
|
||||
v_t, uncond_v_t = v_all[:B], v_all[B:]
|
||||
|
||||
# CFG combination
|
||||
v_t = self.cfg_alpha * v_t + (1 - self.cfg_alpha) * uncond_v_t
|
||||
|
||||
x = x + v_t * dt
|
||||
|
||||
# --- Quantize ---
|
||||
sampled = torch.clamp(x, -1, 1)
|
||||
scaled_x = ((sampled + 1) / 2) * (self.acoustic_embeddings_levels - 1)
|
||||
output_codes = scaled_x.round().long()
|
||||
output_codes[~should_decode] = self.empty_audio_token_id
|
||||
acoustic_codes = output_codes + len(AudioSpecialTokens)
|
||||
|
||||
# --- Combine semantic + acoustic ---
|
||||
audio_codes = torch.cat([semantic_code, acoustic_codes], dim=1) # (B, 1 + n_acoustic)
|
||||
|
||||
# --- Compute fake_eos ---
|
||||
fake_eos = torch.where(
|
||||
audio_codes[:, 0] == self.end_audio_token_id,
|
||||
self.fake_eos_one,
|
||||
self.fake_eos_zero,
|
||||
)
|
||||
|
||||
return fake_eos, audio_codes
|
||||
|
||||
def _capture_graph_for_size(
|
||||
self,
|
||||
size: int,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
hidden_dim: int,
|
||||
):
|
||||
"""Capture a CUDA graph for a specific batch size."""
|
||||
static_input = torch.zeros(size, hidden_dim, device=device, dtype=dtype)
|
||||
static_noise = torch.randn(size, self.n_acoustic_codebook, device=device, dtype=dtype)
|
||||
|
||||
# Stabilizing eager run
|
||||
with torch.no_grad():
|
||||
_ = self._forward_cudagraph_compatible(static_input, noise=static_noise)
|
||||
|
||||
torch.cuda.synchronize(device)
|
||||
|
||||
graph = CUDAGraph()
|
||||
with torch.no_grad():
|
||||
with torch.cuda.graph(graph):
|
||||
static_fake_eos, static_audio_codes = self._forward_cudagraph_compatible(
|
||||
static_input, noise=static_noise
|
||||
)
|
||||
|
||||
self.graphs[size] = graph
|
||||
self.static_inputs[size] = static_input
|
||||
self.static_noise[size] = static_noise
|
||||
self.static_fake_eos[size] = static_fake_eos
|
||||
self.static_audio_codes[size] = static_audio_codes
|
||||
|
||||
def _get_padded_size(self, actual_size: int) -> int | None:
|
||||
"""Round up to the nearest captured bucket size."""
|
||||
for size in self.capture_sizes:
|
||||
if actual_size <= size:
|
||||
return size
|
||||
return None
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, dict[str, list[torch.Tensor]] | None]:
|
||||
"""
|
||||
Drop-in replacement for model.compute_mm_logits().
|
||||
|
||||
Falls back to eager execution if:
|
||||
- CUDA graph is not enabled/warmed up
|
||||
- Batch size exceeds largest captured size
|
||||
"""
|
||||
actual_size = hidden_states.shape[0]
|
||||
|
||||
if not self.enabled or not self._warmed_up:
|
||||
return self.model.compute_mm_logits(hidden_states)
|
||||
|
||||
padded_size = self._get_padded_size(actual_size)
|
||||
if padded_size is None or padded_size not in self.graphs:
|
||||
return self.model.compute_mm_logits(hidden_states)
|
||||
|
||||
# Zero static input, then copy actual data
|
||||
self.static_inputs[padded_size].zero_()
|
||||
self.static_inputs[padded_size][:actual_size] = hidden_states
|
||||
|
||||
# Fill noise buffer with fresh random values before replay so the
|
||||
# flow-matching ODE starts from different initial noise each time.
|
||||
self.static_noise[padded_size].normal_()
|
||||
|
||||
# Replay captured graph
|
||||
self.graphs[padded_size].replay()
|
||||
|
||||
# Clone and slice outputs for actual batch size
|
||||
fake_eos = self.static_fake_eos[padded_size][:actual_size].clone()
|
||||
audio_codes = self.static_audio_codes[padded_size][:actual_size].clone()
|
||||
|
||||
# Package into expected format: (fake_eos, {"audio": [list of tensors]})
|
||||
audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0))
|
||||
mm_tokens = {"audio": audio_list}
|
||||
|
||||
return fake_eos, mm_tokens
|
||||
@@ -0,0 +1,372 @@
|
||||
from collections.abc import Iterable
|
||||
from functools import cached_property
|
||||
from pathlib import Path
|
||||
|
||||
import regex as re
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from huggingface_hub import hf_hub_download
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.models.interfaces import SupportsMultiModal
|
||||
from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm.tokenizers import cached_tokenizer_from_config
|
||||
from vllm.v1.outputs import SamplerOutput
|
||||
from vllm.v1.sample.metadata import SamplingMetadata
|
||||
from vllm.v1.sample.sampler import Sampler
|
||||
|
||||
from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
from vllm_omni.model_executor.models.voxtral_tts.cuda_graph_acoustic_transformer_wrapper import (
|
||||
CUDAGraphAcousticTransformerWrapper,
|
||||
)
|
||||
from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation import (
|
||||
VoxtralTTSDummyInputsBuilder,
|
||||
VoxtralTTSMultiModalProcessor,
|
||||
VoxtralTTSProcessingInfo,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def parse_batched_audio_input(input_ids: torch.Tensor, num_codebooks: int) -> tuple[list[torch.Tensor], list[int]]:
|
||||
"""Parse batched input_ids with [ctx_frames, context_length, ...tokens] format.
|
||||
|
||||
Each request in the batch is laid out as:
|
||||
[ctx_frames, context_length, <flat audio tokens>]
|
||||
where the flat audio tokens have length (ctx_frames + context_length) * num_codebooks.
|
||||
|
||||
Returns:
|
||||
all_audio_tokens: list of (num_frames, num_codebooks) tensors per request.
|
||||
all_ctx_frames: list of ctx_frames values per request.
|
||||
"""
|
||||
all_audio_tokens: list[torch.Tensor] = []
|
||||
all_ctx_frames: list[int] = []
|
||||
offset = 0
|
||||
while offset < input_ids.numel():
|
||||
ctx_frames = int(input_ids[offset].item())
|
||||
context_length = int(input_ids[offset + 1].item())
|
||||
offset += 2
|
||||
|
||||
req_input_ids_length = (ctx_frames + context_length) * num_codebooks
|
||||
req_input_ids = input_ids[offset : offset + req_input_ids_length]
|
||||
offset += req_input_ids_length
|
||||
assert req_input_ids.numel() % num_codebooks == 0, (
|
||||
f"Number of elements must be divisible by {num_codebooks} but get {req_input_ids.numel()=}"
|
||||
)
|
||||
num_of_audios = req_input_ids.numel() // num_codebooks
|
||||
audio_tokens = req_input_ids.view(num_of_audios, num_codebooks)
|
||||
all_audio_tokens.append(audio_tokens)
|
||||
all_ctx_frames.append(ctx_frames)
|
||||
return all_audio_tokens, all_ctx_frames
|
||||
|
||||
|
||||
def apply_ctx_frames_cutting(
|
||||
batch_audio_arrays: list[torch.Tensor],
|
||||
all_ctx_frames: list[int],
|
||||
downsample_factor: int,
|
||||
) -> list[torch.Tensor]:
|
||||
"""Cut leading context samples from decoded audio arrays.
|
||||
|
||||
Args:
|
||||
batch_audio_arrays: raw decoded audio arrays, one per request.
|
||||
all_ctx_frames: number of context frames per request.
|
||||
downsample_factor: samples per frame in the audio tokenizer.
|
||||
|
||||
Returns:
|
||||
List of audio arrays with context samples removed.
|
||||
"""
|
||||
result: list[torch.Tensor] = []
|
||||
for audio_array, ctx_frames in zip(batch_audio_arrays, all_ctx_frames):
|
||||
if ctx_frames > 0:
|
||||
new_sample_cut = downsample_factor * ctx_frames
|
||||
audio_array = audio_array[new_sample_cut:]
|
||||
result.append(audio_array)
|
||||
return result
|
||||
|
||||
|
||||
@MULTIMODAL_REGISTRY.register_processor(
|
||||
VoxtralTTSMultiModalProcessor,
|
||||
info=VoxtralTTSProcessingInfo,
|
||||
dummy_inputs=VoxtralTTSDummyInputsBuilder,
|
||||
)
|
||||
class VoxtralTTSForConditionalGeneration(
|
||||
nn.Module,
|
||||
SupportsMultiModal,
|
||||
CustomProcessMixin,
|
||||
):
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
config = vllm_config.model_config.hf_config
|
||||
self.have_multimodal_outputs = True
|
||||
self.has_preprocess = False
|
||||
self.has_postprocess = False
|
||||
self.config = config
|
||||
self.repo_id = vllm_config.model_config.model
|
||||
self.is_hf_model = not Path(self.repo_id).is_dir()
|
||||
self.model_stage = vllm_config.model_config.model_stage
|
||||
if self.model_stage == "audio_generation":
|
||||
self.has_preprocess = True
|
||||
self.has_postprocess = True
|
||||
self.requires_raw_input_tokens = True
|
||||
self.set_custom_preprocess(self.tts_preprocess)
|
||||
self.set_custom_postprocess(self.tts_postprocess)
|
||||
self.audio_generation = init_vllm_registered_model(
|
||||
vllm_config=vllm_config,
|
||||
hf_config=config,
|
||||
prefix=maybe_prefix(prefix, "audio_generation"),
|
||||
architectures=["VoxtralTTSAudioGeneration"],
|
||||
)
|
||||
self.model = self.audio_generation
|
||||
self.audio_tokenizer = None
|
||||
self._cudagraph_acoustic_transformer = None
|
||||
self._vllm_config = vllm_config
|
||||
tokenizer = cached_tokenizer_from_config(vllm_config.model_config)
|
||||
self._audio_token_id = tokenizer.instruct.audio_encoder.special_ids.audio
|
||||
speaker_id = config.audio_config.get("speaker_id", None)
|
||||
if speaker_id and self.is_hf_model:
|
||||
self.voice_to_embedding = {
|
||||
sid: torch.load(hf_hub_download(repo_id=self.repo_id, filename=f"voice_embedding/{sid}.pt"))
|
||||
for sid in speaker_id
|
||||
}
|
||||
logger.info("Available voice embeddings: %s", list(self.voice_to_embedding.keys()))
|
||||
else:
|
||||
self.voice_to_embedding = {}
|
||||
logger.warning("No speaker_id configured in audio_config.No voice embeddings will be available.")
|
||||
elif self.model_stage == "audio_tokenizer":
|
||||
self.requires_raw_input_tokens = True
|
||||
self.audio_generation = None
|
||||
self.audio_tokenizer = init_vllm_registered_model(
|
||||
vllm_config=vllm_config,
|
||||
hf_config=config,
|
||||
prefix=maybe_prefix(prefix, "audio_tokenizer"),
|
||||
architectures=["VoxtralTTSAudioTokenizer"],
|
||||
)
|
||||
self.model = self.audio_tokenizer
|
||||
else:
|
||||
raise ValueError("Invalid model stage")
|
||||
|
||||
# -------------------- CUDA Graph for acoustic transformer --------------------
|
||||
def _enable_acoustic_transformer_cudagraph(self):
|
||||
"""Initialize and capture CUDA graphs for compute_mm_logits."""
|
||||
if self.model_stage != "audio_generation" or not hasattr(self, "_vllm_config"):
|
||||
return
|
||||
|
||||
model_cfg = getattr(self._vllm_config, "model_config", None)
|
||||
if model_cfg and getattr(model_cfg, "enforce_eager", False):
|
||||
logger.info("CUDA Graph for acoustic transformer not enabled: --enforce-eager is set")
|
||||
return
|
||||
|
||||
try:
|
||||
acoustic_transformer = self.model.acoustic_transformer
|
||||
device = next(acoustic_transformer.parameters()).device
|
||||
dtype = next(acoustic_transformer.parameters()).dtype
|
||||
hidden_dim = self.model.config.text_config.hidden_size
|
||||
|
||||
wrapper = CUDAGraphAcousticTransformerWrapper(self.model)
|
||||
wrapper._warmup_and_capture(device, dtype, hidden_dim)
|
||||
self._cudagraph_acoustic_transformer = wrapper
|
||||
logger.info("CUDA Graph for acoustic transformer enabled")
|
||||
except Exception:
|
||||
logger.warning("Failed to enable CUDA Graph for acoustic transformer", exc_info=True)
|
||||
self._cudagraph_acoustic_transformer = None
|
||||
|
||||
# -------------------- Device utilities --------------------
|
||||
@cached_property
|
||||
def sampler(self):
|
||||
if hasattr(self.model, "sampler"):
|
||||
return self.model.sampler
|
||||
return Sampler()
|
||||
|
||||
def tts_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict | None):
|
||||
self.post_process_idx = 0
|
||||
audio_tokens = info_dict.pop("audio", None)
|
||||
if audio_tokens is not None:
|
||||
kwargs = {"audio_tokens": audio_tokens.to(input_ids.device)}
|
||||
multimodal_embeddings = self.model.embed_multimodal(**kwargs)
|
||||
if input_ids[0] == self._audio_token_id:
|
||||
input_embeds = multimodal_embeddings[0]
|
||||
return input_ids, input_embeds, info_dict
|
||||
voice = info_dict.pop("voice", None)
|
||||
if voice is not None:
|
||||
if isinstance(voice, list):
|
||||
voice = voice[0]
|
||||
multimodal_embeddings = self.voice_to_embedding[voice].to(input_ids.device).clone().detach()
|
||||
is_multimodal = input_ids == self._audio_token_id
|
||||
input_embeds = self.embed_input_ids(
|
||||
input_ids=input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal
|
||||
)
|
||||
return input_ids, input_embeds, info_dict
|
||||
return input_ids, input_embeds, info_dict
|
||||
|
||||
def tts_postprocess(self, hidden_states: torch.Tensor, multimodal_outputs: object, **info_dict: object | None):
|
||||
update_dict = {}
|
||||
if isinstance(multimodal_outputs, dict) and "audio" in multimodal_outputs:
|
||||
assert self.post_process_idx < len(multimodal_outputs["audio"]), (
|
||||
f"Expect {self.post_process_idx=} < {len(multimodal_outputs['audio'])=}"
|
||||
)
|
||||
update_dict["audio"] = multimodal_outputs["audio"][self.post_process_idx]
|
||||
self.post_process_idx += 1
|
||||
return update_dict
|
||||
|
||||
def embed_input_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings=None,
|
||||
is_multimodal=None,
|
||||
) -> torch.Tensor:
|
||||
if self.model_stage == "audio_tokenizer":
|
||||
return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.config.text_config.hidden_size)
|
||||
return self.model.embed_input_ids(
|
||||
input_ids=input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal
|
||||
)
|
||||
|
||||
def embed_multimodal(self, **kwargs):
|
||||
# Delegate to generation model for multimodal processing
|
||||
return self.model.embed_multimodal(**kwargs)
|
||||
|
||||
def last_index_of(self, list, value):
|
||||
return len(list) - 1 - list[::-1].index(value)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**kwargs: object,
|
||||
) -> torch.Tensor | IntermediateTensors | OmniOutput:
|
||||
if self.model_stage == "audio_generation":
|
||||
if inputs_embeds is not None:
|
||||
input_ids = None
|
||||
hidden_states = self.model(
|
||||
input_ids,
|
||||
positions,
|
||||
intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds,
|
||||
)
|
||||
return hidden_states
|
||||
|
||||
if self.model_stage == "audio_tokenizer":
|
||||
if (input_ids == 0).all():
|
||||
logger.info("audio_tokenizer: sample run with dummy input")
|
||||
# TODO(chenyo): Move this to dummy_inputs creation
|
||||
num_codebooks = self.audio_tokenizer.num_codebooks
|
||||
audio_tokens = torch.randint(low=2, high=100, size=(116, num_codebooks), dtype=torch.int32)
|
||||
audio_tokens[-1, :] = 0
|
||||
audio_tokens[-1, 0] = 1
|
||||
batch_audio_arrays = self.audio_tokenizer.decode_helper_batch_async([audio_tokens])
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={"audio": batch_audio_arrays},
|
||||
)
|
||||
all_audio_tokens, all_ctx_frames = parse_batched_audio_input(
|
||||
input_ids, num_codebooks=self.audio_tokenizer.num_codebooks
|
||||
)
|
||||
|
||||
# Batch decode all requests at once
|
||||
batch_audio_arrays_raw = self.audio_tokenizer.decode_helper_batch_async(all_audio_tokens)
|
||||
batch_audio_arrays = apply_ctx_frames_cutting(
|
||||
batch_audio_arrays_raw, all_ctx_frames, self.audio_tokenizer.downsample_factor
|
||||
)
|
||||
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={"audio": batch_audio_arrays},
|
||||
)
|
||||
|
||||
def make_omni_output(
|
||||
self, model_outputs: torch.Tensor | OmniOutput | tuple, logits_index: int | None = None, **kwargs
|
||||
) -> OmniOutput:
|
||||
if isinstance(model_outputs, torch.Tensor):
|
||||
if self.model_stage == "audio_generation":
|
||||
hidden_states = model_outputs
|
||||
assert logits_index is not None
|
||||
input_hidden_states = hidden_states[logits_index]
|
||||
if self._cudagraph_acoustic_transformer is not None:
|
||||
fake_eos, multimodal_outputs = self._cudagraph_acoustic_transformer(input_hidden_states)
|
||||
else:
|
||||
fake_eos, multimodal_outputs = self.model.compute_mm_logits(input_hidden_states)
|
||||
hidden_states[logits_index, 0] = fake_eos
|
||||
return OmniOutput(
|
||||
text_hidden_states=hidden_states,
|
||||
multimodal_outputs=multimodal_outputs,
|
||||
)
|
||||
|
||||
raise ValueError(f"Unsupported {self.model_stage} for model_outputs type: {type(model_outputs)}")
|
||||
if isinstance(model_outputs, OmniOutput):
|
||||
if self.model_stage == "audio_tokenizer":
|
||||
return model_outputs
|
||||
|
||||
raise ValueError(f"Unsupported {self.model_stage} for model_outputs type: {type(model_outputs)}")
|
||||
|
||||
raise ValueError(f"Unsupported model_outputs type: {type(model_outputs)}")
|
||||
|
||||
def compute_logits(self, hidden_states: torch.Tensor | OmniOutput) -> torch.Tensor | None:
|
||||
# Handle OmniOutput type
|
||||
if isinstance(hidden_states, OmniOutput):
|
||||
hidden_states = hidden_states.text_hidden_states
|
||||
if self.model_stage == "audio_generation":
|
||||
text_logits = self.model.fake_logits_for_audio_tokens(fake_eos=hidden_states)
|
||||
return text_logits
|
||||
if self.model_stage == "audio_tokenizer":
|
||||
raise ValueError("Invalid model stage")
|
||||
|
||||
def sample(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> SamplerOutput | None:
|
||||
return self.model.sample(logits, sampling_metadata)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
loaded_weights = set()
|
||||
remapping_rules = [
|
||||
(r"^audio_tokenizer\.(.*)$", r"\1"), # noqa: E501
|
||||
(
|
||||
r"^mm_audio_embeddings\.audio_codebook_embeddings\.embeddings\.(weight|bias)",
|
||||
r"audio_token_embedding.embeddings.\1",
|
||||
), # noqa: E501
|
||||
(r"^mm_audio_embeddings\.tok_embeddings\.weight", r"tok_embeddings.weight"), # noqa: E501
|
||||
]
|
||||
|
||||
def llm_weights_generator():
|
||||
nonlocal loaded_weights
|
||||
for name, w in weights:
|
||||
is_audio_tokenizer = name.startswith(
|
||||
"mm_audio_embeddings.audio_codebook_embeddings"
|
||||
) or name.startswith("audio_tokenizer.")
|
||||
|
||||
if is_audio_tokenizer and self.audio_tokenizer is not None:
|
||||
# Remap name only when loading audio_tokenizer (Stage-1).
|
||||
# Yield the original name in Stage-0 so audio_generation.load_weights()
|
||||
# can still filter by prefix using its own remapping rules.
|
||||
remapped = name
|
||||
for pattern, repl in remapping_rules:
|
||||
if re.fullmatch(pattern, remapped):
|
||||
remapped = re.sub(pattern, repl, remapped)
|
||||
name = self.audio_tokenizer.load_weight((remapped, w))
|
||||
loaded_weights.add(f"audio_tokenizer.{name}")
|
||||
continue
|
||||
|
||||
yield (name, w)
|
||||
|
||||
audio_generation_weights = list(llm_weights_generator())
|
||||
if self.audio_generation is not None:
|
||||
for name in self.audio_generation.load_weights(audio_generation_weights):
|
||||
loaded_weights.add(f"audio_generation.{name}")
|
||||
|
||||
# If encoder weights were not in the checkpoint, mark them as
|
||||
# "loaded" so the weight-validation does not fail.
|
||||
# encode_waveforms() will raise at runtime if called without encoder weights.
|
||||
if self.audio_tokenizer is not None and not self.audio_tokenizer._encoder_loaded:
|
||||
for name, _ in self.audio_tokenizer.named_parameters():
|
||||
if name.startswith(self.audio_tokenizer._encoder_weight_prefixes):
|
||||
loaded_weights.add(f"audio_tokenizer.{name}")
|
||||
|
||||
# Capture CUDA graphs for compute_mm_logits after weights are loaded
|
||||
self._enable_acoustic_transformer_cudagraph()
|
||||
|
||||
return loaded_weights
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,105 @@
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm # Use llm stage type to launch OmniLLM
|
||||
runtime:
|
||||
process: true # Run this stage in a separate process
|
||||
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
|
||||
max_batch_size: 32
|
||||
engine_args:
|
||||
model_stage: audio_generation
|
||||
model_arch: VoxtralTTSForConditionalGeneration
|
||||
worker_type: ar
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.8
|
||||
enforce_eager: false
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
engine_output_type: latent
|
||||
enable_prefix_caching: false
|
||||
tokenizer_mode: mistral
|
||||
config_format: mistral
|
||||
load_format: mistral
|
||||
skip_mm_profiling: true
|
||||
enable_chunked_prefill: false
|
||||
max_model_len: 4096
|
||||
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.voxtral_tts.generator2tokenizer_async_chunk
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
is_comprehension: true
|
||||
final_output: false
|
||||
final_output_type: text
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.1
|
||||
- stage_id: 1
|
||||
stage_type: llm # Use llm stage type to launch OmniLLM
|
||||
runtime:
|
||||
process: true
|
||||
devices: "0"
|
||||
max_batch_size: 32
|
||||
engine_args:
|
||||
model_stage: audio_tokenizer
|
||||
model_arch: VoxtralTTSForConditionalGeneration
|
||||
worker_type: generation
|
||||
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
async_scheduling: false
|
||||
gpu_memory_utilization: 0.1
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
skip_mm_profiling: true
|
||||
engine_output_type: audio
|
||||
tokenizer_mode: mistral
|
||||
config_format: mistral
|
||||
load_format: mistral
|
||||
max_num_batched_tokens: 65536
|
||||
max_model_len: 65536
|
||||
engine_input_source: [0]
|
||||
is_comprehension: false
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
tts_args:
|
||||
max_instructions_length: 500
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_p: 0.8
|
||||
top_k: 40
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
|
||||
# Top-level runtime config (concise): default windows and stage edges
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1 # Simplified: trigger downstream only after full upstream completion
|
||||
max_inflight: 1 # Simplified: process serially within each stage
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_chunk_frames_at_begin: 5
|
||||
codec_left_context_frames: 25
|
||||
|
||||
edges:
|
||||
- from: 0 # language_model → acoustic_transformer: trigger only after receiving full input (-1)
|
||||
to: 1
|
||||
window_size: -1
|
||||
@@ -0,0 +1,46 @@
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
runtime: {process: true, devices: '0', max_batch_size: 32}
|
||||
engine_args: {model_stage: audio_generation, model_arch: VoxtralTTSForConditionalGeneration,
|
||||
worker_type: ar, worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker, scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler,
|
||||
gpu_memory_utilization: 0.8, enforce_eager: false, trust_remote_code: true, async_scheduling: true,
|
||||
engine_output_type: latent, enable_prefix_caching: false, tokenizer_mode: mistral,
|
||||
config_format: mistral, skip_mm_profiling: true, enable_chunked_prefill: false,
|
||||
max_model_len: 4096, custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.voxtral_tts.generator2tokenizer_async_chunk}
|
||||
output_connectors: {to_stage_1: connector_of_shared_memory}
|
||||
is_comprehension: true
|
||||
final_output: false
|
||||
final_output_type: text
|
||||
default_sampling_params: {temperature: 0.0, top_p: 1.0, top_k: -1, max_tokens: 2048,
|
||||
seed: 42, detokenize: true, repetition_penalty: 1.1}
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime: {process: true, devices: '0', max_batch_size: 32}
|
||||
engine_args: {model_stage: audio_tokenizer, model_arch: VoxtralTTSForConditionalGeneration,
|
||||
worker_type: generation, worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker,
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler,
|
||||
async_scheduling: false, gpu_memory_utilization: 0.1, enforce_eager: true, trust_remote_code: true,
|
||||
enable_prefix_caching: false, skip_mm_profiling: true, engine_output_type: audio,
|
||||
tokenizer_mode: mistral, config_format: mistral, max_num_batched_tokens: 65536,
|
||||
max_model_len: 65536}
|
||||
engine_input_source: [0]
|
||||
is_comprehension: false
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
input_connectors: {from_stage_0: connector_of_shared_memory}
|
||||
tts_args: {max_instructions_length: 500}
|
||||
default_sampling_params: {temperature: 0.9, top_p: 0.8, top_k: 40, max_tokens: 2048,
|
||||
seed: 42, detokenize: true, repetition_penalty: 1.05}
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults: {window_size: -1, max_inflight: 1}
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra: {shm_threshold_bytes: 65536, codec_streaming: true, connector_get_sleep_s: 0.01,
|
||||
connector_get_max_wait_first_chunk: 3000, connector_get_max_wait: 300, codec_chunk_frames: 25,
|
||||
codec_chunk_frames_at_begin: 5, codec_left_context_frames: 25}
|
||||
edges:
|
||||
- {from: 0, to: 1, window_size: -1}
|
||||
@@ -0,0 +1,110 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.inputs import TextPrompt
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.inputs.data import OmniTokensPrompt
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def generator2tokenizer(
|
||||
stage_list,
|
||||
engine_input_source,
|
||||
prompt: OmniTokensPrompt | TextPrompt = None,
|
||||
requires_multimodal_data: bool = False,
|
||||
):
|
||||
if not engine_input_source:
|
||||
raise ValueError("engine_input_source cannot be empty")
|
||||
source_stage_id = engine_input_source[0]
|
||||
if source_stage_id >= len(stage_list):
|
||||
raise IndexError(f"Invalid stage_id: {source_stage_id}")
|
||||
if stage_list[source_stage_id].engine_outputs is None:
|
||||
raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
|
||||
|
||||
generator_outputs = stage_list[source_stage_id].engine_outputs
|
||||
tokenizer_inputs = []
|
||||
for generator_output in generator_outputs:
|
||||
output = generator_output.outputs[0]
|
||||
audio_tokens = torch.cat(output.multimodal_output["audio"], dim=-1).flatten().detach().cpu().tolist()
|
||||
tokenizer_inputs.append(
|
||||
OmniTokensPrompt(
|
||||
prompt_token_ids=audio_tokens,
|
||||
multi_modal_data=None,
|
||||
mm_processor_kwargs=None,
|
||||
)
|
||||
)
|
||||
|
||||
return tokenizer_inputs
|
||||
|
||||
|
||||
def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None:
|
||||
audio = pooling_output.get("audio")
|
||||
if not isinstance(audio, torch.Tensor) or audio.numel() == 0:
|
||||
return None
|
||||
return audio.flatten()
|
||||
|
||||
|
||||
def generator2tokenizer_async_chunk(
|
||||
transfer_manager: Any,
|
||||
pooling_output: dict[str, Any],
|
||||
request: Any,
|
||||
is_finished: bool = False,
|
||||
) -> dict[str, Any] | None:
|
||||
request_id = request.external_req_id
|
||||
finished = bool(is_finished or request.is_finished())
|
||||
|
||||
if isinstance(pooling_output, dict):
|
||||
frame = _extract_last_frame(pooling_output)
|
||||
if frame is not None:
|
||||
codec_codes = frame.cpu().tolist()
|
||||
transfer_manager.code_prompt_token_ids[request_id].append(codec_codes)
|
||||
elif not finished:
|
||||
# Some steps may not produce pooling_output. Only flush on finish.
|
||||
return None
|
||||
|
||||
connector = getattr(transfer_manager, "connector", None)
|
||||
raw_cfg = getattr(connector, "config", {}) or {}
|
||||
cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
|
||||
chunk_size = int(cfg.get("codec_chunk_frames", 25))
|
||||
chunk_size_at_begin = int(cfg.get("codec_chunk_frames_at_begin", 5))
|
||||
left_context_size = int(cfg.get("codec_left_context_frames", 25))
|
||||
if chunk_size <= 0 or left_context_size < 0:
|
||||
raise ValueError(
|
||||
f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, "
|
||||
f"codec_left_context_frames={left_context_size}"
|
||||
)
|
||||
length = len(transfer_manager.code_prompt_token_ids[request_id])
|
||||
|
||||
# Avoid emitting empty chunks during normal streaming. If the request is
|
||||
# finished and nothing was produced, emit an EOF marker.
|
||||
if length <= 0:
|
||||
if finished:
|
||||
return {
|
||||
"code_predictor_codes": [],
|
||||
"finished": torch.tensor(True, dtype=torch.bool),
|
||||
}
|
||||
return None
|
||||
|
||||
# Use a small chunk size at begin
|
||||
if length <= chunk_size:
|
||||
chunk_size = chunk_size_at_begin
|
||||
|
||||
chunk_length = length % chunk_size
|
||||
|
||||
if chunk_length != 0 and not finished:
|
||||
return None
|
||||
|
||||
context_length = chunk_length if chunk_length != 0 else chunk_size
|
||||
end_index = min(length, left_context_size + context_length)
|
||||
ctx_frames = max(0, int(end_index - context_length))
|
||||
window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
|
||||
|
||||
# Pack context + chunk into codebook-major flat codes for adapter.
|
||||
code_predictor_codes = torch.tensor(window_frames).reshape(-1).tolist()
|
||||
|
||||
return {
|
||||
"code_predictor_codes": [int(ctx_frames)] + [int(context_length)] + code_predictor_codes,
|
||||
"finished": torch.tensor(finished, dtype=torch.bool),
|
||||
}
|
||||
@@ -575,10 +575,15 @@ class GPUARModelRunner(OmniGPUModelRunner):
|
||||
if sub_dict:
|
||||
mm_cpu[k] = sub_dict
|
||||
elif isinstance(v, list):
|
||||
element = v[0]
|
||||
if isinstance(element, torch.Tensor):
|
||||
element = element.detach().to("cpu").contiguous()
|
||||
mm_cpu[k] = element
|
||||
if len(v) == 0:
|
||||
continue
|
||||
cpu_list = []
|
||||
for elem in v:
|
||||
if isinstance(elem, torch.Tensor):
|
||||
cpu_list.append(elem.detach().to("cpu").contiguous())
|
||||
else:
|
||||
cpu_list.append(elem)
|
||||
mm_cpu[k] = cpu_list
|
||||
except Exception as e:
|
||||
logger.error(f"Error in merge multimodal outputs: {e}")
|
||||
|
||||
@@ -597,6 +602,12 @@ class GPUARModelRunner(OmniGPUModelRunner):
|
||||
mm_payload[k] = v[start:end].contiguous()
|
||||
elif isinstance(v, dict):
|
||||
mm_payload[k] = {sk: sv[start:end].contiguous() for sk, sv in v.items()}
|
||||
elif isinstance(v, list):
|
||||
element = v[idx] if idx < len(v) else v[0]
|
||||
# Clone tensors to avoid cross-request aliasing
|
||||
if isinstance(element, torch.Tensor):
|
||||
element = element.clone()
|
||||
mm_payload[k] = element
|
||||
elif isinstance(v, torch.Tensor):
|
||||
# List-derived tensor payloads are request-invariant; clone to
|
||||
# avoid accidental cross-request aliasing on downstream mutation.
|
||||
|
||||
@@ -92,9 +92,11 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
|
||||
# TODO move this model specific logic to a separate class
|
||||
# TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS.
|
||||
self.has_talker_mtp = False
|
||||
talker_mtp = getattr(self.model, "talker_mtp", None)
|
||||
if talker_mtp is not None:
|
||||
self.talker_mtp = talker_mtp # type: ignore[assignment]
|
||||
self.has_talker_mtp = True
|
||||
cudagraph_mode = self.compilation_config.cudagraph_mode
|
||||
assert cudagraph_mode is not None
|
||||
# Only wrap talker_mtp in CUDAGraphWrapper for Omni models that
|
||||
@@ -790,7 +792,7 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
slot_mapping=slot_mappings,
|
||||
),
|
||||
):
|
||||
if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"):
|
||||
if getattr(self.model, "talker", None) is not None and self.has_talker_mtp:
|
||||
num_tokens_padded_talker_mtp = num_tokens_padded
|
||||
if num_tokens_padded_talker_mtp == self.max_num_tokens:
|
||||
num_tokens_padded_talker_mtp = self.talker_mtp_input_ids.gpu.shape[0]
|
||||
@@ -1034,7 +1036,9 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
s, e = start_offset, start_offset + sched_tokens
|
||||
# only consider to store data into update dict.
|
||||
hidden_states_slice = hidden_states[s:e]
|
||||
update_dict = self.model.postprocess(hidden_states_slice, **req_infos)
|
||||
update_dict = self.model.postprocess(
|
||||
hidden_states_slice, multimodal_outputs=multimodal_outputs, **req_infos
|
||||
)
|
||||
self._update_intermediate_buffer(req_id, update_dict)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
@@ -1261,7 +1265,7 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
dtype=req_embeds.dtype,
|
||||
)
|
||||
|
||||
if hasattr(self.model, "talker_mtp") and span_len == 1:
|
||||
if self.has_talker_mtp and span_len == 1:
|
||||
last_talker_hidden, text_step = update_dict.pop("mtp_inputs")
|
||||
decode_slice = slice(len(decode_req_ids), len(decode_req_ids) + 1)
|
||||
self.talker_mtp_input_ids.gpu[decode_slice].copy_(req_input_ids)
|
||||
@@ -1280,7 +1284,7 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
input_ids[s : s + seg_len] = req_input_ids
|
||||
|
||||
# run talker mtp decode
|
||||
if hasattr(self.model, "talker_mtp"):
|
||||
if self.has_talker_mtp:
|
||||
self._talker_mtp_forward(decode_req_ids, inputs_embeds)
|
||||
|
||||
return (
|
||||
@@ -1347,7 +1351,7 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
**model_kwargs_extra,
|
||||
)
|
||||
if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"):
|
||||
model_output = self.model.make_omni_output(model_output, **model_kwargs_extra)
|
||||
model_output = self.model.make_omni_output(model_output, **model_kwargs, **model_kwargs_extra)
|
||||
# Cache model output so later sample_tokens can consume multimodal results.
|
||||
self._omni_last_model_output = model_output
|
||||
return model_output
|
||||
|
||||
Reference in new issue
Block a user