diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 9e5027272..5b8d7e104 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -248,6 +248,47 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate + - label: "TTS · GLM-TTS Test" + timeout_in_minutes: 20 + commands: + - | + timeout 20m bash -c " + export VLLM_LOGGING_LEVEL=DEBUG + pytest -s -v tests/e2e/online_serving/test_glm_tts.py tests/e2e/offline_inference/test_glm_tts.py -m 'advanced_model and cuda' --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 + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate - label: "Omni · Qwen3-Omni Test" timeout_in_minutes: 30 diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index b0637a5d7..d26a579ed 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -58,6 +58,7 @@ th { | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | +| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | | `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ | | `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | | | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | diff --git a/docs/user_guide/examples/offline_inference/text_to_speech.md b/docs/user_guide/examples/offline_inference/text_to_speech.md index 7ad3d2d40..7937b1180 100644 --- a/docs/user_guide/examples/offline_inference/text_to_speech.md +++ b/docs/user_guide/examples/offline_inference/text_to_speech.md @@ -19,6 +19,7 @@ list of supported architectures across all modalities, see | VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | — | continuation (`--ref-audio` + `--ref-text`) | 48 kHz | | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ (`async_chunk: true` default) | — | 22.05 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ (`--streaming`) | — | 44.1 kHz | +| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz | | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design (`--instruct`), language (`--lang`) | 24 kHz | | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants (`--query-type`) | 24 kHz | | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets (`--voice`) | 24 kHz | @@ -158,6 +159,35 @@ Streaming requires `async_chunk: true` in the stage config. --- +## GLM-TTS + +2-stage TTS pipeline (AR + DiT flow-matching) at 24 kHz. Every request requires reference audio and its transcript for zero-shot voice cloning. + +### Quick start +```bash +python examples/offline_inference/text_to_speech/glm_tts/end2end.py \ + --model zai-org/GLM-TTS \ + --text "你好，这是语音合成测试。" \ + --ref-audio /path/to/reference.wav \ + --ref-text "这是参考音频的文本内容。" \ + --output-dir ./output +``` + +### Architecture +``` +Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz) + (Llama-based) (32k vocab) (Flow Matching) +``` + +### Notes +- `--ref-audio` and `--ref-text` are **required** together; GLM-TTS does not support text-only synthesis. +- Reference audio should be 3-10 seconds. +- First run may be slow due to lazy loading of WhisperVQ tokenizer and CampPlus ONNX speaker embedder. +- Default sampling: temperature=1.0, top_k=25, top_p=0.8 (RAS method). +- The `--model` path should point to the repository root (not `llm/` subdirectory). + +--- + ## OmniVoice Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design). @@ -313,6 +343,10 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T ``````py --8<-- "examples/offline_inference/text_to_speech/fish_speech/end2end.py" `````` +??? abstract "glm_tts/end2end.py" + ``````py + --8<-- "examples/offline_inference/text_to_speech/glm_tts/end2end.py" + `````` ??? abstract "omnivoice/end2end.py" ``````py --8<-- "examples/offline_inference/text_to_speech/omnivoice/end2end.py" diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index ff22dd0fd..37bfde7ad 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -19,6 +19,7 @@ For the full list of supported architectures across all modalities, see |---|---|---|---|---|---| | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — | | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ | +| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ | | OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — | | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) | | VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ | @@ -192,6 +193,45 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server --- +## GLM-TTS + +2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`. + +### Launch +```bash +vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 +# or: +bash examples/online_serving/text_to_speech/glm_tts/run_server.sh /path/to/GLM-TTS +``` + +### Sending requests +```bash +# Voice cloning (required) +python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \ + --text "你好，这是语音克隆测试。" \ + --ref-audio file:///path/to/ref.wav \ + --ref-text "这是参考音频的文本内容。" + +# Custom format +python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \ + --text "Hello, this is a voice cloning test." \ + --ref-audio file:///path/to/ref.wav \ + --ref-text "Transcript of the reference audio." \ + --response-format mp3 -o output.mp3 +``` + +### Gradio demo +```bash +bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh +``` + +### Notes +- Output: 24 kHz mono WAV via HiFT vocoder. +- `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds. +- Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer. + +--- + ## OmniVoice Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline. @@ -419,6 +459,22 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts ``````py --8<-- "examples/online_serving/text_to_speech/fish_speech/speech_client.py" `````` +??? abstract "glm_tts/gradio_demo.py" + ``````py + --8<-- "examples/online_serving/text_to_speech/glm_tts/gradio_demo.py" + `````` +??? abstract "glm_tts/openai_speech_client.py" + ``````py + --8<-- "examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py" + `````` +??? abstract "glm_tts/run_gradio_demo.sh" + ``````sh + --8<-- "examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh" + `````` +??? abstract "glm_tts/run_server.sh" + ``````sh + --8<-- "examples/online_serving/text_to_speech/glm_tts/run_server.sh" + `````` ??? abstract "omnivoice/run_server.sh" ``````sh --8<-- "examples/online_serving/text_to_speech/omnivoice/run_server.sh" diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index 722b9eeb3..0bf0efc4e 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -16,6 +16,7 @@ list of supported architectures across all modalities, see |---|---|---|---|---|---|---| | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz | +| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz | | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz | | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | single (AR + codec) | ✓ (required) | ✓ | voice_clone, continuation | 48 kHz | | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz | @@ -92,6 +93,35 @@ Streaming is enabled by default via `async_chunk: true` in `vllm_omni/deploy/cos --- +## GLM-TTS + +2-stage TTS pipeline (AR + DiT flow-matching) at 24 kHz. Every request requires reference audio and its transcript for zero-shot voice cloning. + +### Quick start +```bash +python examples/offline_inference/text_to_speech/glm_tts/end2end.py \ + --model zai-org/GLM-TTS \ + --text "你好，这是语音合成测试。" \ + --ref-audio /path/to/reference.wav \ + --ref-text "这是参考音频的文本内容。" \ + --output-dir ./output +``` + +### Architecture +``` +Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz) + (Llama-based) (32k vocab) (Flow Matching) +``` + +### Notes +- `--ref-audio` and `--ref-text` are **required** together; GLM-TTS does not support text-only synthesis. +- Reference audio should be 3-10 seconds. +- First run may be slow due to lazy loading of WhisperVQ tokenizer and CampPlus ONNX speaker embedder. +- Default sampling: temperature=1.0, top_k=25, top_p=0.8 (RAS method). +- The `--model` path should point to the repository root (not `llm/` subdirectory). + +--- + ## Fish Speech S2 Pro 4B dual-AR text-to-speech model from FishAudio with the DAC codec at 44.1 kHz. diff --git a/examples/offline_inference/text_to_speech/glm_tts/end2end.py b/examples/offline_inference/text_to_speech/glm_tts/end2end.py new file mode 100644 index 000000000..d1f19f8a5 --- /dev/null +++ b/examples/offline_inference/text_to_speech/glm_tts/end2end.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-TTS End-to-End Offline Inference Example. + +GLM-TTS is a two-stage TTS system: + - Stage 0 (AR): Llama-based model generates speech tokens from text + - Stage 1 (DiT): Flow matching model converts speech tokens to audio + +Usage: + # Sync two-stage (default) + python examples/offline_inference/text_to_speech/glm_tts/end2end.py \ + --model /path/to/GLM-TTS \ + --text "你好，这是一个语音合成测试。" \ + --ref-audio /path/to/reference.wav \ + --ref-text "参考音频的转录文本。" \ + --output-dir ./output + + # Async chunk mode (streaming DiT) + python examples/offline_inference/text_to_speech/glm_tts/end2end.py \ + --model /path/to/GLM-TTS --async-chunk \ + --text "你好，这是一个语音合成测试。" \ + --ref-audio /path/to/reference.wav \ + --ref-text "参考音频的转录文本。" \ + --output-dir ./output +""" + +import base64 +import io +import logging +import os +import tempfile +import time +from typing import Any +from urllib.request import urlopen + +import numpy as np +import soundfile as sf +import torch +import yaml + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni import Omni +from vllm_omni.model_executor.models.glm_tts.glm_tts import build_glm_tts_prefill_metadata + +logger = logging.getLogger(__name__) + +DEFAULT_DEPLOY_CONFIG = os.path.join( + os.path.dirname(__file__), + "..", + "..", + "..", + "..", + "vllm_omni", + "deploy", + "glm_tts.yaml", +) +SAMPLE_RATE = 24000 + + +def _load_ref_audio(ref_audio: str) -> tuple[torch.Tensor, int]: + """Load reference audio from file path, URL, or data URI.""" + if ref_audio.startswith(("http://", "https://")): + with urlopen(ref_audio, timeout=60) as response: + audio_obj: Any = io.BytesIO(response.read()) + elif ref_audio.startswith("data:"): + _, _, encoded = ref_audio.partition(",") + audio_obj = io.BytesIO(base64.b64decode(encoded)) + else: + audio_obj = ref_audio + wav_np, sr = sf.read(audio_obj, dtype="float32") + if wav_np.ndim > 1: + wav_np = wav_np.mean(axis=1) + return torch.from_numpy(wav_np), int(sr) + + +def _concat_audio(audio_val: Any) -> np.ndarray: + """Concatenate audio tensors from multimodal output.""" + if isinstance(audio_val, list): + tensors = [torch.as_tensor(t).float().reshape(-1) for t in audio_val if t is not None] + if not tensors: + return np.zeros((0,), dtype=np.float32) + return torch.cat(tensors, dim=-1).cpu().numpy().astype(np.float32, copy=False) + if isinstance(audio_val, torch.Tensor): + return audio_val.float().cpu().numpy().reshape(-1) + return np.asarray(audio_val, dtype=np.float32).reshape(-1) + + +def _extract_sample_rate(audio_mm: dict) -> int: + """Extract sample rate from multimodal output dict.""" + sr_raw = audio_mm.get("sr", SAMPLE_RATE) + if isinstance(sr_raw, list): + sr_raw = sr_raw[-1] if sr_raw else SAMPLE_RATE + if hasattr(sr_raw, "item"): + return int(sr_raw.item()) + return int(sr_raw) + + +def _modify_deploy_config(base_path: str, async_chunk: bool) -> str: + """Build deploy config with explicit sync/async mode and eager execution. + + Mirrors the logic in ``tests/e2e/offline_inference/test_glm_tts.py`` + (``_get_deploy_config``) so that example runs match CI behavior. + """ + with open(base_path) as f: + cfg = yaml.safe_load(f) + cfg["async_chunk"] = async_chunk + for stage in cfg.get("stages", []): + stage["enforce_eager"] = True + if stage.get("stage_id") == 0: + stage["async_scheduling"] = bool(async_chunk) + tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False, prefix="glm_tts_") + yaml.dump(cfg, tmp) + tmp.close() + return tmp.name + + +def main(args): + """Run offline GLM-TTS inference.""" + os.makedirs(args.output_dir, exist_ok=True) + base_deploy_config = args.deploy_config or DEFAULT_DEPLOY_CONFIG + deploy_config_path = _modify_deploy_config(base_deploy_config, args.async_chunk) + + ref_audio_wav, ref_audio_sr = _load_ref_audio(args.ref_audio) + if not args.ref_text: + raise ValueError("GLM-TTS requires --ref-audio and --ref-text.") + + inputs = [ + { + "prompt": args.text, + "multi_modal_data": { + "audio": (ref_audio_wav.float().cpu().numpy(), ref_audio_sr), + }, + "modalities": ["audio"], + "mm_processor_kwargs": {"prompt_text": args.ref_text}, + "additional_information": build_glm_tts_prefill_metadata( + args.model, + args.text, + args.ref_text, + ), + } + ] + + omni = Omni( + model=args.model, + stage_configs_path=deploy_config_path, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + ) + + t_start = time.perf_counter() + outputs = omni.generate(inputs) + elapsed = (time.perf_counter() - t_start) * 1000 + + assert outputs, "No outputs returned" + audio_mm = outputs[0].multimodal_output + assert "audio" in audio_mm, "No audio output found" + + audio = _concat_audio(audio_mm["audio"]) + sr = _extract_sample_rate(audio_mm) + out_path = os.path.join(args.output_dir, "output.wav") + sf.write(out_path, audio, samplerate=sr, format="WAV") + + logger.info("Saved %s (%.2fs @ %dHz)", out_path, len(audio) / sr, sr) + logger.info("Total inference: %.1f ms", elapsed) + + +def parse_args(): + parser = FlexibleArgumentParser(description="GLM-TTS Text-to-Speech Example") + parser.add_argument("--model", type=str, required=True, help="Model path") + parser.add_argument("--text", type=str, default="你好，这是一个语音合成测试。") + parser.add_argument("--output-dir", type=str, default="./output") + parser.add_argument("--ref-audio", type=str, required=True, help="Reference WAV path/URL") + parser.add_argument("--ref-text", type=str, required=True, help="Transcript of ref audio") + parser.add_argument("--deploy-config", type=str, default=None) + parser.add_argument( + "--async-chunk", + action="store_true", + default=False, + help="Enable async_chunk mode (streaming DiT). Default: sync two-stage.", + ) + parser.add_argument("--log-stats", action="store_true") + parser.add_argument("--stage-init-timeout", type=int, default=600) + return parser.parse_args() + + +if __name__ == "__main__": + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + ) + main(parse_args()) diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 75b85580d..922156d4c 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -15,6 +15,7 @@ For the full list of supported architectures across all modalities, see | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo | |---|---|---|---|---|---| | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ | +| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ | | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | — (caption-controlled) | — | caption fields (`instructions`) | — | | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ | | OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — | @@ -94,6 +95,45 @@ For full request-shape documentation (all parameters, response formats, error co --- +## GLM-TTS + +2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`. + +### Launch +```bash +vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 +# or: +bash examples/online_serving/text_to_speech/glm_tts/run_server.sh /path/to/GLM-TTS +``` + +### Sending requests +```bash +# Voice cloning (required) +python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \ + --text "你好，这是语音克隆测试。" \ + --ref-audio file:///path/to/ref.wav \ + --ref-text "这是参考音频的文本内容。" + +# Custom format +python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \ + --text "Hello, this is a voice cloning test." \ + --ref-audio file:///path/to/ref.wav \ + --ref-text "Transcript of the reference audio." \ + --response-format mp3 -o output.mp3 +``` + +### Gradio demo +```bash +bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh +``` + +### Notes +- Output: 24 kHz mono WAV via HiFT vocoder. +- `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds. +- Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer. + +--- + ## Fish Speech S2 Pro 4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec. diff --git a/examples/online_serving/text_to_speech/glm_tts/gradio_demo.py b/examples/online_serving/text_to_speech/glm_tts/gradio_demo.py new file mode 100644 index 000000000..12e5d4b11 --- /dev/null +++ b/examples/online_serving/text_to_speech/glm_tts/gradio_demo.py @@ -0,0 +1,291 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Gradio demo for GLM-TTS online serving via /v1/audio/speech API. + +Supports: + - Voice-cloned text-to-speech synthesis from reference audio (upload or URL) + - Streaming (progressive PCM chunks) and non-streaming modes + +Usage: + # Start the server first (see run_server.sh), then: + python gradio_demo.py --api-base http://localhost:8091 + + # Or use run_gradio_demo.sh to start both server and demo together. +""" + +import argparse +import base64 +import io + +try: + import gradio as gr +except ImportError: + raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None +import httpx +import numpy as np +import soundfile as sf + +PCM_SAMPLE_RATE = 24000 +DEFAULT_API_BASE = "http://localhost:8091" + + +def encode_audio_to_base64(audio_data: tuple) -> str: + """Encode Gradio audio input (sample_rate, numpy_array) to base64 data URL.""" + sample_rate, audio_np = audio_data + if audio_np.dtype != np.int16: + if audio_np.dtype in (np.float32, np.float64): + audio_np = np.clip(audio_np, -1.0, 1.0) + audio_np = (audio_np * 32767).astype(np.int16) + else: + audio_np = audio_np.astype(np.int16) + buf = io.BytesIO() + sf.write(buf, audio_np, sample_rate, format="WAV") + wav_b64 = base64.b64encode(buf.getvalue()).decode("utf-8") + return f"data:audio/wav;base64,{wav_b64}" + + +def build_payload( + text: str, + ref_audio: tuple | None, + ref_audio_url: str, + ref_text: str, + response_format: str = "wav", + stream: bool = False, + max_new_tokens: int = 4096, +) -> dict: + """Build the /v1/audio/speech request payload.""" + if not text or not text.strip(): + raise gr.Error("Please enter text to synthesize.") + + payload: dict = { + "input": text.strip(), + "voice": "default", + "response_format": "pcm" if stream else response_format, + "stream": stream, + "max_new_tokens": max_new_tokens, + } + + # Voice cloning: ref_audio takes priority over URL. + ref_url_stripped = ref_audio_url.strip() if ref_audio_url else "" + if ref_audio is not None: + payload["ref_audio"] = encode_audio_to_base64(ref_audio) + elif ref_url_stripped: + payload["ref_audio"] = ref_url_stripped + else: + raise gr.Error("GLM-TTS requires a reference audio file or URL.") + + if not ref_text or not ref_text.strip(): + raise gr.Error("GLM-TTS requires a transcript of the reference audio.") + payload["ref_text"] = ref_text.strip() + + return payload + + +def generate_speech(api_base: str, text: str, ref_audio, ref_audio_url, ref_text, response_format, max_new_tokens): + """Non-streaming: call /v1/audio/speech and return full audio.""" + payload = build_payload( + text, ref_audio, ref_audio_url, ref_text, response_format, stream=False, max_new_tokens=max_new_tokens + ) + + try: + with httpx.Client(timeout=300.0) as client: + resp = client.post( + f"{api_base}/v1/audio/speech", + json=payload, + headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"}, + ) + except httpx.TimeoutException: + raise gr.Error("Request timed out. The server may be busy.") + except httpx.ConnectError: + raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") + + if resp.status_code != 200: + raise gr.Error(f"Server error ({resp.status_code}): {resp.text}") + + content_type = resp.headers.get("content-type", "") + if "application/json" in content_type: + try: + raise gr.Error(f"Server error: {resp.json()}") + except ValueError: + pass + + try: + if response_format == "pcm": + audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0 + return (PCM_SAMPLE_RATE, audio_np) + audio_np, sample_rate = sf.read(io.BytesIO(resp.content)) + if audio_np.ndim > 1: + audio_np = audio_np[:, 0] + return (sample_rate, audio_np.astype(np.float32)) + except Exception as e: + raise gr.Error(f"Failed to decode audio: {e}") + + +def generate_speech_stream( + api_base, text, ref_audio, ref_audio_url, ref_text, response_format, max_new_tokens, chunk_seconds=0.5 +): + """Streaming: yield PCM chunks for progressive playback.""" + payload = build_payload( + text, ref_audio, ref_audio_url, ref_text, response_format, stream=True, max_new_tokens=max_new_tokens + ) + + min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds) + prebuffer_chunks = 2 + pending: list[np.ndarray] = [] + pending_len = 0 + chunks_yielded = 0 + prebuffer: list[np.ndarray] = [] + + try: + leftover = b"" + with httpx.Client(timeout=300.0) as client: + with client.stream( + "POST", + f"{api_base}/v1/audio/speech", + json=payload, + headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"}, + ) as resp: + if resp.status_code != 200: + resp.read() + raise gr.Error(f"Server error ({resp.status_code}): {resp.text}") + for chunk in resp.iter_bytes(): + if not chunk: + continue + raw = leftover + chunk + usable = len(raw) - (len(raw) % 2) + leftover = raw[usable:] + if usable == 0: + continue + float_samples = np.frombuffer(raw[:usable], dtype=np.int16).copy().astype(np.float32) / 32767.0 + pending.append(float_samples) + pending_len += len(float_samples) + if pending_len >= min_chunk_samples: + audio_chunk = np.concatenate(pending) + pending.clear() + pending_len = 0 + if chunks_yielded < prebuffer_chunks: + prebuffer.append(audio_chunk) + chunks_yielded += 1 + if chunks_yielded == prebuffer_chunks: + yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer)) + prebuffer.clear() + else: + yield (PCM_SAMPLE_RATE, audio_chunk) + remaining = prebuffer + pending + if remaining: + yield (PCM_SAMPLE_RATE, np.concatenate(remaining)) + except httpx.TimeoutException: + raise gr.Error("Request timed out. The server may be busy.") + except httpx.ConnectError: + raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") + + +def on_stream_change(stream: bool): + """Lock format to PCM when streaming.""" + if stream: + return gr.update(value="pcm", interactive=False) + return gr.update(value="wav", interactive=True) + + +def build_interface(api_base: str, stream_chunk_seconds: float = 0.5): + """Build the Gradio interface.""" + with gr.Blocks(title="GLM-TTS Demo") as demo: + gr.Markdown("# GLM-TTS - Voice Cloning") + gr.Markdown(f"**Server:** `{api_base}` | **Model:** GLM-TTS | **Output:** 24kHz") + + with gr.Row(): + with gr.Column(scale=3): + text_input = gr.Textbox( + label="Text to Synthesize", + placeholder="Enter text here (supports Chinese and English)...", + lines=4, + ) + + with gr.Accordion("Reference Voice (required)", open=True): + gr.Markdown("Upload or link a short reference audio (3-10s) and provide its transcript.") + ref_audio = gr.Audio( + label="Reference Audio", + type="numpy", + sources=["upload", "microphone"], + ) + ref_audio_url = gr.Textbox( + label="Reference Audio URL (alternative to upload)", + placeholder="https://example.com/reference.wav", + lines=1, + ) + ref_text = gr.Textbox( + label="Reference Audio Transcript (required for cloning)", + placeholder="Exact transcript of the reference audio...", + lines=2, + ) + + with gr.Row(): + response_format = gr.Dropdown( + choices=["wav", "mp3", "flac", "pcm"], + value="wav", + label="Audio Format", + scale=1, + ) + stream_checkbox = gr.Checkbox( + label="Stream output", + value=False, + info="Progressive PCM streaming", + scale=1, + ) + max_new_tokens = gr.Slider( + minimum=256, + maximum=8192, + value=4096, + step=256, + label="Max New Tokens", + scale=1, + ) + + generate_btn = gr.Button("Generate Speech", variant="primary", size="lg") + + with gr.Column(scale=2): + audio_output = gr.Audio( + label="Generated Audio", + interactive=False, + streaming=True, + autoplay=True, + ) + gr.Markdown( + "### About\n" + "- **GLM-TTS** by Zhipu AI: Llama-based AR + Flow Matching DiT\n" + "- **Voice cloning**: upload 3-10s reference + transcript\n" + "- **Streaming**: real-time PCM output\n" + "- **24kHz** output via Vocos vocoder" + ) + + stream_checkbox.change(fn=on_stream_change, inputs=[stream_checkbox], outputs=[response_format]) + + all_inputs = [text_input, ref_audio, ref_audio_url, ref_text, response_format, max_new_tokens] + + def dispatch(stream_enabled, *args): + if stream_enabled: + yield from generate_speech_stream(api_base, *args, chunk_seconds=stream_chunk_seconds) + else: + yield generate_speech(api_base, *args) + + generate_btn.click(fn=dispatch, inputs=[stream_checkbox] + all_inputs, outputs=[audio_output]) + demo.queue() + return demo + + +def main(): + parser = argparse.ArgumentParser(description="Gradio demo for GLM-TTS") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})") + parser.add_argument("--host", default="0.0.0.0", help="Gradio host (default: 0.0.0.0)") + parser.add_argument("--port", type=int, default=7860, help="Gradio port (default: 7860)") + parser.add_argument("--share", action="store_true", help="Share publicly") + parser.add_argument("--stream-chunk-seconds", type=float, default=0.5, help="Seconds per streaming chunk") + args = parser.parse_args() + + print(f"Connecting to vLLM server at: {args.api_base}") + demo = build_interface(args.api_base, stream_chunk_seconds=args.stream_chunk_seconds) + demo.launch(server_name=args.host, server_port=args.port, share=args.share) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py b/examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py new file mode 100644 index 000000000..a0ff181c3 --- /dev/null +++ b/examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""OpenAI-compatible client for GLM-TTS via /v1/audio/speech endpoint. + +GLM-TTS is a two-stage TTS system (AR + DiT) that generates audio from text +conditioned on reference speech. Each request requires ref_audio + ref_text. + +Usage: + # Voice cloning + python openai_speech_client.py --text "你好" --ref-audio file:///path/to/ref.wav --ref-text "参考文本" + + # Streaming response, for async_chunk server mode + python openai_speech_client.py --text "你好" --stream --ref-audio file:///path/to/ref.wav --ref-text "参考文本" + + # Specify output format + python openai_speech_client.py --text "你好" --ref-audio file:///path/to/ref.wav \ + --ref-text "参考文本" --response-format mp3 -o output.mp3 +""" + +import argparse + +import httpx + +# Default server configuration +DEFAULT_API_BASE = "http://localhost:8091" +DEFAULT_API_KEY = "EMPTY" + + +def run_tts_generation(args) -> None: + """Run TTS generation via OpenAI-compatible /v1/audio/speech API.""" + if not args.ref_audio or not args.ref_text: + raise ValueError("GLM-TTS requires --ref-audio and --ref-text for voice cloning.") + + payload = { + "model": args.model, + "voice": "default", + "input": args.text, + "response_format": args.response_format, + "stream": bool(args.stream), + "ref_audio": args.ref_audio, + "ref_text": args.ref_text, + } + if args.max_new_tokens: + payload["max_new_tokens"] = args.max_new_tokens + + print(f"Model: {args.model}") + print(f"Text: {args.text}") + print(f"Voice cloning: ref_audio={args.ref_audio}, ref_text={args.ref_text}") + print(f"Stream: {args.stream}") + print("Generating audio...") + + api_url = f"{args.api_base}/v1/audio/speech" + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {args.api_key}", + } + + if args.stream: + output_path = args.output or "tts_output.pcm" + with httpx.Client(timeout=300.0) as client, open(output_path, "wb") as f: + with client.stream("POST", api_url, json=payload, headers=headers) as response: + if response.status_code != 200: + print(f"Error: {response.status_code}") + response.read() + print(response.text) + return + for chunk in response.iter_bytes(): + f.write(chunk) + print(f"Streaming audio saved to: {output_path}") + else: + with httpx.Client(timeout=300.0) as client: + response = client.post(api_url, json=payload, headers=headers) + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.text) + return + try: + text = response.content.decode("utf-8") + if text.startswith('{"error"'): + print(f"Error: {text}") + return + except UnicodeDecodeError: + pass + output_path = args.output or f"tts_output.{args.response_format}" + with open(output_path, "wb") as f: + f.write(response.content) + print(f"Audio saved to: {output_path}") + + +def parse_args(): + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="OpenAI-compatible client for GLM-TTS via /v1/audio/speech", + ) + + # Server configuration + parser.add_argument( + "--api-base", + type=str, + default=DEFAULT_API_BASE, + help=f"API base URL (default: {DEFAULT_API_BASE})", + ) + parser.add_argument( + "--api-key", + type=str, + default=DEFAULT_API_KEY, + help="API key (default: EMPTY)", + ) + parser.add_argument( + "--model", + "-m", + type=str, + default="glm-tts", + help="Model name/path", + ) + + # Input text + parser.add_argument( + "--text", + type=str, + required=True, + help="Text to synthesize", + ) + + # Generation parameters + parser.add_argument( + "--max-new-tokens", + type=int, + default=None, + help="Maximum new tokens to generate (default: model default)", + ) + + # Output + parser.add_argument( + "--response-format", + type=str, + default="wav", + choices=["wav", "mp3", "flac", "pcm", "aac", "opus"], + help="Audio output format (default: wav)", + ) + parser.add_argument( + "--stream", + action="store_true", + help="Request a streaming audio response (use with async_chunk server mode).", + ) + parser.add_argument( + "--output", + "-o", + type=str, + default=None, + help="Output audio file path (default: tts_output.)", + ) + + # Voice cloning parameters + parser.add_argument( + "--ref-audio", + type=str, + default=None, + help="Reference audio URL, file:// URI, or base64 data URL for voice cloning", + ) + parser.add_argument( + "--ref-text", + type=str, + default=None, + help="Transcript of the reference audio (required with --ref-audio)", + ) + + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + run_tts_generation(args) diff --git a/examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh b/examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh new file mode 100644 index 000000000..75cecdf9e --- /dev/null +++ b/examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh @@ -0,0 +1,48 @@ +#!/bin/bash +# Launch GLM-TTS server + Gradio demo together. +# +# Usage: +# ./run_gradio_demo.sh +# CUDA_VISIBLE_DEVICES=0 PORT=8091 GRADIO_PORT=7860 ./run_gradio_demo.sh + +set -e + +MODEL="${MODEL:-zai-org/GLM-TTS}" +PORT="${PORT:-8091}" +GRADIO_PORT="${GRADIO_PORT:-7860}" +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)" + +echo "Starting GLM-TTS server (port $PORT)..." +FLASHINFER_DISABLE_VERSION_CHECK=1 \ +vllm-omni serve "$MODEL" \ + --deploy-config "$REPO_ROOT/vllm_omni/deploy/glm_tts.yaml" \ + --host 0.0.0.0 \ + --port "$PORT" \ + --gpu-memory-utilization 0.9 \ + --trust-remote-code \ + --enforce-eager \ + --omni & +SERVER_PID=$! + +cleanup() { + echo "Stopping server (PID $SERVER_PID)..." + kill $SERVER_PID 2>/dev/null + wait $SERVER_PID 2>/dev/null +} +trap cleanup EXIT + +# Wait for server to be ready. +echo "Waiting for server to start..." +for i in $(seq 1 120); do + if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then + echo "Server ready." + break + fi + sleep 2 +done + +echo "Starting Gradio demo (port $GRADIO_PORT)..." +python "$SCRIPT_DIR/gradio_demo.py" \ + --api-base "http://localhost:$PORT" \ + --port "$GRADIO_PORT" diff --git a/examples/online_serving/text_to_speech/glm_tts/run_server.sh b/examples/online_serving/text_to_speech/glm_tts/run_server.sh new file mode 100644 index 000000000..f5cc3a1f8 --- /dev/null +++ b/examples/online_serving/text_to_speech/glm_tts/run_server.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# Launch vLLM-Omni server for GLM-TTS models +# +# Usage: +# ./run_server.sh # Default model path, async_chunk mode +# ./run_server.sh /path/to/GLM-TTS # Custom model path, async_chunk mode +# ./run_server.sh /path/to/GLM-TTS sync # Sync two-stage mode +# +# NOTE: The model path should point to the repo ROOT (not llm/ subdirectory). +# model_subdir/tokenizer_subdir in the pipeline config resolve subdirectories. + +set -e + +MODEL="${1:-zai-org/GLM-TTS}" +MODE="${2:-async}" + +EXTRA_ARGS=() +case "$MODE" in + async|async_chunk) + ;; + sync|no_async_chunk) + EXTRA_ARGS+=("--no-async-chunk") + ;; + *) + echo "Unknown mode: $MODE (expected async or sync)" >&2 + exit 1 + ;; +esac + +echo "Starting GLM-TTS server with model: $MODEL (mode: $MODE)" + +vllm-omni serve "$MODEL" \ + --deploy-config vllm_omni/deploy/glm_tts.yaml \ + --host 0.0.0.0 \ + --port 8091 \ + --trust-remote-code \ + --omni \ + "${EXTRA_ARGS[@]}" diff --git a/recipes/README.md b/recipes/README.md index 8ddb72943..d2e1eeff3 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -43,6 +43,7 @@ recipes/ | [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 | | [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) | | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB | +| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB | Within a single recipe file, include different hardware support sections such as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like diff --git a/recipes/zai-org/GLM-TTS.md b/recipes/zai-org/GLM-TTS.md new file mode 100644 index 000000000..2bf1a202a --- /dev/null +++ b/recipes/zai-org/GLM-TTS.md @@ -0,0 +1,105 @@ +# GLM-TTS for Chinese/English TTS on 1x GPU + +## Summary + +- Vendor: zai-org +- Model: `zai-org/GLM-TTS` +- Task: Zero-shot voice-cloned text-to-speech synthesis +- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API +- Maintainer: Community + +## When to use this recipe + +Use this recipe to serve GLM-TTS as a two-stage TTS system (AR + DiT +flow-matching) for Chinese and English speech synthesis. Every request is +conditioned on reference audio and its transcript. + +## References + +- Upstream or canonical docs: + [zai-org/GLM-TTS on HuggingFace](https://huggingface.co/zai-org/GLM-TTS) +- GitHub repository: + [zai-org/GLM-TTS](https://github.com/zai-org/GLM-TTS) +- Related example under `examples/`: + [`examples/online_serving/text_to_speech/README.md#glm-tts`](../../examples/online_serving/text_to_speech/README.md#glm-tts) +- Offline inference example: + [`examples/offline_inference/text_to_speech/README.md#glm-tts`](../../examples/offline_inference/text_to_speech/README.md#glm-tts) + +## Hardware Support + +### GPU + +### 1x A40 48GB + +#### Environment + +- OS: Linux +- Python: 3.10+ +- Driver / runtime: NVIDIA CUDA environment with A40 48GB GPU +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Command + +Start the server from the repository root: + +```bash +vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 +``` + +Async chunking is enabled by default in the bundled deployment config. For +the sync (non-streaming) path: + +```bash +vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk +``` + +Use a custom deploy config for advanced cases: + +```bash +vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \ + --deploy-config /path/to/your_glm_tts_overrides.yaml +``` + +#### Verification + +Run the bundled OpenAI-compatible client with reference audio: + +```bash +python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \ + --text "你好，这是一个语音合成测试。" \ + --ref-audio file:///path/to/ref.wav \ + --ref-text "这是参考音频的文本内容。" +``` + +For a quick API smoke test: + +```bash +curl http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "zai-org/GLM-TTS", + "input": "你好，这是一个语音合成测试。", + "response_format": "wav", + "ref_audio": "file:///path/to/ref.wav", + "ref_text": "这是参考音频的文本内容。" + }' --output test.wav +``` + +Voice cloning with reference audio: + +```bash +python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \ + --text "你好，这是语音克隆测试。" \ + --ref-audio file:///path/to/ref.wav \ + --ref-text "这是参考音频的文本内容。" +``` + +#### Notes + +- Hardware scope: the default bundled config is CUDA-only and verified on 1x A40 48GB (~16.6 GiB peak); fits 24GB cards. Split stages across GPUs for higher concurrency. +- Memory usage: ~18-20GB total (AR ~10GB, DiT ~8GB); both stages share GPU 0 by default. +- Audio output: 24kHz mono WAV via HiFT vocoder (Vocos2D 32kHz fallback with resampling). +- Key flags: `--omni` is required; `--trust-remote-code` is needed for the GLM-TTS phoneme tokenizer; the DiT stage enables bucketed model-internal CUDA graphs with eager fallback. +- Voice cloning: requires `ref_audio` + `ref_text` together. Reference audio should be 3-10 seconds. Feature extraction (WhisperVQ tokenizer, CampPlus ONNX, mel) runs on the model side. +- Known limitations: First request may be slow due to lazy model loading (WhisperVQ, CampPlus ONNX). Warm-cache RTF is approximately 0.6-0.7x on A40. diff --git a/tests/assets/glm_tts/jiayan_zh.wav b/tests/assets/glm_tts/jiayan_zh.wav new file mode 100644 index 000000000..1dd528747 Binary files /dev/null and b/tests/assets/glm_tts/jiayan_zh.wav differ diff --git a/tests/e2e/offline_inference/test_glm_tts.py b/tests/e2e/offline_inference/test_glm_tts.py new file mode 100644 index 000000000..0b902c6bc --- /dev/null +++ b/tests/e2e/offline_inference/test_glm_tts.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E Offline inference tests for GLM-TTS voice cloning. + +GLM-TTS is a zero-shot voice cloning model that always requires ref_audio. +There is no text-only / non-clone inference path in the official implementation. +""" + +import os +import tempfile + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +import numpy as np +import pytest +import soundfile as sf + +from tests.helpers.assertions import assert_audio_speech_response +from tests.helpers.mark import hardware_test +from tests.helpers.media import convert_audio_file_to_text, get_asset_path +from tests.helpers.runtime import OmniResponse +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config +from vllm_omni.model_executor.models.glm_tts.glm_tts import build_glm_tts_prefill_metadata + +MODEL = os.environ.get("GLM_TTS_MODEL_PATH", "zai-org/GLM-TTS") +REF_TEXT = "他当时还跟线下其他的站姐吵架，然后，打架进局子了。" + +DEPLOY_CONFIG = get_deploy_config_path("glm_tts.yaml") + +REFERENCE_PROMPT_WAV_PATH = get_asset_path("glm_tts/jiayan_zh.wav") + +ASYNC_CHUNK_MODES = [ + pytest.param(False, id="sync"), + pytest.param(True, id="async_chunk"), +] + + +def _load_ref_audio() -> tuple[np.ndarray, int]: + audio, sr = sf.read(str(REFERENCE_PROMPT_WAV_PATH), dtype="float32", always_2d=False) + if isinstance(audio, np.ndarray) and audio.ndim > 1: + audio = np.mean(audio, axis=-1) + return np.asarray(audio, dtype=np.float32), int(sr) + + +def _concat_audio(audio_val) -> np.ndarray: + import torch + + if isinstance(audio_val, list): + tensors = [torch.as_tensor(t).float().reshape(-1) for t in audio_val if t is not None] + if not tensors: + return np.zeros((0,), dtype=np.float32) + return torch.cat(tensors, dim=-1).cpu().numpy().astype(np.float32, copy=False) + if isinstance(audio_val, torch.Tensor): + return audio_val.float().cpu().numpy().reshape(-1) + return np.asarray(audio_val, dtype=np.float32).reshape(-1) + + +def _extract_sample_rate(audio_mm: dict) -> int: + sr_raw = audio_mm.get("sr", 24000) + if isinstance(sr_raw, list): + sr_raw = sr_raw[-1] if sr_raw else 24000 + if hasattr(sr_raw, "item"): + return int(sr_raw.item()) + return int(sr_raw) + + +def _get_deploy_config(*, async_chunk: bool) -> str: + """Build deploy config with explicit sync/async mode and eager execution.""" + return modify_stage_config( + DEPLOY_CONFIG, + updates={ + "async_chunk": async_chunk, + "stages": { + 0: { + "enforce_eager": True, + "async_scheduling": bool(async_chunk), + }, + 1: { + "enforce_eager": True, + }, + }, + }, + ) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("async_chunk", ASYNC_CHUNK_MODES) +def test_offline_voice_clone_zh(async_chunk: bool) -> None: + """ + Test voice cloning offline inference. + Deploy Setting: glm_tts.yaml with explicit sync/async mode, enforce_eager + Input Modal: text + ref_audio + ref_text + Output Modal: audio + + Uses the official jiayan_zh.wav reference audio from the upstream + GLM-TTS repository (real human speech matching REF_TEXT). + """ + synth_text = "今天天气真不错，适合出去散散步。" + prompt_audio = _load_ref_audio() + from tests.helpers.runtime import OmniRunner + + with OmniRunner( + MODEL, + stage_configs_path=_get_deploy_config(async_chunk=async_chunk), + stage_init_timeout=600, + ) as omni_runner: + outputs = omni_runner.omni.generate( + [ + { + "prompt": synth_text, + "multi_modal_data": {"audio": prompt_audio}, + "modalities": ["audio"], + "mm_processor_kwargs": {"prompt_text": REF_TEXT}, + "additional_information": build_glm_tts_prefill_metadata( + MODEL, + synth_text, + REF_TEXT, + trust_remote_code=True, + ), + } + ], + omni_runner.get_default_sampling_params_list(), + ) + + assert outputs, "No outputs returned" + audio_mm = outputs[0].multimodal_output + assert "audio" in audio_mm, "No audio output found" + audio = _concat_audio(audio_mm["audio"]) + assert audio.size > 0, "Generated audio is empty" + + with tempfile.NamedTemporaryFile(suffix=".wav") as out_file: + sf.write(out_file.name, audio, samplerate=_extract_sample_rate(audio_mm), format="WAV") + transcript = convert_audio_file_to_text(out_file.name) + audio_bytes = open(out_file.name, "rb").read() + + assert_audio_speech_response( + OmniResponse( + success=True, + audio_bytes=audio_bytes, + audio_content=transcript, + audio_format="audio/wav", + ), + {"input": synth_text, "response_format": "wav"}, + run_level="advanced_model", + ) diff --git a/tests/e2e/online_serving/test_glm_tts.py b/tests/e2e/online_serving/test_glm_tts.py new file mode 100644 index 000000000..7c6ad9764 --- /dev/null +++ b/tests/e2e/online_serving/test_glm_tts.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E Online tests for GLM-TTS model with text input and audio output. + +These tests verify the /v1/audio/speech endpoint works correctly with +the GLM-TTS two-stage pipeline (AR + DiT). +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.media import load_test_audio_data_url +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = os.environ.get("GLM_TTS_MODEL_PATH", "zai-org/GLM-TTS") +REF_TEXT = "他当时还跟线下其他的站姐吵架，然后，打架进局子了。" + +DEPLOY_CONFIG = get_deploy_config_path("glm_tts.yaml") + +# Official GLM-TTS reference audio and transcript. Vendored under tests/assets/ +# so Buildkite does not depend on raw.githubusercontent.com during E2E. +REF_AUDIO_URL = load_test_audio_data_url("glm_tts/jiayan_zh.wav") + + +SYNC_EXTRA_ARGS = [ + "--trust-remote-code", + "--disable-log-stats", + "--no-async-chunk", +] + +ASYNC_EXTRA_ARGS = [ + "--trust-remote-code", + "--disable-log-stats", +] + +tts_sync_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=DEPLOY_CONFIG, + server_args=SYNC_EXTRA_ARGS, + stage_init_timeout=600, + ), + id="glm_tts_sync", + ) +] + +tts_async_chunk_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=DEPLOY_CONFIG, + server_args=ASYNC_EXTRA_ARGS, + stage_init_timeout=600, + ), + id="glm_tts_async_chunk", + ) +] + + +# GLM-TTS is a zero-shot voice cloning model (always requires ref_audio). +# There is no text-only / non-clone path — the official implementation +# asserts prompt_feat.shape[1] != 0 in token2wav_with_cache and every +# example prompt includes a reference audio. CosyVoice3 follows the same +# pattern: all tests are voice_clone_*. + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_sync_server_params, indirect=True) +def test_voice_clone_zh_sync(omni_server, openai_client) -> None: + """ + Test voice cloning TTS with Chinese text. + Deploy Setting: glm_tts.yaml with ``--no-async-chunk`` (sync two-stage) + Input Modal: text + ref_audio + ref_text + Output Modal: audio + Input Setting: stream=False + + Uses the official jiayan_zh.wav reference audio from the upstream + GLM-TTS repository (real human speech matching REF_TEXT). + """ + request_config = { + "model": omni_server.model, + "input": "今天天气真不错，适合出去散散步。", + "stream": False, + "response_format": "wav", + "ref_audio": REF_AUDIO_URL, + "ref_text": REF_TEXT, + } + openai_client.send_audio_speech_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True) +def test_voice_clone_zh_async_chunk(omni_server, openai_client) -> None: + """ + Test voice cloning TTS with Chinese text via async_chunk streaming. + Deploy Setting: glm_tts.yaml default ``async_chunk: true`` + Input Modal: text + ref_audio + ref_text + Output Modal: audio + Input Setting: stream=True + """ + request_config = { + "model": omni_server.model, + "input": "今天天气真不错，适合出去散散步。", + "stream": True, + "response_format": "wav", + "ref_audio": REF_AUDIO_URL, + "ref_text": REF_TEXT, + } + openai_client.send_audio_speech_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_sync_server_params, indirect=True) +def test_models_endpoint(omni_server, openai_client) -> None: + """Test the /v1/models endpoint returns loaded model.""" + models = openai_client.client.models.list() + assert len(models.data) > 0 diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 2440fe586..872410aaa 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -2335,6 +2335,84 @@ class TestCosyVoice3Serving: cosyvoice3_server._build_cosyvoice3_prompt.assert_awaited_once() +# ---- GLM-TTS Serving Tests ---- + + +@pytest.fixture +def glm_tts_server(mocker: MockerFixture): + mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set()) + mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None) + + mock_engine_client = mocker.MagicMock() + mock_engine_client.errored = False + mock_engine_client.model_config = mocker.MagicMock( + model="zai-org/GLM-TTS", + hf_config=SimpleNamespace(min_token_text_ratio=2, max_token_text_ratio=20), + ) + mock_engine_client.default_sampling_params_list = [ + SimpleNamespace(max_tokens=2048, min_tokens=None, extra_args=None) + ] + mock_engine_client.tts_batch_max_items = 32 + mock_engine_client.generate = mocker.MagicMock(return_value="generator") + mock_engine_client.stage_configs = [ + SimpleNamespace( + engine_args=SimpleNamespace(model_stage="glm_tts"), + tts_args={}, + ) + ] + + mock_models = mocker.MagicMock() + mock_models.is_base_model.return_value = True + + return OmniOpenAIServingSpeech( + engine_client=mock_engine_client, + models=mock_models, + request_logger=mocker.MagicMock(), + ) + + +class TestGLMTTSServing: + def test_validate_glm_tts_requires_ref_audio(self, glm_tts_server): + request = OpenAICreateSpeechRequest(input="Hello", ref_text="Reference transcript") + error = glm_tts_server._validate_glm_tts_request(request) + assert error is not None + assert "ref_audio" in error + + def test_validate_glm_tts_requires_ref_text(self, glm_tts_server): + request = OpenAICreateSpeechRequest(input="Hello", ref_audio="data:audio/wav;base64,abc") + error = glm_tts_server._validate_glm_tts_request(request) + assert error is not None + assert "ref_text" in error + + def test_estimate_glm_tts_target_text_len_uses_tokenizer_tokens( + self, + glm_tts_server, + mocker: MockerFixture, + ): + class FakeTokenizer: + def encode(self, text): + assert text == "normalized target" + return [10, 20, 30] + + mocker.patch( + "vllm_omni.model_executor.models.glm_tts.glm_tts.resolve_glm_tts_tokenizer_path", + return_value="resolved-tokenizer", + ) + load_tokenizer = mocker.patch( + "vllm_omni.model_executor.models.glm_tts.glm_tts.load_glm_tts_tokenizer", + return_value=FakeTokenizer(), + ) + mocker.patch( + "vllm_omni.model_executor.models.glm_tts.text_frontend.GLMTTSTextFrontend.text_normalize", + return_value="normalized target", + ) + + text_token_len = glm_tts_server._estimate_glm_tts_text_token_len("abcdef") + + assert text_token_len == 3 + load_tokenizer.assert_called_once() + + class TestTTSAsyncOffloading: """Tests for event-loop-safe offloading of blocking TTS operations.""" diff --git a/tests/model_executor/models/test_glm_tts_prompt_len.py b/tests/model_executor/models/test_glm_tts_prompt_len.py new file mode 100644 index 000000000..df12078b4 --- /dev/null +++ b/tests/model_executor/models/test_glm_tts_prompt_len.py @@ -0,0 +1,351 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.common.nucleus_ras_sampling import ras_sample_one +from vllm_omni.model_executor.models.glm_tts import glm_tts as glm_tts_module +from vllm_omni.model_executor.models.glm_tts.glm_tts import GLMTTSForConditionalGeneration +from vllm_omni.model_executor.models.glm_tts.text_frontend import GLMTTSTextFrontend + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_glm_tts_ras_fallback_matches_official_random_resample() -> None: + sampled = ras_sample_one( + torch.tensor([0.0, float("-inf")]), + decoded_tokens=[0] * 10, + top_p=1.0, + top_k=1, + win_size=10, + tau_r=0.1, + generator=None, + ) + + assert sampled == 0 + + +def test_glm_tts_text_frontend_applies_official_symbol_replacements() -> None: + frontend = GLMTTSTextFrontend() + + assert frontend.pre_replace("3-2 咯，") == "3减2 喽，" + assert frontend.post_replace("α②∈A/B~C²") == "阿尔法二属于A每B到C平方" + + +def test_glm_tts_call_hf_processor_builds_prompt_text_text_boa( + monkeypatch: pytest.MonkeyPatch, +) -> None: + processor = object.__new__(glm_tts_module.GLMTTSMultiModalProcessor) + processor.info = SimpleNamespace( + ctx=SimpleNamespace( + get_hf_config=lambda *_args, **_kwargs: SimpleNamespace(), + model_config=SimpleNamespace(model="fake-model"), + ) + ) + processor.text_frontend = object() + processor.special_ids = {"boa": 99} + processor.speech_tokenizer = object() + processor.processor_device = torch.device("cpu") + processor.campplus_session = object() + + monkeypatch.setattr( + processor, + "_ensure_cached_runtime_components", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + processor, + "_encode_text", + lambda text: torch.tensor([[11, 12]]) if text.endswith(" ") else torch.tensor([[21, 22, 23]]), + ) + monkeypatch.setattr(processor, "_get_audio", lambda _mm_data: (torch.zeros(8), 24000)) + monkeypatch.setattr( + glm_tts_module, + "_normalize_glm_tts_processor_text", + lambda _frontend, text, add_trailing_space=False: text + (" " if add_trailing_space else ""), + ) + monkeypatch.setattr(glm_tts_module, "extract_prompt_speech_token", lambda *_args, **_kwargs: [7, 8, 9, 10]) + monkeypatch.setattr(glm_tts_module, "extract_prompt_feat", lambda *_args, **_kwargs: torch.ones(4, 3)) + monkeypatch.setattr(glm_tts_module, "extract_spk_embedding", lambda *_args, **_kwargs: [0.1, 0.2, 0.3]) + + out = processor._call_hf_processor( + prompt="target text", + mm_data={"audio": object()}, + mm_kwargs={"prompt_text": "reference text"}, + tok_kwargs={}, + ) + + assert out["input_ids"].tolist() == [[11, 12, 21, 22, 23, 99]] + assert out["prompt_speech_token"].tolist() == [[7, 8, 9, 10]] + assert out["prompt_speech_token_len"][0].item() == 4 + assert out["glm_tts_text_token_len"][0].item() == 3 + + +def test_glm_tts_prompt_updates_append_single_audio_suffix() -> None: + processor = object.__new__(glm_tts_module.GLMTTSMultiModalProcessor) + out_mm_kwargs = { + "audio": [{"prompt_speech_token_len": SimpleNamespace(data=[torch.tensor([5], dtype=torch.long)])}] + } + + updates = processor._get_prompt_updates( + mm_items=object(), + hf_processor_mm_kwargs={}, + out_mm_kwargs=out_mm_kwargs, + ) + + assert len(updates) == 1 + assert updates[0].modality == "audio" + assert updates[0].insertion(0) == [1, 1, 1, 1, 1] + + +def test_glm_tts_dummy_inputs_forbid_dummy_wav(monkeypatch: pytest.MonkeyPatch) -> None: + builder = object.__new__(glm_tts_module.GLMTTSDummyInputsBuilder) + monkeypatch.setattr( + builder, + "_get_dummy_audios", + lambda *_args, **_kwargs: pytest.fail("GLM-TTS dummy inputs must not fabricate wav audio"), + ) + + assert builder.get_dummy_mm_data(seq_len=16, mm_counts={"audio": 1}) == {} + + +def test_glm_tts_processing_info_declares_audio_token_budget() -> None: + info = object.__new__(glm_tts_module.GLMTTSMultiModalProcessingInfo) + + assert info.get_mm_max_tokens_per_item(seq_len=4096, mm_counts={"audio": 1}) == {"audio": 1024} + assert info.get_mm_max_tokens_per_item(seq_len=512, mm_counts={"audio": 1}) == {"audio": 512} + assert info.get_mm_max_tokens_per_item(seq_len=4096, mm_counts={}) == {} + + +def test_glm_tts_text_only_processor_skips_voice_clone(monkeypatch: pytest.MonkeyPatch) -> None: + class FakeTokenizer: + def encode(self, text: str) -> list[int]: + return [11, 12, 13] + + processor = object.__new__(glm_tts_module.GLMTTSMultiModalProcessor) + processor.info = SimpleNamespace( + ctx=SimpleNamespace( + get_hf_config=lambda *_args, **_kwargs: SimpleNamespace(), + model_config=SimpleNamespace(model="fake-model", tokenizer=None, trust_remote_code=True), + ) + ) + + monkeypatch.setattr(glm_tts_module, "resolve_glm_tts_model_dir", lambda *_args, **_kwargs: "/fake-root") + monkeypatch.setattr(glm_tts_module, "load_glm_tts_tokenizer", lambda *_args, **_kwargs: FakeTokenizer()) + monkeypatch.setattr(glm_tts_module, "get_glm_tts_special_token_ids", lambda _tokenizer: {"boa": 99}) + monkeypatch.setattr(glm_tts_module, "GLMTTSTextFrontend", lambda: object()) + monkeypatch.setattr( + glm_tts_module, + "_normalize_glm_tts_processor_text", + lambda _frontend, text, add_trailing_space=False: text, + ) + monkeypatch.setattr( + glm_tts_module, + "load_voice_clone_frontend", + lambda *_args, **_kwargs: pytest.fail("text-only dummy path should not load voice clone frontend"), + ) + + out = processor._call_hf_processor(prompt="target text", mm_data={}, mm_kwargs={}, tok_kwargs={}) + + assert out["input_ids"].tolist() == [[11, 12, 13, 99]] + assert out["prompt_speech_token"].shape == (1, 1024) + assert out["prompt_speech_token_len"][0].item() == 1024 + assert out["glm_tts_prompt_text_token_len"][0].item() == 0 + assert out["glm_tts_text_token_len"][0].item() == 3 + + fields = processor._get_mm_fields_config(out, {}) + assert {key: value.modality for key, value in fields.items()} == { + "prompt_speech_token": "audio", + "prompt_speech_token_len": "audio", + "glm_tts_prompt_text_token_len": "audio", + "glm_tts_text_token_len": "audio", + } + + updates = processor._get_prompt_updates( + mm_items=object(), + hf_processor_mm_kwargs={}, + out_mm_kwargs={ + "audio": [{"prompt_speech_token_len": SimpleNamespace(data=[out["prompt_speech_token_len"][0]])}] + }, + ) + assert len(updates) == 1 + assert len(updates[0].insertion(0)) == 1024 + + +def test_glm_tts_embed_replaces_multimodal_positions_per_request() -> None: + class FakeModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.marker = torch.nn.Parameter(torch.zeros(())) + + def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: + return input_ids.to(torch.float32).reshape(-1, 1) + + model = object.__new__(GLMTTSForConditionalGeneration) + torch.nn.Module.__init__(model) + model.model_stage = "glm_tts" + model.model = FakeModel() + + input_ids = torch.tensor([101, 102, 1, 1, 201, 1, 1, 1]) + is_multimodal = torch.tensor([False, False, True, True, False, True, True, True]) + mm_embeds = [ + torch.tensor([[1001.0], [1002.0]]), + torch.tensor([[2001.0], [2002.0], [2003.0]]), + ] + + out = model.embed_input_ids(input_ids, multimodal_embeddings=mm_embeds, is_multimodal=is_multimodal) + + assert out.squeeze(-1).tolist() == [ + 101.0, + 102.0, + 1001.0, + 1002.0, + 201.0, + 2001.0, + 2002.0, + 2003.0, + ] + + +def test_glm_tts_embed_multimodal_keeps_all_batched_requests() -> None: + class FakeModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.marker = torch.nn.Parameter(torch.zeros(())) + + def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: + return input_ids.to(torch.float32).reshape(*input_ids.shape, 1) + + model = object.__new__(GLMTTSForConditionalGeneration) + torch.nn.Module.__init__(model) + model.model_stage = "glm_tts" + model.model = FakeModel() + model._ats = 1000 + + out = model.embed_multimodal( + prompt_speech_token=[ + torch.tensor([[1, 2]]), + torch.tensor([[3, 4, 5]]), + ] + ) + + assert out is not None + assert len(out) == 2 + assert out[0].squeeze(-1).tolist() == [1001.0, 1002.0] + assert out[1].squeeze(-1).tolist() == [1003.0, 1004.0, 1005.0] + + +def test_glm_tts_embed_multimodal_returns_one_item_for_single_request() -> None: + class FakeModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.marker = torch.nn.Parameter(torch.zeros(())) + + def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: + return input_ids.to(torch.float32).reshape(*input_ids.shape, 1) + + model = object.__new__(GLMTTSForConditionalGeneration) + torch.nn.Module.__init__(model) + model.model_stage = "glm_tts" + model.model = FakeModel() + model._ats = 1000 + + out = model.embed_multimodal(prompt_speech_token=torch.tensor([[1, 2, 3]])) + + assert out is not None + assert len(out) == 1 + assert out[0].squeeze(-1).tolist() == [1001.0, 1002.0, 1003.0] + + +def test_glm_tts_embed_multimodal_treats_scalar_tensor_list_as_one_item() -> None: + class FakeModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.marker = torch.nn.Parameter(torch.zeros(())) + + def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: + return input_ids.to(torch.float32).reshape(*input_ids.shape, 1) + + model = object.__new__(GLMTTSForConditionalGeneration) + torch.nn.Module.__init__(model) + model.model_stage = "glm_tts" + model.model = FakeModel() + model._ats = 1000 + + out = model.embed_multimodal(prompt_speech_token=[torch.tensor(1), torch.tensor(2)]) + + assert out is not None + assert len(out) == 1 + assert out[0].squeeze(-1).tolist() == [1001.0, 1002.0] + + +def test_glm_tts_preprocess_treats_one_token_text_tail_as_prefill() -> None: + model = object.__new__(GLMTTSForConditionalGeneration) + model._pad = 0 + model._ats = 1000 + model._ate = 2000 + model._eoa = 3000 + model._model_dtype = lambda: torch.float32 + + input_ids, input_embeds, info = model.preprocess( + torch.tensor([123]), + torch.tensor([[1.0, 2.0]]), + glm_tts_mm_prefill_done=True, + ) + + assert input_ids.tolist() == [0] + assert input_embeds.tolist() == [[1.0, 2.0]] + assert info["glm_tts_mm_prefill_done"] is True + assert info["speech_tokens"].tolist() == [[-1]] + + +def test_glm_tts_preprocess_initializes_generation_bounds_before_sampling() -> None: + model = object.__new__(GLMTTSForConditionalGeneration) + model._pad = 0 + model._ats = 1000 + model._ate = 2000 + model._eoa = 3000 + model._model_dtype = lambda: torch.float32 + model.config = SimpleNamespace(min_token_text_ratio=2.0, max_token_text_ratio=20.0) + + _, _, info = model.preprocess( + torch.tensor([11, 12]), + torch.tensor([[1.0], [2.0]]), + glm_tts_text_token_len=[torch.tensor([27])], + ) + + assert info["glm_tts_text_token_len"][0].item() == 27 + + +def test_glm_tts_make_omni_output_emits_prompt_speech_len_once() -> None: + model = object.__new__(GLMTTSForConditionalGeneration) + + out = model.make_omni_output( + torch.randn(2, 4), + model_intermediate_buffer=[{}], + prompt_speech_token_len=[torch.tensor([9])], + glm_tts_text_token_len=[torch.tensor([5])], + ) + + assert out.multimodal_outputs is not None + assert out.multimodal_outputs["prompt_speech_token_len"][0].item() == 9 + assert out.multimodal_outputs["glm_tts_text_token_len"][0].item() == 5 + + +def test_glm_tts_postprocess_caches_prompt_speech_token_len() -> None: + model = object.__new__(GLMTTSForConditionalGeneration) + model.config = SimpleNamespace(min_token_text_ratio=2.0, max_token_text_ratio=20.0) + + update = model.postprocess( + torch.randn(2, 4), + multimodal_outputs={ + "prompt_speech_token_len": [torch.tensor([13])], + "glm_tts_text_token_len": [torch.tensor([6])], + }, + ) + + assert update["prompt_speech_token_len"].item() == 13 + assert update["glm_tts_text_token_len"].item() == 6 diff --git a/tests/model_executor/stage_input_processors/test_glm_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_glm_tts_async_chunk.py new file mode 100644 index 000000000..19bdaf7fe --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_glm_tts_async_chunk.py @@ -0,0 +1,278 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.engine.serialization import serialize_additional_information +from vllm_omni.model_executor.models.glm_tts.glm_tts_dit_wrapper import ( + split_request_ids, + valid_speech_tokens, +) +from vllm_omni.model_executor.stage_input_processors.glm_tts import ( + ar_to_dit, + ar_to_dit_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Request: + external_req_id = "req-1" + additional_information = None + + def is_finished(self) -> bool: + return False + + +def _transfer_manager() -> SimpleNamespace: + connector = SimpleNamespace( + config={ + "extra": { + "codec_chunk_frames": 25, + "codec_left_context_frames": 25, + } + } + ) + return SimpleNamespace(connector=connector) + + +def _progressive_transfer_manager() -> SimpleNamespace: + connector = SimpleNamespace( + config={ + "extra": { + "codec_chunk_frames": [25, 50, 200], + "codec_left_context_frames": 25, + } + } + ) + return SimpleNamespace(connector=connector) + + +def _source_output(mm: dict, request_id: str = "req-0") -> SimpleNamespace: + return SimpleNamespace( + request_id=request_id, + outputs=[ + SimpleNamespace( + multimodal_output=mm, + ) + ], + ) + + +def _audio_codes(payload) -> list[int]: + assert payload.codes is not None + assert payload.codes.audio is not None + return payload.codes.audio.tolist() + + +def test_ar_to_dit_filters_invalid_tokens_and_preserves_conditioning() -> None: + prompt_token = torch.tensor([[1, 2, 3]]) + prompt_feat = torch.randn(1, 4, 80) + embedding = torch.randn(1, 192) + source_outputs = [ + _source_output( + { + "speech_tokens": torch.tensor([[-1, 0, 7, -1, 32767]]), + "prompt_token": prompt_token, + "prompt_feat": prompt_feat, + "embedding": embedding, + } + ) + ] + + outputs = ar_to_dit(source_outputs) + + assert len(outputs) == 1 + assert outputs[0]["prompt_token_ids"] == [0, 7, 32767] + additional_info = outputs[0]["additional_information"] + assert additional_info["speech_tokens"] == [0, 7, 32767] + # Voice clone data stored in nested embed dict (matches async_chunk path) + embed = additional_info["embed"] + assert embed["speech_token"] is prompt_token + assert embed["speech_feat"] is prompt_feat + assert embed["embedding"] is embedding + # Meta dict enables block-causal attention in DiT wrapper + meta = additional_info["meta"] + assert meta["req_id"] == ["req-0"] + assert meta["left_context_size"] == 0 + assert meta["stream_finished"] is True + private = additional_info["kv_metadata"]["glm_tts"] + assert private["block_pattern"] == [3] # len([0, 7, 32767]) + assert private["chunk_sizes_history"] == [3] + + +def test_dit_explicit_empty_speech_tokens_override_scheduler_placeholder() -> None: + placeholder = torch.tensor([0], dtype=torch.long) + + token = valid_speech_tokens([], device=placeholder.device, fallback=placeholder) + + assert token.numel() == 0 + + +def test_ar_to_dit_async_chunk_restores_serialized_conditioning_payload() -> None: + transfer = _transfer_manager() + request = _Request() + prompt_token = torch.tensor([[1, 2, 3]]) + prompt_feat = torch.randn(1, 4, 80) + embedding = torch.randn(1, 192) + request.additional_information = serialize_additional_information( + { + "prompt_speech_token": prompt_token, + "prompt_feat": prompt_feat, + "embedding": embedding, + } + ) + + payload = ar_to_dit_async_chunk( + transfer, + {"speech_tokens": torch.tensor([[0]])}, + request, + is_finished=False, + ) + + assert payload is None + state = transfer.request_payload["req-1"]["_glm_tts_async_state"] + prompt_payload = state["prompt_payload"] + assert torch.equal(prompt_payload["prompt_speech_token"], prompt_token) + assert torch.equal(prompt_payload["prompt_feat"], prompt_feat) + assert torch.equal(prompt_payload["embedding"], embedding) + + +def test_split_request_ids_uses_seq_token_counts() -> None: + ids = torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.long) + + split = split_request_ids(ids, seq_token_counts=[2, 3, 1]) + + assert [part.tolist() for part in split] == [[1, 2], [3, 4, 5], [6]] + + +def test_ar_to_dit_async_chunk_emits_cumulative_prefixes() -> None: + transfer = _transfer_manager() + request = _Request() + payloads = [] + + for token in range(75): + payload = ar_to_dit_async_chunk( + transfer, + {"speech_tokens": torch.tensor([[token]])}, + request, + is_finished=False, + ) + if payload is not None: + payloads.append(payload) + + assert len(payloads) == 3 + assert _audio_codes(payloads[0]) == list(range(25)) + assert payloads[0].meta.left_context_size == 0 + assert payloads[0].kv_metadata["glm_tts"]["block_pattern"] == [25] + assert _audio_codes(payloads[1]) == list(range(50)) + assert payloads[1].meta.left_context_size == 25 + assert payloads[1].kv_metadata["glm_tts"]["block_pattern"] == [25] + assert _audio_codes(payloads[2]) == list(range(75)) + assert payloads[2].meta.left_context_size == 50 + assert payloads[2].kv_metadata["glm_tts"]["block_pattern"] == [25] + + +def test_ar_to_dit_async_chunk_keeps_progressive_block_pattern_fixed() -> None: + transfer = _progressive_transfer_manager() + request = _Request() + payloads = [] + + for token in range(30): + payload = ar_to_dit_async_chunk( + transfer, + {"speech_tokens": torch.tensor([[token]])}, + request, + is_finished=False, + ) + if payload is not None: + payloads.append(payload) + + final_payload = ar_to_dit_async_chunk( + transfer, + None, + request, + is_finished=True, + ) + + assert len(payloads) == 1 + assert payloads[0].kv_metadata["glm_tts"]["chunk_sizes_history"] == [25] + assert payloads[0].kv_metadata["glm_tts"]["block_pattern"] == [25, 50, 200] + assert final_payload is not None + assert final_payload.kv_metadata["glm_tts"]["chunk_sizes_history"] == [25, 5] + assert final_payload.kv_metadata["glm_tts"]["block_pattern"] == [25, 50, 200] + + +def test_ar_to_dit_async_chunk_terminal_boundary_is_metadata_only() -> None: + transfer = _transfer_manager() + request = _Request() + payloads = [] + + for token in range(50): + payload = ar_to_dit_async_chunk( + transfer, + {"speech_tokens": torch.tensor([[token]])}, + request, + is_finished=False, + ) + if payload is not None: + payloads.append(payload) + + final_payload = ar_to_dit_async_chunk( + transfer, + None, + request, + is_finished=True, + ) + + assert len(payloads) == 2 + assert final_payload is not None + assert _audio_codes(final_payload) == [] + assert final_payload.meta.left_context_size == 50 + assert final_payload.meta.req_id == ["req-1"] + assert bool(final_payload.meta.stream_finished.item()) + + +def test_ar_to_dit_async_chunk_finished_flushes_partial_unsent_tail() -> None: + transfer = _transfer_manager() + request = _Request() + payloads = [] + + for token in range(30): + payload = ar_to_dit_async_chunk( + transfer, + {"speech_tokens": torch.tensor([[token]])}, + request, + is_finished=False, + ) + if payload is not None: + payloads.append(payload) + + final_payload = ar_to_dit_async_chunk( + transfer, + None, + request, + is_finished=True, + ) + + assert len(payloads) == 1 + assert _audio_codes(payloads[0]) == list(range(25)) + assert payloads[0].meta.left_context_size == 0 + assert final_payload is not None + assert _audio_codes(final_payload) == list(range(30)) + assert final_payload.meta.left_context_size == 25 + assert final_payload.meta.req_id == ["req-1"] + assert bool(final_payload.meta.finished.item()) + assert bool(final_payload.meta.stream_finished.item()) + assert ( + ar_to_dit_async_chunk( + transfer, + None, + request, + is_finished=True, + ) + is None + ) diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 555f35e17..1474e325c 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -97,6 +97,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = { "vllm_omni.model_executor.models.voxtral_tts.pipeline", "VOXTRAL_TTS_PIPELINE", ), + "glm_tts": ( + "vllm_omni.model_executor.models.glm_tts.pipeline", + "GLM_TTS_PIPELINE", + ), "fish_qwen3_omni": ( "vllm_omni.model_executor.models.fish_speech.pipeline", "FISH_SPEECH_PIPELINE", diff --git a/vllm_omni/deploy/glm_tts.yaml b/vllm_omni/deploy/glm_tts.yaml new file mode 100644 index 000000000..1600b6152 --- /dev/null +++ b/vllm_omni/deploy/glm_tts.yaml @@ -0,0 +1,83 @@ +# GLM-TTS deploy: 2-stage AR -> DiT flow-matching. +# +# Default mode is async-chunk streaming through SharedMemoryConnector. +# Stage 0 (AR LLM) streams speech token chunks to Stage 1 (DiT + vocoder) +# via the shared-memory connector. Stage 1 runs as LLM_GENERATION type +# to support chunk_transfer_adapter consumption. +# +# Stage 1 keeps vLLM eager scheduling but enables the GLM-TTS bucketed DiT +# CUDA graph wrapper via hf_overrides.use_dit_cuda_graphs. Graph capture falls +# back to eager if a dynamic shape exceeds the captured buckets. +# +# Verified on 1x A40 48GB (both stages share GPU 0). +# Memory footprint ~16.6 GiB peak; fits 24GB cards with the settings below. +async_chunk: true +trust_remote_code: true + +connectors: + glm_tts_shm: + name: SharedMemoryConnector + extra: + codec_streaming: true + connector_get_sleep_s: 0.002 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + # Progressive chunk sizes matching original GLM-TTS defaults. + codec_chunk_frames: [25, 50, 200] + codec_left_context_frames: 25 + crossfade_sec: 0.1 + +stages: + - stage_id: 0 + max_num_seqs: 1 + gpu_memory_utilization: 0.55 + # BF16 for AR decode: memory-bandwidth bound on single-batch decode, + # bf16 halves HBM reads per step (~1.7x throughput on L4/A40). + dtype: bfloat16 + enforce_eager: false + async_scheduling: true + enable_prefix_caching: false + max_num_batched_tokens: 512 + max_model_len: 2048 + devices: "0" + output_connectors: + to_stage_1: glm_tts_shm + default_sampling_params: + temperature: 1.0 + top_k: 25 + top_p: 0.8 + max_tokens: 2048 + min_tokens: 1 + seed: 42 + detokenize: false + hf_overrides: + sample_method: ras + ras_top_p: 0.8 + ras_top_k: 25 + ras_win_size: 10 + ras_tau_r: 0.1 + min_token_text_ratio: 2.0 + max_token_text_ratio: 20.0 + skip_mm_profiling: true + + - stage_id: 1 + max_num_seqs: 1 + gpu_memory_utilization: 0.25 + enforce_eager: true + max_model_len: 8192 + max_num_batched_tokens: 2048 + # Force float32 for the flow-matching DiT stage. The Euler ODE solver + # (10 steps × 2 forward passes each) accumulates rounding errors in half + # precision that manifest as audible hiss on H20/H100. float32 matches + # the reference implementation and eliminates the noise floor. + dtype: float32 + devices: "0" + enable_prefix_caching: false + input_connectors: + from_stage_0: glm_tts_shm + default_sampling_params: + max_tokens: 2048 + hf_overrides: + use_dit_cuda_graphs: true + disable_hybrid_kv_cache_manager: true + skip_mm_profiling: true diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 12a5a450a..c46222899 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -20,6 +20,7 @@ logger = init_logger(__name__) # Used when auto-injecting hf_overrides for models with missing config.json. _ARCH_TO_MODEL_TYPE: dict[str, str] = { "CosyVoice3Model": "cosyvoice3", + "GLMTTSForConditionalGeneration": "glm_tts", "OmniVoiceModel": "omnivoice", "VoxCPM2TalkerForConditionalGeneration": "voxcpm2", } @@ -27,6 +28,7 @@ _ARCH_TO_MODEL_TYPE: dict[str, str] = { # Maps model architecture names to tokenizer subfolder paths within HF repos. _TOKENIZER_SUBFOLDER_MAP: dict[str, str] = { "CosyVoice3Model": "CosyVoice-BlankEN", + "GLMTTSForConditionalGeneration": "vq32k-phoneme-tokenizer", } @@ -37,6 +39,10 @@ def _register_omni_hf_configs() -> None: from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( Qwen3TTSConfig, ) + from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config + from vllm_omni.transformers_utils.configs.glm_tts import GLMTTSConfig + from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig + from vllm_omni.transformers_utils.configs.voxcpm2 import VoxCPM2Config 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 @@ -51,6 +57,10 @@ def _register_omni_hf_configs() -> None: for model_type, config_cls in [ ("qwen3_tts", Qwen3TTSConfig), + ("cosyvoice3", CosyVoice3Config), + ("glm_tts", GLMTTSConfig), + ("omnivoice", OmniVoiceConfig), + ("voxcpm2", VoxCPM2Config), ]: try: AutoConfig.register(model_type, config_cls) @@ -298,6 +308,7 @@ class OmniEngineArgs(EngineArgs): model_path, allow_patterns=[ f"{subfolder}/tokenizer*", + f"{subfolder}/tokenization_*", f"{subfolder}/special_tokens*", f"{subfolder}/vocab*", f"{subfolder}/merges*", diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 3b7a2e747..10036978e 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -68,6 +68,7 @@ _COVO_AUDIO_MODEL_STAGES = {"fused_thinker_talker"} _VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"} _MING_TTS_MODEL_STAGES = {"ming_tts"} _MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"} +_GLM_TTS_MODEL_STAGES = {"glm_tts"} _TTS_MODEL_STAGES: set[str] = ( _VOXTRAL_TTS_MODEL_STAGES | _QWEN3_TTS_MODEL_STAGES @@ -78,6 +79,7 @@ _TTS_MODEL_STAGES: set[str] = ( | _VOXCPM2_TTS_MODEL_STAGES | _MING_TTS_MODEL_STAGES | _MOSS_TTS_MODEL_STAGES + | _GLM_TTS_MODEL_STAGES ) _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {"fish_tts", "qwen3_tts", "voxtral_tts", "cosyvoice3", "voxcpm2"} _TTS_LANGUAGES: set[str] = { @@ -335,6 +337,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Determine TTS model type or None self._tts_model_type = self._detect_tts_model_type() + # GLM-TTS lazy-cached resources (populated on first GLM-TTS request) + self._glm_tts_text_tokenizer: object | None = None + self._glm_tts_text_frontend: object | None = None + # Cache TTS configuration values (computed once, reused per request) self._max_instructions_length = self._compute_max_instructions_length() @@ -499,6 +505,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return "ming_flash_omni_tts" if model_stage in _MOSS_TTS_MODEL_STAGES: return "moss_tts_nano" + if model_stage in _GLM_TTS_MODEL_STAGES: + return "glm_tts" return None def _compute_max_instructions_length(self) -> int: @@ -528,6 +536,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # is no static speaker list to surface here. return set() try: + if self._tts_model_type == "glm_tts": + return set() if self._tts_model_type == "voxcpm2": return {"default"} if self._tts_model_type == "voxtral_tts": @@ -785,12 +795,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): voice_lower = request.voice.lower() if voice_lower not in self.uploaded_speakers: - if self._tts_model_type in ("cosyvoice3", "fish_tts", "omnivoice", "moss_tts_nano"): + if self._tts_model_type in ("cosyvoice3", "fish_tts", "omnivoice", "moss_tts_nano", "glm_tts"): label = { "cosyvoice3": "CosyVoice3", "fish_tts": "Fish Speech", "omnivoice": "OmniVoice", "moss_tts_nano": "MOSS-TTS-Nano", + "glm_tts": "GLM-TTS", }.get(self._tts_model_type, self._tts_model_type) return ( f"Unknown voice '{request.voice}'. {label} has no " @@ -1135,6 +1146,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return self._validate_ming_tts_request(request) if self._tts_model_type == "moss_tts_nano": return self._validate_moss_tts_request(request) + if self._tts_model_type == "glm_tts": + return self._validate_glm_tts_request(request) return self._validate_qwen_tts_request(request) def _voxcpm2_encode(self, text: str) -> list[int]: @@ -1383,6 +1396,26 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): params["prompt_audio_array"] = [[wav_list, sr]] return params + def _validate_glm_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: + """Validate GLM-TTS request — requires ref_audio for voice cloning.""" + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + + if request.ref_audio is None: + return "GLM-TTS requires 'ref_audio' for zero-shot voice cloning" + fmt_err = self._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + if not request.ref_text or not request.ref_text.strip(): + return "GLM-TTS voice cloning requires 'ref_text' (transcript of the reference audio)" + + 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 >= {_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_fish_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate Fish Speech request parameters. Returns error message or None.""" if not request.input or not request.input.strip(): @@ -1898,6 +1931,105 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ) return sampling_params_list + # ---- GLM-TTS helpers ---- + + async def _build_glm_tts_prompt( + self, + request: OpenAICreateSpeechRequest, + *, + has_inline_ref_audio: bool = False, + ) -> dict[str, Any]: + """Build prompt for GLM-TTS. + + Uses the multimodal processor path (same as CosyVoice3): + - prompt: synthesis text + - multi_modal_data["audio"]: (wav_samples, sr) reference audio + - mm_processor_kwargs["prompt_text"]: reference text transcript + + AR preprocess() builds [PromptText | Text | BOA | PromptSpeechTokens + ATS]. + DiT receives prompt_token, prompt_feat, embedding for conditioning. + """ + # Voice cloning requires ref_audio + ref_text + if request.ref_audio is not None and request.ref_text: + wav_samples, sr = await self._resolve_ref_audio(request.ref_audio) + audio_data = (np.asarray(wav_samples, dtype=np.float32), int(sr)) + + mm_kwargs: dict[str, Any] = { + "prompt_text": request.ref_text, + } + if request.voice: + voice_lower = request.voice.lower() + if voice_lower in self.uploaded_speakers and not has_inline_ref_audio: + mm_kwargs["voice_name"] = voice_lower + mm_kwargs["voice_created_at"] = self._voice_created_at(voice_lower) + + return { + "prompt": request.input, + "multi_modal_data": { + "audio": audio_data, + }, + "mm_processor_kwargs": mm_kwargs, + "additional_information": self._build_glm_tts_prefill_metadata( + request.input, + request.ref_text, + ), + } + + raise ValueError("GLM-TTS requires ref_audio and ref_text for voice cloning.") + + def _glm_tts_text_tokenizer_and_frontend(self): + from vllm_omni.model_executor.models.glm_tts.glm_tts import ( + load_glm_tts_tokenizer, + resolve_glm_tts_tokenizer_path, + ) + from vllm_omni.model_executor.models.glm_tts.text_frontend import GLMTTSTextFrontend + + cached = self._glm_tts_text_tokenizer + if cached is None: + model_name_or_path = self.engine_client.model_config.model + tokenizer_path = getattr(self.engine_client.model_config, "tokenizer", None) + if tokenizer_path is None: + tokenizer_path = resolve_glm_tts_tokenizer_path(model_name_or_path) + cached = load_glm_tts_tokenizer( + tokenizer_path, + model_name_or_path=model_name_or_path, + trust_remote_code=bool(getattr(self.engine_client.model_config, "trust_remote_code", False)), + ) + self._glm_tts_text_tokenizer = cached + + frontend = self._glm_tts_text_frontend + if frontend is None: + frontend = GLMTTSTextFrontend() + self._glm_tts_text_frontend = frontend + return cached, frontend + + def _estimate_glm_tts_text_token_len(self, text: str | None, *, add_trailing_space: bool = False) -> int: + """Estimate GLM-TTS normalized text length with the model tokenizer.""" + cached, frontend = self._glm_tts_text_tokenizer_and_frontend() + text = text or "" + normalized = frontend.text_normalize(text) or text + normalized = normalized.strip() + if add_trailing_space and normalized: + normalized = f"{normalized} " + return max(1, len(cached.encode(normalized))) + + def _build_glm_tts_prefill_metadata(self, text: str, prompt_text: str | None) -> dict[str, Any]: + """Build GLM-TTS processor length metadata for additional_information. + + The model preprocess hook runs before postprocess can mirror MM kwargs + into additional_information, so these scalar fields must originate from + the request payload rather than runner-level mm_features. + """ + text_len = self._estimate_glm_tts_text_token_len(text) + prompt_text_len = ( + self._estimate_glm_tts_text_token_len(prompt_text, add_trailing_space=True) if prompt_text else 0 + ) + return { + "glm_tts_text_token_len": [text_len], + "glm_tts_prompt_text_token_len": [prompt_text_len], + "input_len": [prompt_text_len + text_len + 1], + } + # ---- Ming-flash-omni standalone-talker (TTS) helpers ---- def _build_ming_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: @@ -1963,7 +2095,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Resolve uploaded voice for non-Qwen3 models. # Qwen3 TTS has its own uploaded voice handling in _build_tts_params(). has_inline_ref_audio = request.ref_audio is not None - if self._tts_model_type in ("fish_tts", "cosyvoice3", "moss_tts_nano"): + if self._tts_model_type in ("fish_tts", "cosyvoice3", "moss_tts_nano", "glm_tts"): err = self._apply_uploaded_speaker(request) if err: raise ValueError(err) @@ -2041,6 +2173,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): elif self._tts_model_type == "cosyvoice3": prompt = await self._build_cosyvoice3_prompt(request, has_inline_ref_audio=has_inline_ref_audio) tts_params = {} + elif self._tts_model_type == "glm_tts": + prompt = await self._build_glm_tts_prompt(request, has_inline_ref_audio=has_inline_ref_audio) + tts_params = {} elif self._tts_model_type == "ming_flash_omni_tts": prompt = self._build_ming_prompt(request) tts_params = {} @@ -2112,6 +2247,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): model_type = "ming_flash_omni_tts" elif self._tts_model_type == "moss_tts_nano": model_type = "moss_tts_nano" + elif self._tts_model_type == "glm_tts": + model_type = "glm_tts" elif self._is_tts: model_type = tts_params.get("task_type", ["unknown"])[0] else: @@ -2129,6 +2266,57 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if self._tts_model_type == "cosyvoice3" and sampling_params_list: sampling_params_list = self._apply_cosyvoice3_dynamic_tokens(sampling_params_list, request) + # GLM-TTS: set dynamic min/max tokens based on text length. + if self._tts_model_type == "glm_tts" and sampling_params_list: + import copy + + sampling_params_list = copy.deepcopy(sampling_params_list) + glm_metadata = prompt.get("additional_information") if isinstance(prompt, dict) else None + text_len_value = None + if isinstance(glm_metadata, dict): + text_len_value = glm_metadata.get("glm_tts_text_token_len") + if isinstance(text_len_value, list) and text_len_value: + text_len_value = text_len_value[0] + text_token_len = ( + int(text_len_value) + if text_len_value is not None + else self._estimate_glm_tts_text_token_len(request.input) + ) + hf_cfg = self.model_config.hf_config + min_ratio = getattr(hf_cfg, "min_token_text_ratio", 2) + max_ratio = getattr(hf_cfg, "max_token_text_ratio", 20) + stage_min_tokens = getattr(sampling_params_list[0], "min_tokens", None) + stage_max_tokens = getattr(sampling_params_list[0], "max_tokens", None) + cap_candidates = [int(cap) for cap in (stage_max_tokens, request.max_new_tokens) if cap is not None] + hard_cap = min(cap_candidates) if cap_candidates else None + + min_tokens = max(1, int(text_token_len * min_ratio)) + if stage_min_tokens is not None: + min_tokens = max(min_tokens, int(stage_min_tokens)) + if hard_cap is not None: + min_tokens = min(min_tokens, hard_cap) + + max_tokens = max(min_tokens, int(text_token_len * max_ratio)) + if hard_cap is not None: + max_tokens = min(max_tokens, hard_cap) + sampling_params_list[0].min_tokens = min_tokens + sampling_params_list[0].max_tokens = max_tokens + seed = getattr(request, "seed", None) + if seed is not None: + sampling_params_list[0].seed = seed + logger.info( + "GLM-TTS dynamic tokens: text_tokens=%d, min_ratio=%s, max_ratio=%s, " + "stage_min=%s, stage_max=%s, request_max=%s, min_tokens=%d, max_tokens=%d", + text_token_len, + min_ratio, + max_ratio, + stage_min_tokens, + stage_max_tokens, + request.max_new_tokens, + min_tokens, + max_tokens, + ) + # Apply model-specific extra parameters if request.extra_params is not None and sampling_params_list: if not isinstance(request.extra_params, dict): diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index de7eb5f4c..86a0d8182 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -329,11 +329,16 @@ def resolve_model_config_path(model: str) -> str: except Exception as e: raise ValueError(f"Failed to read config.json for model: {model}. Error: {e}") from e else: - raise ValueError( - f"Could not determine model_type for model: {model}. " - f"Model is not in standard transformers format and does not have model_index.json. " - f"Please ensure the model has proper configuration files with 'model_type' field" - ) + # No config.json at repo root (e.g. GLM-TTS stores configs in + # subdirectories only). Try matching against registered deploy + # YAML filenames before giving up. + model_type = _try_resolve_omni_model_type(model) + if model_type is None: + raise ValueError( + f"Could not determine model_type for model: {model}. " + f"Model is not in standard transformers format and does not have model_index.json. " + f"Please ensure the model has proper configuration files with 'model_type' field" + ) default_config_path = current_omni_platform.get_default_stage_config_path() if model_type in _DIFFUSERS_CLASS_TO_CONFIG: diff --git a/vllm_omni/model_executor/models/common/nucleus_ras_sampling.py b/vllm_omni/model_executor/models/common/nucleus_ras_sampling.py new file mode 100644 index 000000000..34573ae7d --- /dev/null +++ b/vllm_omni/model_executor/models/common/nucleus_ras_sampling.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared TTS sampling primitives: nucleus (top-p/top-k) and RAS. + +Used by CosyVoice3 and GLM-TTS (and any future TTS model with RAS). +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch + + +def multinomial_sample( + probs: torch.Tensor, + generator: torch.Generator | None = None, +) -> torch.Tensor: + """Draw one sample from a categorical distribution.""" + return torch.multinomial(probs, 1, replacement=True, generator=generator).reshape(()) + + +def nucleus_sample_one( + weighted_scores: torch.Tensor, + *, + top_p: float, + top_k: int, + generator: torch.Generator | None = None, +) -> int: + """Sample one token using nucleus (top-p + top-k) filtering. + + ``weighted_scores`` should be log-softmax-ed logits (for RAS callers) + or raw logits (softmax is applied internally). + + Vectorized via ``torch.cumsum`` — no Python loop over the vocabulary. + """ + # Apply top-k before sort. GLM-TTS and CosyVoice3 usually run with + # top_k=25, so sorting the whole audio vocabulary on every AR step is + # avoidable work in the hottest path. Keep probabilities normalized over + # the full vocabulary so the top-p cumulative mass matches the old path. + if top_k > 0 and int(top_k) < int(weighted_scores.numel()): + top_scores, sorted_idx = torch.topk( + weighted_scores, + k=int(top_k), + dim=0, + largest=True, + sorted=True, + ) + sorted_prob = (top_scores - torch.logsumexp(top_scores, dim=0)).exp() + else: + probs = weighted_scores.softmax(dim=0) + sorted_prob, sorted_idx = probs.sort(descending=True, stable=True) + + # Apply top-p (nucleus) filtering: keep the smallest prefix whose + # cumulative probability exceeds top_p, always including the first token. + # Use float32 for cumulative sum to match the original Python float64 + # accumulation loop. Low-precision (fp16/bf16) cumsum over ~32K vocab + # causes rounding drift that shifts the top-p boundary, changing which + # tokens are kept and ultimately producing different AR sequences. + sorted_prob_f32 = sorted_prob.float() + cum_prob = sorted_prob_f32.cumsum(dim=0) + # Mask: include tokens where cumsum *before* this token is still < top_p + mask = (cum_prob - sorted_prob_f32) < top_p + # Safety: ensure top token is always kept. With sorted descending + # probabilities, mask[0] is always True when top_p > 0 (since the + # cumsum-before-self of the first element is 0 < top_p). Setting it + # unconditionally avoids a GPU→CPU sync from the old `mask.any()` check. + mask[0] = True + + kept_probs = torch.nan_to_num(sorted_prob[mask], nan=0.0, posinf=0.0, neginf=0.0).clamp_min_(0.0) + kept_indices = sorted_idx[mask] + # ``torch.multinomial`` accepts unnormalized weights, but not an all-zero + # vector. Handle the extreme underflow case on-device by forcing the top + # retained token to be selectable without adding another CPU sync. + kept_probs[0] = torch.where( + kept_probs.sum() > 0, + kept_probs[0], + torch.ones_like(kept_probs[0]), + ) + + sample_pos = multinomial_sample(kept_probs, generator=generator) + return int(kept_indices[sample_pos].item()) + + +def ras_sample_one( + weighted_scores: torch.Tensor, + decoded_tokens: Sequence[int], + *, + top_p: float, + top_k: int, + win_size: int, + tau_r: float, + generator: torch.Generator | None = None, +) -> int: + """Repetition-Aware Sampling following GLM-TTS/CosyVoice. + + If the nucleus-sampled token appears too often in the recent window, + upstream samples once from the original full distribution instead of + masking the repeated token. This matters for GLM-TTS because masking the + dominant repeated token can artificially raise EOA probability right after + the minimum length guard is lifted. + """ + + def _random_sample_one() -> int: + return int(multinomial_sample(weighted_scores.softmax(dim=0), generator=generator).item()) + + top_id = nucleus_sample_one( + weighted_scores, + top_p=top_p, + top_k=top_k, + generator=generator, + ) + if win_size > 0 and decoded_tokens: + recent = decoded_tokens[-win_size:] + rep_num = sum(1 for token in recent if int(token) == top_id) + if rep_num >= win_size * tau_r: + top_id = _random_sample_one() + return top_id diff --git a/vllm_omni/model_executor/models/common/whisper_vq.py b/vllm_omni/model_executor/models/common/whisper_vq.py new file mode 100644 index 000000000..9d7af76b5 --- /dev/null +++ b/vllm_omni/model_executor/models/common/whisper_vq.py @@ -0,0 +1,291 @@ +# Copyright 2022 The OpenAI Authors and The HuggingFace Inc. team. All rights reserved. +# 2025 Zhipu AI Inc (authors: CogAudio Group Members) +# SPDX-License-Identifier: Apache-2.0 +"""WhisperVQEncoder: HF WhisperEncoder + VQ codebook + inter-layer pooling. + +Built on standard ``WhisperConfig``. VQ-specific parameters are patched onto +the config object after ``from_pretrained`` so that callers never need a +separate config class. GLM-TTS (and future VQ-based TTS models) only need:: + + from transformers import WhisperConfig + from vllm_omni.model_executor.models.common.whisper_vq import WhisperVQEncoder + + cfg = WhisperConfig.from_pretrained(checkpoint_dir) + cfg.pooling_kernel_size = None # passed through from checkpoint + cfg.pooling_type = "max" # or set by caller + cfg.pooling_position = 0 + cfg.quantize_vocab_size = 32768 + cfg.quantize_position = 16 + cfg.quantize_encoder_only = True + model = WhisperVQEncoder(cfg) +""" + +from __future__ import annotations + +import math +import re +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F +from torch import Tensor, nn +from transformers import WhisperConfig +from transformers.modeling_outputs import BaseModelOutput +from transformers.models.whisper.modeling_whisper import WhisperEncoder + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +@dataclass +class QuantizedBaseModelOutput(BaseModelOutput): + quantized_token_ids: torch.LongTensor | None = None + + +def vector_quantize(inputs: Tensor, codebook: Tensor) -> tuple[Tensor, Tensor, Tensor]: + """Nearest-neighbour codebook lookup.""" + embedding_size = codebook.size(1) + inputs_flatten = inputs.reshape(-1, embedding_size) + codebook_sqr = torch.sum(codebook**2, dim=1) + inputs_sqr = torch.sum(inputs_flatten**2, dim=1, keepdim=True) + distances = torch.addmm( + codebook_sqr + inputs_sqr, + inputs_flatten, + codebook.t(), + alpha=-2.0, + beta=1.0, + ) + _, indices_flatten = torch.min(distances, dim=1) + codes_flatten = torch.index_select(codebook, dim=0, index=indices_flatten) + codes = codes_flatten.view_as(inputs) + return codes, indices_flatten, distances + + +_LEGACY_LAYER_REMAP = ( + (re.compile(r"^layers\.(\d+)\.attn\.query\."), r"layers.\1.self_attn.q_proj."), + (re.compile(r"^layers\.(\d+)\.attn\.key\."), r"layers.\1.self_attn.k_proj."), + (re.compile(r"^layers\.(\d+)\.attn\.value\."), r"layers.\1.self_attn.v_proj."), + (re.compile(r"^layers\.(\d+)\.attn\.out\."), r"layers.\1.self_attn.out_proj."), + (re.compile(r"^layers\.(\d+)\.attn_ln\."), r"layers.\1.self_attn_layer_norm."), + (re.compile(r"^layers\.(\d+)\.mlp_ln\."), r"layers.\1.final_layer_norm."), + (re.compile(r"^layers\.(\d+)\.mlp\.0\."), r"layers.\1.fc1."), + (re.compile(r"^layers\.(\d+)\.mlp\.2\."), r"layers.\1.fc2."), +) + + +def remap_legacy_whisper_vq_state_dict(state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + """Map CogAudio/GLM-TTS fork parameter names onto HF WhisperEncoder names.""" + remapped: dict[str, Tensor] = {} + for name, tensor in state_dict.items(): + new_name = name + if new_name == "embed_positions": + new_name = "embed_positions.weight" + for pattern, replacement in _LEGACY_LAYER_REMAP: + new_name = pattern.sub(replacement, new_name) + remapped[new_name] = tensor + return remapped + + +# --------------------------------------------------------------------------- +# main encoder class +# --------------------------------------------------------------------------- + + +class WhisperVQEncoder(WhisperEncoder): + """HF Whisper encoder with optional VQ codebook and pooling. + + Uses a standard ``WhisperConfig`` with the following VQ-specific attrs + patched on by the caller (or loaded from the checkpoint's config.json): + + * ``pooling_kernel_size`` -- int | None + * ``pooling_type`` -- "max" | "avg" + * ``pooling_position`` -- int (0-based layer index) + * ``quantize_vocab_size`` -- int | None + * ``quantize_position`` -- int (0-based layer index) + * ``quantize_encoder_only`` -- bool + """ + + def __init__(self, config: WhisperConfig): + super().__init__(config) + embed_dim = config.d_model + max_source_positions = config.max_source_positions + + # Truncate layers when quantize_encoder_only is set. + qpos = int(getattr(config, "quantize_position", 0) or 0) + if getattr(config, "quantize_encoder_only", False): + self.layers = nn.ModuleList(list(self.layers[:qpos])) + self.layer_norm = None + else: + qpos = max(qpos, int(config.encoder_layers)) # never exceed total + + # Optional inter-layer pooling. + self.pooling_layer: nn.Module | None = None + ksize = getattr(config, "pooling_kernel_size", None) + if ksize is not None: + ptype = str(getattr(config, "pooling_type", "max") or "max") + self.pooling_layer = { + "max": nn.MaxPool1d(kernel_size=int(ksize)), + "avg": nn.AvgPool1d(kernel_size=int(ksize)), + }[ptype] + + # Optional vector-quantization codebook. + self.codebook: nn.Embedding | None = None + self.embed_positions2: nn.Embedding | None = None + vq_sz = getattr(config, "quantize_vocab_size", None) + if vq_sz is not None: + self.codebook = nn.Embedding(int(vq_sz), embed_dim) + pos2_len = max_source_positions + if ksize is not None: + pos2_len = math.ceil(max_source_positions / int(ksize)) + self.embed_positions2 = nn.Embedding(pos2_len, embed_dim) + + # cached for forward() + self._pooling_position = int(getattr(config, "pooling_position", 0) or 0) + self._quantize_position = qpos if vq_sz is not None else -1 + + # ------------------------------------------------------------------ + # forward / helpers + # ------------------------------------------------------------------ + + @property + def device(self) -> torch.device: # type: ignore[override] + return self.conv1.weight.device + + @property + def dtype(self) -> torch.dtype: # type: ignore[override] + return self.conv1.weight.dtype + + def load_state_dict( + self, + state_dict: dict[str, Tensor], + strict: bool = True, + assign: bool = False, + ): # type: ignore[override] + remapped = remap_legacy_whisper_vq_state_dict(state_dict) + try: + return super().load_state_dict(remapped, strict=strict, assign=assign) + except TypeError: + return super().load_state_dict(remapped, strict=strict) + + def forward( + self, + input_features: Tensor, + attention_mask: Tensor | None = None, + **_: Any, + ) -> QuantizedBaseModelOutput: + batch_size, _, _ = input_features.shape + hidden_states = F.gelu(self.conv1(input_features)) + hidden_states = F.gelu(self.conv2(hidden_states)) + hidden_states = hidden_states.permute(0, 2, 1) + seq_len = int(hidden_states.shape[1]) + + pos_embed = self._position_embeddings(seq_len, hidden_states.device, hidden_states.dtype) + hidden_states = hidden_states + pos_embed + + valid_mask = self._downsample_attention_mask(attention_mask, batch_size, seq_len, hidden_states.device) + hidden_states = hidden_states.masked_fill(~valid_mask.unsqueeze(-1), 0) + + quantized_token_ids = None + for idx, layer in enumerate(self.layers): + layer_mask = self._make_layer_attention_mask(valid_mask, hidden_states.dtype) + layer_outputs = layer( + hidden_states, + attention_mask=layer_mask, + layer_head_mask=None, + output_attentions=False, + ) + hidden_states = layer_outputs[0] if isinstance(layer_outputs, tuple) else layer_outputs + hidden_states = hidden_states.masked_fill(~valid_mask.unsqueeze(-1), 0) + + if idx + 1 == self._pooling_position and self.pooling_layer is not None: + hidden_states, valid_mask = self._apply_pooling(hidden_states, valid_mask) + + if idx + 1 == self._quantize_position and self.codebook is not None: + hidden_states, quantized_token_ids = self._apply_vq(hidden_states, valid_mask) + + if self.layer_norm is not None: + hidden_states = self.layer_norm(hidden_states) + hidden_states = hidden_states.masked_fill(~valid_mask.unsqueeze(-1), 0) + + return QuantizedBaseModelOutput( + last_hidden_state=hidden_states, + quantized_token_ids=quantized_token_ids, + ) + + def _position_embeddings(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> Tensor: + ep = self.embed_positions + if isinstance(ep, nn.Embedding): + positions = torch.arange(seq_len, device=device) + return ep(positions).to(dtype=dtype).unsqueeze(0) + return ep[:seq_len].to(device=device, dtype=dtype).unsqueeze(0) + + def _downsample_attention_mask( + self, + attention_mask: Tensor | None, + batch_size: int, + seq_len: int, + device: torch.device, + ) -> Tensor: + if attention_mask is None: + return torch.ones(batch_size, seq_len, dtype=torch.bool, device=device) + conv_stride = int(self.conv1.stride[0]) * int(self.conv2.stride[0]) + downsampled = attention_mask[:, ::conv_stride].to(device=device, dtype=torch.bool) + if downsampled.shape[1] < seq_len: + pad = seq_len - int(downsampled.shape[1]) + downsampled = F.pad(downsampled, (0, pad), value=False) + return downsampled[:, :seq_len] + + @staticmethod + def _make_layer_attention_mask(valid_mask: Tensor, dtype: torch.dtype) -> Tensor: + min_value = torch.finfo(dtype).min + additive = torch.zeros(valid_mask.shape, dtype=dtype, device=valid_mask.device) + additive = additive.masked_fill(~valid_mask, min_value) + return additive[:, None, None, :] + + def _apply_pooling(self, hidden_states: Tensor, valid_mask: Tensor) -> tuple[Tensor, Tensor]: + k = int(self.config.pooling_kernel_size) # type: ignore[attr-defined] + pooled_segments: list[Tensor] = [] + pooled_masks: list[Tensor] = [] + for batch_idx in range(int(hidden_states.shape[0])): + length = int(valid_mask[batch_idx].sum().item()) + segment = hidden_states[batch_idx, :length] + if segment.numel() == 0: + pooled = hidden_states.new_zeros((0, hidden_states.shape[-1])) + else: + segment_t = segment.unsqueeze(0).transpose(1, 2) + if int(segment_t.shape[-1]) % k != 0: + segment_t = F.pad(segment_t, (0, k - int(segment_t.shape[-1]) % k)) + pooled = self.pooling_layer(segment_t).transpose(1, 2).squeeze(0) # type: ignore[misc] + pooled_segments.append(pooled) + pooled_masks.append(torch.ones(pooled.shape[0], dtype=torch.bool, device=hidden_states.device)) + + max_len = max((int(seg.shape[0]) for seg in pooled_segments), default=0) + output = hidden_states.new_zeros((hidden_states.shape[0], max_len, hidden_states.shape[-1])) + output_mask = torch.zeros(hidden_states.shape[0], max_len, dtype=torch.bool, device=hidden_states.device) + for i, segment in enumerate(pooled_segments): + if segment.shape[0] > 0: + output[i, : segment.shape[0]] = segment + output_mask[i, : segment.shape[0]] = pooled_masks[i] + return output, output_mask + + def _apply_vq(self, hidden_states: Tensor, valid_mask: Tensor) -> tuple[Tensor, Tensor | None]: + output = hidden_states.clone() + token_ids = torch.full( + (hidden_states.shape[0], hidden_states.shape[1]), + -1, + dtype=torch.long, + device=hidden_states.device, + ) + for batch_idx in range(int(hidden_states.shape[0])): + length = int(valid_mask[batch_idx].sum().item()) + if length <= 0: + continue + segment = hidden_states[batch_idx : batch_idx + 1, :length] + quantized, indices, _ = vector_quantize(segment, self.codebook.weight) # type: ignore[union-attr] + if self.embed_positions2 is not None: + quantized = quantized + self.embed_positions2.weight[:length].to(dtype=quantized.dtype).unsqueeze(0) + output[batch_idx, :length] = quantized.squeeze(0) + token_ids[batch_idx, :length] = indices.reshape(-1) + return output.masked_fill(~valid_mask.unsqueeze(-1), 0), token_ids diff --git a/vllm_omni/model_executor/models/glm_tts/__init__.py b/vllm_omni/model_executor/models/glm_tts/__init__.py new file mode 100644 index 000000000..9618227ca --- /dev/null +++ b/vllm_omni/model_executor/models/glm_tts/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-TTS model support for vLLM-Omni. + +GLM-TTS is a two-stage text-to-speech system: + - Stage 0 (LLM): Llama-based AR model generates speech tokens from text + - Stage 1 (DiT): Flow matching model converts speech tokens to mel-spectrogram + +Reference: https://github.com/zai-org/GLM-TTS +""" + +from vllm_omni.model_executor.models.glm_tts.glm_tts import ( + GLMTTSForConditionalGeneration, +) + +__all__ = [ + "GLMTTSForConditionalGeneration", +] diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts.py b/vllm_omni/model_executor/models/glm_tts/glm_tts.py new file mode 100644 index 000000000..6e14840e2 --- /dev/null +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts.py @@ -0,0 +1,1423 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-TTS AR Model (Stage 0): Text → Speech Tokens. + +Based on Llama architecture, generates speech token sequences from input text. +Analogous to Fish Speech Slow AR and Qwen3-TTS Talker models. +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import replace +from typing import Any + +import torch +import torch.nn as nn +from transformers import AutoTokenizer +from transformers.feature_extraction_utils import BatchFeature +from transformers.utils.hub import cached_file +from vllm.config import VllmConfig +from vllm.config.multimodal import BaseDummyOptions +from vllm.distributed import get_pp_group +from vllm.inputs import MultiModalDataDict +from vllm.logger import init_logger +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.interfaces import SupportsMultiModal +from vllm.model_executor.models.llama import LlamaModel +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + WeightsMapper, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems +from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + ProcessorInputs, + PromptIndexTargets, + PromptInsertion, + PromptUpdate, +) +from vllm.sequence import IntermediateTensors +from vllm.v1.outputs import SamplerOutput +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.sample.sampler import Sampler + +from vllm_omni.model_executor.models.common.nucleus_ras_sampling import ras_sample_one as _ras_sample_one +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.transformers_utils.configs.glm_tts import GLMTTSConfig +from vllm_omni.utils.speaker_cache import get_speaker_cache + +from .text_frontend import GLMTTSTextFrontend +from .voice_clone import ( + extract_prompt_feat, + extract_prompt_speech_token, + extract_spk_embedding, + load_voice_clone_frontend, +) + +logger = init_logger(__name__) + +_GLM_TTS_DEFAULT_REPO_ID = "zai-org/GLM-TTS" +_GLM_TTS_TOKENIZER_SUBDIR = "vq32k-phoneme-tokenizer" +_GLM_TTS_MAX_PROMPT_SPEECH_TOKENS = 1024 + + +def _req_float(param: torch.Tensor | None, req_idx: int, default: float) -> float: + if param is None or param.numel() == 0: + return default + index = min(req_idx, int(param.numel()) - 1) + return float(param.reshape(-1)[index].item()) + + +def _tensor_param_values( + param: torch.Tensor | None, + count: int, + default: float, +) -> list[float]: + if param is None or param.numel() == 0: + return [default] * count + values = param.reshape(-1).detach().cpu().tolist() + if not values: + return [default] * count + if len(values) < count: + values.extend([values[-1]] * (count - len(values))) + return [float(value) for value in values[:count]] + + +def resolve_glm_tts_tokenizer_path(model_name_or_path: Any) -> str: + """Resolve vq32k-phoneme-tokenizer directory to a local path. + + For local model dirs, scans for the tokenizer subfolder on disk. + For HF hub names, downloads the tokenizer subfolder via + ``snapshot_download`` (same pattern as CosyVoice3's + ``_ensure_cached_runtime_components`` and + ``arg_utils._TOKENIZER_SUBFOLDER_MAP``). + """ + model_path = os.fspath(model_name_or_path) + if os.path.exists(model_path): + for candidate in [ + os.path.join(model_path, _GLM_TTS_TOKENIZER_SUBDIR), + os.path.join(os.path.dirname(model_path), _GLM_TTS_TOKENIZER_SUBDIR), + ]: + if os.path.isdir(candidate): + return candidate + return model_path + + from huggingface_hub import snapshot_download + + local_dir = snapshot_download( + model_path, + allow_patterns=[ + f"{_GLM_TTS_TOKENIZER_SUBDIR}/tokenizer*", + f"{_GLM_TTS_TOKENIZER_SUBDIR}/tokenization_*", + f"{_GLM_TTS_TOKENIZER_SUBDIR}/special_tokens*", + f"{_GLM_TTS_TOKENIZER_SUBDIR}/vocab*", + f"{_GLM_TTS_TOKENIZER_SUBDIR}/merges*", + f"{_GLM_TTS_TOKENIZER_SUBDIR}/added_tokens*", + ], + ) + candidate = os.path.join(local_dir, _GLM_TTS_TOKENIZER_SUBDIR) + if os.path.isdir(candidate): + return candidate + return model_path + + +def resolve_glm_tts_model_dir( + model_name_or_path: Any, + *, + tokenizer_path: Any | None = None, + required_files: Sequence[str] = (), + optional_files: Sequence[str] = (), +) -> str: + """Resolve GLM-TTS repo root: local dir → tokenizer parent → cached_file → snapshot.""" + model_path = os.fspath(model_name_or_path) + if os.path.isdir(model_path): + return os.path.abspath(model_path) + + def _has_files(root: str) -> bool: + has_req = all(os.path.isfile(os.path.join(root, *f.split("/"))) for f in required_files) + has_opt = not optional_files or any(os.path.isfile(os.path.join(root, *f.split("/"))) for f in optional_files) + return has_req and has_opt + + # Try tokenizer parent directory + if tokenizer_path: + tp = os.path.abspath(os.path.normpath(os.fspath(tokenizer_path))) + root = os.path.dirname(tp) if os.path.basename(tp) == _GLM_TTS_TOKENIZER_SUBDIR else tp + if os.path.isdir(root) and _has_files(root): + return root + + # Try cached_file to locate individual files without full snapshot + for filename in list(required_files) + list(optional_files): + try: + resolved = cached_file(model_name_or_path, filename) + if resolved: + root = os.path.abspath(resolved) + for _ in filename.split("/"): + root = os.path.dirname(root) + if _has_files(root): + return root + except Exception: + pass + + from huggingface_hub import snapshot_download + + return snapshot_download(model_name_or_path) + + +def _first_glm_tts_value(value: Any) -> Any: + return value[0] if isinstance(value, list) and value else value + + +def _glm_tts_int_value(value: Any) -> int | None: + """Extract a scalar integer from vLLM multimodal/runtime wrappers.""" + value = _first_glm_tts_value(value) + if isinstance(value, torch.Tensor): + if value.numel() == 0: + return None + return int(value.reshape(-1)[0].item()) + if isinstance(value, (int, float)): + return int(value) + return None + + +def _decode_glm_tts_audio_data(value: Any) -> tuple[torch.Tensor | None, int | None]: + """Extract (waveform_tensor, sample_rate) from multimodal input.""" + if isinstance(value, list) and len(value) == 1: + value = value[0] + if isinstance(value, torch.Tensor): + return value, None + if isinstance(value, (tuple, list)) and len(value) == 2 and isinstance(value[1], (int, float)): + return torch.as_tensor(value[0]), int(value[1]) + if hasattr(value, "shape"): # numpy array + return torch.as_tensor(value), None + return None, None + + +_glm_tts_tokenizer_cache: dict[str, Any] = {} + + +def load_glm_tts_tokenizer( + tokenizer_path: Any, + *, + model_name_or_path: Any | None = None, + trust_remote_code: bool = False, + **kwargs: Any, +) -> Any: + """Load GLM-TTS phoneme tokenizer (ChatGLM4Tokenizer, slow-only). + + Expects a resolved local directory (use ``resolve_glm_tts_tokenizer_path`` + to handle HF hub names first). Results are cached by *tokenizer_path*. + """ + cache_key = str(tokenizer_path) + cached = _glm_tts_tokenizer_cache.get(cache_key) + if cached is not None: + return cached + + tokenizer = AutoTokenizer.from_pretrained( + tokenizer_path, + use_fast=False, + trust_remote_code=trust_remote_code, + **kwargs, + ) + _glm_tts_tokenizer_cache[cache_key] = tokenizer + return tokenizer + + +def get_glm_tts_special_token_ids(tokenizer: Any) -> dict[str, int]: + """Resolve GLM-TTS special IDs from the phoneme/audio tokenizer.""" + special_tokens = { + "ats": "<|audio_0|>", + "ate": "<|audio_32767|>", + "boa": "<|begin_of_audio|>", + "eoa": "<|user|>", + "pad": "<|endoftext|>", + } + + result: dict[str, int] = {} + for key, token_str in special_tokens.items(): + token_ids = tokenizer.encode(token_str, add_special_tokens=False) + if len(token_ids) != 1: + raise ValueError(f"Token '{key}' ({token_str}) should encode to single ID, got: {token_ids}") + result[key] = int(token_ids[0]) + return result + + +def resolve_glm_tts_campplus_path(model_dir: str) -> str | None: + """Resolve ``campplus.onnx`` for GLM-TTS voice cloning.""" + local = os.path.join(model_dir, "frontend", "campplus.onnx") + if os.path.isfile(local): + return local + + try: + resolved = cached_file("FunAudioLLM/CosyVoice-300M", "campplus.onnx") + if resolved is not None: + logger.info("Resolved campplus.onnx from FunAudioLLM/CosyVoice-300M: %s", resolved) + return resolved + except Exception: + logger.debug("cached_file could not fetch campplus.onnx", exc_info=True) + + logger.warning( + "campplus.onnx not found locally or could not be downloaded. " + "Voice cloning speaker embedding will not be available.", + ) + return None + + +def _normalize_glm_tts_processor_text( + frontend: GLMTTSTextFrontend, + value: str | None, + *, + add_trailing_space: bool = False, +) -> str: + if value is None: + return "" + normalized = frontend.text_normalize(value) + normalized = (normalized or value).strip() + if add_trailing_space and normalized: + normalized = f"{normalized} " + return normalized + + +_glm_tts_text_frontend: GLMTTSTextFrontend | None = None + + +def _get_text_frontend() -> GLMTTSTextFrontend: + global _glm_tts_text_frontend + if _glm_tts_text_frontend is None: + _glm_tts_text_frontend = GLMTTSTextFrontend() + return _glm_tts_text_frontend + + +def build_glm_tts_prefill_metadata( + model_name_or_path: Any, + text: str, + prompt_text: str | None, + *, + tokenizer_path: Any | None = None, + trust_remote_code: bool = False, +) -> dict[str, list[int]]: + """Build request-local GLM-TTS length metadata for AR prefill. + + These scalars are mirrored into ``model_intermediate_buffer`` so the model + preprocess hook can initialize request state from the metadata alone. + """ + if tokenizer_path is None: + tokenizer_path = resolve_glm_tts_tokenizer_path(model_name_or_path) + tokenizer = load_glm_tts_tokenizer( + tokenizer_path, + model_name_or_path=model_name_or_path, + trust_remote_code=trust_remote_code, + ) + frontend = _get_text_frontend() + text_len = max(1, len(tokenizer.encode(_normalize_glm_tts_processor_text(frontend, text)))) + prompt_text_len = 0 + if prompt_text: + prompt_text_len = max( + 1, + len( + tokenizer.encode( + _normalize_glm_tts_processor_text( + frontend, + prompt_text, + add_trailing_space=True, + ) + ) + ), + ) + return { + "glm_tts_text_token_len": [text_len], + "glm_tts_prompt_text_token_len": [prompt_text_len], + "input_len": [prompt_text_len + text_len + 1], + } + + +class GLMTTSMultiModalProcessingInfo(BaseProcessingInfo): + def get_hf_config(self): + return self.ctx.get_hf_config(GLMTTSConfig) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"audio": 1} + + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> Mapping[str, int] | None: + if mm_counts.get("audio", 0) <= 0: + return {} + return {"audio": min(seq_len, _GLM_TTS_MAX_PROMPT_SPEECH_TOKENS)} + + def get_data_parser(self): + return MultiModalDataParser( + target_sr=24000, + expected_hidden_size=self._get_expected_hidden_size(), + ) + + +class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessingInfo]): + """GLM-TTS voice-clone processor. + + Unlike CosyVoice3, GLM-TTS prompt speech tokens are normal Llama vocab IDs + (``<|audio_N|>``). The processor exposes them as multimodal embeddings for + the AR prompt and also carries WhisperVQ/CampPlus outputs to the AR->DiT + handoff. + """ + + def _ensure_cached_runtime_components( + self, + model_dir: str, + config: GLMTTSConfig, + *, + load_voice_clone: bool = True, + ) -> None: + # Skip if already loaded for this model_dir + mode + cache_key = (model_dir, load_voice_clone) + if getattr(self, "_loaded_cache_key", None) == cache_key: + return + + tokenizer_path = getattr(self.info.ctx.model_config, "tokenizer", None) + required = ("speech_tokenizer/config.json", "speech_tokenizer/model.safetensors") if load_voice_clone else () + model_dir = resolve_glm_tts_model_dir( + model_dir, + tokenizer_path=tokenizer_path, + required_files=required, + ) + + # Load text tokenizer + frontend (always needed) + if getattr(self, "_text_model_dir", None) != model_dir: + tok_path = tokenizer_path + if not tok_path or not os.path.isdir(os.fspath(tok_path)): + tok_path = os.path.join(model_dir, _GLM_TTS_TOKENIZER_SUBDIR) + if not os.path.isdir(tok_path): + tok_path = model_dir + trust_remote_code = bool(getattr(self.info.ctx.model_config, "trust_remote_code", False)) + self.tokenizer = load_glm_tts_tokenizer( + tok_path, + model_name_or_path=self.info.ctx.model_config.model, + trust_remote_code=trust_remote_code, + ) + self.special_ids = get_glm_tts_special_token_ids(self.tokenizer) + self.text_frontend = GLMTTSTextFrontend() + self._text_model_dir = model_dir + + # Load voice clone components (speech tokenizer + CampPlus) + if load_voice_clone: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.processor_device = device + campplus_path = getattr(self, "_campplus_path", None) or resolve_glm_tts_campplus_path(model_dir) + self._campplus_path = campplus_path + self.speech_tokenizer, self.campplus_session = load_voice_clone_frontend( + model_dir, + device, + speech_tokenizer_cache=getattr(self, "speech_tokenizer", None), + campplus_cache=getattr(self, "campplus_session", None), + campplus_path=campplus_path, + ) + + if not hasattr(self, "_speaker_cache"): + self._speaker_cache = get_speaker_cache() + + self._loaded_cache_key = cache_key + + def _encode_text(self, text: str) -> torch.Tensor: + token_ids = self.tokenizer.encode(text) + return torch.tensor([token_ids], dtype=torch.long) + + def _get_audio(self, mm_data: Mapping[str, object]) -> tuple[torch.Tensor | None, int | None]: + audio = mm_data.get("audio") + if audio is None: + audios = mm_data.get("audios") + if isinstance(audios, (list, tuple)) and audios: + # Match CosyVoice3's compatibility behavior and accept the + # first item from the legacy plural field name. + audio = audios[0] + else: + audio = audios + wav, sr = _decode_glm_tts_audio_data(audio) + if wav is not None and wav.ndim > 1: + wav = wav.float() + wav = wav.mean(dim=0) if wav.shape[0] <= wav.shape[-1] else wav.mean(dim=-1) + return wav, int(sr or 24000) if wav is not None else sr + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + config = self.info.ctx.get_hf_config() + model_dir = self.info.ctx.model_config.model + wav, sr = self._get_audio(mm_data) + if wav is None: + audio_supplied = mm_data.get("audio") is not None or mm_data.get("audios") is not None + assert not audio_supplied, ( + "GLM-TTS received an audio multimodal payload that could not be decoded. " + "Runtime voice cloning must provide decodable multi_modal_data['audio']; " + "the text-only processor path is reserved for profiling/cache." + ) + self._ensure_cached_runtime_components(model_dir, config, load_voice_clone=False) + normalized_text = _normalize_glm_tts_processor_text(self.text_frontend, prompt) + text_ids = self._encode_text(normalized_text) + dummy_audio_len = min( + _GLM_TTS_MAX_PROMPT_SPEECH_TOKENS, + max(1, int(getattr(config, "max_position_embeddings", _GLM_TTS_MAX_PROMPT_SPEECH_TOKENS))), + ) + boa_tensor = torch.tensor([[int(self.special_ids["boa"])]], dtype=torch.long) + input_ids = torch.cat([text_ids, boa_tensor], dim=1) + # vLLM profiling/cache setup may call the processor without an + # audio item. Do not fabricate a dummy wav here: for GLM-TTS that + # would exercise the real voice-clone frontend during profiling. + # This shape-only branch is not a supported runtime inference mode. + return BatchFeature( + { + "input_ids": input_ids, + "prompt_speech_token": torch.zeros((1, dummy_audio_len), dtype=torch.long), + "prompt_speech_token_len": [torch.tensor([dummy_audio_len], dtype=torch.long)], + "glm_tts_prompt_text_token_len": [torch.tensor([0], dtype=torch.long)], + "glm_tts_text_token_len": [torch.tensor([int(text_ids.shape[1])], dtype=torch.long)], + } + ) + + self._ensure_cached_runtime_components(model_dir, config) + normalized_text = _normalize_glm_tts_processor_text(self.text_frontend, prompt) + text_ids = self._encode_text(normalized_text) + prompt_text = mm_kwargs.get("prompt_text") + if not isinstance(prompt_text, str) or not prompt_text.strip(): + raise ValueError("GLM-TTS voice cloning requires mm_processor_kwargs['prompt_text'].") + + normalized_prompt_text = _normalize_glm_tts_processor_text( + self.text_frontend, + prompt_text, + add_trailing_space=True, + ) + prompt_text_ids = self._encode_text(normalized_prompt_text) + + # Speaker cache: skip WhisperVQ + mel + CampPlus on hit + voice_name = mm_kwargs.get("voice_name") + cache_key = None + if voice_name and isinstance(voice_name, str): + cache_key = self._speaker_cache.make_cache_key( + voice_name, + model_type="glm_tts", + created_at=int(mm_kwargs.get("voice_created_at") or 0), + ) + cached = self._speaker_cache.get(cache_key) + if cached is not None: + boa_tensor = torch.tensor([[int(self.special_ids["boa"])]], dtype=torch.long) + input_ids = torch.cat([prompt_text_ids, text_ids, boa_tensor], dim=1) + return BatchFeature( + { + "input_ids": input_ids, + "prompt_speech_token": cached["prompt_speech_token"].clone(), + "prompt_speech_token_len": [cached["prompt_speech_token_len"].clone()], + "glm_tts_prompt_text_token_len": [ + torch.tensor([int(prompt_text_ids.shape[1])], dtype=torch.long) + ], + "prompt_feat": cached["prompt_feat"].clone(), + "embedding": cached["embedding"].clone(), + "glm_tts_text_token_len": [torch.tensor([int(text_ids.shape[1])], dtype=torch.long)], + } + ) + + prompt_speech_token = extract_prompt_speech_token(wav, int(sr or 24000), self.speech_tokenizer) + if not prompt_speech_token: + raise RuntimeError("GLM-TTS failed to extract WhisperVQ prompt speech tokens from ref_audio.") + + prompt_speech_token_tensor = torch.tensor([prompt_speech_token], dtype=torch.long) + boa_tensor = torch.tensor([[int(self.special_ids["boa"])]], dtype=torch.long) + input_ids = torch.cat([prompt_text_ids, text_ids, boa_tensor], dim=1) + logger.info( + "GLM-TTS processor prompt: prompt_text_tokens=%d text_tokens=%d " + "prompt_speech_tokens=%d input_tokens_before_audio=%d expected_prefill_tokens=%d", + int(prompt_text_ids.shape[1]), + int(text_ids.shape[1]), + len(prompt_speech_token), + int(input_ids.shape[1]), + int(input_ids.shape[1]) + len(prompt_speech_token), + ) + + prompt_feat = extract_prompt_feat(wav, int(sr or 24000), self.processor_device) + if prompt_feat is None: + raise RuntimeError("GLM-TTS failed to extract prompt mel features from ref_audio.") + embedding = extract_spk_embedding(wav, int(sr or 24000), self.campplus_session) + if embedding is None: + raise RuntimeError("GLM-TTS failed to extract CampPlus speaker embedding from ref_audio.") + + prompt_speech_token_len = torch.tensor([len(prompt_speech_token)], dtype=torch.long) + prompt_feat_cpu = prompt_feat.detach().to("cpu").unsqueeze(0).contiguous() + embedding_tensor = torch.tensor([embedding], dtype=torch.float32) + + if cache_key is not None: + self._speaker_cache.put( + cache_key, + { + "prompt_speech_token": prompt_speech_token_tensor.detach().cpu(), + "prompt_speech_token_len": prompt_speech_token_len.detach().cpu(), + "prompt_feat": prompt_feat_cpu.detach().cpu(), + "embedding": embedding_tensor.detach().cpu(), + }, + ) + + return BatchFeature( + { + "input_ids": input_ids, + "prompt_speech_token": prompt_speech_token_tensor, + "prompt_speech_token_len": [prompt_speech_token_len], + "glm_tts_prompt_text_token_len": [torch.tensor([int(prompt_text_ids.shape[1])], dtype=torch.long)], + "prompt_feat": prompt_feat_cpu, + "embedding": embedding_tensor, + "glm_tts_text_token_len": [torch.tensor([int(text_ids.shape[1])], dtype=torch.long)], + } + ) + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + return { + key: MultiModalFieldConfig.batched("audio") + for key in ( + "prompt_speech_token", + "prompt_speech_token_len", + "glm_tts_prompt_text_token_len", + "prompt_feat", + "embedding", + "glm_tts_text_token_len", + ) + if key in hf_inputs + } + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return False + + def _cached_apply_hf_processor(self, inputs: ProcessorInputs, timing_ctx: Any): + # GLM-TTS builds the actual AR prompt from both the request text and + # mm_processor_kwargs["prompt_text"]. The base cache path separates text + # processing from audio processing, which drops the reference transcript + # and BOA token from the final prompt. + return self._apply_hf_processor(inputs, timing_ctx) + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + def insertion_end(item_idx: int) -> list[int]: + audio_items = out_mm_kwargs["audio"] + item = audio_items[item_idx] if item_idx < len(audio_items) else audio_items[0] + token_len = item["prompt_speech_token_len"].data[0].item() + return [1] * int(token_len) + + return [ + PromptInsertion( + modality="audio", + target=PromptIndexTargets.end(), + insertion=insertion_end, + ) + ] + + +class GLMTTSDummyInputsBuilder(BaseDummyInputsBuilder[GLMTTSMultiModalProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return "This is a test of the GLM-TTS voice cloning system." + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions] | None = None, + ) -> MultiModalDataDict: + return {} + + def get_dummy_processor_inputs( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions] | None = None, + ) -> ProcessorInputs: + inputs = super().get_dummy_processor_inputs(seq_len, mm_counts, mm_options) + inputs.hf_processor_mm_kwargs = {"prompt_text": "This is the reference voice."} + return inputs + + +@MULTIMODAL_REGISTRY.register_processor( + GLMTTSMultiModalProcessor, + info=GLMTTSMultiModalProcessingInfo, + dummy_inputs=GLMTTSDummyInputsBuilder, +) +class GLMTTSForConditionalGeneration(nn.Module, SupportsMultiModal): + """vLLM model for GLM-TTS. + + Handles both stages via model_stage branching: + - ``glm_tts`` (Stage 0): Text → Speech tokens (LLM AR, Llama backbone). + - ``glm_tts_dit`` (Stage 1): Speech tokens → Audio (DiT flow-matching + vocoder). + + Attributes: + have_multimodal_outputs: Signals scheduler to collect multimodal outputs. + has_preprocess: Model has preprocess hook for input preparation (stage 0 only). + has_postprocess: Model has postprocess hook for hidden state caching (stage 0 only). + """ + + supports_multimodal_raw_input_only = True + supports_multimodal = True + requires_raw_input_tokens = True + prefer_model_sampler = True + _sampling_eps = 1e-5 + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "llama_embedding.": "model.embed_tokens.", + "llama.model.": "model.", + "llama.": "model.", + } + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.vllm_config = vllm_config + self.model_path = vllm_config.model_config.model + self.model_dir = self.model_path + self.model_stage = getattr(vllm_config.model_config, "model_stage", "glm_tts") + + # Load configuration + config: GLMTTSConfig = vllm_config.model_config.hf_config # type: ignore[assignment] + self.config = config + + # ---- Stage 1 (DiT): delegate to GLMTTSDiTForGeneration ---- + if self.model_stage == "glm_tts_dit": + from .glm_tts_dit_wrapper import GLMTTSDiTForGeneration + + self._dit_gen = GLMTTSDiTForGeneration(vllm_config=vllm_config, prefix=prefix) + # Expose the DiT as self.model for weight loading routing + self.model = self._dit_gen + self.have_multimodal_outputs = True + self.enable_update_additional_information = True + # DiT stage does not use preprocess/postprocess/sample + self.has_preprocess = False + self.has_postprocess = False + # DiT loads all weights internally from flow/flow.pt etc. + # Point DefaultModelLoader to flow/flow.pt (a .pt file) so it can + # find and load it without triggering the vLLM#39699 bug where + # subdirectory safetensors patterns cause pt_weights_iterator + # to be used on .safetensors files. load_weights() ignores the + # yielded weights and loads flow/flow.pt itself. + self.allow_patterns_overrides = ["flow/flow.pt"] + self.fall_back_to_pt_during_load = False + return + sample_method = str(getattr(config, "sample_method", "ras")).lower() + if sample_method != "ras": + raise ValueError(f"Unsupported GLM-TTS sample_method={sample_method!r}; expected 'ras'.") + + # Resolve repo root to local path (CosyVoice3 pattern). + # Model weights are in llm/ subdirectory; tokenizer and other resources + # are siblings of llm/ under the repo root. + self.model_dir = resolve_glm_tts_model_dir( + self.model_dir, + tokenizer_path=getattr(vllm_config.model_config, "tokenizer", None), + ) + + # Stage 0 weights live under llm/model-*.safetensors and are loaded + # explicitly by _iter_llm_safetensors() during load_weights(). + # + # We still point DefaultModelLoader at flow/flow.pt as a bootstrap + # sentinel because the current upstream loader path does not reliably + # target subdirectory safetensors patterns here without hitting the + # vLLM#39699 iterator bug. Keep the real stage-0 source of truth in + # _iter_llm_safetensors(); this override is only to get through the + # generic loader bootstrap step. + self.allow_patterns_overrides = ["flow/flow.pt"] + self.fall_back_to_pt_during_load = False + + # Load tokenizer for special token ID resolution. + # Prefer the tokenizer path auto-detected by arg_utils + # (_TOKENIZER_SUBFOLDER_MAP). Fall back to vq32k-phoneme-tokenizer/ + # under model_dir. + tokenizer_path = getattr(vllm_config.model_config, "tokenizer", None) + if not (tokenizer_path and os.path.isdir(tokenizer_path)): + tokenizer_path = os.path.join(self.model_dir, _GLM_TTS_TOKENIZER_SUBDIR) + if not os.path.exists(tokenizer_path): + tokenizer_path = self.model_dir + trust_remote_code = bool(getattr(vllm_config.model_config, "trust_remote_code", False)) + self._tokenizer = load_glm_tts_tokenizer( + tokenizer_path, + model_name_or_path=self.model_path, + trust_remote_code=trust_remote_code, + ) + special_ids = get_glm_tts_special_token_ids(self._tokenizer) + self._ats = special_ids["ats"] + self._ate = special_ids["ate"] + self._boa = special_ids["boa"] + self._eoa = special_ids["eoa"] + self._pad = special_ids["pad"] + self._bos = self._tokenizer.bos_token_id or self._pad + + logger.info( + "GLM-TTS token IDs: ATS=%d, ATE=%d, BOA=%d, EOA=%d, PAD=%d, vocab_size=%d", + self._ats, + self._ate, + self._boa, + self._eoa, + self._pad, + config.vocab_size, + ) + + # Validate special token sanity with runtime exceptions (asserts can be + # stripped under python -O). + if self._ate - self._ats != 32767: + raise ValueError(f"Audio token range should be 32768, got {self._ate - self._ats + 1}") + if self._boa >= self._ats: + raise ValueError(f"BOA={self._boa} should be < ATS={self._ats} (BOA is text token)") + + # Validate vocab_size covers all special tokens + max_token = max(self._ats, self._ate, self._boa, self._eoa, self._pad) + if max_token >= config.vocab_size: + raise ValueError( + f"vocab_size ({config.vocab_size}) must be > max token ID ({max_token}). " + f"Check model's config.json has correct vocab_size." + ) + + # Update config with dynamic token IDs so vLLM uses correct eos_token + # This enables proper stop detection when EOA is sampled + config.eos_token_id = self._eoa + config.eoa_token_id = self._eoa + config.audio_token_start = self._ats + config.audio_token_end = self._ate + config.boa_token_id = self._boa + config.pad_token_id = self._pad + config.bos_token_id = self._bos + + # Model flags for AR scheduler + self.have_multimodal_outputs = True + self.has_preprocess = True + self.has_postprocess = True + # GLM-TTS async chunk transfer consumes speech_tokens plus voice-clone + self.gpu_resident_buffer_keys: set[tuple[str, str]] = {("last_hidden", "last")} + + self.model = LlamaModel(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) + + # LM head for speech token prediction + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors + + # Runtime validation: dynamic EOA must match the hardcoded pipeline + # stop_token_ids constant (59253). A mismatch means the upstream + # tokenizer vocabulary has changed and pipeline.py needs updating. + _PIPELINE_EOA = 59253 + if self._eoa != _PIPELINE_EOA: + logger.warning( + "GLM-TTS EOA token mismatch: tokenizer resolved %d but " + "pipeline.py hardcodes stop_token_ids=[%d]. Update " + "pipeline.py to match the current checkpoint.", + self._eoa, + _PIPELINE_EOA, + ) + + # RAS sampling config — stored as attributes for sample() + self._ras_win_size = int(getattr(config, "ras_win_size", 10)) + self._ras_tau_r = float(getattr(config, "ras_tau_r", 0.1)) + self._ras_top_p = float(getattr(config, "ras_top_p", 0.8)) + self._ras_top_k = int(getattr(config, "ras_top_k", 25)) + + def _model_dtype(self) -> torch.dtype: + """Return the active parameter dtype for locally-created embeddings.""" + try: + return next(self.model.parameters()).dtype + except StopIteration: + dtype = getattr(self.vllm_config.model_config, "dtype", None) + if isinstance(dtype, torch.dtype): + return dtype + if isinstance(dtype, str): + _aliases = {"bf16": "bfloat16", "fp16": "float16", "half": "float16", "fp32": "float32"} + resolved = getattr(torch, _aliases.get(dtype.lower(), dtype.lower()), None) + if resolved is not None: + return resolved + return torch.get_default_dtype() + + def _resolve_prefill_target_text_len( + self, + span_len: int, + info_dict: Mapping[str, Any], + ) -> int | None: + """Resolve the target-text token count before the first AR sample. + + The request builder mirrors the processor's length metadata into + model_intermediate_buffer before the first prefill. Prefer lengths + inferred from that metadata, then fall back to the explicit target + scalar. Never silently use a prompt-text length as the target length. + """ + + provided_text_len = _glm_tts_int_value(info_dict.get("glm_tts_text_token_len")) + prompt_text_len = _glm_tts_int_value(info_dict.get("glm_tts_prompt_text_token_len")) + prompt_speech_len = _glm_tts_int_value(info_dict.get("prompt_speech_token_len")) + input_len = _glm_tts_int_value(info_dict.get("input_len")) + + inferred_text_len: int | None = None + if input_len is not None and prompt_text_len is not None: + inferred_text_len = input_len - prompt_text_len - 1 # strip prompt text and BOA. + elif prompt_text_len is not None and prompt_speech_len is not None: + inferred_text_len = span_len - prompt_text_len - prompt_speech_len - 1 + + if inferred_text_len is not None and inferred_text_len > 0: + if provided_text_len is not None and provided_text_len != inferred_text_len: + logger.warning( + "GLM-TTS target text token length mismatch: processor=%d inferred=%d; using inferred length.", + provided_text_len, + inferred_text_len, + ) + return inferred_text_len + + return provided_text_len if provided_text_len is not None and provided_text_len > 0 else None + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: Any | None = None, + is_multimodal: Any | None = None, + **kwargs: Any, + ) -> torch.Tensor: + if self.model_stage == "glm_tts_dit": + return self._dit_gen.embed_input_ids(input_ids, **kwargs) + embed_tokens = self.model.embed_tokens(input_ids) + if multimodal_embeddings is None or is_multimodal is None: + return embed_tokens + + mask = torch.as_tensor(is_multimodal, device=embed_tokens.device, dtype=torch.bool).reshape(-1) + if not bool(mask.any()): + return embed_tokens + + if isinstance(multimodal_embeddings, (list, tuple)): + tensors = [ + torch.as_tensor(item, device=embed_tokens.device, dtype=embed_tokens.dtype) + for item in multimodal_embeddings + if item is not None + ] + if not tensors: + return embed_tokens + mm_embeds = torch.cat([item.reshape(-1, item.shape[-1]) for item in tensors], dim=0) + else: + mm_embeds = multimodal_embeddings + if mm_embeds is None: + return embed_tokens + mm_embeds = torch.as_tensor(mm_embeds, device=embed_tokens.device, dtype=embed_tokens.dtype) + if mm_embeds.ndim == 3 and int(mm_embeds.shape[0]) == 1: + mm_embeds = mm_embeds.squeeze(0) + if mm_embeds.ndim != 2: + shape = tuple(mm_embeds.shape) + raise ValueError(f"GLM-TTS multimodal embeddings should be 2D, got shape={shape}") + + flat_tokens = embed_tokens.reshape(-1, embed_tokens.shape[-1]) + if int(flat_tokens.shape[0]) != int(mask.numel()): + raise ValueError( + "GLM-TTS multimodal mask/token length mismatch: " + f"mask={int(mask.numel())}, tokens={int(flat_tokens.shape[0])}" + ) + mm_len = int(mask.sum().item()) + if int(mm_embeds.shape[0]) < mm_len: + raise ValueError( + "GLM-TTS multimodal embedding length mismatch: " + f"embeddings={int(mm_embeds.shape[0])}, placeholders={mm_len}" + ) + + flat_tokens = flat_tokens.clone() + flat_tokens[mask] = mm_embeds[:mm_len] + return flat_tokens.reshape_as(embed_tokens) + + def embed_multimodal(self, **kwargs: Any) -> list[torch.Tensor] | None: + if self.model_stage != "glm_tts": + return None + prompt_speech_token = kwargs.get("prompt_speech_token") + if prompt_speech_token is None: + return None + + def embed_one(value: Any) -> torch.Tensor | None: + if value is None: + return None + speech_token = torch.as_tensor(value, device=next(self.model.parameters()).device) + if speech_token.numel() == 0: + return None + speech_token = speech_token.to(dtype=torch.long) + if speech_token.ndim == 0: + speech_token = speech_token.reshape(1, 1) + elif speech_token.ndim == 1: + speech_token = speech_token.unsqueeze(0) + elif speech_token.ndim > 2: + speech_token = speech_token.reshape(int(speech_token.shape[0]), -1) + if int(speech_token.min().item()) >= self._ats: + speech_ids = speech_token + else: + speech_ids = speech_token + self._ats + speech_embeds = self.model.embed_tokens(speech_ids) + return speech_embeds.reshape(-1, speech_embeds.shape[-1]) + + def is_flat_token_sequence(value: list[Any] | tuple[Any, ...]) -> bool: + if not value: + return False + try: + return all(torch.as_tensor(item).ndim == 0 for item in value) + except (TypeError, ValueError): + return False + + if isinstance(prompt_speech_token, (list, tuple)): + if is_flat_token_sequence(prompt_speech_token): + one_embed = embed_one(prompt_speech_token) + return [one_embed] if one_embed is not None else None + embeds = [embed_one(item) for item in prompt_speech_token] + embeds = [item for item in embeds if item is not None] + if not embeds: + return None + return embeds + + one_embed = embed_one(prompt_speech_token) + return [one_embed] if one_embed is not None else None + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor | IntermediateTensors | OmniOutput: + if self.model_stage == "glm_tts_dit": + return self._dit_gen.forward(input_ids, positions, **kwargs) + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + + def compute_logits( + self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None + ) -> torch.Tensor | None: + if self.model_stage == "glm_tts_dit": + return self._dit_gen.compute_logits(hidden_states, sampling_metadata) + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None: + return None + logits = self.logits_processor(self.lm_head, hidden_states) + if logits is None: + return None + + return logits + + def _glm_tts_ras_enabled(self, sampling_metadata: SamplingMetadata) -> bool: + if self.model_stage != "glm_tts": + return False + if sampling_metadata.max_num_logprobs is not None: + return False + if sampling_metadata.temperature is None: + return False + if bool(sampling_metadata.bad_words_token_ids): + return False + if torch.any(sampling_metadata.frequency_penalties != 0): + return False + if torch.any(sampling_metadata.presence_penalties != 0): + return False + return True + + def _apply_allowed_mask(self, logits: torch.Tensor) -> torch.Tensor: + """Restrict logits to valid audio tokens [ATS..ATE] ∪ {EOA}.""" + allowed_mask = getattr(self, "_ar_allowed_mask", None) + if allowed_mask is None or allowed_mask.device != logits.device or allowed_mask.shape[0] != logits.shape[-1]: + allowed_mask = torch.full((logits.shape[-1],), float("-inf"), device=logits.device) + allowed_mask[self._ats : self._ate + 1] = 0.0 + allowed_mask[self._eoa] = 0.0 + self._ar_allowed_mask = allowed_mask + return logits + allowed_mask + + def _select_allowed_audio_logits(self, logits: torch.Tensor) -> torch.Tensor: + """Return compact logits for [audio tokens] plus EOA. + + The RAS path samples only from GLM-TTS audio tokens and EOA. Keeping a + compact view avoids top-k/logsumexp over the full LLM vocabulary on + every AR step; this is equivalent to sampling after ``_apply_allowed_mask``. + """ + audio_logits = logits[:, self._ats : self._ate + 1] + if self._ats <= self._eoa <= self._ate: + return audio_logits + return torch.cat((audio_logits, logits[:, self._eoa : self._eoa + 1]), dim=-1) + + def _allowed_local_to_token_id(self, local_id: int) -> int: + audio_vocab = self._ate - self._ats + 1 + return self._ats + local_id if local_id < audio_vocab else self._eoa + + def _recent_allowed_local_ids(self, decoded_tokens: Sequence[int], win_size: int) -> list[int]: + if win_size <= 0 or not decoded_tokens: + return [] + audio_vocab = self._ate - self._ats + 1 + recent: list[int] = [] + for token in decoded_tokens[-win_size:]: + token_id = int(token) + if self._ats <= token_id <= self._ate: + recent.append(token_id - self._ats) + elif token_id == self._eoa: + recent.append(audio_vocab) + else: + recent.append(-1) + return recent + + def sample( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> SamplerOutput | None: + """RAS sampler following CosyVoice3 pattern. + + Uses vLLM Sampler for logits processing (logit_bias_state handles + min_tokens/max_tokens/stop_token_ids). When RAS is enabled, applies + per-request nucleus+repetition-aware sampling; otherwise falls back + to standard vLLM sampling. + """ + if logits is None or logits.numel() == 0: + return None + if self.model_stage != "glm_tts": + return None + + sampler = getattr(self, "_ar_sampler", None) + if sampler is None: + sampler = Sampler() + self._ar_sampler = sampler + + if not self._glm_tts_ras_enabled(sampling_metadata): + logits = self._apply_allowed_mask(logits) + return sampler(logits=logits, sampling_metadata=sampling_metadata) + + logits = logits.to(torch.float32) + sampling_for_processors = replace(sampling_metadata, no_penalties=True) + logits = sampler.apply_logits_processors(logits, sampling_for_processors, predict_bonus_token=False) + allowed_logits = self._select_allowed_audio_logits(logits) + + default_top_p = self._ras_top_p + default_top_k = self._ras_top_k + ws = self._ras_win_size + tr = self._ras_tau_r + + num_reqs = int(allowed_logits.shape[0]) + t_data = sampling_metadata.temperature + tp_data = sampling_metadata.top_p + tk_data = sampling_metadata.top_k + cache_key = ( + t_data.data_ptr() if t_data is not None else 0, + tp_data.data_ptr() if tp_data is not None else 0, + tk_data.data_ptr() if tk_data is not None else 0, + num_reqs, + ) + _sp_cache = getattr(self, "_sampling_param_cache", None) + if _sp_cache is not None and _sp_cache[0] == cache_key: + temperatures, top_ps, top_ks = _sp_cache[1] + else: + temperatures = _tensor_param_values(t_data, num_reqs, 1.0) + top_ps = _tensor_param_values(tp_data, num_reqs, default_top_p) + top_ks = _tensor_param_values(tk_data, num_reqs, default_top_k) + self._sampling_param_cache = (cache_key, (temperatures, top_ps, top_ks)) + + sampled_ids: list[int] = [] + for req_idx in range(num_reqs): + row_logits = allowed_logits[req_idx] + temperature = temperatures[req_idx] + if temperature < self._sampling_eps: + local_id = int(torch.argmax(row_logits).item()) + sampled_ids.append(self._allowed_local_to_token_id(local_id)) + continue + + top_p = top_ps[req_idx] + top_k = int(top_ks[req_idx]) + generator = sampling_metadata.generators.get(req_idx) + weighted_scores = row_logits / max(temperature, self._sampling_eps) + decoded_tokens = ( + sampling_metadata.output_token_ids[req_idx] if req_idx < len(sampling_metadata.output_token_ids) else [] + ) + local_id = _ras_sample_one( + weighted_scores, + self._recent_allowed_local_ids(decoded_tokens, ws), + top_p=top_p, + top_k=top_k, + win_size=ws, + tau_r=tr, + generator=generator, + ) + sampled_ids.append(self._allowed_local_to_token_id(local_id)) + + sampled = torch.tensor(sampled_ids, device=logits.device, dtype=torch.int32) + return SamplerOutput(sampled_token_ids=sampled.unsqueeze(-1), logprobs_tensors=None) + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + """Package hidden states, speech tokens, and voice clone data into OmniOutput. + + Streaming contract: **delta**. Each decode step emits exactly one + speech token (or a prefill placeholder). The engine's output + processor concatenates per-step deltas into the final tensor. + """ + if isinstance(model_outputs, OmniOutput): + return model_outputs + + hidden = model_outputs + info_dicts = kwargs.get("model_intermediate_buffer") or [] + if isinstance(info_dicts, dict): + info_dicts = [info_dicts] + + speech_tokens_list: list[torch.Tensor] = [] + multimodal_extras: dict[str, Any] = {} + mm_conditioning_keys = ( + "prompt_speech_token", + "prompt_speech_token_len", + "glm_tts_prompt_text_token_len", + "prompt_feat", + "embedding", + "glm_tts_text_token_len", + ) + info_for_flags = next((info for info in info_dicts if isinstance(info, dict)), None) + if info_for_flags is None or not info_for_flags.get("_voice_clone_emitted"): + copied_from_kwargs = False + for key in mm_conditioning_keys: + val = kwargs.get(key) + if val is not None: + multimodal_extras[key] = val + copied_from_kwargs = True + if copied_from_kwargs and info_for_flags is not None: + info_for_flags["_voice_clone_emitted"] = True + + for info in info_dicts: + if not isinstance(info, dict): + continue + tokens = info.get("speech_tokens") + if isinstance(tokens, torch.Tensor) and tokens.numel() > 0: + speech_tokens_list.append(tokens) + # Propagate voice cloning features from preprocess info_update. + # IMPORTANT: Only emit once — the output processor accumulates + # every decode step and concatenates tensors, which would corrupt + # constant data (e.g. embedding [192] → [N*192]). + if not info.get("_voice_clone_emitted"): + for key in mm_conditioning_keys: + val = info.get(key) + if val is not None and key not in multimodal_extras: + multimodal_extras[key] = val + if any(info.get(k) is not None for k in mm_conditioning_keys): + info["_voice_clone_emitted"] = True + + if not speech_tokens_list: + return OmniOutput(text_hidden_states=hidden, multimodal_outputs=multimodal_extras) + + speech_tokens = torch.cat(speech_tokens_list, dim=0) + span_len = int(speech_tokens.shape[0]) + hidden = hidden[:span_len] + multimodal_extras["speech_tokens"] = speech_tokens + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs=multimodal_extras, + ) + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Prepare inputs for GLM-TTS AR model. + + GLM-TTS only supports the multimodal processor path: + text prompt + ``multi_modal_data["audio"]`` + ``mm_processor_kwargs["prompt_text"]``. + Legacy placeholder prompts via ``additional_information`` are rejected. + """ + if info_dict.get("additional_information") is not None: + raise ValueError( + "GLM-TTS no longer accepts legacy additional_information prompts; " + "use prompt + multi_modal_data['audio'] + mm_processor_kwargs['prompt_text']." + ) + + span_len = int(input_ids.shape[0]) + logger.debug("preprocess: span_len=%d, input_ids.shape=%s", span_len, input_ids.shape) + if span_len <= 0: + return input_ids, input_embeds if input_embeds is not None else self.embed_input_ids(input_ids), {} + + device = input_ids.device + + if isinstance(info_dict.get("text"), list) and info_dict["text"]: + raise ValueError( + "GLM-TTS no longer accepts legacy text/additional_information prefill payloads; " + "use the multimodal processor path." + ) + + mm_prefill_done = bool(info_dict.get("glm_tts_mm_prefill_done")) + sampled_token: int | None = None + if input_embeds is not None and span_len == 1: + sampled_token = int(input_ids[0].item()) + one_token_prefill_tail = ( + input_embeds is not None + and span_len == 1 + and mm_prefill_done + and sampled_token is not None + and sampled_token != self._eoa + and not (self._ats <= sampled_token <= self._ate) + ) + is_prefill_span = span_len > 1 or not mm_prefill_done or one_token_prefill_tail + if input_embeds is None and is_prefill_span: + raise ValueError("Missing GLM-TTS multimodal input embeddings.") + if input_embeds is not None and is_prefill_span: + input_ids_out = input_ids.clone() + input_ids_out[:] = self._pad + info_update: dict[str, Any] = { + "glm_tts_mm_prefill_done": True, + "speech_tokens": torch.full((span_len, 1), -1, device=device, dtype=torch.long), + } + # One-token prefill tails re-enter the prefill branch but + # text_token_len was already resolved during the initial prefill. + if not one_token_prefill_tail: + text_token_len = self._resolve_prefill_target_text_len(span_len, info_dict) + if text_token_len is not None: + info_update["glm_tts_text_token_len"] = torch.tensor( + [text_token_len], device=device, dtype=torch.long + ) + prompt_text_len = info_dict.get("glm_tts_prompt_text_token_len") + if prompt_text_len is not None: + info_update["glm_tts_prompt_text_token_len"] = prompt_text_len + prompt_speech_len = info_dict.get("prompt_speech_token_len") + if prompt_speech_len is None: + input_len = _glm_tts_int_value(info_dict.get("input_len")) + if input_len is not None: + prompt_speech_len = torch.tensor( + [max(0, span_len - input_len)], device=device, dtype=torch.long + ) + if prompt_speech_len is not None: + info_update["prompt_speech_token_len"] = prompt_speech_len + return input_ids_out, input_embeds.to(dtype=self._model_dtype()), info_update + + # Decode: span_len == 1 + # Standard autoregressive decode - use input_ids directly + if input_embeds is not None and int(input_embeds.shape[0]) == 1: + inputs_embeds_out = input_embeds.reshape(1, -1).to(dtype=self._model_dtype()) + else: + inputs_embeds_out = self.embed_input_ids(input_ids.reshape(1, 1).to(torch.long)) + inputs_embeds_out = inputs_embeds_out.reshape(1, -1).to(dtype=self._model_dtype()) + + # Convert sampled token to speech token (relative to ATS) + # -1 = invalid/EOA, valid range = [0, ATE-ATS] + if sampled_token is None: + sampled_token = int(input_ids[0].item()) + if self._ats <= sampled_token <= self._ate: + speech_token = sampled_token - self._ats + else: + # EOA or other non-audio token → mark as invalid (-1) + speech_token = -1 + logger.debug("GLM-TTS decode: non-audio token %d (EOA=%d)", sampled_token, self._eoa) + speech_tokens = torch.tensor([[speech_token]], device=device, dtype=torch.long) + + info_update = {"speech_tokens": speech_tokens} + return input_ids, inputs_embeds_out, info_update + + def postprocess(self, hidden_states: torch.Tensor, **kwargs: Any) -> dict[str, Any]: + """Cache last hidden state for next decode step.""" + if hidden_states.numel() == 0: + return {} + last = hidden_states[-1, :].detach() + update: dict[str, Any] = {"last_hidden": {"last": last}} + + multimodal_outputs = kwargs.get("multimodal_outputs") + if isinstance(multimodal_outputs, dict): + copied_conditioning = False + for key in ("prompt_speech_token", "prompt_speech_token_len", "prompt_feat", "embedding"): + val = multimodal_outputs.get(key) + if val is not None: + update[key] = _first_glm_tts_value(val) + copied_conditioning = True + + prompt_text_len = multimodal_outputs.get("glm_tts_prompt_text_token_len") + if prompt_text_len is not None: + update["glm_tts_prompt_text_token_len"] = _first_glm_tts_value(prompt_text_len) + + text_token_len = _glm_tts_int_value(multimodal_outputs.get("glm_tts_text_token_len")) + if text_token_len is not None: + update["glm_tts_text_token_len"] = torch.tensor([int(text_token_len)], dtype=torch.long) + + if copied_conditioning: + update["_voice_clone_emitted"] = True + + return update + + def _iter_llm_safetensors(self) -> Iterable[tuple[str, torch.Tensor]]: + """Yield (name, tensor) pairs from llm/model-*.safetensors.""" + import glob as glob_module + + from safetensors.torch import load_file + + llm_dir = os.path.join(self.model_dir, "llm") + if os.path.isdir(llm_dir): + sf_files = sorted(glob_module.glob(os.path.join(llm_dir, "model-*.safetensors"))) + if sf_files: + for sf_path in sf_files: + yield from load_file(sf_path, device="cpu").items() + return + + from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific + + model_root = download_weights_from_hf_specific( + self.model_path, + self.vllm_config.load_config.download_dir, + allow_patterns=["llm/model-*.safetensors"], + ) + llm_dir = os.path.join(model_root, "llm") + sf_files = sorted(glob_module.glob(os.path.join(llm_dir, "model-*.safetensors"))) + if not sf_files: + raise RuntimeError(f"No LLM safetensors found under {model_root}. Expected llm/model-*.safetensors.") + for sf_path in sf_files: + yield from load_file(sf_path, device="cpu").items() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights from checkpoint. + + Stage 0 (glm_tts): HuggingFace Llama-format checkpoint from llm/ subdir. + Stage 1 (glm_tts_dit): DiT flow.pt + vocoder. + """ + if self.model_stage == "glm_tts_dit": + return self._dit_gen.load_weights(weights) + + def _glm_tts_weights() -> Iterable[tuple[str, torch.Tensor]]: + for name, loaded_weight in self._iter_llm_safetensors(): + if "rotary_emb.inv_freq" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + continue + yield name, loaded_weight + + loader = AutoWeightsLoader(self) + loaded_params = loader.load_weights(_glm_tts_weights(), mapper=self.hf_to_vllm_mapper) + + params_dict = dict(self.named_parameters()) + # Handle tie_word_embeddings: copy embed_tokens to lm_head if not loaded. + lm_head_key = "lm_head.weight" + embed_key = "model.embed_tokens.weight" + if lm_head_key not in loaded_params and embed_key in loaded_params: + if lm_head_key in params_dict and embed_key in params_dict: + lm_head_param = params_dict[lm_head_key] + embed_param = params_dict[embed_key] + weight_loader = getattr(lm_head_param, "weight_loader", default_weight_loader) + weight_loader(lm_head_param, embed_param.data) + loaded_params.add(lm_head_key) + logger.info("Tied lm_head.weight to embed_tokens.weight") + + logger.info("Loaded %d weights for GLMTTSForConditionalGeneration", len(loaded_params)) + return loaded_params diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py new file mode 100644 index 000000000..5cda6ec45 --- /dev/null +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py @@ -0,0 +1,650 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-TTS DiT (Diffusion Transformer) Model. + +Flow Matching model that converts speech tokens to mel-spectrogram. +Ported from: https://github.com/zai-org/GLM-TTS/blob/main/flow/dit.py +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb + +from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention + + +def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor: float = 1.0): + """Precompute Rotary Positional Embeddings (RoPE).""" + theta *= theta_rescale_factor ** (dim / (dim - 2)) + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + t = torch.arange(end, device=freqs.device) + freqs = torch.outer(t, freqs).float() + freqs_cos = torch.cos(freqs) + freqs_sin = torch.sin(freqs) + return torch.cat([freqs_cos, freqs_sin], dim=-1) + + +def get_pos_embed_indices(start, length, max_pos, scale=1.0): + """Get position embedding indices.""" + scale = scale * torch.ones_like(start, dtype=torch.float32) + pos = ( + start.unsqueeze(1) + + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long() + ) + pos = torch.where(pos < max_pos, pos, max_pos - 1) + return pos + + +class SinusPositionEmbedding(nn.Module): + """Sinusoidal position embedding for timesteps.""" + + def __init__(self, dim: int): + super().__init__() + self.dim = dim + + def forward(self, x: torch.Tensor, scale: float = 1000) -> torch.Tensor: + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) + emb = scale * x.unsqueeze(1) * emb.unsqueeze(0) + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + +class TimestepEmbedding(nn.Module): + """Timestep embedding with MLP.""" + + def __init__(self, dim: int, freq_embed_dim: int = 256): + super().__init__() + self.time_embed = SinusPositionEmbedding(freq_embed_dim) + self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + + def forward(self, timestep: torch.Tensor) -> torch.Tensor: + time_hidden = self.time_embed(timestep) + # Cast to model weight dtype (bfloat16) not input dtype (may be float32) + weight_dtype = self.time_mlp[0].weight.dtype + time_hidden = time_hidden.to(weight_dtype) + return self.time_mlp(time_hidden) + + +class GRN(nn.Module): + """Global Response Normalization layer.""" + + def __init__(self, dim: int): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gx = torch.norm(x, p=2, dim=1, keepdim=True) + nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-6) + return self.gamma * (x * nx) + self.beta + x + + +class ConvNeXtV2Block(nn.Module): + """ConvNeXt-V2 Block for text embedding.""" + + def __init__(self, dim: int, intermediate_dim: int, dilation: int = 1): + super().__init__() + padding = (dilation * (7 - 1)) // 2 + self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation) + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, intermediate_dim) + self.act = nn.GELU() + self.grn = GRN(intermediate_dim) + self.pwconv2 = nn.Linear(intermediate_dim, dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = x.transpose(1, 2) + x = self.dwconv(x) + x = x.transpose(1, 2) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + return residual + x + + +class ConvPositionEmbedding(nn.Module): + """Convolutional position embedding.""" + + def __init__(self, dim: int, kernel_size: int = 31, groups: int = 16): + super().__init__() + if kernel_size % 2 == 0: + raise ValueError(f"kernel_size must be odd, got {kernel_size}") + self.conv1d = nn.Sequential( + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2), + nn.Mish(), + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2), + nn.Mish(), + ) + + def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: + if mask is not None: + mask = mask[..., None] + x = x.masked_fill(~mask, 0.0) + + x = rearrange(x, "b n d -> b d n") + x = self.conv1d(x) + out = rearrange(x, "b d n -> b n d") + + if mask is not None: + out = out.masked_fill(~mask, 0.0) + return out + + +class TextEmbedding(nn.Module): + """Text/speech token embedding with optional ConvNeXt modeling.""" + + def __init__( + self, + text_num_embeds: int, + output_dim: int, + conv_layers: int = 0, + conv_mult: int = 2, + length_align: str = "fill", + ): + super().__init__() + self.pad_id = int(text_num_embeds) + self.text_embed = nn.Embedding(text_num_embeds + 1, output_dim, padding_idx=self.pad_id) + self.length_align = length_align + + if conv_layers > 0: + self.extra_modeling = True + self.precompute_max_pos = 4096 + self.register_buffer( + "freqs_cis", precompute_freqs_cis(output_dim, self.precompute_max_pos), persistent=False + ) + self.text_blocks = nn.Sequential( + *[ConvNeXtV2Block(output_dim, output_dim * conv_mult) for _ in range(conv_layers)] + ) + else: + self.extra_modeling = False + + def forward( + self, + text: torch.Tensor, + aim_seq_len: int, + text_lens: torch.Tensor | None = None, + ) -> torch.Tensor: + batch, text_len = text.shape[0], text.shape[1] + if self.length_align == "fill": + text = text[:, :aim_seq_len] + text = F.pad(text, (0, aim_seq_len - min(text_len, aim_seq_len)), value=self.pad_id) + elif self.length_align == "interpolate_token": + if text_lens is None: + text = ( + F.interpolate( + text.unsqueeze(1).float(), + size=aim_seq_len, + mode="nearest", + ) + .squeeze(1) + .long() + ) + else: + text_lens = text_lens.to(device=text.device, dtype=torch.long) + text_lens = torch.clamp(text_lens, min=1, max=text_len) + positions = torch.arange(aim_seq_len, device=text.device, dtype=torch.long) + src_idx = positions.unsqueeze(0) * text_lens.unsqueeze(1) // int(aim_seq_len) + src_idx = torch.minimum(src_idx, (text_lens - 1).unsqueeze(1)) + text = torch.gather(text, dim=1, index=src_idx) + elif self.length_align != "interpolate_feature": + raise ValueError(f"Unsupported length_align={self.length_align!r}") + + hidden = self.text_embed(text) + + if self.length_align == "interpolate_feature": + hidden = F.interpolate( + hidden.permute(0, 2, 1), + size=aim_seq_len, + mode="nearest", + ).permute(0, 2, 1) + + if self.extra_modeling: + # Get model dtype from text_blocks weights + model_dtype = self.text_blocks[0].dwconv.weight.dtype + hidden = hidden.to(model_dtype) + + batch_start = torch.zeros((batch,), dtype=torch.long, device=text.device) + pos_idx = get_pos_embed_indices(batch_start, aim_seq_len, max_pos=self.precompute_max_pos) + text_pos_embed = self.freqs_cis[pos_idx].to(model_dtype) + hidden = hidden + text_pos_embed + hidden = self.text_blocks(hidden) + + return hidden + + +class EmbeddingConcater(nn.Module): + """Concatenates and projects noisy audio, condition, and text embeddings.""" + + def __init__(self, in_dim: int, out_dim: int): + super().__init__() + self.proj = nn.Linear(in_dim, out_dim) + self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim) + + def forward( + self, + x: torch.Tensor, + cond: torch.Tensor, + text_embed: torch.Tensor, + drop_audio_cond: bool = False, + ) -> torch.Tensor: + if drop_audio_cond: + cond = torch.zeros_like(cond) + x = self.proj(torch.cat((x, cond, text_embed), dim=-1)) + x = self.conv_pos_embed(x) + x + return x + + +class AdaLayerNormZero(nn.Module): + """Adaptive LayerNorm with modulation from time/condition embedding.""" + + def __init__(self, dim: int, additional_dim: int = 0): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(dim + additional_dim, dim * 6) + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x: torch.Tensor, emb: torch.Tensor): + emb = self.linear(self.silu(emb)) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1) + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +class AdaLayerNormZeroFinal(nn.Module): + """Adaptive LayerNorm for final layer.""" + + def __init__(self, dim: int, additional_dim: int = 0): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(dim + additional_dim, dim * 2) + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: + emb = self.linear(self.silu(emb)) + scale, shift = torch.chunk(emb, 2, dim=1) + x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return x + + +class FeedForward(nn.Module): + """Feed-forward network with GELU activation. + + Matches original GLM-TTS structure with nested Sequential for checkpoint compatibility. + Weight paths: ff.ff.0.0.* (first Linear), ff.ff.2.* (second Linear) + """ + + def __init__(self, dim: int, dim_out: int | None = None, mult: int = 4, dropout: float = 0.0): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + # Nested structure to match original checkpoint paths + project_in = nn.Sequential( + nn.Linear(dim, inner_dim), + nn.GELU(approximate="tanh"), + ) + self.ff = nn.Sequential( + project_in, + nn.Dropout(dropout), + nn.Linear(inner_dim, dim_out), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.ff(x) + + +class DiTAttention(nn.Module): + """Attention module using diffusion infrastructure.""" + + def __init__(self, dim: int, heads: int = 8, dim_head: int = 64, dropout: float = 0.0): + super().__init__() + self.dim = dim + self.heads = heads + self.dim_head = dim_head + self.inner_dim = dim_head * heads + self.scale = 1.0 / math.sqrt(dim_head) + + self.to_q = nn.Linear(dim, self.inner_dim) + self.to_k = nn.Linear(dim, self.inner_dim) + self.to_v = nn.Linear(dim, self.inner_dim) + self.to_out = nn.Sequential(nn.Linear(self.inner_dim, dim), nn.Dropout(dropout)) + + self.attn = DiffusionAttention( + num_heads=heads, + head_size=dim_head, + softmax_scale=self.scale, + causal=False, + ) + + def forward( + self, + x: torch.Tensor, + padding_mask: torch.Tensor | None = None, + rope=None, + attn_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + batch_size, seq_len = x.shape[0], x.shape[1] + + query = self.to_q(x) + key = self.to_k(x) + value = self.to_v(x) + + # Apply RoPE + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + # Reshape for attention + query = query.view(batch_size, seq_len, self.heads, self.dim_head) + key = key.view(batch_size, seq_len, self.heads, self.dim_head) + value = value.view(batch_size, seq_len, self.heads, self.dim_head) + + # Block-causal mask requires arbitrary 2D attention mask which + # FlashAttention does not support. Fall back to PyTorch SDPA + # (matches official GLM-TTS flow/modules.py:272). + if attn_mask is not None: + q = query.permute(0, 2, 1, 3) # [B, H, T, D] + k = key.permute(0, 2, 1, 3) + v = value.permute(0, 2, 1, 3) + out = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_mask, + dropout_p=0.0, + is_causal=False, + ) + out = out.permute(0, 2, 1, 3) # [B, T, H, D] + else: + out = self.attn(query, key, value, attn_metadata=None) + out = out.view(batch_size, seq_len, self.inner_dim) + out = out.to(query.dtype) + out = self.to_out(out) + + # Apply padding mask + if padding_mask is not None: + padding_mask = rearrange(padding_mask, "b 1 n -> b n 1").bool() + out = out.masked_fill(~padding_mask, 0.0) + + return out + + +class DiTBlock(nn.Module): + """DiT block with AdaLayerNorm modulation.""" + + def __init__( + self, + dim: int, + heads: int, + dim_head: int, + ff_mult: int = 4, + dropout: float = 0.1, + additional_condition_dim: int = 0, + ): + super().__init__() + self.attn_norm = AdaLayerNormZero(dim, additional_condition_dim) + self.attn = DiTAttention(dim=dim, heads=heads, dim_head=dim_head, dropout=dropout) + self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout) + + def forward( + self, + x: torch.Tensor, + t: torch.Tensor, + padding_mask: torch.Tensor | None = None, + rope=None, + attn_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + # Pre-norm & modulation for attention + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) + + # Attention + attn_output = self.attn(x=norm, padding_mask=padding_mask, rope=rope, attn_mask=attn_mask) + x = x + gate_msa.unsqueeze(1) * attn_output + + # FeedForward + norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ff_output = self.ff(norm) + x = x + gate_mlp.unsqueeze(1) * ff_output + + return x + + +class GLMTTSDiT(nn.Module): + """GLM-TTS Diffusion Transformer. + + Converts speech tokens to mel-spectrogram via flow matching. + """ + + _repeated_blocks = ["DiTBlock"] + _layerwise_offload_blocks_attrs = ["transformer_blocks"] + + def __init__( + self, + *, + trans_dim: int = 768, + depth: int = 18, + heads: int = 12, + dim_head: int = 64, + ff_mult: int = 2, + dropout: float = 0.1, + mel_dim: int = 80, + text_vocab_size: int = 100000, + text_emb_dim: int = 512, + conv_layers: int = 4, + condition_dim: int = 160, # mel_dim + spk_embed_dim + spkr_emb_adaln: bool = False, + spkr_dim: int = 192, + use_wavlm_emb: bool = False, + long_skip_connection: bool = False, + ): + super().__init__() + + self.time_embed = TimestepEmbedding(trans_dim) + self.text_emb_layer = TextEmbedding( + text_vocab_size, + text_emb_dim, + conv_layers=conv_layers, + length_align="interpolate_token", + ) + + in_dim = mel_dim + text_emb_dim + condition_dim + self.emb_concator = EmbeddingConcater(in_dim, trans_dim) + + self.rotary_embed = RotaryEmbedding(dim_head) + + self.dim = trans_dim + self.depth = depth + self.num_heads = heads + self.spkr_emb_adaln = spkr_emb_adaln + + # Frame rates for block-causal mask token→mel conversion + self._input_frame_rate: float = 25.0 + self._mel_framerate: float = 50.0 + + # Speaker embedding dimension for AdaLN + if spkr_emb_adaln: + additional_dim = spkr_dim + 256 if use_wavlm_emb else spkr_dim + else: + additional_dim = 0 + + self.transformer_blocks = nn.ModuleList( + [ + DiTBlock( + dim=trans_dim, + heads=heads, + dim_head=dim_head, + ff_mult=ff_mult, + dropout=dropout, + additional_condition_dim=additional_dim, + ) + for _ in range(depth) + ] + ) + + self.long_skip_connection = nn.Linear(trans_dim * 2, trans_dim, bias=False) if long_skip_connection else None + + self.norm_out = AdaLayerNormZeroFinal(trans_dim, additional_dim=additional_dim) + self.proj_out = nn.Linear(trans_dim, mel_dim) + + # ------------------------------------------------------------------ + # Block-causal attention mask (ported from GLM-TTS block_mask_util) + # ------------------------------------------------------------------ + + def _token_to_mel_size(self, token_count: int) -> int: + """Convert token block size to mel-frame block size for streaming mask. + + Matches upstream GLM-TTS ``flow/dit.py`` block-mask conversion. + This is intentionally separate from normal ``feat_len`` calculation, + which still follows ``input_frame_rate`` and ``mel_framerate``. + """ + return int(token_count / 12.5 * 22050 / 256) + + def _create_block_causal_mask( + self, + block_pattern: list[int], + seq_len: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + """Build block-causal attention mask [1, 1, T, T]. + + Ported from GLM-TTS utils/block_mask_util.py ``create()``. + Each position can attend to its own block plus a look-ahead into + the next block (50% for small blocks, 66% for large blocks). + + Vectorized via ``torch.searchsorted`` — no Python loop over T. + """ + import itertools + + mel_blocks = [self._token_to_mel_size(b) for b in block_pattern] + # Extend with copies of the last block (enough for any seq_len) + extended = mel_blocks + [mel_blocks[-1]] * ((seq_len // max(mel_blocks[-1], 1)) + 2) + accum = list(itertools.accumulate(extended)) + + # Vectorized: map each position to its block index via searchsorted + accum_t = torch.tensor(accum, dtype=torch.long) + extended_t = torch.tensor(extended, dtype=torch.long) + positions = torch.arange(seq_len, dtype=torch.long) + + # searchsorted(accum, positions, right=True) gives the block index + # for each position (accum is strictly increasing). + block_idx = torch.searchsorted(accum_t, positions, right=True) + + # Per-position: block start, block size, position within block + block_start = torch.where(block_idx > 0, accum_t[block_idx - 1], torch.zeros_like(block_idx)) + cur_block_size = extended_t[block_idx] + idx_in_block = positions - block_start + + # Look-ahead: positions near the end of a block can see into the next + look_future = cur_block_size - 1 - idx_in_block + min_future = torch.where( + cur_block_size <= 172, + cur_block_size // 2, + (cur_block_size.float() / 1.5).long(), + ) + delta = torch.clamp(min_future - look_future, min=0) + aim_len = torch.clamp(block_start + cur_block_size + delta, max=seq_len) + + # Build mask: position i can attend to columns [0, aim_len[i]) + col_idx = torch.arange(seq_len, device=device) + mask = col_idx.unsqueeze(0) < aim_len.unsqueeze(1).to(device) + + # [1, 1, T, T] for broadcast across batch and heads + return mask.unsqueeze(0).unsqueeze(0) + + def forward( + self, + middle_point: torch.Tensor, + condition: torch.Tensor, + text: torch.Tensor, + time_step: torch.Tensor, + padding_mask: torch.Tensor, + spkr_emb: torch.Tensor | None = None, + text_lens: torch.Tensor | None = None, + is_causal: bool = False, + attn_mask: torch.Tensor | None = None, + block_pattern: list[int] | None = None, + ) -> torch.Tensor: + """Forward pass. + + Args: + middle_point: Noised audio tensor [B, T, mel_dim] + condition: Masked audio condition [B, T, condition_dim] + text: Speech token tensor [B, T_text] + time_step: Timestep tensor [B] + padding_mask: Padding mask [B, T] + spkr_emb: Speaker embedding [B, spkr_dim] + is_causal: Use block-causal attention + attn_mask: Custom attention mask (overrides is_causal) + block_pattern: Token-level block sizes for block-causal mask + + Returns: + Predicted mel-spectrogram [B, T, mel_dim] + """ + # Cast floating inputs to model dtype (weights may be bfloat16) + model_dtype = self.proj_out.weight.dtype + middle_point = middle_point.to(model_dtype) + condition = condition.to(model_dtype) + time_step = time_step.to(model_dtype) + if spkr_emb is not None: + spkr_emb = spkr_emb.to(model_dtype) + + _, seq_len = middle_point.shape[0], middle_point.shape[1] + + # Time embedding + time_emb = self.time_embed(time_step) + + # Concatenate speaker embedding if using AdaLN + if self.spkr_emb_adaln and spkr_emb is not None: + time_emb = torch.cat([time_emb, spkr_emb], dim=-1) + + # Text embedding - ensure cast to model dtype + text_embed = self.text_emb_layer(text, seq_len, text_lens=text_lens).to(model_dtype) + + # Input projection + x = self.emb_concator(middle_point, condition, text_embed, drop_audio_cond=False) + + # Rotary embeddings + rope = self.rotary_embed.forward_from_seq_len(seq_len) + + # Build block-causal mask when is_causal and no explicit attn_mask + if attn_mask is None and is_causal and block_pattern is not None: + attn_mask = self._create_block_causal_mask( + block_pattern, + seq_len, + middle_point.device, + model_dtype, + ) + + # Long skip connection + if self.long_skip_connection is not None: + residual = x + + # Transformer blocks + for block in self.transformer_blocks: + x = block(x, time_emb, padding_mask=padding_mask.unsqueeze(1), rope=rope, attn_mask=attn_mask) + + # Long skip connection + if self.long_skip_connection is not None: + x = self.long_skip_connection(torch.cat((x, residual), dim=-1)) + + # Output projection + x = self.norm_out(x, time_emb) + output = self.proj_out(x) + + return output diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts_dit_wrapper.py b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit_wrapper.py new file mode 100644 index 000000000..a7c26af30 --- /dev/null +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit_wrapper.py @@ -0,0 +1,1301 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-TTS DiT wrapper: LLM_GENERATION-compatible flow-matching + vocoder.""" + +from __future__ import annotations + +import dataclasses +import json +import math +import os +import threading +from collections.abc import Iterable +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +from torch.cuda import CUDAGraph +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.platforms import current_platform + +from vllm_omni.model_executor.models.glm_tts.glm_tts import ( + resolve_glm_tts_model_dir, +) +from vllm_omni.model_executor.models.glm_tts.vocoder import ( + load_vocoder, + mel_to_audio, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + + +@dataclasses.dataclass +class _StreamRequestState: + """Per-request streaming state (replaces 5 separate dicts).""" + + diff_cache: dict[int | str, Any] = dataclasses.field(default_factory=dict) + prompt_payload: dict[str, torch.Tensor] = dataclasses.field(default_factory=dict) + audio_tail: torch.Tensor | None = None + wav_pointer: int = 0 + fade_tail: torch.Tensor | None = None + + +_GLM_TTS_OFFICIAL_FLOW_FALLBACK: dict[str, Any] = { + "spkr_emb_adaLN": True, + "speech_token_cfg": False, + "remove_spkr_concat_condition": True, + "mel_dim": 80, + "mel_framerate": 50, + "input_frame_rate": 25, +} + +_GLM_TTS_RUNTIME_FLOW_DEFAULTS: dict[str, Any] = { + "sample_rate": 24000, + "n_timesteps": 10, + "t_scheduler": "cosine", + "inference_cfg_rate": 0.7, + "spk_embed_dim": 80, + "trans_dim": 768, + "depth": 18, + "heads": 12, +} + + +def _load_glm_tts_flow_config(flow_dir: str) -> dict[str, Any]: + """Load GLM-TTS flow config from JSON/YAML, falling back to official defaults.""" + config: dict[str, Any] = dict(_GLM_TTS_OFFICIAL_FLOW_FALLBACK) + + # Try config.json first (standard HuggingFace layout) + json_path = os.path.join(flow_dir, "config.json") + if os.path.exists(json_path): + try: + with open(json_path) as f: + raw = json.load(f) + flow = raw.get("flow", raw) if isinstance(raw, dict) else {} + if isinstance(flow, dict): + config.update(flow) + except Exception as e: + logger.warning("Failed to parse GLM-TTS flow config.json: %s", e) + + # Try config.yaml (HyperPyYAML — strip tags, only extract kwargs) + yaml_path = os.path.join(flow_dir, "config.yaml") + if os.path.exists(yaml_path): + try: + import yaml + + class _TagStripper(yaml.SafeLoader): + pass + + _TagStripper.add_multi_constructor( + "!", + lambda loader, _tag, node: ( + loader.construct_mapping(node, deep=True) + if isinstance(node, yaml.MappingNode) + else loader.construct_sequence(node, deep=True) + if isinstance(node, yaml.SequenceNode) + else loader.construct_scalar(node) + ), + ) + with open(yaml_path) as f: + raw = yaml.load(f, Loader=_TagStripper) + flow = raw.get("flow", raw) if isinstance(raw, dict) else {} + if isinstance(flow, dict): + config.update(flow) + except Exception as e: + logger.warning("Failed to parse GLM-TTS flow config.yaml: %s", e) + + return config + + +def make_pad_mask(lengths: torch.Tensor, max_len: int | None = None) -> torch.Tensor: + if max_len is None: + max_len = int(lengths.max().item()) + seq_range = torch.arange(0, max_len, device=lengths.device) + return seq_range.unsqueeze(0).expand(lengths.shape[0], max_len) >= lengths.unsqueeze(1) + + +def _unwrap(value: object) -> object: + """Unwrap single-element lists (common runtime_info wrapping).""" + return value[0] if isinstance(value, list) and value else value + + +def nested_get(info: dict[str, Any], section: str, key: str) -> Any: + nested = info.get(section) + return nested[key] if isinstance(nested, dict) and key in nested else info.get(key) + + +def _glm_tts_private_info(info: dict[str, Any]) -> dict[str, Any]: + private = info.get("glm_tts") + if isinstance(private, dict): + return private + kv_metadata = info.get("kv_metadata") + if isinstance(kv_metadata, dict): + private = kv_metadata.get("glm_tts") + if isinstance(private, dict): + return private + return {} + + +def as_tensor(value: object) -> torch.Tensor | None: + value = _unwrap(value) + return value if isinstance(value, torch.Tensor) else None + + +def as_bool(value: object) -> bool: + value = _unwrap(value) + if isinstance(value, torch.Tensor) and value.numel() > 0: + return bool(value.reshape(-1)[0].item()) + return bool(value) if value is not None else False + + +def as_str(value: object) -> str | None: + value = _unwrap(value) + return str(value) if value is not None else None + + +def _as_scalar(value: object, typ: type, default: int | float = 0): + """Extract scalar int/float from runtime_info value.""" + value = _unwrap(value) + if isinstance(value, torch.Tensor) and value.numel() > 0: + return typ(value.reshape(-1)[0].item()) + if value is None: + return default + try: + return typ(value) + except (TypeError, ValueError): + return default + + +def as_int(value: object, default: int = 0) -> int: + return _as_scalar(value, int, default) + + +def as_float(value: object, default: float = 0.0) -> float: + return _as_scalar(value, float, default) + + +def as_int_list(value: object) -> list[int] | None: + if value is None: + return None + if isinstance(value, torch.Tensor): + return [int(x) for x in value.reshape(-1).tolist()] if value.numel() else [] + if isinstance(value, list): + return [int(x.item() if hasattr(x, "item") else x) for x in value] + return None + + +def valid_speech_tokens( + value: object, + *, + device: torch.device, + fallback: torch.Tensor | None = None, +) -> torch.Tensor: + """Return GLM-TTS speech tokens, treating an explicit empty list as empty.""" + if value is None: + token = fallback if fallback is not None else torch.empty(0, device=device, dtype=torch.long) + elif isinstance(value, torch.Tensor): + token = value.to(device=device, dtype=torch.long) + elif isinstance(value, np.ndarray): + token = torch.as_tensor(value, device=device, dtype=torch.long) + elif isinstance(value, list): + if len(value) == 1 and isinstance(value[0], torch.Tensor): + token = value[0].to(device=device, dtype=torch.long) + elif value and all(isinstance(item, torch.Tensor) for item in value): + token = torch.cat([item.reshape(-1).to(device=device, dtype=torch.long) for item in value]) + else: + token = torch.tensor(value, device=device, dtype=torch.long) + else: + token = fallback if fallback is not None else torch.empty(0, device=device, dtype=torch.long) + + token = token.reshape(-1).to(device=device, dtype=torch.long) + return token[(token >= 0) & (token < 32768)] + + +def split_request_ids(ids: torch.Tensor, seq_token_counts: list[int] | None = None) -> list[torch.Tensor]: + if seq_token_counts is not None: + boundaries, total = [0], ids.numel() + for count in seq_token_counts: + boundaries.append(boundaries[-1] + int(count)) + return [ids[boundaries[i] : min(boundaries[i + 1], total)] for i in range(len(seq_token_counts))] + try: + from vllm.forward_context import get_forward_context, is_forward_context_available + + if is_forward_context_available(): + slices = get_forward_context().ubatch_slices + if slices is not None and len(slices) > 1 and not any(hasattr(s, "token_slice") for s in slices): + boundaries = [0] + for s in slices: + boundaries.append(boundaries[-1] + int(s)) + return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)] + except ImportError: + pass + return [ids] + + +class GLMTTSDiTForGeneration(nn.Module): + """GLM-TTS DiT flow-matching stage wrapped as LLM_GENERATION.""" + + have_multimodal_outputs = True + enable_update_additional_information = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + self.config = config + self.model_stage = "glm_tts_dit" + self.max_num_seqs = int(getattr(vllm_config.scheduler_config, "max_num_seqs", 1)) + self._use_dit_cuda_graphs = bool(getattr(config, "use_dit_cuda_graphs", False)) + self._dit_cudagraph: CUDAGraphGLMTTSDiTWrapper | None = None + self._codec_chunk_frames, self._codec_left_context_frames = self._connector_chunk_config(vllm_config) + + model_path = vllm_config.model_config.model + model_path = resolve_glm_tts_model_dir( + model_path, + tokenizer_path=getattr(vllm_config.model_config, "tokenizer", None), + required_files=("flow/flow.pt",), + optional_files=("hift/hift.pt", "vocos2d/generator_jit.ckpt"), + ) + self.model_dir = model_path + flow_dir = os.path.join(self.model_dir, "flow") + self.hf_flow_config = _load_glm_tts_flow_config(flow_dir) + logger.info("GLM-TTS flow config: %s", self.hf_flow_config) + + self.mel_dim = self.hf_flow_config.get("mel_dim", 80) + self.input_frame_rate = self.hf_flow_config.get("input_frame_rate", 25.0) + self.mel_framerate = self.hf_flow_config.get("mel_framerate", 50) + self.sample_rate = self.hf_flow_config.get("sample_rate", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["sample_rate"]) + + self.n_timesteps = self.hf_flow_config.get("n_timesteps", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["n_timesteps"]) + config_n_ts = getattr(config, "n_timesteps", None) + if config_n_ts is not None: + self.n_timesteps = int(config_n_ts) + self.t_scheduler = self.hf_flow_config.get("t_scheduler", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["t_scheduler"]) + self.inference_cfg_rate = self.hf_flow_config.get( + "inference_cfg_rate", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["inference_cfg_rate"] + ) + self.speech_token_cfg = self.hf_flow_config.get("speech_token_cfg", False) + + self.spk_embed_dim = self.hf_flow_config.get("spk_embed_dim", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["spk_embed_dim"]) + self.spkr_emb_adaln = self.hf_flow_config.get("spkr_emb_adaLN", True) + self.remove_spkr_concat_condition = self.hf_flow_config.get("remove_spkr_concat_condition", True) + + from vllm_omni.model_executor.models.glm_tts.glm_tts_dit import GLMTTSDiT + + trans_dim = self.hf_flow_config.get("trans_dim", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["trans_dim"]) + depth = self.hf_flow_config.get("depth", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["depth"]) + heads = self.hf_flow_config.get("heads", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["heads"]) + condition_dim = self.mel_dim if self.remove_spkr_concat_condition else self.mel_dim + self.spk_embed_dim + + self.dit = GLMTTSDiT( + trans_dim=trans_dim, + depth=depth, + heads=heads, + dim_head=64, + ff_mult=2, + dropout=0.0, + mel_dim=self.mel_dim, + text_vocab_size=100000, + text_emb_dim=512, + conv_layers=4, + condition_dim=condition_dim, + spkr_emb_adaln=self.spkr_emb_adaln, + spkr_dim=192, + use_wavlm_emb=False, + ) + self._model_dtype = self._resolve_model_dtype(vllm_config) + if self._model_dtype is not None: + self.dit.to(dtype=self._model_dtype) + + if not self.remove_spkr_concat_condition: + self.spk_embed_affine_layer = nn.Linear(192, self.spk_embed_dim) + if self._model_dtype is not None: + self.spk_embed_affine_layer.to(dtype=self._model_dtype) + + self._vocoder = None + self._vocoder_sample_rate: int | None = None + + self._stream_cache_lock = threading.Lock() + self._stream_state_by_req: dict[str, _StreamRequestState] = {} + # Track recently-seen req_ids so stale entries (aborted requests) + # can be evicted without modifying the shared model runner. + self._stream_last_seen: dict[str, int] = {} + self._stream_forward_count: int = 0 + + self._flow_weights_path = os.path.join(self.model_dir, "flow", "flow.pt") + self._hift_weights_path = os.path.join(self.model_dir, "hift", "hift.pt") + self._vocos_jit_path = os.path.join(self.model_dir, "vocos2d", "generator_jit.ckpt") + + @staticmethod + def _connector_chunk_config(vllm_config: VllmConfig) -> tuple[int, int]: + cc = getattr(vllm_config.model_config, "stage_connector_config", None) + extra = (cc or {}).get("extra") if isinstance(cc, dict) else getattr(cc, "extra", None) + if not isinstance(extra, dict): + return 25, 25 + cf = extra.get("codec_chunk_frames", 25) + return (int(cf[0]) if isinstance(cf, list) else int(cf)), int(extra.get("codec_left_context_frames", 25)) + + @staticmethod + def _resolve_model_dtype(vllm_config: VllmConfig) -> torch.dtype | None: + dtype = getattr(vllm_config.model_config, "dtype", None) + if isinstance(dtype, torch.dtype): + return dtype + if isinstance(dtype, str): + _aliases = {"bf16": "bfloat16", "fp16": "float16", "half": "float16", "fp32": "float32"} + return getattr(torch, _aliases.get(dtype.lower(), dtype.lower()), None) + return None + + def _ensure_vocoder_loaded(self) -> None: + if self._vocoder is not None: + return + + device = next(self.dit.parameters()).device + self._vocoder, self._vocoder_sample_rate = load_vocoder(self.model_dir, device, self.sample_rate) + if self._vocoder is None: + raise RuntimeError( + "GLM-TTS stage 1 could not load a vocoder. " + f"Expected either {self._hift_weights_path} or {self._vocos_jit_path}." + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + **_: Any, + ) -> torch.Tensor: + """Return zero embeddings (DiT does not use token embeddings).""" + if input_ids.dim() != 1: + raise ValueError(f"GLM-TTS DiT input_ids must be 1D, got shape={tuple(input_ids.shape)}") + hidden = int(self.config.hidden_size) + return torch.zeros( + (input_ids.shape[0], hidden), + device=input_ids.device, + dtype=self._model_dtype or torch.float32, + ) + + def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None: + """DiT does not produce logits for sampling.""" + return None + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + **kwargs: Any, + ) -> OmniOutput: + """Run DiT flow-matching + vocoder on speech tokens.""" + # Evict stale streaming caches (aborted requests whose stream_finished + # was never sent). Entries not seen for 64 forward calls are dropped. + self._evict_stale_stream_caches() + + runtime_info = kwargs.get("model_intermediate_buffer") or [] + + seq_token_counts = kwargs.get("seq_token_counts") + flat_ids = input_ids.reshape(-1).to(dtype=torch.long) + if seq_token_counts is None and self.max_num_seqs > 1: + raise RuntimeError( + "GLM-TTS DiT stage requires seq_token_counts when max_num_seqs > 1; " + "otherwise concatenated speech tokens cannot be split per request." + ) + request_ids_list = split_request_ids(flat_ids, seq_token_counts) + + num_reqs = max(1, len(request_ids_list)) + sample_rate = torch.tensor(int(self.sample_rate), dtype=torch.int32) + empty_audio = torch.zeros((0,), dtype=torch.float32, device=input_ids.device) + audios: list[torch.Tensor] = [empty_audio] * num_reqs + srs: list[torch.Tensor] = [sample_rate] * num_reqs + + if not isinstance(runtime_info, list): + runtime_info = [] + + for idx, req_ids in enumerate(request_ids_list): + info = runtime_info[idx] if idx < len(runtime_info) and isinstance(runtime_info[idx], dict) else {} + req_id = as_str(nested_get(info, "meta", "req_id")) if info else None + if req_id is not None: + with self._stream_cache_lock: + self._stream_last_seen[req_id] = self._stream_forward_count + stream_finished_raw = nested_get(info, "meta", "stream_finished") if info else None + if stream_finished_raw is None and info: + stream_finished_raw = info.get("stream_finished") + stream_finished = as_bool(stream_finished_raw) if info else False + prompt_token = as_tensor(nested_get(info, "embed", "speech_token")) if info else None + if prompt_token is None and info: + prompt_token = as_tensor(info.get("prompt_token")) # backward compat + if prompt_token is None and info: + prompt_token = as_tensor(info.get("prompt_speech_token")) # backward compat + prompt_feat = as_tensor(nested_get(info, "embed", "speech_feat")) if info else None + if prompt_feat is None and info: + prompt_feat = as_tensor(info.get("prompt_feat")) # backward compat + embedding = as_tensor(nested_get(info, "embed", "embedding")) if info else None + if embedding is None and info: + embedding = as_tensor(info.get("embedding")) # backward compat + + speech_tokens_raw = info.get("speech_tokens") if info else None + has_meta = isinstance(info.get("meta"), dict) if info else False + uses_streaming = bool(info) and ( + has_meta + or "stream_finished" in info + or "token_offset" in info + or "left_context_size" in info + or "glm_tts" in info + or "kv_metadata" in info + ) + if uses_streaming and req_id is not None: + with self._stream_cache_lock: + state = self._stream_state_by_req.setdefault(req_id, _StreamRequestState()) + prompt_payload = state.prompt_payload + if prompt_token is not None: + prompt_payload["prompt_speech_token"] = prompt_token + elif "prompt_speech_token" in prompt_payload: + prompt_token = prompt_payload["prompt_speech_token"] + if prompt_feat is not None: + prompt_payload["prompt_feat"] = prompt_feat + elif "prompt_feat" in prompt_payload: + prompt_feat = prompt_payload["prompt_feat"] + if embedding is not None: + prompt_payload["embedding"] = embedding + elif "embedding" in prompt_payload: + embedding = prompt_payload["embedding"] + + fallback_token = req_ids.reshape(-1).to(device=input_ids.device, dtype=torch.long) + if uses_streaming: + token = valid_speech_tokens(None, device=input_ids.device, fallback=fallback_token) + else: + token = valid_speech_tokens( + speech_tokens_raw, + device=input_ids.device, + fallback=fallback_token, + ) + + if token.numel() == 0: + if stream_finished and req_id is not None: + with self._stream_cache_lock: + self._stream_state_by_req.pop(req_id, None) + self._stream_last_seen.pop(req_id, None) + audios[idx] = empty_audio + continue + + if uses_streaming: + private_info = _glm_tts_private_info(info) + token_offset_raw = nested_get(info, "meta", "left_context_size") + if token_offset_raw is None: + token_offset_raw = info.get("left_context_size") + if token_offset_raw is None: + token_offset_raw = nested_get(info, "meta", "token_offset") + if token_offset_raw is None: + token_offset_raw = info.get("token_offset") + token_offset = max( + 0, + as_int(token_offset_raw, 0), + ) + + chunk_sizes_raw = private_info.get("chunk_sizes_history") + if chunk_sizes_raw is None: + chunk_sizes_raw = nested_get(info, "meta", "chunk_sizes_history") + if chunk_sizes_raw is None: + chunk_sizes_raw = info.get("chunk_sizes_history") + block_pattern_raw = private_info.get("block_pattern") + if block_pattern_raw is None: + block_pattern_raw = nested_get(info, "meta", "block_pattern") + if block_pattern_raw is None: + block_pattern_raw = info.get("block_pattern") + crossfade_raw = private_info.get("crossfade_sec") + if crossfade_raw is None: + crossfade_raw = nested_get(info, "meta", "crossfade_sec") + if crossfade_raw is None: + crossfade_raw = info.get("crossfade_sec") + chunk_sizes_history = as_int_list(chunk_sizes_raw) + block_pattern = as_int_list(block_pattern_raw) + crossfade_sec = as_float(crossfade_raw, 0.1) + + audio = self._forward_full( + speech_tokens=token, + prompt_token=prompt_token, + prompt_feat=prompt_feat, + embedding=embedding, + token_offset=token_offset, + req_id=req_id, + stream_finished=stream_finished, + chunk_sizes_history=chunk_sizes_history, + block_pattern=block_pattern, + crossfade_sec=crossfade_sec, + ) + if stream_finished and req_id is not None: + with self._stream_cache_lock: + self._stream_state_by_req.pop(req_id, None) + self._stream_last_seen.pop(req_id, None) + else: + audio = self._forward_full( + speech_tokens=token, + prompt_token=prompt_token, + prompt_feat=prompt_feat, + embedding=embedding, + ) + + audios[idx] = audio.reshape(-1).to(dtype=torch.float32) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"audio": audios, "sr": srs}, + ) + + @staticmethod + def _to_2d(val: torch.Tensor | np.ndarray | None, device: torch.device, dtype: torch.dtype) -> torch.Tensor | None: + """Convert numpy/tensor to 2D+ tensor on device, or return None.""" + if val is None: + return None + if isinstance(val, np.ndarray): + val = torch.from_numpy(val) + val = val.to(device=device, dtype=dtype) + if val.ndim == 1: + val = val.unsqueeze(0) + return val + + @staticmethod + def _validate_voice_clone(prompt_token, prompt_feat, embedding) -> None: + """Raise if voice clone payload is incomplete.""" + fields = {"prompt_speech_token": prompt_token, "prompt_feat": prompt_feat, "embedding": embedding} + present = [k for k, v in fields.items() if v is not None] + missing = [k for k, v in fields.items() if v is None] + if present and missing: + raise RuntimeError(f"GLM-TTS voice clone payload incomplete: present={present}, missing={missing}.") + + @torch.inference_mode() + def _forward_full( + self, + speech_tokens: torch.Tensor, + prompt_token: torch.Tensor | None = None, + prompt_feat: torch.Tensor | None = None, + embedding: torch.Tensor | None = None, + token_offset: int = 0, + req_id: str | None = None, + stream_finished: bool = False, + chunk_sizes_history: list[int] | None = None, + block_pattern: list[int] | None = None, + crossfade_sec: float = 0.1, + ) -> torch.Tensor: + """DiT flow-matching + vocoder pipeline. Returns audio waveform.""" + device = speech_tokens.device + self._validate_voice_clone(prompt_token, prompt_feat, embedding) + + if speech_tokens.ndim == 1: + speech_tokens = speech_tokens.unsqueeze(0) + generated_token_len = int(speech_tokens.shape[1]) + + has_voice_clone = prompt_token is not None or embedding is not None + prompt_token = self._to_2d(prompt_token, device, torch.long) + if prompt_token is not None: + prompt_token_len = prompt_token.shape[1] + speech_tokens = torch.cat([prompt_token, speech_tokens], dim=1) + else: + prompt_token_len = 0 + has_prompt_audio_context = prompt_token_len > 0 + + token_len = speech_tokens.shape[1] + if token_len == 0: + return torch.zeros(0, device=device) + + text_lens = torch.tensor([token_len], device=device, dtype=torch.long) + feat_len = max(1, int(token_len / self.input_frame_rate * self.mel_framerate)) + + padding_mask = ~make_pad_mask(torch.tensor([feat_len], device=device), max_len=feat_len) + mel_cond = torch.zeros([1, feat_len, self.mel_dim], device=device) + + prompt_feat_len = 0 + prompt_feat = self._to_2d(prompt_feat, device, mel_cond.dtype) + if prompt_feat is not None: + if prompt_feat.ndim == 2: + prompt_feat = prompt_feat.unsqueeze(0) + copy_len = min(prompt_feat.shape[1], feat_len) + mel_cond[:, :copy_len, :] = prompt_feat[:, -copy_len:, :] + prompt_feat_len = copy_len + has_prompt_audio_context = has_prompt_audio_context and prompt_feat_len > 0 + + embedding = self._to_2d(embedding, device, mel_cond.dtype) + spkr_embedding = None + if not self.remove_spkr_concat_condition: + if has_voice_clone and embedding is not None: + spk_embedding_normed = torch.nn.functional.normalize(embedding, dim=1) + spk_embedding = self.spk_embed_affine_layer(spk_embedding_normed) + spk_embedding_expanded = spk_embedding.unsqueeze(1).expand(-1, feat_len, -1) + condition = torch.cat([mel_cond, spk_embedding_expanded], dim=-1) + spkr_embedding = spk_embedding_normed + else: + spk_zeros = torch.zeros([1, feat_len, self.spk_embed_dim], device=device) + condition = torch.cat([mel_cond, spk_zeros], dim=-1) + else: + condition = mel_cond + + if self.spkr_emb_adaln: + if has_voice_clone and embedding is not None: + if spkr_embedding is None: + spkr_embedding = torch.nn.functional.normalize(embedding, dim=1) + else: + spkr_embedding = torch.zeros([1, 192], device=device, dtype=mel_cond.dtype) + + with self._stream_cache_lock: + _rs = self._stream_state_by_req.get(req_id) if req_id is not None else None + last_step_cache = _rs.diff_cache if _rs is not None else None + + dit_block_pattern: list[int] | None = None + is_streaming = req_id is not None and token_offset >= 0 + if is_streaming: + pattern_source = block_pattern or chunk_sizes_history + if pattern_source is not None and len(pattern_source) > 0: + if prompt_token_len > 0: + dit_block_pattern = [prompt_token_len] + pattern_source + else: + dit_block_pattern = list(pattern_source) + + lookahead_mel_len = 0 + if is_streaming and chunk_sizes_history is not None and len(chunk_sizes_history) > 0: + last_chunk_tokens = chunk_sizes_history[-1] + lookahead_tokens = last_chunk_tokens // 2 + lookahead_mel_len = int(lookahead_tokens / self.input_frame_rate * self.mel_framerate) + + # Compute cache_len: only cache the stable region for streaming + stream_cache_len: int | None = None + if req_id is not None: + stable_mel_len = int(generated_token_len / self.input_frame_rate * self.mel_framerate) + override_mel = max(0, stable_mel_len - lookahead_mel_len) + stream_cache_len = min(feat_len, prompt_feat_len + override_mel + 1) + + mel, current_step_cache = self._do_sample( + speech_tokens=speech_tokens, + mel_cond=mel_cond, + condition=condition, + padding_mask=padding_mask, + spkr_embedding=spkr_embedding, + n_timesteps=self.n_timesteps, + last_step_cache=last_step_cache, + block_pattern=dit_block_pattern, + cache_len=stream_cache_len, + text_lens=text_lens, + ) + + if req_id is not None: + current_step_cache["override_len"] = int(stream_cache_len or 0) + with self._stream_cache_lock: + rs = self._stream_state_by_req.setdefault(req_id, _StreamRequestState()) + rs.diff_cache = current_step_cache + + if has_prompt_audio_context: + crop_len = min(prompt_feat_len, max(0, int(mel.shape[1]) - 1)) + if crop_len > 0: + mel = mel[:, crop_len:, :] + + self._ensure_vocoder_loaded() + + audio = mel_to_audio(self._vocoder, mel) + + effective_sr = self._vocoder_sample_rate or self.sample_rate + if effective_sr != self.sample_rate and audio is not None and audio.numel() > 0: + try: + import torchaudio + + audio = torchaudio.functional.resample(audio, effective_sr, self.sample_rate) + except Exception as e: + logger.warning( + "Failed to resample audio from %d to %d: %s", + effective_sr, + self.sample_rate, + e, + ) + + audio = audio.reshape(-1) + if is_streaming and req_id is not None: + audio = self._slice_streaming_audio( + audio, + req_id=req_id, + stream_finished=stream_finished, + lookahead_mel_len=lookahead_mel_len, + crossfade_sec=crossfade_sec, + ) + + return audio + + def _slice_streaming_audio( + self, + audio: torch.Tensor, + *, + req_id: str, + stream_finished: bool, + lookahead_mel_len: int, + crossfade_sec: float, + ) -> torch.Tensor: + """Slice cumulative streaming audio using the official wav pointer. + + Official GLM-TTS runs the flow/vocoder on the cumulative token prefix + every chunk, then drops the waveform samples that were already emitted. + Non-final chunks additionally keep a lookahead plus fade tail for the + next chunk; the final full-prefix pass flushes that tail. + """ + look_back_len = max(0, int(lookahead_mel_len * self.sample_rate / self.mel_framerate)) + fade_len = max(0, int(crossfade_sec * self.sample_rate)) + overlap_len = look_back_len + fade_len + + with self._stream_cache_lock: + _rs = self._stream_state_by_req.get(req_id) + wav_pointer = _rs.wav_pointer if _rs is not None else 0 + fade_tail = _rs.fade_tail if _rs is not None else None + + current = audio[min(wav_pointer, int(audio.numel())) :].clone() + if fade_tail is not None and fade_tail.numel() > 0 and current.numel() > 0: + blend_len = min(int(fade_tail.numel()), int(current.numel())) + fade_in = torch.linspace(0.0, 1.0, blend_len, dtype=current.dtype, device=current.device) + fade_out = 1.0 - fade_in + current[:blend_len] = ( + fade_tail[:blend_len].to(device=current.device, dtype=current.dtype) * fade_out + + current[:blend_len] * fade_in + ) + + if stream_finished: + with self._stream_cache_lock: + _rs = self._stream_state_by_req.get(req_id) + if _rs is not None: + _rs.wav_pointer = 0 + _rs.fade_tail = None + return current + + if overlap_len <= 0 or current.numel() <= overlap_len: + emit = torch.zeros(0, dtype=current.dtype, device=current.device) + next_tail = current.detach().clone() + else: + emit = current[:-overlap_len] + tail_end = -look_back_len if look_back_len > 0 else None + next_tail = current[-overlap_len:tail_end].detach().clone() + + with self._stream_cache_lock: + rs = self._stream_state_by_req.setdefault(req_id, _StreamRequestState()) + rs.wav_pointer = wav_pointer + int(emit.numel()) + rs.fade_tail = next_tail + return emit + + # Maximum number of forward calls before a cache entry is considered stale. + _STALE_THRESHOLD = 64 + + def _evict_stale_stream_caches(self) -> None: + """Evict streaming cache entries for requests not seen recently. + + This is a self-contained fallback: if ``stream_finished`` is never + sent (e.g. client disconnect), the entry would otherwise leak forever. + Called at the top of every ``forward`` so no changes to the shared + model runner are needed. + """ + self._stream_forward_count += 1 + # Only check every 16 forward calls to amortize the lock overhead. + if self._stream_forward_count % 16 != 0: + return + with self._stream_cache_lock: + if not self._stream_last_seen: + return + stale_ids = [ + rid + for rid, last_seen in self._stream_last_seen.items() + if self._stream_forward_count - last_seen > self._STALE_THRESHOLD + ] + for rid in stale_ids: + self._stream_state_by_req.pop(rid, None) + self._stream_last_seen.pop(rid, None) + + def _do_sample( + self, + speech_tokens: torch.Tensor, + mel_cond: torch.Tensor, + condition: torch.Tensor, + padding_mask: torch.Tensor, + spkr_embedding: torch.Tensor | None, + n_timesteps: int, + last_step_cache: dict[int | str, Any] | None = None, + block_pattern: list[int] | None = None, + cache_len: int | None = None, + text_lens: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, dict[int | str, Any]]: + """Euler ODE sampling for flow matching.""" + # CUDA graph path: the captured graph does NOT pass is_causal / + # block_pattern to the DiT, so it always uses bidirectional attention. + # That is correct for non-streaming (full-sequence) inference but + # wrong for streaming where block-causal attention is required. + # Skip the graph when block_pattern is set so we fall through to the + # eager path that correctly forwards these arguments. + use_cudagraph = block_pattern is None + graph_result = ( + self._try_cudagraph_sample( + speech_tokens=speech_tokens, + mel_cond=mel_cond, + condition=condition, + padding_mask=padding_mask, + spkr_embedding=spkr_embedding, + n_timesteps=n_timesteps, + last_step_cache=last_step_cache, + cache_len=cache_len, + text_lens=text_lens, + ) + if use_cudagraph + else None + ) + if graph_result is not None: + return graph_result + + device = speech_tokens.device + # Always accumulate the Euler ODE state in float32 regardless of the + # DiT model dtype. Half-precision rounding errors across 10 steps (×2 + # CFG passes each) manifest as audible hiss on H20/H100. + x = torch.randn(mel_cond.shape, device=device, dtype=torch.float32) + # Cast conditioning tensors to float32 for the ODE loop. + condition_f32 = condition.to(torch.float32) + spkr_embedding_f32 = spkr_embedding.to(torch.float32) if spkr_embedding is not None else None + # Pre-cast invariant conditioning tensors to DiT dtype once (avoids + # repeated .to() calls inside the inner Euler loop). + condition_dit = condition_f32.to(mel_cond.dtype) + spkr_embedding_dit = spkr_embedding_f32.to(mel_cond.dtype) if spkr_embedding_f32 is not None else None + current_step_cache: dict[int | str, Any] = {} + effective_cache_len = min(int(cache_len or 0), int(mel_cond.shape[1])) + should_cache_steps = effective_cache_len > 0 + t_span = torch.linspace(0, 1, n_timesteps + 1, device=device, dtype=torch.float32) + if self.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + + for step in range(n_timesteps): + cache_step = step + 1 + if last_step_cache is not None and cache_step in last_step_cache: + x_cache = last_step_cache[cache_step].get("x") + if isinstance(x_cache, torch.Tensor): + override_len = int(last_step_cache.get("override_len", x_cache.shape[1])) + safe_len = min(int(x.shape[1]), int(x_cache.shape[1]), max(0, override_len)) + if safe_len > 0: + x[:, :safe_len, :] = x_cache[:, :safe_len, :].to(device=device, dtype=torch.float32) + if should_cache_steps: + current_step_cache[cache_step] = {"x": x[:, :effective_cache_len, :].detach().clone()} + + t_current = t_span[step] + dt = t_span[step + 1] - t_current + + use_causal = block_pattern is not None + x_dit = x.to(mel_cond.dtype) + + if self.inference_cfg_rate > 0: + # Batch conditional + unconditional into a single forward + middle_b = x_dit.expand(2, -1, -1) + condition_b = torch.cat([condition_dit, torch.zeros_like(condition_dit)], dim=0) + if self.speech_token_cfg: + text_b = torch.cat([speech_tokens, torch.zeros_like(speech_tokens)], dim=0) + else: + text_b = speech_tokens.expand(2, -1) + time_b = t_current.unsqueeze(0).expand(2) + padding_b = padding_mask.expand(2, -1) + if spkr_embedding_dit is not None and self.spkr_emb_adaln: + spkr_b = torch.cat([spkr_embedding_dit, torch.zeros_like(spkr_embedding_dit)], dim=0) + else: + spkr_b = None + text_lens_b = text_lens.expand(2) if text_lens is not None else None + pred_b = self.dit( + middle_point=middle_b, + condition=condition_b, + text=text_b, + time_step=time_b, + padding_mask=padding_b, + spkr_emb=spkr_b, + is_causal=use_causal, + block_pattern=block_pattern, + text_lens=text_lens_b, + ).to(torch.float32) + dphi_dt, cfg_dphi_dt = pred_b.chunk(2, dim=0) + dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt + else: + dphi_dt = self.dit( + middle_point=x_dit, + condition=condition_dit, + text=speech_tokens, + time_step=t_current.unsqueeze(0), + padding_mask=padding_mask, + spkr_emb=spkr_embedding_dit if spkr_embedding_dit is not None and self.spkr_emb_adaln else None, + is_causal=use_causal, + block_pattern=block_pattern, + text_lens=text_lens, + ).to(torch.float32) + + x = x + dt * dphi_dt + + return x.to(mel_cond.dtype), current_step_cache + + def _try_cudagraph_sample( + self, + *, + speech_tokens: torch.Tensor, + mel_cond: torch.Tensor, + condition: torch.Tensor, + padding_mask: torch.Tensor, + spkr_embedding: torch.Tensor | None, + n_timesteps: int, + last_step_cache: dict[int | str, Any] | None, + cache_len: int | None = None, + text_lens: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, dict[int | str, Any]] | None: + if not self._use_dit_cuda_graphs or mel_cond.device.type != "cuda": + return None + if self._dit_cudagraph is None: + condition_dim = self.mel_dim if self.remove_spkr_concat_condition else self.mel_dim + self.spk_embed_dim + self._dit_cudagraph = CUDAGraphGLMTTSDiTWrapper( + self.dit, + mel_dim=self.mel_dim, + condition_dim=condition_dim, + input_frame_rate=self.input_frame_rate, + mel_framerate=self.mel_framerate, + inference_cfg_rate=self.inference_cfg_rate, + speech_token_cfg=self.speech_token_cfg, + spkr_emb_adaln=self.spkr_emb_adaln, + t_scheduler=self.t_scheduler, + enabled=True, + ) + try: + self._dit_cudagraph.warmup( + device=mel_cond.device, + dtype=mel_cond.dtype, + n_timesteps=n_timesteps, + codec_chunk_frames=self._codec_chunk_frames, + codec_left_context_frames=self._codec_left_context_frames, + ) + except Exception: + logger.warning("Disabling GLM-TTS DiT CUDA graphs after warmup failure", exc_info=True) + self._dit_cudagraph.enabled = False + return None + + try: + return self._dit_cudagraph.sample( + speech_tokens=speech_tokens, + mel_cond=mel_cond, + condition=condition, + padding_mask=padding_mask, + spkr_embedding=spkr_embedding, + n_timesteps=n_timesteps, + last_step_cache=last_step_cache, + cache_len=cache_len, + text_lens=text_lens, + ) + except Exception: + logger.warning("GLM-TTS DiT CUDA graph replay failed; falling back to eager", exc_info=True) + return None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load DiT + vocoder weights from {model_dir}/flow/flow.pt.""" + for _ in weights: + pass + + weight_dict: dict[str, torch.Tensor] = {} + if not os.path.exists(self._flow_weights_path): + raise FileNotFoundError(f"GLM-TTS flow weights not found at {self._flow_weights_path}") + + logger.info("Loading GLM-TTS flow weights from %s", self._flow_weights_path) + checkpoint = torch.load(self._flow_weights_path, map_location="cpu", weights_only=True) + if "model" in checkpoint: + weight_dict = checkpoint["model"] + else: + weight_dict = checkpoint + + dit_state = { + (name.removeprefix("estimator.") if name.startswith("estimator.") else name.removeprefix("dit.")): tensor + for name, tensor in weight_dict.items() + if (name.startswith("estimator.") or name.startswith("dit.")) + and "inv_freq" not in name + and "rotary_embed" not in name + } + if not dit_state: + raise RuntimeError( + "GLM-TTS flow checkpoint did not contain any DiT weights under 'estimator.*' or 'dit.*'." + ) + missing, unexpected = self.dit.load_state_dict(dit_state, strict=False) + missing = [k for k in missing if "inv_freq" not in k] + if missing: + logger.warning("Missing DiT weights: %s", missing[:10]) + if unexpected: + logger.warning("Unexpected DiT weights: %s", unexpected[:10]) + + loaded_spk: set[str] = set() + if hasattr(self, "spk_embed_affine_layer"): + for suffix in ("weight", "bias"): + key = f"spk_embed_affine_layer.{suffix}" + if key in weight_dict: + param = getattr(self.spk_embed_affine_layer, suffix) + param.data.copy_(weight_dict[key].to(device=param.device, dtype=param.dtype)) + loaded_spk.add(key) + self.spk_embed_affine_layer.eval() + self.dit.eval() + self._ensure_vocoder_loaded() + + # Return both prefixed and unprefixed param names for nested module compat + loaded_params: set[str] = set() + for k in dit_state: + loaded_params.update((f"dit.{k}", f"_dit_gen.dit.{k}")) + for k in loaded_spk: + loaded_params.update((k, f"_dit_gen.{k}")) + logger.info("Loaded GLMTTSDiTForGeneration: %d params", len(loaded_params)) + return loaded_params + + +class CUDAGraphGLMTTSDiTWrapper: + """Capture GLM-TTS DiT Euler sampling for fixed mel-frame buckets. + + The wrapper owns static buffers for a single-request DiT stage. Dynamic + requests are right-padded to the nearest captured mel-frame bucket and + trimmed back to the actual mel length after graph replay. + """ + + def __init__( + self, + dit: torch.nn.Module, + *, + mel_dim: int, + condition_dim: int, + input_frame_rate: float, + mel_framerate: float, + inference_cfg_rate: float, + speech_token_cfg: bool, + spkr_emb_adaln: bool, + t_scheduler: str, + enabled: bool = True, + capture_sizes: list[int] | None = None, + ) -> None: + self.dit = dit + self.mel_dim = int(mel_dim) + self.condition_dim = int(condition_dim) + self.input_frame_rate = float(input_frame_rate) + self.mel_framerate = float(mel_framerate) + self.inference_cfg_rate = float(inference_cfg_rate) + self.speech_token_cfg = bool(speech_token_cfg) + self.spkr_emb_adaln = bool(spkr_emb_adaln) + self.t_scheduler = str(t_scheduler) + self.enabled = bool(enabled) + self.capture_sizes = sorted(set(capture_sizes or [])) + + self.graphs: dict[int, CUDAGraph] = {} + self.static_x: dict[int, torch.Tensor] = {} + self.static_condition: dict[int, torch.Tensor] = {} + self.static_text: dict[int, torch.Tensor] = {} + self.static_text_lens: dict[int, torch.Tensor] = {} + self.static_padding_mask: dict[int, torch.Tensor] = {} + self.static_spkr_emb: dict[int, torch.Tensor] = {} + self.static_output: dict[int, torch.Tensor] = {} + self.static_step_cache: dict[int, list[torch.Tensor]] = {} + + self._warmed_up = False + self._n_timesteps: int | None = None + + @staticmethod + def compute_capture_sizes( + *, + codec_chunk_frames: int = 25, + codec_left_context_frames: int = 25, + input_frame_rate: float = 25.0, + mel_framerate: float = 50.0, + max_bucket: int = 1024, + ) -> list[int]: + sizes: set[int] = set() + + def token_to_mel(tokens: int) -> int: + return max(1, int(math.ceil(float(tokens) / input_frame_rate * mel_framerate))) + + if codec_chunk_frames > 0: + sizes.add(token_to_mel(codec_chunk_frames)) + if codec_left_context_frames > 0: + sizes.add(token_to_mel(codec_chunk_frames + codec_left_context_frames)) + + for size in (32, 64, 128, 256, 512, 1024): + if size <= max_bucket: + sizes.add(size) + return sorted(sizes) + + def _get_padded_size(self, actual_feat_len: int) -> int | None: + for size in self.capture_sizes: + if actual_feat_len <= size: + return size + return None + + def _token_bucket_len(self, bucket_feat_len: int) -> int: + return max(1, int(math.ceil(float(bucket_feat_len) * self.input_frame_rate / self.mel_framerate))) + + def warmup( + self, + *, + device: torch.device, + dtype: torch.dtype, + n_timesteps: int, + codec_chunk_frames: int = 25, + codec_left_context_frames: int = 25, + ) -> None: + if not self.enabled or self._warmed_up or device.type != "cuda" or torch.cuda.is_current_stream_capturing(): + return + self._n_timesteps = int(n_timesteps) + if not self.capture_sizes: + self.capture_sizes = self.compute_capture_sizes( + codec_chunk_frames=codec_chunk_frames, + codec_left_context_frames=codec_left_context_frames, + input_frame_rate=self.input_frame_rate, + mel_framerate=self.mel_framerate, + ) + + self.dit.eval() + logger.info("Starting GLM-TTS DiT CUDA graph warmup for sizes: %s", self.capture_sizes) + for size in self.capture_sizes: + try: + self._allocate_static_buffers(size, device, dtype) + self._capture(size, device) + logger.info("Captured GLM-TTS DiT CUDA graph for mel frames=%d", size) + except Exception: + logger.warning("Failed to capture GLM-TTS DiT CUDA graph for size=%d", size, exc_info=True) + self._warmed_up = True + + def _allocate_static_buffers(self, size: int, device: torch.device, dtype: torch.dtype) -> None: + token_len = self._token_bucket_len(size) + self.static_x[size] = torch.zeros(1, size, self.mel_dim, device=device, dtype=dtype) + self.static_condition[size] = torch.zeros(1, size, self.condition_dim, device=device, dtype=dtype) + self.static_text[size] = torch.zeros(1, token_len, device=device, dtype=torch.long) + self.static_text_lens[size] = torch.full((1,), token_len, device=device, dtype=torch.long) + self.static_padding_mask[size] = torch.zeros(1, size, device=device, dtype=torch.bool) + self.static_spkr_emb[size] = torch.zeros(1, 192, device=device, dtype=dtype) + self.static_output[size] = torch.zeros(1, size, self.mel_dim, device=device, dtype=dtype) + self.static_step_cache[size] = [ + torch.zeros(1, size, self.mel_dim, device=device, dtype=dtype) for _ in range(int(self._n_timesteps or 0)) + ] + + def _time_schedule(self, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + n_timesteps = int(self._n_timesteps or 0) + t_span = torch.linspace(0, 1, n_timesteps + 1, device=device, dtype=dtype) + if self.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + return t_span + + def _run_static(self, size: int) -> torch.Tensor: + x = self.static_x[size] + condition = self.static_condition[size] + text = self.static_text[size] + text_lens = self.static_text_lens[size] + padding_mask = self.static_padding_mask[size] + spkr_emb = self.static_spkr_emb[size] if self.spkr_emb_adaln else None + t_span = self._time_schedule(x.device, x.dtype) + + for step in range(int(self._n_timesteps or 0)): + self.static_step_cache[size][step].copy_(x) + t_current = t_span[step] + dt = t_span[step + 1] - t_current + if self.inference_cfg_rate > 0: + middle_b = x.expand(2, -1, -1) + condition_b = torch.cat([condition, torch.zeros_like(condition)], dim=0) + if self.speech_token_cfg: + text_b = torch.cat([text, torch.zeros_like(text)], dim=0) + else: + text_b = text.expand(2, -1) + time_b = t_current.unsqueeze(0).expand(2) + padding_b = padding_mask.expand(2, -1) + text_lens_b = text_lens.expand(2) + if spkr_emb is not None: + spkr_b = torch.cat([spkr_emb, torch.zeros_like(spkr_emb)], dim=0) + else: + spkr_b = None + pred_b = self.dit( + middle_point=middle_b, + condition=condition_b, + text=text_b, + text_lens=text_lens_b, + time_step=time_b, + padding_mask=padding_b, + spkr_emb=spkr_b, + ) + dphi_dt, cfg_dphi_dt = pred_b.chunk(2, dim=0) + dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt + else: + dphi_dt = self.dit( + middle_point=x, + condition=condition, + text=text, + text_lens=text_lens, + time_step=t_current.unsqueeze(0), + padding_mask=padding_mask, + spkr_emb=spkr_emb, + ) + x = x + dt * dphi_dt + self.static_output[size].copy_(x) + return self.static_output[size] + + def _capture(self, size: int, device: torch.device) -> None: + with torch.no_grad(): + self.static_x[size].normal_() + _ = self._run_static(size) + torch.accelerator.synchronize(device) + + graph = CUDAGraph() + with torch.no_grad(): + with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + self._run_static(size) + self.graphs[size] = graph + + def sample( + self, + *, + speech_tokens: torch.Tensor, + mel_cond: torch.Tensor, + condition: torch.Tensor, + padding_mask: torch.Tensor, + spkr_embedding: torch.Tensor | None, + n_timesteps: int, + last_step_cache: dict[int | str, Any] | None = None, + cache_len: int | None = None, + text_lens: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, dict[int | str, Any]] | None: + if ( + not self.enabled + or not self._warmed_up + or int(n_timesteps) != int(self._n_timesteps or -1) + or speech_tokens.shape[0] != 1 + or mel_cond.shape[0] != 1 + or mel_cond.device.type != "cuda" + or torch.cuda.is_current_stream_capturing() + ): + return None + + actual_feat_len = int(mel_cond.shape[1]) + size = self._get_padded_size(actual_feat_len) + if size is None or size not in self.graphs: + return None + + token_bucket_len = int(self.static_text[size].shape[1]) + if int(speech_tokens.shape[1]) > token_bucket_len: + return None + + self.static_x[size].normal_() + self.static_condition[size].zero_() + self.static_text[size].zero_() + self.static_padding_mask[size].zero_() + self.static_spkr_emb[size].zero_() + + # Warm-start: when streaming, pre-fill the stable region of the + # initial noise from the step-1 cache of the previous chunk. + # The graph's Euler loop then evolves from this warm-started state. + # New (lookahead) region keeps random noise — no per-step overlay. + if last_step_cache is not None and 1 in last_step_cache: + cached_x = last_step_cache[1]["x"] + override_len = int(last_step_cache.get("override_len", cached_x.shape[1])) + safe = min(size, int(cached_x.shape[1]), max(0, override_len)) + if safe > 0: + self.static_x[size][:, :safe, :].copy_(cached_x[:, :safe, :].to(self.static_x[size])) + + self.static_condition[size][:, :actual_feat_len, : condition.shape[-1]] = condition + self.static_text[size][:, : speech_tokens.shape[1]] = speech_tokens + actual_text_len = int(speech_tokens.shape[1]) + if text_lens is not None: + self.static_text_lens[size].copy_(text_lens) + else: + self.static_text_lens[size].fill_(actual_text_len) + self.static_padding_mask[size][:, :actual_feat_len] = padding_mask + if spkr_embedding is not None: + self.static_spkr_emb[size].copy_(spkr_embedding.to(self.static_spkr_emb[size])) + + self.graphs[size].replay() + mel = self.static_output[size][:, :actual_feat_len, :].clone() + current_step_cache: dict[int | str, Any] = {} + cache_len = min(int(cache_len or 0), actual_feat_len) + if cache_len > 0: + for idx, cache in enumerate(self.static_step_cache[size], start=1): + current_step_cache[idx] = {"x": cache[:, :cache_len, :].detach().clone()} + return mel, current_step_cache diff --git a/vllm_omni/model_executor/models/glm_tts/pipeline.py b/vllm_omni/model_executor/models/glm_tts/pipeline.py new file mode 100644 index 000000000..f016a993e --- /dev/null +++ b/vllm_omni/model_executor/models/glm_tts/pipeline.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-TTS pipeline: Stage 0 (AR) → Stage 1 (DiT).""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.glm_tts" + +GLM_TTS_PIPELINE = PipelineConfig( + model_type="glm_tts", + model_arch="GLMTTSForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="glm_tts", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + engine_output_type="latent", + async_chunk_process_next_stage_input_func=(f"{_PROC}.ar_to_dit_async_chunk"), + sampling_constraints={ + # GLM-TTS uses 👂 (token string "👂", ID 59253) as + # end-of-audio marker. The ID is resolved dynamically by + # GLMTTSForConditionalGeneration.__init__ from the tokenizer and + # validated at runtime against this hardcoded value. If the + # upstream checkpoint changes the mapping, the model will log a + # warning so the constant here can be updated. + "stop_token_ids": [59253], + }, + ), + StagePipelineConfig( + stage_id=1, + model_stage="glm_tts_dit", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + engine_output_type="latent", + sync_process_input_func=f"{_PROC}.ar_to_dit", + ), + ), +) diff --git a/vllm_omni/model_executor/models/glm_tts/text_frontend.py b/vllm_omni/model_executor/models/glm_tts/text_frontend.py new file mode 100644 index 000000000..8cfc0e5c8 --- /dev/null +++ b/vllm_omni/model_executor/models/glm_tts/text_frontend.py @@ -0,0 +1,393 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lightweight text frontend aligned with the official GLM-TTS preprocessing.""" + +from __future__ import annotations + +import json +import os +import re + +try: + import emoji +except ImportError: + emoji = None # type: ignore[assignment] + +try: + import inflect +except ImportError: + inflect = None # type: ignore[assignment] + +try: + from tn.chinese.normalizer import Normalizer as ZhNormalizer +except ImportError: + ZhNormalizer = None + +try: + from tn.english.normalizer import Normalizer as EnNormalizer +except ImportError: + EnNormalizer = None + + +CHINESE_CHAR_PATTERN = re.compile(r"[\u4e00-\u9fff]+") +PUNCTUATION_CHARS = r"。？！；：、.?!;:，," + + +def contains_chinese(text: str) -> bool: + return bool(CHINESE_CHAR_PATTERN.search(text)) + + +def markdown_norm(markdown_text: str) -> str: + markdown_text = re.sub(r"^(\d+)\. ", r"\1。", markdown_text) + return markdown_text.replace("\\n", "\n") + + +def multi_line_process(plain_text: str) -> str: + lines: list[str] = [] + for line in plain_text.split("\n"): + line = line.strip() + if not line: + continue + if line[-1] not in {".", "!", "?", ";", ":", "：", "。", "！", "？", "；", "，"}: + line = f"{line}{'。' if contains_chinese(line) else '. '}" + lines.append(line) + return "".join(lines) if contains_chinese("".join(lines)) else " ".join(lines) + + +def remove_bracket(text: str, lang: str = "zh") -> str: + brackets_to_remove = [ + ("(", ")"), + ("（", "）"), + ("【", "】"), + ("「", "」"), + ("`", "`"), + ("《", "》"), + ("『", "』"), + ("{", "}"), + ("[", "]"), + ] + if lang != "en": + brackets_to_remove.append(("‘", "’")) + for left, right in brackets_to_remove: + text = text.replace(left, "").replace(right, "") + return text + + +def number_to_chinese(number: int | float | str) -> str: + units = ["", "十", "百", "千", "万", "十万", "百万", "千万", "亿", "十亿", "百亿", "千亿"] + nums = "零一二三四五六七八九" + if isinstance(number, (float, str)): + str_num = str(float(number)) + if "." in str_num: + int_part, decimal_part = str_num.split(".") + if decimal_part.strip("0") == "": + return number_to_chinese(int(int_part)) + chinese_int = number_to_chinese(int(int_part)) if int(int_part) != 0 else "零" + chinese_decimal = "".join(nums[int(d)] for d in decimal_part) + return f"{chinese_int}点{chinese_decimal}" + + number = int(number) + if number == 0: + return "零" + if number < 0: + return "负" + number_to_chinese(abs(number)) + + result: list[str] = [] + unit_position = 0 + zero_flag = False + last_unit = "" + while number > 0: + num = number % 10 + current_unit = units[unit_position] if unit_position < len(units) else "" + if num == 0: + if not zero_flag and result and last_unit not in ["万", "亿"]: + result.append(nums[num]) + zero_flag = True + else: + if current_unit: + result.append(current_unit) + result.append(nums[num]) + zero_flag = False + last_unit = current_unit + unit_position += 1 + number //= 10 + + result.reverse() + text = "".join(result) + text = re.sub(r"^一十", "十", text) + text = re.sub(r"零+", "零", text) + text = re.sub(r"零$", "", text) + return text + + +def tn_scientific_notation(text: str) -> str: + pattern = re.compile(r"(.*?)(-?\d+(?:\.\d+)?)[*x×](\d+(?:\.\d+)?)\^(-?\d+(?:\.\d+)?)(.*?)$") + match = pattern.match(text) + if not match: + return text + prefix, base, multiplier, exponent, suffix = match.groups() + description = f"{number_to_chinese(base)}乘{number_to_chinese(multiplier)}的{number_to_chinese(exponent)}次方" + return f"{prefix}{description}{suffix}" + + +def replace_asterisk_with_multiply(text: str, lang: str) -> str: + if lang == "en": + replace_name = "multiply" + replace_number = "0-9" + else: + replace_name = "乘" + replace_number = "一二三四五六七八九十百千万亿" + + rule0 = rf"(?<=[{replace_number}a-zA-Z])\s*\*\s*(?=[{replace_number}a-zA-Z])" + rule1 = rf"(?<=[{replace_number}a-zA-Z]\))\s*\*\s*(?=[{replace_number}a-zA-Z])" + rule2 = rf"(?<=[{replace_number}a-zA-Z])\s*\*\s*(?=\([{replace_number}a-zA-Z])" + rule3 = rf"(?<=[{replace_number}a-zA-Z]\))\s*\*\s*(?=\([{replace_number}a-zA-Z])" + return re.sub(f"{rule0}|{rule1}|{rule2}|{rule3}", replace_name, text).replace("*", "") + + +def special_replace(text: str) -> str: + text = text.replace("\\", "") + text = f" {text} " + text = re.sub(r"(?<=\W)额(?=\W)", "呃", text) + text = re.sub(r"[~～]+", "。", text) + text = replace_asterisk_with_multiply(text, "zh") + return text.strip() + + +def spell_out_number(text: str, inflect_parser: object | None) -> str: + if inflect_parser is None: + return text + new_text: list[str] = [] + start = None + for idx, char in enumerate(text): + if not char.isdigit(): + if start is not None: + new_text.append(inflect_parser.number_to_words(text[start:idx])) + start = None + new_text.append(char) + elif start is None: + start = idx + if start is not None and start < len(text): + new_text.append(inflect_parser.number_to_words(text[start:])) + return "".join(new_text) + + +def replace_space(text: str) -> str: + alphanumeric_pattern = r"[a-zA-Z0-9]" + punctuation_pattern = r"[.,!?;:]" + text = re.sub(r" +", " ", text) + result = "" + idx = 0 + while idx < len(text): + current_char = text[idx] + if current_char != " ": + result += current_char + idx += 1 + continue + prev_char = text[idx - 1] if idx > 0 else "" + next_char = text[idx + 1] if idx + 1 < len(text) else "" + if re.match(alphanumeric_pattern, prev_char) and re.match(alphanumeric_pattern, next_char): + result += " " + elif re.match(punctuation_pattern, prev_char) and re.match(alphanumeric_pattern, next_char): + result += " " + idx += 1 + return result + + +def normalize_punctuation(text: str, punctuation_chars: str) -> str: + text = replace_space(text) + text = re.sub(rf"([{punctuation_chars}])\1+", r"\1", text) + text = re.sub(rf"([{punctuation_chars}])(?=[{punctuation_chars}])", "", text) + text = text.replace("#", "") + text = text.replace("！", "。") + text = text.replace("!", ".") + return text + + +def ensure_proper_ending(text: str) -> str: + if not text: + return text + if text[-1] in "？?": + return text + if text[-1] in PUNCTUATION_CHARS: + return f"{text[:-1]}{'。' if contains_chinese(text) else '.'}" + return f"{text}{'。' if contains_chinese(text) else '.'}" + + +class GLMTTSTextFrontend: + """Subset of the official GLM-TTS text frontend used by the adapter.""" + + def __init__(self) -> None: + self.inflect_parser = inflect.engine() if inflect is not None else None + self.zh_tn_model = ( + ZhNormalizer( + remove_erhua=False, + full_to_half=True, + remove_interjections=False, + overwrite_cache=True, + ) + if ZhNormalizer is not None + else None + ) + self.en_tn_model = EnNormalizer() if EnNormalizer is not None else None + self._custom_replacements = self._load_custom_replacements() + + @staticmethod + def _load_custom_replacements() -> list[tuple[str, str]]: + custom_replace_path = os.path.join(os.path.dirname(__file__), "configs", "custom_replace.jsonl") + if not os.path.exists(custom_replace_path): + return [] + replacements: list[tuple[str, str]] = [] + with open(custom_replace_path, encoding="utf-8") as f: + for line in f: + if not line.strip(): + continue + item = json.loads(line) + replacements.append((item["origin"], item["new"])) + return replacements + + def pre_replace(self, sentence: str) -> str: + sentence = tn_scientific_notation(sentence) + sentence = re.sub(r"(?<=\d)\s*-\s*(?=\d)", "减", sentence) + sentence = sentence.replace("-", "") + sentence = re.sub(r"咯([" + re.escape(PUNCTUATION_CHARS) + r"])", r"喽\1", sentence) + for origin, new in self._custom_replacements: + sentence = sentence.replace(origin, new) + return sentence + + def post_replace(self, sentence: str) -> str: + sentence = remove_bracket(sentence) + sentence = sentence.replace(" - ", "，") + sentence = sentence.replace("——", "，") + sentence = re.sub(r"[,:：;；、]+", "，", sentence) + sentence = re.sub(r"[.…]+", "。", sentence) + sentence = re.sub(r"[_·]+", "", sentence) + sentence = re.sub(r"""['"‘’“”|]+""", "", sentence) + + symbol_map = { + "†": "，", + "²": "平方", + "³": "立方", + "/": "每", + "~": "到", + "～": "到", + "①": "一", + "②": "二", + "③": "三", + "④": "四", + "⑤": "五", + "⑥": "六", + "⑦": "七", + "⑧": "八", + "⑨": "九", + "⑩": "十", + "α": "阿尔法", + "β": "贝塔", + "γ": "伽玛", + "Γ": "伽玛", + "δ": "德尔塔", + "Δ": "德尔塔", + "△": "德尔塔", + "ε": "艾普西龙", + "ζ": "捷塔", + "η": "依塔", + "θ": "西塔", + "Θ": "西塔", + "ι": "艾欧塔", + "κ": "喀帕", + "λ": "拉姆达", + "Λ": "拉姆达", + "μ": "缪", + "ν": "拗", + "ξ": "克西", + "Ξ": "克西", + "ο": "欧米克伦", + "π": "派", + "Π": "派", + "ρ": "肉", + "ς": "西格玛", + "Σ": "西格玛", + "σ": "西格玛", + "τ": "套", + "υ": "宇普西龙", + "φ": "服艾", + "Φ": "服艾", + "χ": "器", + "ψ": "普赛", + "Ψ": "普赛", + "ω": "欧米伽", + "Ω": "欧米伽", + "□": "方框", + ">": "大于", + "<": "小于", + "∈": "属于", + "∉": "不属于", + "∪": "并", + "∩": "交", + "⊥": "垂直", + "∥": "平行", + "≠": "不等于", + "∵": "因为", + "∴": "所以", + "∅": "空集", + "⊂": "真包含于", + "⊃": "包含", + "⊆": "包含于", + "⊇": "真包含", + "⊄": "不属于", + "⊅": "非超集", + "⊈": "不属于", + "⊉": "非超集", + } + for old, new in symbol_map.items(): + sentence = sentence.replace(old, new) + return sentence + + def text_normalize(self, text: str | None) -> str | None: + if text is None: + return None + text = self._preprocess_text(text) + if contains_chinese(text): + text = self._normalize_chinese_text(text).lower() + else: + text = self._normalize_english_text(text) + text = normalize_punctuation(text, PUNCTUATION_CHARS) + return ensure_proper_ending(text) + + def _preprocess_text(self, text: str) -> str: + text = markdown_norm(text) + text = multi_line_process(text) + if emoji is not None: + text = emoji.replace_emoji(text, replace="") + return re.sub(r"(?<=[a-zA-Z])-(?=[a-zA-Z])", " ", text) + + def _normalize_chinese_text(self, text: str) -> str: + text = self.pre_replace(text) + if self.zh_tn_model is not None: + text = self.zh_tn_model.normalize(text) + text = special_replace(text) + return self.post_replace(text).strip() + + def _normalize_english_text(self, text: str) -> str: + text = text.replace("'", "’") + if self.en_tn_model is not None: + text = self.en_tn_model.normalize(text) + text = remove_bracket(text, "en") + text = replace_asterisk_with_multiply(text, "en") + text = text.replace("—", " ") + text = text.replace("’", "'") + text = spell_out_number(text, self.inflect_parser) + text = re.sub(r"\s+", " ", text) + keep_punctuation = r"\.,!\?'\:;" + pattern = rf"[^\w\s{keep_punctuation}]" + text = re.sub(pattern, "", text) + text = text.lower() + text = re.sub(r"\.+", ".", text) + text = re.sub(r"\,+", ",", text) + text = re.sub(r"!+", "!", text) + text = re.sub(r"\?+", "?", text) + text = re.sub(r"'+", "'", text) + text = re.sub(r":+", ":", text) + text = re.sub(r";+", ";", text) + text = re.sub(r"\s*([.,?!':;])\s*", r"\1 ", text) + return text.strip() diff --git a/vllm_omni/model_executor/models/glm_tts/vocoder.py b/vllm_omni/model_executor/models/glm_tts/vocoder.py new file mode 100644 index 000000000..77c61c0de --- /dev/null +++ b/vllm_omni/model_executor/models/glm_tts/vocoder.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Vocoder loading and mel-to-audio conversion for GLM-TTS. + +Supports: + - HiFT (24kHz) — reuses CosyVoice3's vendored HiFTGenerator + - Vocos2D JIT (32kHz) from TorchScript checkpoint + +Extracted from glm_tts_dit_wrapper.py to keep file sizes under 800 lines. +""" + +from __future__ import annotations + +import os +from typing import Any + +import torch +import torch.nn as nn +from vllm.logger import init_logger + +try: + from torch.nn.utils.parametrizations import weight_norm +except ImportError: + from torch.nn.utils import weight_norm + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# HiFT vocoder (24kHz) — reuses CosyVoice3's vendored HiFTGenerator +# --------------------------------------------------------------------------- + + +class ConvRNNF0Predictor(nn.Module): + """Non-causal F0 predictor for GLM-TTS HiFT vocoder. + + GLM-TTS ships a non-causal HiFT checkpoint whose F0 predictor uses + standard ``nn.Conv1d`` layers (kernel_size=3, padding=1). CosyVoice3's + vendored hifigan only includes the *causal* variant + (``CausalConvRNNF0Predictor``), so we provide the non-causal version + here. The network structure and state_dict keys are identical to the + original ``cosyvoice.hifigan_cosy2.f0_predictor.ConvRNNF0Predictor``. + """ + + def __init__( + self, + num_class: int = 1, + in_channels: int = 80, + cond_channels: int = 512, + ): + super().__init__() + self.num_class = num_class + self.condnet = nn.Sequential( + weight_norm(nn.Conv1d(in_channels, cond_channels, kernel_size=3, padding=1)), + nn.ELU(), + weight_norm(nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)), + nn.ELU(), + weight_norm(nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)), + nn.ELU(), + weight_norm(nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)), + nn.ELU(), + weight_norm(nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)), + nn.ELU(), + ) + self.classifier = nn.Linear(in_features=cond_channels, out_features=self.num_class) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.condnet(x) + x = x.transpose(1, 2) + return torch.abs(self.classifier(x).squeeze(-1)) + + +class HiFTWrapper: + """Thin wrapper around CosyVoice3's vendored HiFTGenerator.""" + + def __init__(self, state_dict: dict[str, Any], device: torch.device): + self.device = device + self.sample_rate = 24000 + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( + HiFTGenerator, + ) + + f0_predictor = ConvRNNF0Predictor(num_class=1, in_channels=80, cond_channels=512) + self.model = HiFTGenerator( + in_channels=80, + base_channels=512, + nb_harmonics=8, + sampling_rate=24000, + nsf_alpha=0.1, + nsf_sigma=0.003, + nsf_voiced_threshold=10, + upsample_rates=[8, 5, 3], + upsample_kernel_sizes=[16, 11, 7], + istft_params={"n_fft": 16, "hop_len": 4}, + resblock_kernel_sizes=[3, 7, 11], + resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]], + source_resblock_kernel_sizes=[7, 7, 11], + source_resblock_dilation_sizes=[ + [1, 3, 5], + [1, 3, 5], + [1, 3, 5], + ], + lrelu_slope=0.1, + audio_limit=0.99, + f0_predictor=f0_predictor, + ).to(device) + self.model.load_state_dict(state_dict) + self.model.eval() + self.model.float() + + def __call__(self, mel: torch.Tensor) -> torch.Tensor: + mel = mel.to(self.device) + with torch.no_grad(): + audio, _ = self.model.inference(mel) + return audio + + +def load_hift(ckpt_path: str, device: torch.device) -> HiFTWrapper: + """Load HiFT vocoder from checkpoint.""" + state_dict = torch.load(ckpt_path, map_location=device, weights_only=True) + if isinstance(state_dict, dict): + for key in ("state_dict", "model", "generator"): + nested = state_dict.get(key) + if isinstance(nested, dict): + state_dict = nested + break + return HiFTWrapper(state_dict, device) + + +# --------------------------------------------------------------------------- +# Vocos2D JIT vocoder (32kHz) +# --------------------------------------------------------------------------- + +_MEL_LOGDIFF = -7.847762537473608 + + +class Vocos2DWrapper: + """Thin wrapper around Vocos2D TorchScript model.""" + + def __init__(self, ckpt_path: str, device: torch.device): + self.device = device + self.gen_model = torch.jit.load(ckpt_path, map_location=device) + self.gen_model.eval() + self._mel_logdiff = _MEL_LOGDIFF + + def __call__(self, mel: torch.Tensor) -> torch.Tensor: + if mel.size(-1) == 1: + mel = torch.cat([mel, mel], dim=-1) + xs_mel = mel.transpose(-1, -2) + self._mel_logdiff + xs = self.gen_model(xs_mel.to(self.device)) + if xs.ndim == 2: + xs = xs.unsqueeze(1) + return xs + + +def load_vocos2d_jit(ckpt_path: str, device: torch.device) -> Vocos2DWrapper: + """Load Vocos2D JIT vocoder from TorchScript checkpoint.""" + return Vocos2DWrapper(ckpt_path, device) + + +# --------------------------------------------------------------------------- +# Unified loader +# --------------------------------------------------------------------------- + + +def load_vocoder( + model_root: str, + device: torch.device, + sample_rate: int = 24000, +) -> tuple[Any | None, int]: + """Try to load the best available vocoder. + + Returns: + (vocoder, actual_sample_rate) or (None, sample_rate) on failure. + """ + # Try HiFT first for 24kHz + if sample_rate == 24000: + hift_path = os.path.join(model_root, "hift", "hift.pt") + if os.path.isfile(hift_path): + try: + vocoder = load_hift(hift_path, device) + logger.info("Loaded HiFT vocoder from %s", hift_path) + return vocoder, 24000 + except Exception as e: + logger.warning("Failed to load HiFT: %s", e) + + # Try Vocos2D JIT + vocos_jit_path = os.path.join(model_root, "vocos2d", "generator_jit.ckpt") + if os.path.isfile(vocos_jit_path): + try: + vocoder = load_vocos2d_jit(vocos_jit_path, device) + logger.info("Loaded Vocos2D JIT vocoder (32kHz)") + return vocoder, 32000 + except Exception as e: + logger.warning("Failed to load Vocos2D JIT: %s", e) + + logger.warning("No vocoder available") + return None, sample_rate + + +def mel_to_audio( + vocoder: Any | None, + mel: torch.Tensor, +) -> torch.Tensor: + """Convert mel-spectrogram to audio waveform. + + Args: + vocoder: Vocoder instance (HiFTWrapper, Vocos2DWrapper, or None). + mel: Mel-spectrogram [B, T, mel_dim]. + + Returns: + Audio waveform [B, 1, samples]. + """ + # Transpose to [B, mel_dim, T] for vocoder. + # HiFT / Vocos2D are precision-sensitive; always feed float32. + mel = mel.to(torch.float32).transpose(1, 2) + + if vocoder is not None: + if callable(vocoder) and not hasattr(vocoder, "decode"): + audio = vocoder(mel) + else: + audio = vocoder.decode(mel) + if audio.ndim == 2: + audio = audio.unsqueeze(1) + else: + raise RuntimeError("GLM-TTS vocoder is not loaded.") + + return audio diff --git a/vllm_omni/model_executor/models/glm_tts/voice_clone.py b/vllm_omni/model_executor/models/glm_tts/voice_clone.py new file mode 100644 index 000000000..d4d32f4f6 --- /dev/null +++ b/vllm_omni/model_executor/models/glm_tts/voice_clone.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-TTS voice cloning frontend: speech tokenizer, speaker embedding, mel features. + +Lazy-loads external models (WhisperVQ, CampPlus ONNX) and provides extraction +helpers consumed by the GLM-TTS multimodal processor. +""" + +from __future__ import annotations + +import os + +import torch +import torchaudio.compliance.kaldi as kaldi +import torchaudio.transforms as T +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_MEL_BASIS_CACHE: dict[tuple[int, int, int, int, int, str], torch.Tensor] = {} +_HANN_WINDOW_CACHE: dict[tuple[int, str], torch.Tensor] = {} + + +def _spectral_normalize_torch(magnitudes: torch.Tensor) -> torch.Tensor: + return torch.log(torch.clamp(magnitudes, min=1e-5)) + + +def _get_glm_tts_mel_basis( + *, + device: torch.device, + sample_rate: int, + n_fft: int, + n_mels: int, + f_min: int, + f_max: int, +) -> torch.Tensor: + """Compute or retrieve cached GLM-TTS mel filter bank. + + Uses torchaudio melscale_fbanks with Slaney parameters (norm="slaney", + mel_scale="slaney") to match the official GLM-TTS ``utils/audio.py`` + librosa.filters.mel output. The previous HTK / norm=None config + produced a numerically different filter bank that shifted the + prompt_feat distribution and degraded voice-clone conditioning. + """ + cache_key = (sample_rate, n_fft, n_mels, f_min, f_max, str(device)) + mel_basis = _MEL_BASIS_CACHE.get(cache_key) + if mel_basis is not None: + return mel_basis + + import torchaudio.functional as AF + + mel_basis = ( + AF.melscale_fbanks( + n_freqs=(n_fft // 2) + 1, + f_min=float(f_min), + f_max=float(f_max), + n_mels=n_mels, + sample_rate=sample_rate, + norm="slaney", + mel_scale="slaney", + ) + .transpose(0, 1) + .contiguous() + .float() + .to(device) + ) + + _MEL_BASIS_CACHE[cache_key] = mel_basis + return mel_basis + + +def _extract_glm_tts_mel_feature( + audio: torch.Tensor, + *, + sample_rate: int, + n_fft: int, + win_length: int, + hop_length: int, + n_mels: int, + f_min: int, + f_max: int, +) -> torch.Tensor: + """Extract time-major log-mel spectrogram features from audio waveform.""" + device = audio.device + mel_basis = _get_glm_tts_mel_basis( + device=device, + sample_rate=sample_rate, + n_fft=n_fft, + n_mels=n_mels, + f_min=f_min, + f_max=f_max, + ) + + hann_key = (win_length, str(device)) + window = _HANN_WINDOW_CACHE.get(hann_key) + if window is None: + window = torch.hann_window(win_length, device=device) + _HANN_WINDOW_CACHE[hann_key] = window + + if audio.ndim == 1: + audio = audio.unsqueeze(0) + + pad = int((n_fft - hop_length) / 2) + audio = torch.nn.functional.pad( + audio.unsqueeze(1), + (pad, pad), + mode="reflect", + ).squeeze(1) + + spec = torch.stft( + audio, + n_fft, + hop_length=hop_length, + win_length=win_length, + window=window, + center=False, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + spec = torch.sqrt(spec.real.pow(2) + spec.imag.pow(2) + 1e-9) + spec = torch.matmul(mel_basis, spec) + spec = _spectral_normalize_torch(spec) + return spec.squeeze(0).transpose(0, 1).contiguous() + + +def load_voice_clone_frontend( + model_root: str, + model_device: torch.device, + *, + speech_tokenizer_cache: tuple | None, + campplus_cache: object | None, + campplus_path: str | None = None, +) -> tuple[tuple | None, object | None]: + """Lazy-load voice cloning frontend models. + + Args: + campplus_path: Pre-resolved path to ``campplus.onnx``. When provided + this takes priority over the default ``model_root/frontend/`` + lookup. The caller (``GLMTTSForConditionalGeneration``) resolves + the path via ``transformers.utils.hub.cached_file`` during + ``__init__`` so the file is already downloaded. + + Returns: + ``(speech_tokenizer, campplus_session)`` — guaranteed non-None + on success. Raises ``RuntimeError`` when a required component + cannot be loaded. + """ + speech_tokenizer = speech_tokenizer_cache + campplus_session = campplus_cache + + if speech_tokenizer is None: + speech_tokenizer_paths = ( + os.path.join(model_root, "speech_tokenizer"), + os.path.join(model_root, "ckpt", "speech_tokenizer"), + ) + speech_tokenizer_path = next((path for path in speech_tokenizer_paths if os.path.isdir(path)), None) + if speech_tokenizer_path is None: + raise RuntimeError( + f"GLM-TTS speech tokenizer directory not found at {speech_tokenizer_paths[0]} " + f"or {speech_tokenizer_paths[1]}. Voice cloning requires the WhisperVQ " + f"speech tokenizer from the model snapshot." + ) + try: + from safetensors.torch import load_file + from transformers import AutoFeatureExtractor, WhisperConfig + + from vllm_omni.model_executor.models.common.whisper_vq import ( + WhisperVQEncoder, + ) + + _config = WhisperConfig.from_pretrained(speech_tokenizer_path) + _config.quantize_encoder_only = True + _model = WhisperVQEncoder(_config) + _sd = load_file(os.path.join(speech_tokenizer_path, "model.safetensors")) + _prefix, _cleaned = "model.encoder.", {} + for k, v in _sd.items(): + _cleaned[k[len(_prefix) :] if k.startswith(_prefix) else k] = v + _model.load_state_dict(_cleaned, strict=False) + _model = _model.to(model_device).eval() + _fe = AutoFeatureExtractor.from_pretrained(speech_tokenizer_path) + speech_tokenizer = (_model, _fe) + logger.info("Loaded GLM-TTS speech tokenizer from %s", speech_tokenizer_path) + except Exception as exc: + raise RuntimeError(f"Failed to load GLM-TTS speech tokenizer from {speech_tokenizer_path}") from exc + + if campplus_session is None: + cp = campplus_path or os.path.join(model_root, "frontend", "campplus.onnx") + if not os.path.isfile(cp): + raise RuntimeError( + f"GLM-TTS campplus.onnx not found at {cp}. " + f"Voice cloning requires the CampPlus speaker embedding model." + ) + try: + import onnxruntime + + cuda_avail = "CUDAExecutionProvider" in onnxruntime.get_available_providers() + cuda_ok = model_device.type == "cuda" and cuda_avail + providers = (["CUDAExecutionProvider"] if cuda_ok else []) + ["CPUExecutionProvider"] + campplus_session = onnxruntime.InferenceSession(cp, providers=providers) + logger.info("Loaded GLM-TTS campplus ONNX from %s", cp) + except Exception as exc: + raise RuntimeError(f"Failed to load GLM-TTS campplus ONNX from {cp}") from exc + + return speech_tokenizer, campplus_session + + +def extract_prompt_speech_token( + ref_audio_wav: torch.Tensor, + ref_audio_sr: int, + speech_tokenizer: tuple, +) -> list[int] | None: + """Extract prompt speech tokens from reference audio using WhisperVQ.""" + + model, feature_extractor = speech_tokenizer + device = model.device + + audio = ref_audio_wav.float().to(device) + if ref_audio_sr != 16000: + resampler = T.Resample(orig_freq=ref_audio_sr, new_freq=16000).to(device) + audio = resampler(audio) + + if audio.ndim > 1: + audio = audio[0] + + audio_np = audio.cpu().numpy() + pooling_kernel_size = getattr(model.config, "pooling_kernel_size", 1) + stride = model.conv1.stride[0] * model.conv2.stride[0] * pooling_kernel_size * feature_extractor.hop_length + + all_tokens: list[int] = [] + time_step = 0 + while time_step * 16000 < audio_np.shape[0]: + segment = audio_np[time_step * 16000 : (time_step + 30) * 16000] + features = feature_extractor( + [segment], + sampling_rate=16000, + return_attention_mask=True, + return_tensors="pt", + padding="longest", + pad_to_multiple_of=stride, + ).to(device) + with torch.no_grad(): + outputs = model(**features) + tokens = outputs.quantized_token_ids + attn = features.attention_mask[:, :: model.conv1.stride[0] * model.conv2.stride[0]][:, ::pooling_kernel_size] + all_tokens.extend(tokens[0][attn[0].bool()].tolist()) + time_step += 30 + + return all_tokens if all_tokens else None + + +def extract_spk_embedding( + ref_audio_wav: torch.Tensor, + ref_audio_sr: int, + campplus_session: object, +) -> list[float] | None: + """Extract speaker embedding from reference audio using CampPlus ONNX.""" + + audio = ref_audio_wav.float().cpu() + if ref_audio_sr != 16000: + resampler = T.Resample(orig_freq=ref_audio_sr, new_freq=16000) + audio = resampler(audio) + if audio.ndim > 1: + audio = audio[0] + + feat = kaldi.fbank(audio.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000) + feat = feat - feat.mean(dim=0, keepdim=True) + + input_name = campplus_session.get_inputs()[0].name + embedding = campplus_session.run(None, {input_name: feat.unsqueeze(dim=0).cpu().numpy()})[0].flatten().tolist() + return embedding + + +def extract_prompt_feat( + ref_audio_wav: torch.Tensor, + ref_audio_sr: int, + model_device: torch.device, +) -> torch.Tensor | None: + """Extract mel features from reference audio for DiT conditioning.""" + audio = ref_audio_wav.float().to(model_device) + if ref_audio_sr != 24000: + resampler = T.Resample(orig_freq=ref_audio_sr, new_freq=24000).to(model_device) + audio = resampler(audio) + if audio.ndim > 1: + audio = audio[0] + + with torch.no_grad(): + feat = _extract_glm_tts_mel_feature( + audio, + sample_rate=24000, + n_fft=1920, + win_length=1920, + hop_length=480, + n_mels=80, + f_min=0, + f_max=8000, + ) + return feat # [T_mel, mel_dim] diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 4f29c0ae3..4c49afc36 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -129,6 +129,12 @@ _OMNI_MODELS = { "glm_image_ar", "GlmImageForConditionalGeneration", ), + ## glm_tts + "GLMTTSForConditionalGeneration": ( + "glm_tts", + "glm_tts", + "GLMTTSForConditionalGeneration", + ), "OmniBagelForConditionalGeneration": ( "bagel", "bagel", diff --git a/vllm_omni/model_executor/stage_input_processors/glm_tts.py b/vllm_omni/model_executor/stage_input_processors/glm_tts.py new file mode 100644 index 000000000..caafbf86c --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/glm_tts.py @@ -0,0 +1,362 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage input processor for GLM-TTS: AR → DiT Pipeline. + +Supports both sync (non-streaming) and async_chunk (streaming) modes. +Adapted for LLM_GENERATION execution type on stage 1 (DiT). +""" + +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.data_entry_keys import ( + CodesStruct, + EmbeddingsStruct, + MetaStruct, + OmniPayloadStruct, +) +from vllm_omni.engine.serialization import deserialize_additional_information + +logger = init_logger(__name__) + +_GLM_TTS_PRIVATE_KEY = "glm_tts" + + +def _copy_voice_clone_payload( + src: dict[str, Any], + dst: dict[str, Any], + *, + to_cpu: bool = False, + skip_existing: bool = False, +) -> None: + if not skip_existing or "prompt_speech_token" not in dst: + prompt = src.get("prompt_speech_token") + if prompt is None: + prompt = src.get("prompt_token") + if prompt is not None: + if to_cpu: + prompt = _to_cpu_tensor(prompt) + if prompt is not None: + dst["prompt_speech_token"] = prompt + + for key in ("prompt_feat", "embedding"): + if skip_existing and key in dst: + continue + val = src.get(key) + if val is not None: + if to_cpu: + val = _to_cpu_tensor(val) + if val is not None: + dst[key] = val + + +def _build_voice_clone_embed_struct(payload: dict[str, Any]) -> EmbeddingsStruct | None: + prompt_token = payload.get("prompt_speech_token") + if prompt_token is None: + prompt_token = payload.get("prompt_token") + prompt_feat = payload.get("prompt_feat") + embedding = payload.get("embedding") + if prompt_token is None and prompt_feat is None and embedding is None: + return None + return EmbeddingsStruct( + speech_token=prompt_token, + speech_feat=prompt_feat, + embedding=embedding, + ) + + +# --------------------------------------------------------------------------- +# Sync processor: collect all speech tokens from AR, pass to DiT +# --------------------------------------------------------------------------- + + +def ar_to_dit( + source_outputs: list[Any], + prompt: Any = None, + _requires_multimodal_data: bool = False, +) -> list[Any]: + """Non-streaming: collect speech tokens from AR, pass to DiT. + + Propagates voice cloning data (prompt_token, prompt_feat, embedding) + from the AR model's multimodal_output to the DiT stage using the same + nested ``embed`` dict format as the async_chunk path. This ensures + ``nested_get(info, "embed", ...)`` in the DiT wrapper reads all three + fields without relying on backward-compat flat-key fallbacks. + """ + from vllm_omni.inputs.data import OmniTokensPrompt + + dit_inputs: list[OmniTokensPrompt] = [] + + for output in source_outputs: + mm = output.outputs[0].multimodal_output + speech_tokens = mm.get("speech_tokens") + + # Extract valid token list (filter -1 placeholder markers from prefill) + if isinstance(speech_tokens, torch.Tensor): + valid = speech_tokens.to(torch.long).reshape(-1) + token_list = valid[valid >= 0].cpu().tolist() + elif speech_tokens is not None: + token_list = [t for t in speech_tokens if t >= 0] + else: + token_list = [] + + if not token_list: + logger.warning("No valid speech tokens for request %s", output.request_id) + + additional_info: dict[str, Any] = {"speech_tokens": token_list} + + # Build nested embed dict (unified with async_chunk path). + # _build_voice_clone_embed_struct reads prompt_speech_token/prompt_feat/ + # embedding from *mm* and returns an EmbeddingsStruct; we convert to a + # plain dict so flatten_payload/unflatten_payload round-trips correctly + # (key "embed" is in _NESTED_KEYS). + embed = _build_voice_clone_embed_struct(mm) + if embed is not None: + additional_info["embed"] = { + "speech_token": embed.speech_token, + "speech_feat": embed.speech_feat, + "embedding": embed.embedding, + } + + # Populate meta dict so DiT wrapper uses block-causal attention. + # Without meta, forward() sets uses_streaming=False → bidirectional + # attention → garbage audio. The full token sequence is treated as + # a single "chunk" with stream_finished=True. + n_tokens = len(token_list) if token_list else 1 + additional_info["meta"] = { + "req_id": [output.request_id], + "left_context_size": 0, + "stream_finished": True, + } + additional_info["kv_metadata"] = { + _GLM_TTS_PRIVATE_KEY: { + "chunk_sizes_history": [n_tokens], + "block_pattern": [n_tokens], + } + } + + dit_inputs.append( + OmniTokensPrompt( + # LLM_GENERATION scheduler needs at least 1 token + prompt_token_ids=token_list or [0], + additional_information=additional_info, + ) + ) + + return dit_inputs + + +# --------------------------------------------------------------------------- +# Helper: extract last speech token from AR model pooling output +# --------------------------------------------------------------------------- + + +def _extract_last_speech_token(pooling_output: dict[str, Any]) -> int | None: + """Extract the last valid speech token from AR model output. + + GLM-TTS AR produces one speech token per decode step. + Returns the token ID (relative to ATS, i.e. 0-based), or None. + """ + speech_tokens = pooling_output.get("speech_tokens") + if not isinstance(speech_tokens, torch.Tensor) or speech_tokens.numel() == 0: + return None + token_val = int(speech_tokens.reshape(-1).to(torch.long)[-1].item()) + # -1 = invalid/EOA marker + if token_val < 0: + return None + return token_val + + +def _to_cpu_tensor(value: Any) -> torch.Tensor | None: + """Convert value to CPU tensor if possible.""" + if isinstance(value, torch.Tensor): + return value.detach().cpu() + if isinstance(value, list): + if not value: + return None + if isinstance(value[0], torch.Tensor): + return value[0].detach().cpu() + return None + + +# --------------------------------------------------------------------------- +# Async streaming processor: emit speech token chunks as AR produces them +# --------------------------------------------------------------------------- + + +def ar_to_dit_async_chunk( + transfer_manager: Any, + pooling_output: dict[str, Any] | None, + request: Any, + is_finished: bool = False, +) -> OmniPayloadStruct | None: + """Async streaming processor: emit speech token chunks as AR produces them. + + Accumulates per-step speech tokens and emits fixed-size cumulative prefixes + for GLM-TTS flow-cache streaming. + + Follows the CosyVoice3 talker2code2wav_async_chunk transfer pattern: + - Per-request state tracking (seen tokens, prompt sent flag) + - First chunk carries voice clone conditioning payload + - Each chunk includes left_context_size, stream_finished, req_id + - code_predictor_codes for chunk_transfer_adapter consumption + + Unlike codec left-context decoders, official GLM-TTS sends the cumulative + AR token prefix to the flow stage on every chunk and reuses diffusion + latents internally. Here left_context_size means the stable token prefix + that has already been emitted, so DiT can crop regenerated audio after + sampling. + + GLM-TTS produces single-token-per-step (no multi-codebook), so each entry + in code_prompt_token_ids is a plain int, not a list of codebook values. + """ + request_id = getattr(request, "external_req_id", None) or getattr(request, "request_id", None) + if request_id is None: + raise ValueError("GLM-TTS async chunk request is missing request id") + finished = bool(is_finished or request.is_finished()) + + # Read connector chunk config (supports progressive list or single int) + 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_frames_cfg = cfg.get("codec_chunk_frames", 25) + if isinstance(chunk_frames_cfg, list): + progressive_chunk_sizes = [int(c) for c in chunk_frames_cfg] + else: + progressive_chunk_sizes = [int(chunk_frames_cfg)] + left_context_size_config = int(cfg.get("codec_left_context_frames", 25)) + crossfade_sec = float(cfg.get("crossfade_sec", 0.1)) + + if not progressive_chunk_sizes or any(c <= 0 for c in progressive_chunk_sizes) or left_context_size_config < 0: + raise ValueError( + f"Invalid codec chunk config: codec_chunk_frames={chunk_frames_cfg}, " + f"codec_left_context_frames={left_context_size_config}" + ) + + # Initialize per-request state (like CosyVoice3) + request_payload = getattr(transfer_manager, "request_payload", None) + if request_payload is None: + request_payload = {} + transfer_manager.request_payload = request_payload + code_prompt_token_ids = getattr(transfer_manager, "code_prompt_token_ids", None) + if code_prompt_token_ids is None: + code_prompt_token_ids = {} + transfer_manager.code_prompt_token_ids = code_prompt_token_ids + code_prompt_token_ids.setdefault(request_id, []) + request_state = request_payload.get(request_id) + if not isinstance(request_state, dict) or "_glm_tts_async_state" not in request_state: + # Extract voice clone conditioning from request additional_information + info = deserialize_additional_information(getattr(request, "additional_information", None)) + prompt_payload: dict[str, Any] = {} + _copy_voice_clone_payload(info, prompt_payload, to_cpu=True) + + # Also try to extract from pooling_output (first call) + if isinstance(pooling_output, dict): + _copy_voice_clone_payload(pooling_output, prompt_payload, to_cpu=True, skip_existing=True) + + request_state = { + "_glm_tts_async_state": { + "seen_len": 0, + "sent_prompt": False, + "emitted_chunks": 0, + "emitted_token_len": 0, + "terminal_sent": False, + "prompt_payload": prompt_payload, + "chunk_sizes_history": [], + "block_pattern": progressive_chunk_sizes, + } + } + request_payload[request_id] = request_state + + state = request_state["_glm_tts_async_state"] + if state.get("terminal_sent", False): + return None + + # Helper: get voice clone embed (once per request) + def _pop_embed() -> EmbeddingsStruct | None: + if state.get("sent_prompt", False): + return None + state["sent_prompt"] = True + return _build_voice_clone_embed_struct(state.get("prompt_payload", {})) + + # Helper: build terminal or chunk payload + def _make_payload( + codes: list[int], + *, + tok_offset: int, + is_terminal: bool, + ) -> OmniPayloadStruct: + return OmniPayloadStruct( + codes=CodesStruct( + audio=(torch.tensor(codes, dtype=torch.long) if codes else torch.empty(0, dtype=torch.long)) + ), + meta=MetaStruct( + finished=torch.tensor(is_terminal, dtype=torch.bool), + stream_finished=torch.tensor(is_terminal, dtype=torch.bool), + req_id=[request_id], + left_context_size=tok_offset, + ), + kv_metadata={ + _GLM_TTS_PRIVATE_KEY: { + "chunk_sizes_history": list(state.get("chunk_sizes_history", [])), + "block_pattern": list(state.get("block_pattern", progressive_chunk_sizes)), + "crossfade_sec": crossfade_sec, + } + }, + embed=_pop_embed(), + ) + + # Accumulate new speech token from this step + if isinstance(pooling_output, dict): + token = _extract_last_speech_token(pooling_output) + if token is not None: + code_prompt_token_ids[request_id].append(token) + elif not finished: + return None + + token_frames = code_prompt_token_ids[request_id] + length = len(token_frames) + emitted_token_len = int(state.get("emitted_token_len", 0)) + + # Terminal: no tokens, or finished on/before chunk boundary + if finished and length <= emitted_token_len: + state["terminal_sent"] = True + # Boundary finish is metadata-only: the cumulative prefix has already + # been emitted, so resending it would make the DiT/vocoder redo work + # and can duplicate audio. + return _make_payload([], tok_offset=emitted_token_len, is_terminal=True) + + if length <= 0: + return None + + # Progressive chunk size: 25 → 50 → 200 (official GLM-TTS pattern) + chunk_count = int(state.get("emitted_chunks", 0)) + idx = min(chunk_count, len(progressive_chunk_sizes) - 1) + current_chunk_size = progressive_chunk_sizes[idx] + + # Check if enough tokens accumulated for next chunk + if not finished and (length - emitted_token_len) < current_chunk_size: + return None + + # Determine end index for cumulative prefix + if emitted_token_len == 0: + end_index = min(length, current_chunk_size) + token_offset = 0 + else: + end_index = length if finished else min(length, emitted_token_len + current_chunk_size) + token_offset = emitted_token_len + + # Track chunk sizes for cache slicing + chunk_sizes_history: list[int] = list(state.get("chunk_sizes_history", [])) + chunk_sizes_history.append(end_index - emitted_token_len) + state["chunk_sizes_history"] = chunk_sizes_history + state["emitted_chunks"] = chunk_count + 1 + + if finished: + state["terminal_sent"] = True + else: + state["emitted_token_len"] = max(emitted_token_len, end_index) + + return _make_payload(list(token_frames[:end_index]), tok_offset=token_offset, is_terminal=finished) diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 708d9c2b6..1a806cc46 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -17,6 +17,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "FishSpeechConfig": "vllm_omni.transformers_utils.configs.fish_speech", "FishSpeechSlowARConfig": "vllm_omni.transformers_utils.configs.fish_speech", "FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech", + "GLMTTSConfig": "vllm_omni.transformers_utils.configs.glm_tts", "VoxCPM2Config": "vllm_omni.transformers_utils.configs.voxcpm2", "VoxtralTTSConfig": "vllm_omni.transformers_utils.configs.voxtral_tts", "CosyVoice3Config": "vllm_omni.transformers_utils.configs.cosyvoice3", @@ -36,6 +37,7 @@ __all__ = [ "FishSpeechConfig", "FishSpeechSlowARConfig", "FishSpeechFastARConfig", + "GLMTTSConfig", "VoxCPM2Config", "VoxtralTTSConfig", "CosyVoice3Config", @@ -65,6 +67,7 @@ def __dir__(): # run as soon as `vllm_omni.transformers_utils.configs` is imported. from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import glm_tts as _glm_tts # noqa: F401, E402 from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import ming_flash_omni as _ming_flash_omni # noqa: F401, E402 from vllm_omni.transformers_utils.configs import omnivoice as _omnivoice # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/glm_tts.py b/vllm_omni/transformers_utils/configs/glm_tts.py new file mode 100644 index 000000000..9a7f8dfbf --- /dev/null +++ b/vllm_omni/transformers_utils/configs/glm_tts.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-TTS config registration with transformers AutoConfig. + +Registers GLMTTSConfig (model_type="glm_tts") so that +``AutoConfig.from_pretrained("path/to/glm-tts")`` returns the correct config class. + +Note: GLM-TTS uses a Llama backbone, but we register a custom config +to handle the special token IDs and flow model parameters. +""" + +from __future__ import annotations + +from typing import Any + +from transformers import AutoConfig, PretrainedConfig + + +class GLMTTSConfig(PretrainedConfig): + """Llama-based AR model for text-to-speech token generation. + + Special token IDs are loaded dynamically from the tokenizer at init time. + """ + + model_type: str = "glm_tts" + + def __init__( + self, + vocab_size: int = 98304, + hidden_size: int = 2048, + intermediate_size: int = 6144, + num_hidden_layers: int = 28, + num_attention_heads: int = 16, + num_key_value_heads: int = 4, + hidden_act: str = "silu", + max_position_embeddings: int = 8192, + rms_norm_eps: float = 1e-5, + rope_theta: float = 10000.0, + audio_token_start: int = -1, + audio_token_end: int = -1, + boa_token_id: int = -1, + eoa_token_id: int = -1, + pad_token_id: int = -1, + bos_token_id: int = -1, + speech_token_vocab_size: int = 32768, + speech_token_dim: int = 512, + spk_embed_dim: int = 192, + mel_dim: int = 80, + input_frame_rate: float = 25.0, + mel_framerate: int = 50, + max_token_text_ratio: float = 20.0, + min_token_text_ratio: float = 2.0, + sample_method: str = "ras", + ras_top_p: float = 0.8, + ras_top_k: int = 25, + ras_win_size: int = 10, + ras_tau_r: float = 0.1, + **kwargs: Any, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.rms_norm_eps = rms_norm_eps + self.rope_theta = rope_theta + self.audio_token_start = audio_token_start + self.audio_token_end = audio_token_end + self.boa_token_id = boa_token_id + self.eoa_token_id = eoa_token_id + self.speech_token_vocab_size = speech_token_vocab_size + self.speech_token_dim = speech_token_dim + self.spk_embed_dim = spk_embed_dim + self.mel_dim = mel_dim + self.input_frame_rate = input_frame_rate + self.mel_framerate = mel_framerate + self.max_token_text_ratio = max_token_text_ratio + self.min_token_text_ratio = min_token_text_ratio + self.sample_method = sample_method + self.ras_top_p = ras_top_p + self.ras_top_k = ras_top_k + self.ras_win_size = ras_win_size + self.ras_tau_r = ras_tau_r + super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, **kwargs) + + +AutoConfig.register("glm_tts", GLMTTSConfig) + +__all__ = [ + "GLMTTSConfig", +]