# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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 collections import Counter from typing import TypedDict, cast 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.model_executor.models.cosyvoice3.runtime import ( cosyvoice3_batch_flow_debug, cosyvoice3_batch_flow_profile, ) from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config logger = init_logger(__name__) class StreamingFlowItem(TypedDict, total=False): """One entry of the batched-streaming item list passed to forward_streaming_batch.""" index: int req_id: str | None stream_finished: bool token: torch.Tensor prompt_token: torch.Tensor prompt_feat: torch.Tensor embedding: torch.Tensor cache_state: dict[str, torch.Tensor] | None token_offset_tokens: int finalize: bool 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) # Must cover decode()'s own causal receptive field, not just the F0 # margin; window_len=48 already passes # test_incremental_hift_bounded_window_is_close, so 64 has headroom. self._hift_window_len = 64 @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, token_lens: torch.Tensor | None = None, prompt_token_lens: torch.Tensor | None = None, prompt_feat_lens: torch.Tensor | None = None, ) -> 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) batch_size = int(token.shape[0]) token_len = ( token_lens.to(device=device, dtype=torch.int32) if token_lens is not None else torch.full((batch_size,), token.shape[1], device=device, dtype=torch.int32) ) prompt_token_len = ( prompt_token_lens.to(device=device, dtype=torch.int32) if prompt_token_lens is not None else torch.full((batch_size,), prompt_token.shape[1], device=device, dtype=torch.int32) ) prompt_feat_len = ( prompt_feat_lens.to(device=device, dtype=torch.int32) if prompt_feat_lens is not None else torch.full((batch_size,), 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 def _stream_hift_from_feat( self, feat: torch.Tensor, *, cache_state: dict[str, torch.Tensor] | None = None, finalize: bool = False, ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: 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") cached_offset = 0 if not cache_state else int(cache_state.get("mel_offset", 0)) window_len = self._hift_window_len samples_per_mel = int(np.prod(self.hift.upsample_rates) * self.hift.istft_params["hop_len"]) trim = int(self.hift.f0_predictor.condnet[0].causal_padding) f0_margin_frames = int(self.hift.f0_predictor.left_context_frames) f0_margin = 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) total_hist = cached_mel.shape[-1] overlap = min(window_len + trim, total_hist) window_mel = torch.cat([cached_mel[..., -overlap:], chunk_mel], dim=-1) window_offset = cached_offset + total_hist - overlap uv_offset = window_offset * samples_per_mel f0_margin = min(f0_margin_frames, total_hist - overlap) f0_input_mel = ( torch.cat( [cached_mel[..., total_hist - overlap - f0_margin : total_hist - overlap], window_mel], dim=-1 ) if f0_margin > 0 else window_mel ) else: window_mel = chunk_mel f0_input_mel = chunk_mel window_offset = 0 uv_offset = 0 phase_acc = None if not cache_state else cache_state.get("phase_acc") if phase_acc is not None: phase_acc = phase_acc.to(device=chunk_mel.device, dtype=chunk_mel.dtype) if window_mel.shape[-1] == 0: tts_speech = torch.zeros((chunk_mel.shape[0], 1, 0), device=chunk_mel.device, dtype=chunk_mel.dtype) new_phase_acc = phase_acc else: next_overlap = min(window_len + trim, window_mel.shape[-1]) # On finalize the F0 predictor consumes the whole window (it does not # hold back `trim` frames), so trim must not inflate the phase-carry # index -- it would point past the end of a short window. carry_trim = 0 if finalize else trim tts_speech, _, new_phase_acc = self.hift.inference( speech_feat=f0_input_mel, finalize=finalize, phase_acc=phase_acc, uv_offset=uv_offset, next_overlap=next_overlap, trim=carry_trim, f0_margin=f0_margin, ) tts_speech = tts_speech.reshape(tts_speech.shape[0], -1) new_samples = chunk_mel.shape[-1] * samples_per_mel if finalize: frames_withheld_per_chunk = trim + self.hift.conv_pre_look_right + 1 released = frames_withheld_per_chunk * samples_per_mel emitted_speech = tts_speech[:, -(new_samples + released) :] else: emitted_speech = tts_speech[:, -new_samples:] if new_samples > 0 else tts_speech[:, :0] if finalize: return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), None new_state = { "mel": f0_input_mel.detach().cpu().contiguous(), "mel_offset": window_offset - f0_margin, "phase_acc": new_phase_acc.detach().cpu().contiguous() if new_phase_acc is not None else None, } return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), new_state @torch.inference_mode() def forward_streaming_batch( self, items: list[StreamingFlowItem], *, n_timesteps: int = 10, ) -> list[tuple[torch.Tensor, dict[str, torch.Tensor] | None]]: """Batch the flow-matching mel path, then run HiFT per request. Items are grouped by prompt condition shape and finalization state. Codec tokens may have different lengths; those are padded within the group and passed to the flow as per-row token lengths. """ results: list[tuple[torch.Tensor, dict[str, torch.Tensor] | None] | None] = [None] * len(items) groups: dict[tuple[int, int, int, bool], list[tuple[int, StreamingFlowItem]]] = {} for index, item in enumerate(items): assert isinstance(item["token"], torch.Tensor) assert isinstance(item["prompt_token"], torch.Tensor) assert isinstance(item["prompt_feat"], torch.Tensor) assert isinstance(item["embedding"], torch.Tensor) key = ( int(item["prompt_token"].shape[1]), int(item["prompt_feat"].shape[1]), int(item["embedding"].shape[1]), bool(item.get("finalize", False)), ) groups.setdefault(key, []).append((index, item)) if cosyvoice3_batch_flow_debug(): group_summary = {key: len(group) for key, group in groups.items()} group_size_distribution = Counter(group_summary.values()) batchable_items = sum(size for size in group_summary.values() if size > 1) logger.info( "CosyVoice3 code2wav debug: forward_streaming_batch items=%d " "groups=%s group_size_distribution=%s batchable_items=%d", len(items), group_summary, dict(sorted(group_size_distribution.items())), batchable_items, ) for _key, group in groups.items(): if len(group) == 1: index, item = group[0] result = self.forward_streaming( token=item["token"], prompt_token=item["prompt_token"], prompt_feat=item["prompt_feat"], embedding=item["embedding"], cache_state=item.get("cache_state"), n_timesteps=n_timesteps, token_offset_tokens=int(item.get("token_offset_tokens", 0)), finalize=bool(item.get("finalize", False)), ) results[index] = result continue token_tensors = [item["token"] for _, item in group] token_lens = torch.tensor( [int(token.shape[1]) for token in token_tensors], dtype=torch.int32, ) max_token_len = int(token_lens.max().item()) padded_tokens = [] for token in token_tensors: if int(token.shape[1]) == max_token_len: padded_tokens.append(token) else: pad = torch.zeros( (token.shape[0], max_token_len - int(token.shape[1])), device=token.device, dtype=token.dtype, ) padded_tokens.append(torch.cat([token, pad], dim=1)) tokens = torch.cat(padded_tokens, dim=0) prompt_tokens = torch.cat([item["prompt_token"] for _, item in group], dim=0) prompt_feats = torch.cat([item["prompt_feat"] for _, item in group], dim=0) embeddings = torch.cat([item["embedding"] for _, item in group], dim=0) prompt_token_lens = torch.full((len(group),), prompt_tokens.shape[1], dtype=torch.int32) prompt_feat_lens = torch.full((len(group),), prompt_feats.shape[1], dtype=torch.int32) finalize = bool(group[0][1].get("finalize", False)) with cosyvoice3_batch_flow_profile(f"cosyvoice3_flow_batch_b{len(group)}_t{tokens.shape[1]}"): feat = self._forward_mel( token=tokens, prompt_token=prompt_tokens, prompt_feat=prompt_feats, embedding=embeddings, n_timesteps=n_timesteps, token_offset_tokens=0, streaming=True, finalize=finalize, token_lens=token_lens, prompt_token_lens=prompt_token_lens, prompt_feat_lens=prompt_feat_lens, ) for row, (index, item) in enumerate(group): trim_mel = max(0, int(item.get("token_offset_tokens", 0))) * int(self.token_mel_ratio) valid_tokens = int(token_lens[row].item()) if not finalize: valid_tokens = max(0, valid_tokens - int(self.flow_model.pre_lookahead_len)) valid_mel = valid_tokens * int(self.token_mel_ratio) row_feat = feat[row : row + 1, :, :valid_mel] if trim_mel > 0: row_feat = row_feat[:, :, trim_mel:] results[index] = self._stream_hift_from_feat( row_feat, cache_state=item.get("cache_state"), finalize=finalize, ) assert all(result is not None for result in results), "every streaming item must produce exactly one result" return cast(list[tuple[torch.Tensor, dict[str, torch.Tensor] | None]], results) @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, ) return self._stream_hift_from_feat(feat, cache_state=cache_state, finalize=finalize) @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}")