264 lines
10 KiB
Python
264 lines
10 KiB
Python
# Copyright (c) 2025 Zhipu AI Inc (authors: CogAudio Group Members)
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""
|
|
Flow Matching Wrapper for Streaming Inference.
|
|
Manages the DiT model initialization and the ODE sampling process.
|
|
"""
|
|
|
|
import torch
|
|
from torch.nn import functional as F
|
|
from cosyvoice.utils.mask import make_pad_mask
|
|
from flow.dit import DiT
|
|
|
|
|
|
class Flow(torch.nn.Module):
|
|
def __init__(self,
|
|
# Speech token configuration
|
|
speech_token_dim: int = 512,
|
|
vocab_size: int = 100000,
|
|
input_frame_rate: float = 12.5,
|
|
# Speaker embedding configuration
|
|
spk_embed_dim: int = 80,
|
|
mel_dim: int = 80,
|
|
big_model: bool = False,
|
|
spkr_emb_adaLN: bool = False,
|
|
# Flags and options
|
|
calc_prompt_mel_loss: bool = True,
|
|
speech_token_cfg: bool = True,
|
|
use_wavlm_emb: bool = False,
|
|
remove_spkr_concat_condition: bool = False,
|
|
t_scheduler: str = 'cosine',
|
|
loss_type: str = "l2",
|
|
use_mdt: bool = False,
|
|
mel_framerate: int = 86,
|
|
):
|
|
super().__init__()
|
|
self.mel_dim = mel_dim
|
|
self.input_frame_rate = input_frame_rate
|
|
self.mel_framerate = mel_framerate
|
|
|
|
# Configuration Flags
|
|
self.calc_prompt_mel_loss = calc_prompt_mel_loss
|
|
self.use_wavlm_emb = use_wavlm_emb
|
|
self.remove_spkr_concat_condition = remove_spkr_concat_condition
|
|
self.spkr_emb_adaLN = spkr_emb_adaLN
|
|
self.speech_token_cfg = speech_token_cfg
|
|
self.loss_type = loss_type
|
|
self.use_mdt = use_mdt
|
|
|
|
# Scheduler and CFG settings
|
|
self.t_scheduler = t_scheduler
|
|
self.training_cfg_rate = 0.2
|
|
self.inference_cfg_rate = 0.7
|
|
self.sigma_min = 1e-06
|
|
|
|
# Initialize Speaker Embedding Layer (Project to match DiT condition dim)
|
|
if not remove_spkr_concat_condition:
|
|
# 192 is the expected input dimension of the raw speaker embedding
|
|
self.spk_embed_affine_layer = torch.nn.Linear(192, spk_embed_dim)
|
|
|
|
# DiT Model Configuration
|
|
dit_config = {
|
|
"trans_dim": 1024 if big_model else 768,
|
|
"depth": 22 if big_model else 18,
|
|
"heads": 16 if big_model else 12,
|
|
"ff_mult": 2,
|
|
"text_emb_dim": speech_token_dim,
|
|
"conv_layers": 4,
|
|
"mel_dim": mel_dim,
|
|
"text_vocab_size": vocab_size,
|
|
"condition_dim": mel_dim if remove_spkr_concat_condition else mel_dim + spk_embed_dim,
|
|
"spkr_emb_adaLN": spkr_emb_adaLN,
|
|
"wav_lm_emb": use_wavlm_emb,
|
|
}
|
|
|
|
# Use 'model' instead of 'estimator' for clarity, though estimator is fine
|
|
self.estimator = DiT(**dit_config)
|
|
|
|
@torch.inference_mode()
|
|
def inference_with_cache(self,
|
|
token,
|
|
prompt_token,
|
|
prompt_feat,
|
|
embedding,
|
|
n_timesteps=10,
|
|
last_step_cache=None,
|
|
wavlm_emb_bt=None,
|
|
is_causal=False,
|
|
block_pattern=None
|
|
):
|
|
"""
|
|
Streaming inference method that supports KV-caching via last_step_cache.
|
|
"""
|
|
assert token.shape[0] == 1, "Batch size must be 1 for streaming inference."
|
|
device = token.device
|
|
token_len = token.shape[1]
|
|
|
|
# Handle Prompt Concatenation
|
|
if prompt_token is not None and prompt_feat is not None:
|
|
token = torch.cat([prompt_token, token], dim=1)
|
|
prompt_token_len = prompt_token.shape[1]
|
|
token_len += prompt_token_len
|
|
prompt_feat_len = prompt_feat.shape[1]
|
|
else:
|
|
prompt_feat_len = 0
|
|
|
|
# Calculate feature length based on frame rates
|
|
feat_len = int(token_len / self.input_frame_rate * self.mel_framerate)
|
|
|
|
# Prepare masks and conditions
|
|
padding_mask_bt = (~make_pad_mask(torch.LongTensor([feat_len]))).to(device)
|
|
mel_cond_btd = torch.zeros([1, feat_len, self.mel_dim]).to(device)
|
|
|
|
# Align prompt features
|
|
if prompt_feat is not None:
|
|
# Ensure we don't exceed dimensions when copying prompt features
|
|
copy_len = min(prompt_feat.shape[1], mel_cond_btd.shape[1])
|
|
mel_cond_btd[:, :copy_len, :] = prompt_feat[:, -copy_len:, :]
|
|
|
|
# Process Speaker Embedding
|
|
spkr_embedding_normed = F.normalize(embedding, dim=1)
|
|
|
|
if not self.remove_spkr_concat_condition:
|
|
spkr_embedding = self.spk_embed_affine_layer(spkr_embedding_normed)
|
|
# Expand speaker embedding to match time dimension
|
|
spkr_embedding_expanded = spkr_embedding.unsqueeze(1).expand(-1, mel_cond_btd.shape[1], -1)
|
|
# Concatenate mel-condition and speaker-condition
|
|
condition_btd = torch.cat([mel_cond_btd, spkr_embedding_expanded], dim=-1)
|
|
else:
|
|
condition_btd = mel_cond_btd
|
|
|
|
# Handle WavLM Embedding
|
|
if self.use_wavlm_emb and wavlm_emb_bt is not None:
|
|
emb2 = F.normalize(wavlm_emb_bt.to(device), dim=1)
|
|
spkr_embedding_normed = torch.cat([spkr_embedding_normed, emb2], dim=1)
|
|
|
|
# Run Sampling
|
|
result_btd, current_step2cache = self.do_sample(
|
|
token,
|
|
mel_cond_btd,
|
|
condition_btd,
|
|
padding_mask_bt,
|
|
spkr_embedding_normed,
|
|
is_causal,
|
|
block_pattern,
|
|
n_timesteps,
|
|
last_step_cache
|
|
)
|
|
|
|
# Remove the prompt part from the result
|
|
if prompt_feat is not None:
|
|
result_btd = result_btd[:, prompt_feat_len:, :]
|
|
|
|
# Permute to (Batch, Dimension, Time) for output convention
|
|
result_bdt = result_btd.permute(0, 2, 1)
|
|
|
|
return result_bdt, current_step2cache
|
|
|
|
def do_sample(self,
|
|
speech_token_bt,
|
|
mel_cond_btd,
|
|
condition_btd,
|
|
padding_mask_bt,
|
|
spkr_embedding_normed,
|
|
is_causal,
|
|
block_pattern,
|
|
n_timesteps,
|
|
last_step_cache):
|
|
"""
|
|
Executes the sampling process using a manual Euler method loop to support
|
|
step-by-step caching for streaming.
|
|
"""
|
|
current_step2cache = {}
|
|
|
|
# Initial noise
|
|
x = torch.randn_like(mel_cond_btd)
|
|
device = speech_token_bt.device
|
|
|
|
# Time scheduler setup
|
|
t_span = torch.linspace(0, 1, n_timesteps + 1, device=device)
|
|
if self.t_scheduler == 'cosine':
|
|
t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
|
|
# Note: If other schedulers are needed, add elif blocks here.
|
|
|
|
t_current = t_span[0]
|
|
dt = t_span[1] - t_span[0]
|
|
|
|
sol = []
|
|
|
|
# Iterative Denoising (Euler method)
|
|
for step in range(1, len(t_span)):
|
|
# Apply cache if available (Overwriting the beginning of the sequence)
|
|
if last_step_cache is not None:
|
|
x_cache = last_step_cache[step]['x']
|
|
override_len = last_step_cache.get('override_len', x_cache.shape[-1])
|
|
|
|
# Safety check for dimensions
|
|
safe_len = min(x.shape[1], override_len)
|
|
x[:, :safe_len, :] = x_cache[:, :safe_len, :]
|
|
|
|
# Cache current input for the next streaming chunk
|
|
current_step2cache[step] = {
|
|
"x": x.clone().detach(),
|
|
}
|
|
|
|
# 1. Forward pass (Conditional)
|
|
dphi_dt = self.estimator(
|
|
middle_point_btd=x,
|
|
condition_btd=condition_btd,
|
|
text=speech_token_bt,
|
|
time_step_1d=t_current[None],
|
|
padding_mask_bt=padding_mask_bt,
|
|
spkr_emb_bd=spkr_embedding_normed,
|
|
is_causal=is_causal,
|
|
block_pattern=block_pattern
|
|
)
|
|
|
|
# 2. Classifier-Free Guidance (CFG)
|
|
if self.inference_cfg_rate > 0:
|
|
# Prepare unconditional input
|
|
if self.speech_token_cfg:
|
|
text_for_no_condition = torch.zeros_like(speech_token_bt)
|
|
else:
|
|
text_for_no_condition = speech_token_bt
|
|
|
|
# Forward pass (Unconditional)
|
|
cfg_dphi_dt = self.estimator(
|
|
middle_point_btd=x,
|
|
condition_btd=torch.zeros_like(condition_btd),
|
|
text=text_for_no_condition,
|
|
time_step_1d=t_current[None],
|
|
padding_mask_bt=padding_mask_bt,
|
|
spkr_emb_bd=torch.zeros_like(spkr_embedding_normed),
|
|
is_causal=is_causal,
|
|
block_pattern=block_pattern
|
|
)
|
|
|
|
# Combine conditional and unconditional outputs
|
|
dphi_dt = ((1.0 + self.inference_cfg_rate) * dphi_dt -
|
|
self.inference_cfg_rate * cfg_dphi_dt)
|
|
|
|
# Euler Step
|
|
x = x + dt * dphi_dt
|
|
|
|
# Update time
|
|
t_current = t_current + dt
|
|
sol.append(x)
|
|
|
|
# Update dt for the next step
|
|
if step < len(t_span) - 1:
|
|
dt = t_span[step + 1] - t_current
|
|
|
|
result_btd = sol[-1]
|
|
return result_btd, current_step2cache |