diff --git a/examples/offline_inference/text_to_audio/glm_tts_example.py b/examples/offline_inference/text_to_audio/glm_tts_example.py new file mode 100644 index 00000000000..45041e8a4b3 --- /dev/null +++ b/examples/offline_inference/text_to_audio/glm_tts_example.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Example script for text-to-speech generation using GLM-TTS. + +This script demonstrates how to generate speech from text using +the GLM-TTS model with vLLM-Omni. + +GLM-TTS is a two-stage model: +- Stage 0: LLM (Llama-based) generates speech tokens from text +- Stage 1: DiT (Flow matching) generates mel-spectrogram from speech tokens + +Usage: + python glm_tts_example.py --text "Hello, this is a test of GLM-TTS." + python glm_tts_example.py --text "Welcome to the future of text to speech." --num_inference_steps 16 +""" + +import argparse +import os +import time + +import numpy as np +import torch +from vllm import SamplingParams + +from vllm_omni.entrypoints.omni import Omni + +# Path to stage config (relative to repo root) +STAGE_CONFIG_PATH = "vllm_omni/model_executor/stage_configs/glm_tts.yaml" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Generate speech with GLM-TTS.") + parser.add_argument( + "--model", + default="zai-org/GLM-TTS", + help="GLM-TTS model name or local path.", + ) + parser.add_argument( + "--stage-config", + default=None, + help="Path to stage config YAML file. If not provided, uses default GLM-TTS config.", + ) + parser.add_argument( + "--text", + default="Hello, this is a test of GLM-TTS text to speech synthesis.", + help="Text to convert to speech.", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Random seed for deterministic results.", + ) + parser.add_argument( + "--num_inference_steps", + type=int, + default=32, + help="Number of denoising steps for the flow matching sampler.", + ) + parser.add_argument( + "--output", + type=str, + default="glm_tts_output.wav", + help="Path to save the generated audio (WAV format).", + ) + parser.add_argument( + "--sample_rate", + type=int, + default=24000, + help="Sample rate for output audio.", + ) + return parser.parse_args() + + +def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 24000): + """Save audio data to a WAV file.""" + try: + import soundfile as sf + + sf.write(output_path, audio_data, sample_rate) + print(f"✓ Audio saved to: {output_path}") + except ImportError: + print("Warning: soundfile not installed. Install with: pip install soundfile") + # Fallback to scipy + try: + from scipy.io import wavfile + + wavfile.write(output_path, sample_rate, audio_data) + print(f"✓ Audio saved to: {output_path}") + except ImportError: + print("Error: Neither soundfile nor scipy is installed.") + print("Install with: pip install soundfile") + + +def main(): + args = parse_args() + + # Resolve stage config path + stage_config = args.stage_config + if stage_config is None: + # Try to find the default config relative to this script + script_dir = os.path.dirname(os.path.abspath(__file__)) + repo_root = os.path.dirname(os.path.dirname(os.path.dirname(script_dir))) + stage_config = os.path.join(repo_root, STAGE_CONFIG_PATH) + if not os.path.exists(stage_config): + print(f"Error: Stage config not found at {stage_config}") + print("Please provide --stage-config path or run from repository root.") + return + + print(f"Loading GLM-TTS model: {args.model}") + print(f"Using stage config: {stage_config}") + start_time = time.time() + + # Initialize the Omni engine with GLM-TTS stage config + # Note: stage_configs_path tells Omni to use the YAML config instead of + # auto-detecting the model format + omni = Omni( + model=args.model, + stage_configs_path=stage_config, + ) + + load_time = time.time() - start_time + print(f"✓ Model loaded in {load_time:.2f}s") + + print(f"\nGenerating speech for: '{args.text}'") + print(f"Parameters:") + print(f" - Inference steps: {args.num_inference_steps}") + print(f" - Seed: {args.seed}") + + gen_start = time.time() + + # Build prompt for GLM-TTS LLM stage + # The LLM will generate speech tokens which are then processed by the DiT stage + prompt = args.text + + # Sampling parameters for LLM stage (speech token generation) + llm_sampling_params = SamplingParams( + temperature=0.9, + top_p=0.8, + top_k=40, + max_tokens=2048, + seed=args.seed, + detokenize=False, + repetition_penalty=1.05, + stop_token_ids=[151330], # GLM_TTS_EOA_TOKEN_ID + ) + + # Sampling parameters for DiT stage (audio generation) + dit_sampling_params = SamplingParams( + temperature=0.0, + max_tokens=1, # DiT doesn't use token generation + seed=args.seed, + ) + + sampling_params_list = [llm_sampling_params, dit_sampling_params] + + # Generate audio + outputs = list(omni.generate( + [{"prompt": prompt}], + sampling_params_list, + py_generator=True, + )) + + gen_time = time.time() - gen_start + print(f"✓ Audio generated in {gen_time:.2f}s") + + # Extract audio from outputs + for stage_output in outputs: + if stage_output.final_output_type == "audio": + for req_output in stage_output.request_output: + if hasattr(req_output, "images") and len(req_output.images) > 0: + audio = req_output.images[0] + + # Save audio + if isinstance(audio, np.ndarray): + # Ensure audio is in the correct format for saving + if audio.ndim == 3: # (batch, channels, samples) + audio = audio[0] # Take first batch + if audio.ndim == 2: # (channels, samples) + audio = audio.T # Transpose to (samples, channels) + + save_audio(audio, args.output, args.sample_rate) + + # Print audio stats + duration_actual = audio.shape[0] / args.sample_rate + print(f"\nAudio Statistics:") + print(f" - Shape: {audio.shape}") + print(f" - Duration: {duration_actual:.2f}s") + print(f" - Sample rate: {args.sample_rate} Hz") + print(f" - Min/Max: {audio.min():.3f} / {audio.max():.3f}") + else: + print(f"Error: Unexpected audio format: {type(audio)}") + else: + print("Error: No audio data in output") + + total_time = time.time() - start_time + print(f"\n✓ Total time: {total_time:.2f}s") + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/offline_inference/test_glm_tts_model.py b/tests/e2e/offline_inference/test_glm_tts_model.py new file mode 100644 index 00000000000..00a2f69eb75 --- /dev/null +++ b/tests/e2e/offline_inference/test_glm_tts_model.py @@ -0,0 +1,162 @@ +import os +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch + +from vllm_omni.outputs import OmniRequestOutput + +# ruff: noqa: E402 +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vllm_omni import Omni + +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" + +# GLM-TTS model from HuggingFace +# Note: For CI testing, a smaller random-weights model could be created +models = ["zai-org/GLM-TTS"] + + +@pytest.mark.skip(reason="GLM-TTS model requires full download; enable when model available locally") +@pytest.mark.parametrize("model_name", models) +def test_glm_tts_model(model_name: str): + """Test GLM-TTS text-to-speech generation.""" + m = Omni(model=model_name) + + # GLM-TTS parameters + audio_duration_s = 5.0 # 5 second audio + + # Generate speech tokens placeholder (in production, LLM generates these) + # For testing, we provide mock tokens + speech_tokens = torch.randint(0, 10000, (1, 100), dtype=torch.long) + speaker_embedding = torch.randn(1, 192) + + outputs = m.generate( + "Hello, this is a test of GLM-TTS text to speech synthesis.", + num_inference_steps=8, # Minimal steps for speed + guidance_scale=1.0, + generator=torch.Generator("cuda").manual_seed(42), + num_outputs_per_prompt=1, + extra={ + "audio_duration_s": audio_duration_s, + "speech_tokens": speech_tokens, + "speaker_embedding": speaker_embedding, + }, + ) + + # Verify output structure + assert outputs is not None + first_output = outputs[0] + assert hasattr(first_output, "request_output") and first_output.request_output + + req_out = first_output.request_output[0] + assert isinstance(req_out, OmniRequestOutput) + assert hasattr(req_out, "images") and len(req_out.images) >= 1 + + # For TTS, the "images" field contains audio numpy arrays + audio = req_out.images[0] + assert isinstance(audio, np.ndarray) + # audio shape: (batch, channels, samples) or (channels, samples) + assert audio.ndim >= 2 + + +@pytest.mark.skip(reason="Unit test for pipeline components") +def test_glm_tts_dit_model(): + """Test GLM-TTS DiT model forward pass.""" + from vllm_omni.diffusion.models.glm_tts.glm_tts_dit import GLMTTSDiTModel + + # Create model with default config + model = GLMTTSDiTModel( + hidden_size=256, # Small for testing + num_attention_heads=4, + num_hidden_layers=2, + head_dim=64, + mel_dim=80, + speech_token_dim=128, + speech_token_vocab_size=1000, + speaker_embed_dim=64, + ).cuda() + + # Test forward pass + batch_size = 2 + seq_len = 100 + token_len = 50 + + noisy_mel = torch.randn(batch_size, seq_len, 80).cuda() + timestep = torch.rand(batch_size).cuda() + speech_tokens = torch.randint(0, 1000, (batch_size, token_len)).cuda() + speaker_embedding = torch.randn(batch_size, 64).cuda() + + output = model( + noisy_mel=noisy_mel, + timestep=timestep, + speech_tokens=speech_tokens, + speaker_embedding=speaker_embedding, + ) + + assert output.shape == (batch_size, seq_len, 80) + + +@pytest.mark.skip(reason="Unit test for stage input processor") +def test_glm_tts_stage_input_processor(): + """Test GLM-TTS stage input processor.""" + from vllm_omni.model_executor.stage_input_processors.glm_tts import ( + GLM_TTS_AUDIO_TOKEN_END, + GLM_TTS_AUDIO_TOKEN_START, + extract_speech_tokens, + ) + + # Test speech token extraction + # Create mock token IDs with some audio tokens + token_ids = [ + 100, # Non-audio token + GLM_TTS_AUDIO_TOKEN_START, # First audio token (should become 0) + GLM_TTS_AUDIO_TOKEN_START + 100, # Audio token (should become 100) + 200, # Non-audio token + GLM_TTS_AUDIO_TOKEN_END, # Last audio token + ] + + speech_tokens = extract_speech_tokens(token_ids) + + # Should have 3 audio tokens (normalized to 0-based) + assert len(speech_tokens) == 3 + assert speech_tokens[0] == 0 # First audio token + assert speech_tokens[1] == 100 + assert speech_tokens[2] == GLM_TTS_AUDIO_TOKEN_END - GLM_TTS_AUDIO_TOKEN_START + + +@pytest.mark.skip(reason="Integration test for two-stage pipeline") +def test_glm_tts_two_stage_pipeline(): + """Test GLM-TTS two-stage pipeline (LLM + DiT).""" + # This test requires: + # 1. GLM-TTS LLM model weights + # 2. GLM-TTS DiT model weights + # 3. Stage config file + + from vllm_omni import Omni + + # Load with two-stage config + m = Omni( + model="zai-org/GLM-TTS", + stage_config="vllm_omni/model_executor/stage_configs/glm_tts.yaml", + ) + + outputs = m.generate( + "Hello, this is a test of GLM-TTS.", + num_inference_steps=8, + seed=42, + ) + + assert outputs is not None + first_output = outputs[0] + assert hasattr(first_output, "request_output") + + # Audio should be in images field (X2S pattern) + req_out = first_output.request_output[0] + audio = req_out.images[0] + assert isinstance(audio, np.ndarray) diff --git a/vllm_omni/diffusion/models/glm_tts/__init__.py b/vllm_omni/diffusion/models/glm_tts/__init__.py new file mode 100644 index 00000000000..d12b1f56f75 --- /dev/null +++ b/vllm_omni/diffusion/models/glm_tts/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""GLM-TTS model support for vLLM-Omni.""" + +from vllm_omni.diffusion.models.glm_tts.glm_tts_dit import ( + GLMTTSDiTModel, +) +from vllm_omni.diffusion.models.glm_tts.pipeline_glm_tts import ( + GLMTTSPipeline, + get_glm_tts_post_process_func, +) + +__all__ = [ + "GLMTTSDiTModel", + "GLMTTSPipeline", + "get_glm_tts_post_process_func", +] diff --git a/vllm_omni/diffusion/models/glm_tts/glm_tts_dit.py b/vllm_omni/diffusion/models/glm_tts/glm_tts_dit.py new file mode 100644 index 00000000000..5cb71b4718e --- /dev/null +++ b/vllm_omni/diffusion/models/glm_tts/glm_tts_dit.py @@ -0,0 +1,519 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +GLM-TTS DiT (Flow Matching) Model for vLLM-Omni. + +This module implements the flow matching model that converts speech tokens +to mel-spectrograms, based on the GLM-TTS architecture. +""" + +import math +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.data import OmniDiffusionConfig + +logger = init_logger(__name__) + + +def apply_rotary_emb_glm_tts( + hidden_states: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], +) -> torch.Tensor: + """ + Apply rotary embeddings to input tensors for GLM-TTS. + + Args: + hidden_states: Input tensor of shape [B, S, H, D] where D is head_dim + freqs_cis: Tuple of (cos, sin) frequency tensors of shape [S, rotary_dim] + + Returns: + Tensor with rotary embeddings applied. + """ + cos, sin = freqs_cis # [S, rotary_dim] + rotary_dim = cos.shape[-1] + + # Rotate only the first rotary_dim entries; leave the rest unchanged + x_rot = hidden_states[..., :rotary_dim] + x_pass = hidden_states[..., rotary_dim:] + + cos = cos[None, :, None, :] # [1, S, 1, rotary_dim] + sin = sin[None, :, None, :] # [1, S, 1, rotary_dim] + + # [B, S, H, rotary_dim] -> [B, S, H, 2, rotary_dim//2] -> two halves + x_real, x_imag = x_rot.reshape(*x_rot.shape[:-1], 2, rotary_dim // 2).unbind(-2) + x_rotated = torch.cat([-x_imag, x_real], dim=-1) + + x_rot = (x_rot.float() * cos + x_rotated.float() * sin).to(hidden_states.dtype) + return torch.cat([x_rot, x_pass], dim=-1) + + +class GLMTTSGaussianFourierProjection(nn.Module): + """Gaussian Fourier embeddings for timestep conditioning. + + Matches diffusers pattern with flip_sin_to_cos=True. + """ + + def __init__(self, embedding_size: int = 256, scale: float = 1.0): + super().__init__() + self.weight = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x shape: [batch] or [batch, 1] + # Output: [batch, embedding_size * 2] + x_proj = 2 * math.pi * x[:, None] @ self.weight[None, :] + # flip_sin_to_cos=True means cos comes first + return torch.cat([torch.cos(x_proj), torch.sin(x_proj)], dim=-1) + + +class GLMTTSSelfAttention(nn.Module): + """ + Optimized self-attention for GLM-TTS using vLLM layers. + + Uses full attention (all heads for Q, K, V). + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout: float = 0.0, + ): + super().__init__() + + self.dim = dim + self.num_heads = num_attention_heads + self.head_dim = attention_head_dim + self.inner_dim = num_attention_heads * attention_head_dim + + # All projections use inner_dim for output + self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=False) + self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=False) + self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=False) + + # Output projection + self.to_out = nn.ModuleList( + [ + ReplicatedLinear(self.inner_dim, dim, bias=False), + nn.Dropout(dropout), + ] + ) + + # Full attention (no GQA for self-attention) + self.attn = Attention( + num_heads=num_attention_heads, + head_size=attention_head_dim, + softmax_scale=1.0 / (attention_head_dim**0.5), + causal=False, + num_kv_heads=num_attention_heads, # Same as query heads + ) + + def forward( + self, + hidden_states: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape + + # Projections - all output inner_dim + query, _ = self.to_q(hidden_states) + key, _ = self.to_k(hidden_states) + value, _ = self.to_v(hidden_states) + + # Reshape for multi-head attention (all use full heads) + query = query.view(batch_size, seq_len, self.num_heads, self.head_dim) + key = key.view(batch_size, seq_len, self.num_heads, self.head_dim) + value = value.view(batch_size, seq_len, self.num_heads, self.head_dim) + + # Apply rotary embeddings + if rotary_emb is not None: + query = apply_rotary_emb_glm_tts(query, rotary_emb) + key = apply_rotary_emb_glm_tts(key, rotary_emb) + + # Compute attention + hidden_states = self.attn(query, key, value) + hidden_states = hidden_states.view(batch_size, seq_len, self.inner_dim) + + # Output projection + hidden_states, _ = self.to_out[0](hidden_states) + hidden_states = self.to_out[1](hidden_states) + + return hidden_states + + +class SwiGLU(nn.Module): + """SwiGLU activation - matches diffusers structure.""" + + def __init__(self, dim_in: int, dim_out: int, bias: bool = True): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias) + self.activation = nn.SiLU() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.proj(hidden_states) + hidden_states, gate = hidden_states.chunk(2, dim=-1) + return hidden_states * self.activation(gate) + + +class GLMTTSFeedForward(nn.Module): + """ + Feed-forward network with SwiGLU activation for GLM-TTS. + Matches diffusers FeedForward structure with activation_fn="swiglu". + """ + + def __init__(self, dim: int, inner_dim: int, bias: bool = True): + super().__init__() + self.net = nn.Sequential( + SwiGLU(dim, inner_dim, bias=bias), + nn.Dropout(0.0), + nn.Linear(inner_dim, dim, bias=bias), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.net(hidden_states) + + +class GLMTTSAdaLayerNorm(nn.Module): + """Adaptive LayerNorm with timestep conditioning.""" + + def __init__(self, dim: int): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 6) + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward( + self, + hidden_states: torch.Tensor, + timestep_emb: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Apply adaptive layer normalization. + + Returns: + Tuple of (normed_hidden, gate_msa, shift_mlp, scale_mlp, gate_mlp) + """ + emb = self.linear(self.silu(timestep_emb)) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1) + + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) + + return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +class GLMTTSAdaLayerNormFinal(nn.Module): + """Final adaptive LayerNorm (no MLP modulation needed).""" + + def __init__(self, dim: int): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 2) + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward( + self, + hidden_states: torch.Tensor, + timestep_emb: torch.Tensor, + ) -> torch.Tensor: + emb = self.linear(self.silu(timestep_emb)) + scale, shift = emb.chunk(2, dim=-1) + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + return hidden_states + + +class GLMTTSDiTBlock(nn.Module): + """ + GLM-TTS DiT block with self-attention and FFN. + Uses adaptive layer norm for timestep conditioning. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + ff_mult: int = 4, + ): + super().__init__() + + # Self-attention with adaptive layer norm + self.attn_norm = GLMTTSAdaLayerNorm(dim) + self.attn = GLMTTSSelfAttention( + dim=dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + ) + + # Feed-forward with layer norm + self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = GLMTTSFeedForward(dim, inner_dim=dim * ff_mult) + + def forward( + self, + hidden_states: torch.Tensor, + timestep_emb: torch.Tensor, + rotary_embedding: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + # Self-attention with adaptive norm + norm_hidden, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(hidden_states, timestep_emb) + attn_out = self.attn(norm_hidden, rotary_emb=rotary_embedding, attention_mask=attention_mask) + hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_out + + # Feed-forward with adaptive norm + norm_hidden = self.ff_norm(hidden_states) + norm_hidden = norm_hidden * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1) + ff_out = self.ff(norm_hidden) + hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_out + + return hidden_states + + +class GLMTTSDiTModel(nn.Module): + """ + Optimized GLM-TTS DiT model using vLLM layers. + + This model implements flow matching for speech token to mel-spectrogram generation. + + Architecture: + - Input: Speech tokens + noisy mel + speaker embedding + - Preprocess: Project and combine inputs + - Transformer blocks with adaptive layer norm + - Output: Predicted velocity for flow matching + + Args: + od_config: OmniDiffusion configuration object + sample_size: Maximum mel spectrogram length + mel_dim: Mel spectrogram dimension (typically 80) + num_layers: Number of transformer blocks + attention_head_dim: Dimension per attention head + num_attention_heads: Number of attention heads + speech_token_vocab_size: Vocabulary size for speech tokens + speech_token_dim: Embedding dimension for speech tokens + speaker_embed_dim: Dimension of speaker embeddings + time_proj_dim: Time projection dimension + """ + + def __init__( + self, + od_config: OmniDiffusionConfig | None = None, + sample_size: int = 2048, + mel_dim: int = 80, + num_layers: int = 22, + attention_head_dim: int = 64, + num_attention_heads: int = 16, + speech_token_vocab_size: int = 100000, + speech_token_dim: int = 512, + speaker_embed_dim: int = 192, + time_proj_dim: int = 256, + ): + super().__init__() + + self.sample_size = sample_size + self.mel_dim = mel_dim + self.num_layers = num_layers + self.attention_head_dim = attention_head_dim + self.num_attention_heads = num_attention_heads + + # inner_dim is the transformer hidden dimension + self.inner_dim = num_attention_heads * attention_head_dim + + # Store config for compatibility (like Stable Audio) + self.config = type( + "Config", + (), + { + "sample_size": sample_size, + "mel_dim": mel_dim, + "num_layers": num_layers, + "attention_head_dim": attention_head_dim, + "num_attention_heads": num_attention_heads, + "speech_token_vocab_size": speech_token_vocab_size, + "speech_token_dim": speech_token_dim, + "speaker_embed_dim": speaker_embed_dim, + "time_proj_dim": time_proj_dim, + }, + )() + + # Time projection (Gaussian Fourier features) + self.time_proj = GLMTTSGaussianFourierProjection(embedding_size=time_proj_dim // 2) + + # Timestep projection: time_proj_dim -> inner_dim + self.timestep_proj = nn.Sequential( + nn.Linear(time_proj_dim, self.inner_dim, bias=True), + nn.SiLU(), + nn.Linear(self.inner_dim, self.inner_dim, bias=True), + ) + + # Speech token embedding + self.speech_token_embed = nn.Embedding(speech_token_vocab_size + 1, speech_token_dim) + + # Speaker embedding projection + self.speaker_proj = nn.Sequential( + nn.Linear(speaker_embed_dim, self.inner_dim, bias=False), + nn.SiLU(), + nn.Linear(self.inner_dim, self.inner_dim, bias=False), + ) + + # Input projection: mel + speech_token_dim + inner_dim (speaker) -> inner_dim + self.proj_in = nn.Linear(mel_dim + speech_token_dim + self.inner_dim, self.inner_dim, bias=False) + + # Rotary embedding (computed on the fly based on sequence length) + self.rotary_embed_dim = attention_head_dim // 2 + + # Transformer blocks + self.transformer_blocks = nn.ModuleList( + [ + GLMTTSDiTBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + ) + for _ in range(num_layers) + ] + ) + + # Final norm and output projection + self.norm_out = GLMTTSAdaLayerNormFinal(self.inner_dim) + self.proj_out = nn.Linear(self.inner_dim, mel_dim, bias=False) + + @property + def dtype(self) -> torch.dtype: + """Return the dtype of the model parameters.""" + return next(self.parameters()).dtype + + def _get_rotary_embedding( + self, + seq_len: int, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compute rotary embeddings for the given sequence length.""" + dim = self.rotary_embed_dim + inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).float() / dim)) + t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype) + freqs = torch.outer(t, inv_freq) + cos = freqs.cos().to(dtype) + sin = freqs.sin().to(dtype) + # Double the dimension by concatenating + cos = torch.cat([cos, cos], dim=-1) + sin = torch.cat([sin, sin], dim=-1) + return cos, sin + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + speech_tokens: torch.Tensor, + speaker_embedding: torch.Tensor, + rotary_embedding: tuple[torch.Tensor, torch.Tensor] | None = None, + return_dict: bool = True, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor | Transformer2DModelOutput: + """ + Forward pass of the GLM-TTS DiT model. + + Args: + hidden_states: Noisy mel-spectrogram [B, T, mel_dim] + timestep: Timestep tensor [B] or [1] + speech_tokens: Speech token indices [B, T_tokens] + speaker_embedding: Speaker embedding [B, speaker_dim] + rotary_embedding: Precomputed rotary embeddings (cos, sin) + return_dict: Whether to return a dataclass or tuple + attention_mask: Optional attention mask + + Returns: + Predicted velocity for flow matching [B, T, mel_dim] + """ + batch_size, seq_len, _ = hidden_states.shape + device = hidden_states.device + + # Time embedding + time_hidden_states = self.timestep_proj(self.time_proj(timestep.to(self.dtype))) + + # Speech token embedding - interpolate to match mel length + token_emb = self.speech_token_embed(speech_tokens) # [B, T_tokens, D] + if token_emb.shape[1] != seq_len: + token_emb = token_emb.transpose(1, 2) # [B, D, T_tokens] + token_emb = F.interpolate(token_emb, size=seq_len, mode="linear", align_corners=False) + token_emb = token_emb.transpose(1, 2) # [B, T, D] + + # Speaker embedding projection and expand to sequence + spk_emb = self.speaker_proj(speaker_embedding) # [B, inner_dim] + spk_emb = spk_emb.unsqueeze(1).expand(-1, seq_len, -1) # [B, T, inner_dim] + + # Concatenate inputs and project + combined = torch.cat([hidden_states, token_emb, spk_emb], dim=-1) + hidden_states = self.proj_in(combined) + + # Compute rotary embeddings if not provided + if rotary_embedding is None: + rotary_embedding = self._get_rotary_embedding(seq_len, device, hidden_states.dtype) + + # Transformer blocks + for block in self.transformer_blocks: + hidden_states = block( + hidden_states, + time_hidden_states, + rotary_embedding=rotary_embedding, + attention_mask=attention_mask, + ) + + # Final norm and output projection + hidden_states = self.norm_out(hidden_states, time_hidden_states) + output = self.proj_out(hidden_states) + + if return_dict: + return Transformer2DModelOutput(sample=output) + return (output,) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """ + Load weights from a pretrained model. + + Maps diffusers/HF weight names to our module structure. + + Returns: + Set of parameter names that were successfully loaded. + """ + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + # Weight name mapping from HF/diffusers to our implementation + name_mapping = { + # Timestep projection + "timestep_proj.linear_1.weight": "timestep_proj.0.weight", + "timestep_proj.linear_1.bias": "timestep_proj.0.bias", + "timestep_proj.linear_2.weight": "timestep_proj.2.weight", + "timestep_proj.linear_2.bias": "timestep_proj.2.bias", + # Speaker projection + "speaker_proj.linear_1.weight": "speaker_proj.0.weight", + "speaker_proj.linear_2.weight": "speaker_proj.2.weight", + } + + for name, loaded_weight in weights: + # Apply name mapping if needed + mapped_name = name_mapping.get(name, name) + + if mapped_name in params_dict: + param = params_dict[mapped_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(mapped_name) + else: + logger.debug(f"Skipping weight {name} - not found in model") + + return loaded_params diff --git a/vllm_omni/diffusion/models/glm_tts/pipeline_glm_tts.py b/vllm_omni/diffusion/models/glm_tts/pipeline_glm_tts.py new file mode 100644 index 00000000000..608257bd10a --- /dev/null +++ b/vllm_omni/diffusion/models/glm_tts/pipeline_glm_tts.py @@ -0,0 +1,471 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +GLM-TTS Pipeline for vLLM-Omni. + +This module provides text-to-speech generation using the GLM-TTS model, +integrated with the vLLM-Omni diffusion framework. + +GLM-TTS uses a two-stage architecture: +1. LLM (Llama-based): Converts text to speech tokens +2. Flow Matching: Converts speech tokens to mel-spectrograms +3. Vocoder: Converts mel-spectrograms to waveforms +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable + +import torch +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from vllm.logger import init_logger +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.glm_tts.glm_tts_dit import GLMTTSDiTModel +from vllm_omni.diffusion.request import OmniDiffusionRequest + +logger = init_logger(__name__) + + +def get_glm_tts_post_process_func( + od_config: OmniDiffusionConfig, +): + """ + Create post-processing function for GLM-TTS output. + + Converts raw audio tensor to numpy array for saving. + """ + + def post_process_func( + audio: torch.Tensor, + output_type: str = "np", + ): + if output_type == "latent": + return audio + if output_type == "pt": + return audio + # Convert to numpy + audio_np = audio.cpu().float().numpy() + return audio_np + + return post_process_func + + +class GLMTTSPipeline(nn.Module): + """ + Pipeline for text-to-speech generation using GLM-TTS. + + This pipeline generates audio from text prompts using the GLM-TTS model, + integrated with vLLM-Omni's diffusion framework. + + The pipeline consists of: + 1. LLM: Generates speech tokens from text (loaded externally or provided) + 2. Flow Model (DiT): Converts speech tokens to mel-spectrograms + 3. Vocoder: Converts mel-spectrograms to waveforms + + Args: + od_config: OmniDiffusion configuration object + prefix: Weight prefix for loading (default: "") + """ + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + + self.device = get_local_device() + dtype = getattr(od_config, "dtype", torch.float16) + + model = od_config.model + local_files_only = os.path.exists(model) + + # Set up weights sources for the flow model (transformer) + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="flow", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ), + ] + + # Initialize our custom transformer (weights loaded via load_weights) + self.transformer = GLMTTSDiTModel(od_config=od_config) + + # Try to load vocoder if available + self._load_vocoder(model, local_files_only, dtype) + + # Model parameters + self.sample_rate = 22050 # GLM-TTS default sample rate + self.mel_dim = 80 + self.hop_length = 256 + + # Compute rotary embedding dimension + self.rotary_embed_dim = self.transformer.config.attention_head_dim // 2 + + # Cache backend (set by worker if needed) + self._cache_backend = None + + # Properties for generation tracking + self._guidance_scale = None + self._num_timesteps = None + self._current_timestep = None + + def _load_vocoder( + self, + model: str, + local_files_only: bool, + dtype: torch.dtype, + ): + """Load vocoder from model path if available.""" + self.vocoder = None + self.vocoder_type = None + + # Try to load Vocos vocoder + try: + vocos_path = os.path.join(model, "vocos2d") if local_files_only else model + if os.path.exists(vocos_path) or not local_files_only: + try: + from vocos import Vocos + + self.vocoder = Vocos.from_pretrained( + vocos_path if local_files_only else model, + subfolder="vocos2d" if not local_files_only else None, + ) + self.vocoder = self.vocoder.to(self.device) + self.vocoder_type = "vocos" + logger.info("Loaded Vocos vocoder") + except Exception as e: + logger.debug(f"Could not load Vocos vocoder: {e}") + except Exception as e: + logger.debug(f"Vocoder loading failed: {e}") + + if self.vocoder is None: + logger.warning("No vocoder loaded. Output will be mel-spectrograms. Install vocos for waveform generation.") + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale is not None and self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + def check_inputs( + self, + prompt: str | list[str] | None, + speech_tokens: torch.Tensor | None, + audio_duration_s: float, + from_stage_processor: bool = False, + ): + """Validate input parameters.""" + if audio_duration_s <= 0: + raise ValueError(f"`audio_duration_s` must be positive, got {audio_duration_s}") + + # When receiving from stage processor, speech_tokens will be provided + if not from_stage_processor and prompt is None and speech_tokens is None: + raise ValueError("Provide either `prompt` or `speech_tokens` in extra params. Cannot leave both undefined.") + + def encode_text_to_tokens( + self, + prompt: str | list[str], + batch_size: int, + audio_duration_s: float, + ) -> torch.Tensor: + """ + Convert text to speech tokens. + + Note: This is a placeholder. In production, this should use the + GLM-TTS LLM to generate speech tokens from text. + + Args: + prompt: Input text to synthesize + batch_size: Batch size + audio_duration_s: Target audio duration + + Returns: + Speech token indices [B, T_tokens] + """ + # Estimate token length based on duration + # GLM-TTS uses ~12.5 tokens per second + tokens_per_second = 12.5 + token_length = int(audio_duration_s * tokens_per_second) + + logger.warning( + "LLM not loaded. Using placeholder speech tokens. " + "For proper synthesis, provide speech_tokens in request.extra." + ) + + return torch.randint( + 0, + self.transformer.config.speech_token_vocab_size, + (batch_size, token_length), + device=self.device, + dtype=torch.long, + ) + + def get_speaker_embedding( + self, + batch_size: int, + speaker_embedding: torch.Tensor | None = None, + ) -> torch.Tensor: + """ + Get or generate speaker embedding. + + Args: + batch_size: Batch size + speaker_embedding: Pre-computed speaker embedding + + Returns: + Speaker embedding [B, speaker_dim] + """ + speaker_dim = self.transformer.config.speaker_embed_dim + + if speaker_embedding is not None: + return speaker_embedding.to(self.device, dtype=self.transformer.dtype) + + # Return random speaker embedding as placeholder + logger.warning( + "No speaker embedding provided. Using random embedding. " + "For voice cloning, provide speaker_embedding in request.extra." + ) + return torch.randn( + batch_size, + speaker_dim, + device=self.device, + dtype=self.transformer.dtype, + ) + + def prepare_latents( + self, + batch_size: int, + mel_length: int, + dtype: torch.dtype, + device: torch.device, + generator: torch.Generator | None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + """Prepare initial latent noise for flow matching.""" + shape = (batch_size, mel_length, self.mel_dim) + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device) + + return latents + + def _get_rotary_embedding( + self, + seq_len: int, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compute rotary embeddings for the given sequence length.""" + dim = self.rotary_embed_dim + inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).float() / dim)) + t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype) + freqs = torch.outer(t, inv_freq) + cos = freqs.cos().to(dtype) + sin = freqs.sin().to(dtype) + # Double the dimension by concatenating + cos = torch.cat([cos, cos], dim=-1) + sin = torch.cat([sin, sin], dim=-1) + return cos, sin + + def decode_mel_to_audio( + self, + mel: torch.Tensor, + output_type: str = "np", + ) -> torch.Tensor: + """ + Decode mel-spectrogram to audio waveform. + + Args: + mel: Mel-spectrogram [B, T, mel_dim] or [B, mel_dim, T] + output_type: Output type + + Returns: + Audio waveform + """ + if output_type == "latent": + return mel + + # Ensure mel is in [B, mel_dim, T] format for vocoder + if mel.shape[-1] == self.mel_dim: + mel = mel.transpose(1, 2) + + if self.vocoder is not None: + mel_for_vocoder = mel.to(dtype=torch.float32) + if self.vocoder_type == "vocos": + audio = self.vocoder.decode(mel_for_vocoder) + if audio.dim() == 2: + audio = audio.unsqueeze(1) + else: + audio = mel_for_vocoder + else: + # Return mel if no vocoder + audio = mel + + return audio + + @torch.no_grad() + def forward( + self, + req: OmniDiffusionRequest, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + audio_duration_s: float = 10.0, + num_inference_steps: int = 32, + guidance_scale: float = 1.0, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + output_type: str = "np", + ) -> DiffusionOutput: + """ + Generate audio from text prompt or speech tokens. + + Args: + req: OmniDiffusionRequest containing generation parameters + prompt: Text prompt for audio generation + negative_prompt: Negative prompt (not used in flow matching) + audio_duration_s: Target audio duration in seconds + num_inference_steps: Number of flow matching steps + guidance_scale: Classifier-free guidance scale + generator: Random generator for reproducibility + latents: Pre-generated latents + output_type: Output format ("np", "pt", or "latent") + + Returns: + DiffusionOutput containing generated audio + """ + # Extract from request + prompt = req.prompt if req.prompt is not None else prompt + num_inference_steps = req.num_inference_steps or num_inference_steps + if req.guidance_scale_provided: + guidance_scale = req.guidance_scale + + if generator is None: + generator = req.generator + if generator is None and req.seed is not None: + generator = torch.Generator(device=self.device).manual_seed(req.seed) + + # Check if input comes from stage processor (two-stage pipeline) + # Stage processor provides speech_tokens in additional_information + additional_info = req.extra.get("additional_information", {}) + from_stage_processor = bool(additional_info) + + # Get speech tokens - priority: additional_information > extra > generate + speech_tokens = additional_info.get("speech_tokens", None) + if speech_tokens is None: + speech_tokens = req.extra.get("speech_tokens", None) + + # Get speaker embedding - priority: additional_information > extra > generate + speaker_embedding = additional_info.get("speaker_embedding", None) + if speaker_embedding is None: + speaker_embedding = req.extra.get("speaker_embedding", None) + + audio_duration_s = req.extra.get("audio_duration_s", audio_duration_s) + + # Validate inputs + self.check_inputs(prompt, speech_tokens, audio_duration_s, from_stage_processor) + + # Determine batch size + if speech_tokens is not None and isinstance(speech_tokens, torch.Tensor): + batch_size = speech_tokens.shape[0] + elif prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = 1 + + device = self.device + self._guidance_scale = guidance_scale + + # Generate or use provided speech tokens + if speech_tokens is None: + speech_tokens = self.encode_text_to_tokens(prompt, batch_size, audio_duration_s) + else: + if from_stage_processor: + logger.info("Using speech tokens from LLM stage (two-stage pipeline)") + speech_tokens = speech_tokens.to(device) + # Ensure speech_tokens has batch dimension [B, T] + if speech_tokens.dim() == 1: + speech_tokens = speech_tokens.unsqueeze(0) + + # Get speaker embedding + speaker_embedding = self.get_speaker_embedding(batch_size, speaker_embedding) + + # Calculate mel-spectrogram length + mel_length = int(audio_duration_s * self.sample_rate / self.hop_length) + + # Prepare timesteps (flow matching: 0 -> 1) + timesteps = torch.linspace(0, 1, num_inference_steps + 1, device=device, dtype=self.transformer.dtype) + self._num_timesteps = num_inference_steps + + # Prepare latents (initial noise) + latents = self.prepare_latents( + batch_size, + mel_length, + self.transformer.dtype, + device, + generator, + latents, + ) + + # Compute rotary embeddings + rotary_embedding = self._get_rotary_embedding(mel_length, device, latents.dtype) + + # Flow matching loop (Euler integration) + for i in range(num_inference_steps): + t = timesteps[i] + dt = timesteps[i + 1] - timesteps[i] + self._current_timestep = t + + # Expand timestep for batch + t_batch = t.expand(batch_size) + + # Predict velocity + velocity = self.transformer( + latents, + t_batch, + speech_tokens=speech_tokens, + speaker_embedding=speaker_embedding, + rotary_embedding=rotary_embedding, + return_dict=False, + )[0] + + # Euler step + latents = latents + velocity * dt + + self._current_timestep = None + + # Decode to audio + audio = self.decode_mel_to_audio(latents, output_type) + + return DiffusionOutput(output=audio) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights using AutoWeightsLoader for vLLM integration.""" + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 5edd87a827d..f656faf23b0 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -79,6 +79,11 @@ "pipeline_flux2_klein", "Flux2KleinPipeline", ), + "GLMTTSPipeline": ( + "glm_tts", + "pipeline_glm_tts", + "GLMTTSPipeline", + ), } @@ -127,6 +132,7 @@ def initialize_model( "LongCatImageEditPipeline": "get_longcat_image_post_process_func", "StableDiffusion3Pipeline": "get_sd3_image_post_process_func", "Flux2KleinPipeline": "get_flux2_klein_post_process_func", + "GLMTTSPipeline": "get_glm_tts_post_process_func", } _DIFFUSION_PRE_PROCESS_FUNCS = { diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index a2790cd06e5..88d5d1651fe 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -603,10 +603,14 @@ def _stage_worker( ) try: if stage_type == "diffusion": - engine_args.pop("model_stage") + # Remove parameters that are not valid for OmniDiffusionConfig + engine_args.pop("model_stage", None) + engine_args.pop("max_num_seqs", None) # Added by runtime config, not valid for diffusion stage_engine = OmniDiffusion(**engine_args) else: # Default to LLM engine + # Remove 'model' from engine_args if present to avoid duplicate argument + engine_args.pop("model", None) stage_engine = OmniLLM(model=model, **engine_args) finally: # Release all locks by closing file descriptors diff --git a/vllm_omni/model_executor/stage_configs/glm_tts.yaml b/vllm_omni/model_executor/stage_configs/glm_tts.yaml new file mode 100644 index 00000000000..c41f2b945a4 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/glm_tts.yaml @@ -0,0 +1,74 @@ +# Stage config for running GLM-TTS with two-stage architecture. +# Stage 0: LLM (Llama-based) - Text → Speech tokens +# Stage 1: DiT (Flow matching) - Speech tokens → Audio + +# Model structure on HuggingFace (zai-org/GLM-TTS): +# llm/ - LLM for speech token generation +# flow/ - DiT/Flow matching model for mel generation +# hift/ - HiFi-GAN vocoder +# vocos2d/ - Vocos vocoder +# speech_tokenizer/ - Speech tokenizer + +# The following config is designed for single GPU usage. +stage_args: + - stage_id: 0 + stage_type: llm # Use llm stage type for speech token generation + runtime: + process: true + devices: "0" + max_batch_size: 1 + engine_args: + model: zai-org/GLM-TTS + model_stage: llm # Load from llm/ subdirectory + model_arch: LlamaForCausalLM # GLM-TTS uses Llama architecture + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.4 + enforce_eager: true + trust_remote_code: true + engine_output_type: latent + enable_prefix_caching: false + max_num_batched_tokens: 4096 + is_comprehension: false + final_output: false # Speech tokens are intermediate, not returned to user + final_output_type: text + default_sampling_params: + temperature: 0.9 + top_p: 0.8 + top_k: 40 + max_tokens: 2048 + seed: 42 + detokenize: false + repetition_penalty: 1.05 + stop_token_ids: [151330] # GLM_TTS_EOA_TOKEN_ID + + - stage_id: 1 + stage_type: diffusion # Use diffusion stage type for DiT/flow matching + runtime: + process: true + devices: "0" + max_batch_size: 1 + engine_args: + model: zai-org/GLM-TTS + model_stage: flow # Load from flow/ subdirectory + model_class_name: GLMTTSPipeline + enforce_eager: true + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.glm_tts.llm2dit + final_output: true + final_output_type: audio + default_generation_params: + num_inference_steps: 32 + guidance_scale: 1.0 + +# Runtime config: stage edges define data flow +runtime: + enabled: true + defaults: + window_size: -1 # Trigger downstream only after full upstream completion + max_inflight: 1 # Process serially within each stage + + edges: + - from: 0 # LLM → DiT: trigger after full speech token generation + to: 1 + window_size: -1 diff --git a/vllm_omni/model_executor/stage_input_processors/glm_tts.py b/vllm_omni/model_executor/stage_input_processors/glm_tts.py new file mode 100644 index 00000000000..01b4114c834 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/glm_tts.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Stage input processor for GLM-TTS two-stage pipeline. + +Handles the transition from LLM stage (speech token generation) +to DiT stage (mel-spectrogram generation via flow matching). +""" + +import torch +from vllm.inputs import TextPrompt + +from vllm_omni.inputs.data import OmniTokensPrompt + +# GLM-TTS special tokens +# These define the audio token range and boundaries +GLM_TTS_BOA_TOKEN_ID = 151329 # Beginning of Audio +GLM_TTS_EOA_TOKEN_ID = 151330 # End of Audio +GLM_TTS_AUDIO_TOKEN_START = 151936 # Start of audio token range +GLM_TTS_AUDIO_TOKEN_END = 184703 # End of audio token range (32768 tokens) + + +def extract_speech_tokens(token_ids: list[int]) -> list[int]: + """ + Extract speech tokens from LLM output. + + Filters token IDs to only include valid audio tokens + (between AUDIO_TOKEN_START and AUDIO_TOKEN_END). + + Args: + token_ids: Raw token IDs from LLM output + + Returns: + List of speech token IDs (normalized to 0-based indexing) + """ + speech_tokens = [] + for token_id in token_ids: + if GLM_TTS_AUDIO_TOKEN_START <= token_id <= GLM_TTS_AUDIO_TOKEN_END: + # Normalize to 0-based index for the DiT model + normalized_token = token_id - GLM_TTS_AUDIO_TOKEN_START + speech_tokens.append(normalized_token) + return speech_tokens + + +def llm2dit( + stage_list, + engine_input_source, + prompt: OmniTokensPrompt | TextPrompt = None, + requires_multimodal_data: bool = False, +): + """ + Transform LLM outputs (speech tokens) into DiT stage inputs. + + This processor extracts speech tokens generated by the LLM stage + and packages them for the DiT (flow matching) stage. + + Args: + stage_list: List of all stages in the pipeline + engine_input_source: List of upstream stage IDs to consume from + prompt: Original prompt (may contain multimodal data) + requires_multimodal_data: Whether to pass multimodal data downstream + + Returns: + List of OmniTokensPrompt objects for the DiT stage + """ + if not engine_input_source: + raise ValueError("engine_input_source cannot be empty") + + source_stage_id = engine_input_source[0] + if source_stage_id >= len(stage_list): + raise IndexError(f"Invalid stage_id: {source_stage_id}") + if stage_list[source_stage_id].engine_outputs is None: + raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") + + llm_outputs = stage_list[source_stage_id].engine_outputs + dit_inputs = [] + + if not isinstance(prompt, list): + prompt = [prompt] + + # Extract multimodal data (speaker embeddings, etc.) from original prompts + multi_modal_data = {} + for llm_output, p in zip(llm_outputs, prompt): + if p is not None and isinstance(p, dict): + multi_modal_data[llm_output.request_id] = p.get("multi_modal_data", None) + else: + multi_modal_data[llm_output.request_id] = None + + for i, llm_output in enumerate(llm_outputs): + output = llm_output.outputs[0] + token_ids = output.token_ids + + # Extract speech tokens from LLM output + speech_tokens = extract_speech_tokens(token_ids) + + if not speech_tokens: + # If no speech tokens found, use the raw tokens + # (LLM may not have started generating audio yet) + speech_tokens = token_ids + + # Build additional information for DiT stage + additional_information = { + "speech_tokens": torch.tensor(speech_tokens, dtype=torch.long), + "raw_token_ids": token_ids, + "prompt_token_ids": llm_output.prompt_token_ids, + } + + # Extract speaker embedding if available in multimodal output + if hasattr(output, "multimodal_output") and output.multimodal_output: + if "speaker_embedding" in output.multimodal_output: + additional_information["speaker_embedding"] = output.multimodal_output["speaker_embedding"] + if "latent" in output.multimodal_output: + additional_information["llm_hidden_states"] = output.multimodal_output["latent"] + + dit_inputs.append( + OmniTokensPrompt( + prompt_token_ids=speech_tokens, + additional_information=additional_information, + multi_modal_data=(multi_modal_data.get(llm_output.request_id) if requires_multimodal_data else None), + mm_processor_kwargs=None, + ) + ) + + return dit_inputs