# 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 """ import os from collections import Counter from collections.abc import Mapping from typing import Any, Literal, TypedDict, cast, overload 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.config import get_current_diffusion_config_or_none, set_current_diffusion_config from vllm_omni.diffusion.data import OmniDiffusionConfig 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, cosyvoice3_full_response_enabled, cosyvoice3_packed_inference_enabled, cosyvoice3_packed_streaming_enabled, ) from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config logger = init_logger(__name__) def _build_dit_estimator(estimator_config: Mapping[str, Any]) -> DiT: """Build CosyVoice's embedded DiT with diffusion backend configuration. The code2wav model runs inside an LLM-generation worker, so it does not go through the diffusion model loader that normally installs an ``OmniDiffusionConfig`` during model construction. Install a local default config when needed so ``DIFFUSION_ATTENTION_BACKEND`` and the platform default are resolved by the same path as standalone diffusion models. """ if get_current_diffusion_config_or_none() is not None: return DiT(**estimator_config) diffusion_config = OmniDiffusionConfig() with set_current_diffusion_config(diffusion_config): return DiT(**estimator_config) class StreamingHiFTState(TypedDict, total=False): """Request-owned mel history, its absolute offset, and oscillator phase.""" mel: torch.Tensor mel_offset: int phase_acc: torch.Tensor | None 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: StreamingHiFTState | 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 = _build_dit_estimator(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: StreamingHiFTState | None = None, finalize: bool = False, ) -> tuple[torch.Tensor, StreamingHiFTState | 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 # Clone GPU views so each request owns its cache and does not retain # the full batch allocation or alias a later inference output. new_state = { "mel": ( f0_input_mel.detach().clone() if cosyvoice3_packed_streaming_enabled() else f0_input_mel.detach().cpu().contiguous() ), "mel_offset": window_offset - f0_margin, "phase_acc": ( ( new_phase_acc.detach().clone() if cosyvoice3_packed_streaming_enabled() else 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, StreamingHiFTState | 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. """ if items and cosyvoice3_packed_streaming_enabled(): results = [None] * len(items) for finalize in (False, True): group = [(i, x) for i, x in enumerate(items) if bool(x.get("finalize", False)) == finalize] if group: values = self.forward_batch([x for _, x in group], n_timesteps=n_timesteps, stream_items=True) for (i, _), value in zip(group, values): results[i] = value return results if ( items and cosyvoice3_packed_inference_enabled() and all( item.get("finalize", False) and item.get("cache_state") is None and int(item.get("token_offset_tokens", 0)) == 0 for item in items ) ): return [(speech, None) for speech in self.forward_batch(items, n_timesteps=n_timesteps)] results: list[tuple[torch.Tensor, StreamingHiFTState | 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, StreamingHiFTState | 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: StreamingHiFTState | None = None, n_timesteps: int = 10, token_offset_tokens: int = 0, finalize: bool = False, ) -> tuple[torch.Tensor, StreamingHiFTState | 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 _stream_position_noise(self, width: int, like: torch.Tensor) -> torch.Tensor: """Extend fixed-position noise without changing previously emitted prefixes. Fixed-size CPU blocks preserve the initial experiment's seed-0 noise, isolate the global RNG, and make growth independent of request ordering. """ block_frames = 15000 cached = getattr(self, "_position_noise", None) present = 0 if cached is None else cached.shape[-1] blocks = [] if cached is None else [cached.to(like)] for index in range(present // block_frames, (width + block_frames - 1) // block_frames): generator = torch.Generator(device="cpu").manual_seed(index) blocks.append(torch.randn((1, 80, block_frames), generator=generator).to(like)) if not blocks: return like.new_empty((1, 80, 0)) self._position_noise = torch.cat(blocks, dim=-1) if len(blocks) > 1 else blocks[0] return self._position_noise[..., :width] @overload def forward_batch( self, items: list[StreamingFlowItem], *, n_timesteps: int = 10, stream_items: Literal[False] = False ) -> list[torch.Tensor]: ... @overload def forward_batch( self, items: list[StreamingFlowItem], *, n_timesteps: int = 10, stream_items: Literal[True] ) -> list[tuple[torch.Tensor, StreamingHiFTState | None]]: ... @torch.inference_mode() def forward_batch( self, items: list[StreamingFlowItem], *, n_timesteps: int = 10, stream_items: bool = False ) -> list[torch.Tensor] | list[tuple[torch.Tensor, StreamingHiFTState | None]]: """Batch Flow across different reference lengths and restore request order. Build each prompt/continuation contiguously before right padding. Padding the prompt and continuation separately would insert a gap in conditioning. The default keeps exact-length CPU-F0 HiFT. The opt-in Hopper profile uses packed Flow, batched HiFT, and FP64 GPU F0. """ if not items: return [] if stream_items: if not cosyvoice3_packed_streaming_enabled(): raise ValueError("Packed streaming requires COSYVOICE3_PACKED_STREAMING on Hopper") if any(bool(item.get("finalize", False)) != bool(items[0].get("finalize", False)) for item in items): raise ValueError("Use forward_streaming_batch to group mixed finalization states") if len(items) > 8: ordered = sorted( enumerate(items), key=lambda pair: pair[1]["token"].shape[1] + pair[1]["prompt_token"].shape[1] ) outputs: list[torch.Tensor | tuple[torch.Tensor, StreamingHiFTState | None] | None] = [None] * len(items) for start in range(0, len(ordered), 8): group = ordered[start : start + 8] for (index, _), speech in zip( group, self.forward_batch([item for _, item in group], n_timesteps=n_timesteps, stream_items=stream_items), ): outputs[index] = speech return cast(list[torch.Tensor] | list[tuple[torch.Tensor, StreamingHiFTState | None]], outputs) flow = self.flow_model weight = next(flow.parameters()) mus, conds, speakers, lengths, prompt_lengths = [], [], [], [], [] for item in items: tokens = torch.cat( tuple(item[key].to(device=weight.device, dtype=torch.long) for key in ("prompt_token", "token")), dim=1 ) embedded = flow.input_embedding(tokens.clamp(min=0)) if stream_items and not item.get("finalize", False) and int(flow.pre_lookahead_len) > 0: lookahead = int(flow.pre_lookahead_len) hidden = flow.pre_lookahead_layer(embedded[:, :-lookahead], context=embedded[:, -lookahead:]) else: hidden = flow.pre_lookahead_layer(embedded) mu = hidden.repeat_interleave(self.token_mel_ratio, dim=1) prompt = item["prompt_feat"].to(device=weight.device, dtype=mu.dtype) cond = torch.zeros_like(mu) cond[:, : prompt.shape[1]] = prompt speaker = torch.nn.functional.normalize(item["embedding"].to(weight), dim=1) speakers.append(flow.spk_embed_affine_layer(speaker)) mus.append(mu.transpose(1, 2)) conds.append(cond.transpose(1, 2)) lengths.append(mu.shape[1]) prompt_lengths.append(prompt.shape[1]) width = max(lengths) mu = torch.cat([torch.nn.functional.pad(x, (0, width - x.shape[-1])) for x in mus]) cond = torch.cat([torch.nn.functional.pad(x, (0, width - x.shape[-1])) for x in conds]) mask = ( (torch.arange(width, device=weight.device)[None, :] < torch.tensor(lengths, device=weight.device)[:, None]) .unsqueeze(1) .to(mu) ) if cosyvoice3_packed_inference_enabled(): from .code2wav_core.packed_dit import ( PackedDiT, gather_rows, pack_rows, scatter_rows, solve_flow_euler_packed, ) if not hasattr(self, "_packed_estimator"): self._packed_estimator = PackedDiT(flow.decoder.estimator) self._packed_estimator.compile(mu.dtype) rows = pack_rows(lengths, mu.device) if stream_items: noise = self._stream_position_noise(width, mu).expand(len(items), -1, -1) else: noise = torch.randn(mu.shape, device=mu.device, dtype=mu.dtype) time_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) if flow.decoder.t_scheduler == "cosine": time_span = 1 - torch.cos(time_span * 0.5 * torch.pi) with torch.autocast("cuda", dtype=mu.dtype): packed = solve_flow_euler_packed( self._packed_estimator, gather_rows(noise.transpose(1, 2), rows), time_span, gather_rows(mu.transpose(1, 2), rows), torch.cat(speakers), gather_rows(cond.transpose(1, 2), rows), rows, cfg_rate=flow.decoder.inference_cfg_rate, streaming=stream_items and not items[0].get("finalize", False), modulation_key=(n_timesteps, flow.decoder.t_scheduler, mu.dtype, mu.device), ) feat = scatter_rows(packed, rows, width).transpose(1, 2) else: feat, _ = flow.decoder( mu=mu, mask=mask, spks=torch.cat(speakers), cond=cond, n_timesteps=n_timesteps, streaming=False ) if stream_items: results = [] for row, item in enumerate(items): start = prompt_lengths[row] + max(0, int(item.get("token_offset_tokens", 0))) * self.token_mel_ratio mel = feat[row : row + 1, :, start : lengths[row]].float() results.append( self._stream_hift_from_feat( mel, cache_state=item.get("cache_state"), finalize=bool(item.get("finalize", False)) ) ) return results hift_weight = self.hift.m_source.l_linear.weight if cosyvoice3_full_response_enabled(): mels = [] for row, item in enumerate(items): start = prompt_lengths[row] + max(0, int(item.get("token_offset_tokens", 0))) * self.token_mel_ratio mels.append(feat[row : row + 1, :, start : lengths[row]].float().to(hift_weight)) width = max(mel.shape[-1] for mel in mels) if width == 0: return [mel.new_zeros((1, 1, 0)) for mel in mels] from .code2wav_core.packed_dit import bucketed_width # Recurring shapes let cuDNN/cuFFT reuse plans; rows are already # right padded to the batch maximum and trimmed below. width = bucketed_width(width) padded = torch.cat([torch.nn.functional.pad(mel, (0, width - mel.shape[-1])) for mel in mels]) speech, _, _ = self.hift.inference(speech_feat=padded, finalize=True) stride = int(np.prod(self.hift.upsample_rates) * self.hift.istft_params["hop_len"]) return [speech[row : row + 1, ..., : mel.shape[-1] * stride] for row, mel in enumerate(mels)] results = [] for row, item in enumerate(items): start = prompt_lengths[row] + max(0, int(item.get("token_offset_tokens", 0))) * self.token_mel_ratio mel = feat[row : row + 1, :, start : lengths[row]].float().to(hift_weight) if mel.shape[-1] == 0: speech = mel.new_zeros((1, 1, 0)) else: speech, _, _ = self.hift.inference(speech_feat=mel, finalize=True) results.append(speech) return results @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 """ # 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) # Fold after loading and device placement to avoid recomputing weights # for every chunk. Folding changes state_dict keys, so reloading the # original checkpoint into this instance is unsupported. folded = self.hift.remove_weight_norm() logger.info("Folded %d weight-norm layers in HiFT generator", folded) if folded == 0: logger.warning("HiFT generator had no weight-norm layers to fold; check config drift") # Fold on the same device as mainline F0 inference. Folding on CPU # and moving to GPU later changes the normalized FP32 weights. f0_device = device if device.type == "cuda" and os.getenv("COSYVOICE3_F0_ON_CPU", "0") != "1" else "cpu" self.hift.f0_predictor.to(device=f0_device, dtype=torch.float32) f0_folded = self.hift.f0_predictor.remove_weight_norm() logger.info("Folded %d weight-norm layers in HiFT F0 predictor", f0_folded) if f0_folded == 0: logger.warning("HiFT F0 predictor had no weight-norm layers to fold; check config drift") self.hift.eval() if device.type == "cuda" and cosyvoice3_packed_streaming_enabled(): if os.getenv("COSYVOICE3_CACHED_ISTFT", "0") == "1": self.hift.enable_cached_istft() logger.info("CosyVoice3 streaming HiFT: cached ISTFT enabled") if os.getenv("COSYVOICE3_HIFT_GRAPH", "0") == "1": self.hift.enable_decode_graphs() logger.info("CosyVoice3 streaming HiFT: decode CUDA graphs enabled") logger.info(f"Loaded hift weights from {hift_path}")