# 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 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.transformers_utils.configs.cosyvoice3 import CosyVoice3Config 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_mel( self, token: torch.Tensor, prompt_token: torch.Tensor, prompt_feat: torch.Tensor, embedding: torch.Tensor, n_timesteps: int = 10, token_offset_tokens: int = 0, streaming: bool = True, finalize: bool = False, ) -> torch.Tensor: """Generate mel features via the upstream flow-model inference path.""" flow_weight = next(self.flow_model.parameters()) device = flow_weight.device dtype = flow_weight.dtype token = token.to(device=device, dtype=torch.int32) prompt_token = prompt_token.to(device=device, dtype=torch.int32) prompt_feat = prompt_feat.to(device=device, dtype=dtype) embedding = embedding.to(device=device, dtype=dtype) token_len = torch.tensor([token.shape[1]], device=device, dtype=torch.int32) prompt_token_len = torch.tensor([prompt_token.shape[1]], device=device, dtype=torch.int32) prompt_feat_len = torch.tensor([prompt_feat.shape[1]], device=device, dtype=torch.int32) feat, _ = self.flow_model.inference( token=token, token_len=token_len, prompt_token=prompt_token, prompt_token_len=prompt_token_len, prompt_feat=prompt_feat, prompt_feat_len=prompt_feat_len, embedding=embedding, streaming=streaming, finalize=finalize, n_timesteps=n_timesteps, ) trim_mel = max(0, int(token_offset_tokens)) * int(self.token_mel_ratio) if trim_mel > 0: feat = feat[:, :, trim_mel:] return feat @torch.inference_mode() def forward_streaming( self, token: torch.Tensor, prompt_token: torch.Tensor, prompt_feat: torch.Tensor, embedding: torch.Tensor, *, cache_state: dict[str, torch.Tensor] | None = None, n_timesteps: int = 10, token_offset_tokens: int = 0, finalize: bool = False, ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: """Decode streaming audio using cumulative mel + emitted-speech offset. This mirrors upstream CosyVoice3 streaming semantics more closely than waveform-domain overlap-add: keep a cumulative mel history per request, re-run causal HiFT on the history, and emit only the newly grown speech suffix. That preserves causal look-right handling without double trimming or duplicated overlap at chunk boundaries. """ feat = self._forward_mel( token=token, prompt_token=prompt_token, prompt_feat=prompt_feat, embedding=embedding, n_timesteps=n_timesteps, token_offset_tokens=token_offset_tokens, streaming=True, finalize=finalize, ) hift_weight = self.hift.m_source.l_linear.weight chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype) cached_mel = None if not cache_state else cache_state.get("mel") speech_offset_obj = None if not cache_state else cache_state.get("speech_offset") try: speech_offset = int(speech_offset_obj) if speech_offset_obj is not None else 0 except (TypeError, ValueError): speech_offset = 0 if isinstance(cached_mel, torch.Tensor) and cached_mel.numel() > 0: cached_mel = cached_mel.to(device=chunk_mel.device, dtype=chunk_mel.dtype) tts_mel = torch.cat([cached_mel, chunk_mel], dim=-1) if chunk_mel.numel() > 0 else cached_mel else: tts_mel = chunk_mel if tts_mel.shape[-1] == 0: tts_speech = torch.zeros((chunk_mel.shape[0], 1, 0), device=chunk_mel.device, dtype=chunk_mel.dtype) else: tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize) tts_speech = tts_speech.reshape(tts_speech.shape[0], -1) speech_offset = max(0, min(speech_offset, int(tts_speech.shape[-1]))) emitted_speech = tts_speech[:, speech_offset:] if finalize: return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), None new_state = { "mel": tts_mel.detach().cpu().contiguous(), "speech_offset": int(tts_speech.shape[-1]), } return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), new_state @torch.inference_mode() def forward( self, token: torch.Tensor, prompt_token: torch.Tensor, prompt_feat: torch.Tensor, embedding: torch.Tensor, n_timesteps: int = 10, token_offset_tokens: int = 0, ) -> torch.Tensor: """Generate audio waveform from speech tokens.""" feat = self._forward_mel( token=token, prompt_token=prompt_token, prompt_feat=prompt_feat, embedding=embedding, n_timesteps=n_timesteps, token_offset_tokens=token_offset_tokens, streaming=False, finalize=True, ) # 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, finalize=True) 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}")