# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ CosyVoice3 Code2Wav Stage - Converts speech tokens to audio waveforms. This module contains the code2wav (token-to-waveform) stage which uses: 1. DiT (Diffusion Transformer) with optimized attention backends 2. CFM (Conditional Flow Matching) for mel spectrogram generation 3. HiFiGAN vocoder for waveform synthesis """ from __future__ import annotations import numpy as np import torch import torch.nn as nn from omegaconf import DictConfig from torch.nn import functional as F from vllm.logger import init_logger from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import ( CausalConditionalCFM, CausalMaskedDiffWithDiT, ) from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( CausalConvRNNF0Predictor, CausalHiFTGenerator, ) from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask logger = init_logger(__name__) class CosyVoice3Code2Wav(nn.Module): """CosyVoice3 Code2Wav stage for token-to-waveform conversion. This class encapsulates: - Flow matching decoder with DiT backbone (using diffusion attention) - HiFiGAN vocoder for mel-to-waveform conversion """ def __init__(self, config: CosyVoice3Config): super().__init__() self.config = config # Build flow matching components pre_lookahead_layer = PreLookaheadLayer(**config.flow["pre_lookahead_layer"]) decoder_cfg = config.flow["decoder"] cfm_params = DictConfig(decoder_cfg["cfm_params"]) # DiT estimator using diffusion attention (Flash/Sage/SDPA backends) estimator = DiT(**decoder_cfg["estimator"]) decoder = CausalConditionalCFM( in_channels=decoder_cfg["in_channels"], estimator=estimator, cfm_params=cfm_params, n_spks=decoder_cfg["n_spks"], spk_emb_dim=decoder_cfg["spk_emb_dim"], ) self.flow_model = CausalMaskedDiffWithDiT( input_size=config.flow["input_size"], output_size=config.flow["output_size"], spk_embed_dim=config.flow["spk_embed_dim"], output_type=config.flow["output_type"], vocab_size=config.flow["vocab_size"], input_frame_rate=config.flow["input_frame_rate"], only_mask_loss=config.flow["only_mask_loss"], token_mel_ratio=config.flow["token_mel_ratio"], pre_lookahead_len=config.flow["pre_lookahead_len"], pre_lookahead_layer=pre_lookahead_layer, decoder=decoder, ) # Build HiFiGAN vocoder f0_predictor = CausalConvRNNF0Predictor( num_class=config.hift["f0_predictor"]["num_class"], in_channels=config.hift["f0_predictor"]["in_channels"], cond_channels=config.hift["f0_predictor"]["cond_channels"], ) self.hift = CausalHiFTGenerator( in_channels=config.hift["in_channels"], base_channels=config.hift["base_channels"], nb_harmonics=config.hift["nb_harmonics"], sampling_rate=config.hift["sampling_rate"], nsf_alpha=config.hift["nsf_alpha"], nsf_sigma=config.hift["nsf_sigma"], nsf_voiced_threshold=config.hift["nsf_voiced_threshold"], upsample_rates=config.hift["upsample_rates"], upsample_kernel_sizes=config.hift["upsample_kernel_sizes"], istft_params=config.hift["istft_params"], resblock_kernel_sizes=config.hift["resblock_kernel_sizes"], resblock_dilation_sizes=config.hift["resblock_dilation_sizes"], source_resblock_kernel_sizes=config.hift["source_resblock_kernel_sizes"], source_resblock_dilation_sizes=config.hift["source_resblock_dilation_sizes"], lrelu_slope=config.hift["lrelu_slope"], audio_limit=config.hift["audio_limit"], conv_pre_look_right=config.hift["conv_pre_look_right"], f0_predictor=f0_predictor, ) # Run hift in float32 to avoid dtype mismatches in internal ops self.hift = self.hift.float() # Streaming/chunking parameters self.token_overlap_len = 20 self.mel_overlap_len = int(self.token_overlap_len / self.flow_model.input_frame_rate * 22050 / 256) self.mel_window = np.hamming(2 * self.mel_overlap_len) self.mel_cache_len = 20 self.source_cache_len = int(self.mel_cache_len * 256) self.speech_window = np.hamming(2 * self.source_cache_len) @property def input_frame_rate(self) -> int: """Input frame rate from flow model.""" return self.flow_model.input_frame_rate @property def token_mel_ratio(self) -> int: """Token to mel ratio.""" return self.flow_model.token_mel_ratio @property def output_size(self) -> int: """Output mel dimension.""" return self.flow_model.output_size @property def input_embedding(self) -> nn.Embedding: """Token embedding layer.""" return self.flow_model.input_embedding @property def pre_lookahead_layer(self) -> nn.Module: """Pre-lookahead layer.""" return self.flow_model.pre_lookahead_layer @property def decoder(self) -> nn.Module: """Flow matching decoder.""" return self.flow_model.decoder @property def spk_embed_affine_layer(self) -> nn.Linear: """Speaker embedding affine layer.""" return self.flow_model.spk_embed_affine_layer @torch.inference_mode() def forward( self, token: torch.Tensor, prompt_token: torch.Tensor, prompt_feat: torch.Tensor, embedding: torch.Tensor, n_timesteps: int = 10, ) -> torch.Tensor: """Generate audio waveform from speech tokens. Args: token: Speech tokens from talker stage [batch, seq_len] prompt_token: Prompt speech tokens [batch, prompt_len] prompt_feat: Prompt mel features [batch, feat_len, mel_dim] embedding: Speaker embedding [batch, spk_dim] n_timesteps: Number of diffusion steps Returns: Audio waveform [batch, 1, audio_len] """ device = token.device dtype = next(self.flow_model.parameters()).dtype # Normalize and project speaker embedding embedding = embedding.to(device=device, dtype=dtype) embedding = F.normalize(embedding, dim=1) embedding = self.spk_embed_affine_layer(embedding) # Prepare tokens prompt_token = prompt_token.to(device=device) token_len1, token_len2 = prompt_token.shape[1], token.shape[1] prompt_token_len = torch.tensor([token_len1], device=device, dtype=torch.int32) token_len = torch.tensor([token_len2], device=device, dtype=torch.int32) # Concatenate prompt and target tokens full_token = torch.cat([prompt_token, token], dim=1) full_token_len = prompt_token_len + token_len # Create mask mask = (~make_pad_mask(full_token_len)).unsqueeze(-1).to(embedding) # Token embedding token_emb = self.input_embedding(torch.clamp(full_token, min=0)) * mask # Pre-lookahead processing h = self.pre_lookahead_layer(token_emb) h = h.repeat_interleave(self.token_mel_ratio, dim=1) # Calculate mel lengths mel_len1 = prompt_feat.shape[1] mel_len2 = h.shape[1] - mel_len1 # Build conditioning conds = torch.zeros( [1, mel_len1 + mel_len2, self.output_size], device=device, dtype=h.dtype, ) conds[:, :mel_len1] = prompt_feat conds = conds.transpose(1, 2) # Create mel mask mel_mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h) # Run flow matching decoder feat, _ = self.decoder( mu=h.transpose(1, 2).contiguous(), mask=mel_mask.unsqueeze(1), spks=embedding, cond=conds, n_timesteps=n_timesteps, ) # Extract generated portion (after prompt) feat = feat[:, :, mel_len1:] # Run vocoder hift_weight = self.hift.m_source.l_linear.weight tts_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype) if tts_mel.shape[-1] == 0: tts_speech = torch.zeros( (tts_mel.shape[0], 1, 0), device=tts_mel.device, dtype=tts_mel.dtype, ) else: tts_speech, _ = self.hift.inference(speech_feat=tts_mel) return tts_speech def load_weights(self, model_dir: str, device: torch.device) -> None: """Load flow.pt and hift.pt weights. Args: model_dir: Model directory containing flow.pt and hift.pt device: Device to load weights to """ import os # Load flow weights flow_path = os.path.join(model_dir, "flow.pt") self.flow_model.load_state_dict(torch.load(flow_path, map_location=device), strict=True) self.flow_model.to(device).eval() logger.info(f"Loaded flow weights from {flow_path}") # Load hift weights hift_path = os.path.join(model_dir, "hift.pt") hift_state_dict = { k.replace("generator.", ""): v for k, v in torch.load(hift_path, map_location=device).items() } self.hift.load_state_dict(hift_state_dict, strict=True) self.hift.to(device).eval() logger.info(f"Loaded hift weights from {hift_path}")