[Model] Add GLM-TTS text-to-speech model support (#3141)
Signed-off-by: BeatSeat <wendavid552@gmail.com> Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
This commit is contained in:
BeatSeatandYueqian Lin authored and GitHub committed
1 parent
748470c1fb
commit
28ce618f0e
38 files changed
+7568
-7
No files matched your search
@@ -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
|
||||
|
||||
@@ -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` | ✅︎ | ✅︎ | | |
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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())
|
||||
@@ -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.
|
||||
|
||||
@@ -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()
|
||||
@@ -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.<format>)",
|
||||
)
|
||||
|
||||
# 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)
|
||||
@@ -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"
|
||||
@@ -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[@]}"
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
Binary file not shown.
@@ -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",
|
||||
)
|
||||
@@ -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
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
)
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
@@ -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*",
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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",
|
||||
]
|
||||
File diff suppressed because it is too large.
Load diff
@@ -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
|
||||
File diff suppressed because it is too large.
Load diff
@@ -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",
|
||||
),
|
||||
),
|
||||
)
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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]
|
||||
@@ -129,6 +129,12 @@ _OMNI_MODELS = {
|
||||
"glm_image_ar",
|
||||
"GlmImageForConditionalGeneration",
|
||||
),
|
||||
## glm_tts
|
||||
"GLMTTSForConditionalGeneration": (
|
||||
"glm_tts",
|
||||
"glm_tts",
|
||||
"GLMTTSForConditionalGeneration",
|
||||
),
|
||||
"OmniBagelForConditionalGeneration": (
|
||||
"bagel",
|
||||
"bagel",
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
Reference in new issue
Block a user