""" ein notation: b - batch n - sequence nt - text sequence nw - raw wave length d - dimension """ from __future__ import annotations from random import random import random as _random from typing import Callable, Dict, OrderedDict import math from pathlib import Path import torch import torch.nn.functional as F import torchaudio from torch import nn from torch.nn.utils.rnn import pad_sequence from torchdiffeq import odeint from lemas_tts.model.modules import MelSpec from lemas_tts.model.modules import MIEsitmator, AccentClassifier, grad_reverse from lemas_tts.model.backbones.ecapa_tdnn import ECAPA_TDNN from lemas_tts.model.backbones.prosody_encoder import ProsodyEncoder, extract_fbank_16k from lemas_tts.model.utils import ( default, exists, lens_to_mask, list_str_to_idx, list_str_to_tensor, mask_from_frac_lengths, ) def clip_and_shuffle(mel, mel_len, sample_rate=24000, hop_length=256, ratio=None): """ Randomly clip a mel-spectrogram segment and shuffle 1-second chunks to create an accent-invariant conditioning segment. This is a inference-time utility used by the accent GRL path. Args: mel: [n_mels, T] mel_len: int, original mel length (T) """ frames_per_second = int(sample_rate / hop_length) # ≈ 94 frames / second # ---- 1. Randomly crop 25%~75% of the original length (or ratio * length) ---- total_len = mel_len if not ratio: seg_len = _random.randint(int(0.25 * total_len), int(0.75 * total_len)) else: seg_len = int(total_len * ratio) start = _random.randint(0, max(0, total_len - seg_len)) mel_seg = mel[:, start : start + seg_len] # ---- 2. Split into ~1-second chunks ---- n_chunks = (mel_seg.size(1) + frames_per_second - 1) // frames_per_second chunks = [] for i in range(n_chunks): chunk = mel_seg[:, i * frames_per_second : (i + 1) * frames_per_second] chunks.append(chunk) # ---- 3. Shuffle chunk order ---- _random.shuffle(chunks) shuffled_mel = torch.cat(chunks, dim=1) # ---- 4. Repeat random chunks until reaching original length ---- if shuffled_mel.size(1) < total_len: repeat_chunks = [] while sum(c.size(1) for c in repeat_chunks) < total_len: repeat_chunks.append(_random.choice(chunks)) shuffled_mel = torch.cat([shuffled_mel] + repeat_chunks, dim=1) # ---- 5. Trim to exactly mel_len ---- shuffled_mel = shuffled_mel[:, :total_len] assert shuffled_mel.shape == mel.shape, f"shuffled_mel.shape != mel.shape: {shuffled_mel.shape} != {mel.shape}" return shuffled_mel class CFM(nn.Module): def __init__( self, transformer: nn.Module, sigma=0.0, odeint_kwargs: dict = dict( # atol = 1e-5, # rtol = 1e-5, method="euler" # 'midpoint' ), audio_drop_prob=0.3, text_drop_prob=0.1, num_channels=None, mel_spec_module: nn.Module | None = None, mel_spec_kwargs: dict = dict(), frac_lengths_mask: tuple[float, float] = (0.7, 1.0), vocab_char_map: dict[str:int] | None = None, use_ctc_loss: bool = False, use_spk_enc: bool = False, use_prosody_encoder: bool = False, prosody_cfg_path: str | None = None, prosody_ckpt_path: str | None = None, ): super().__init__() self.frac_lengths_mask = frac_lengths_mask # mel spec self.mel_spec = default(mel_spec_module, MelSpec(**mel_spec_kwargs)) num_channels = default(num_channels, self.mel_spec.n_mel_channels) self.num_channels = num_channels # classifier-free guidance self.audio_drop_prob = audio_drop_prob self.text_drop_prob = text_drop_prob # transformer self.transformer = transformer dim = transformer.dim self.dim = dim # conditional flow related self.sigma = sigma # sampling related self.odeint_kwargs = odeint_kwargs # vocab map for tokenization self.vocab_char_map = vocab_char_map # Prosody encoder (Pretssel ECAPA-TDNN) self.use_prosody_encoder = ( use_prosody_encoder and prosody_cfg_path is not None and prosody_ckpt_path is not None ) if self.use_prosody_encoder: cfg_path = Path(prosody_cfg_path) ckpt_path = Path(prosody_ckpt_path) self.prosody_encoder = ProsodyEncoder(cfg_path, ckpt_path, freeze=True) # 512-d prosody -> mel channel dimension self.prosody_to_mel = nn.Linear(512, self.num_channels) self.prosody_dropout = nn.Dropout(p=0.2) else: self.prosody_encoder = None # Speaker encoder self.use_spk_enc = use_spk_enc if use_spk_enc: self.speaker_encoder = ECAPA_TDNN( self.num_channels, self.dim, channels=[512, 512, 512, 512, 1536], kernel_sizes=[5, 3, 3, 3, 1], dilations=[1, 2, 3, 4, 1], attention_channels=128, res2net_scale=4, se_channels=128, global_context=True, batch_norm=True, ) # self.load_partial_weights(self.speaker_encoder, "/cto_labs/vistring/zhaozhiyuan/outputs/F5-TTS/pretrain/speaker.bin", device="cpu") self.use_ctc_loss = use_ctc_loss if use_ctc_loss: # print("vocab_char_map:", len(vocab_char_map)+1, "dim:", dim, "mel_spec_kwargs:",mel_spec_kwargs) self.ctc = MIEsitmator(len(self.vocab_char_map), self.num_channels, self.dim, dropout=self.text_drop_prob) self.accent_classifier = AccentClassifier(input_dim=self.num_channels, hidden_dim=self.dim, num_accents=12) self.accent_criterion = nn.CrossEntropyLoss() def load_partial_weights(self, model: nn.Module, ckpt_path: str, device="cpu", verbose=True) -> int: """ 仅加载形状匹配的参数，其余跳过。 返回成功加载的参数数量。 """ state_dict = torch.load(ckpt_path, map_location=device) model_dict = model.state_dict() ok_count = 0 new_dict: OrderedDict[str, torch.Tensor] = OrderedDict() for k, v in state_dict.items(): if k in model_dict and v.shape == model_dict[k].shape: new_dict[k] = v ok_count += 1 else: if verbose: print(f"[SKIP] {k} ckpt:{v.shape} model:{model_dict[k].shape if k in model_dict else 'N/A'}") model_dict.update(new_dict) model.load_state_dict(model_dict) if verbose: print(f"=> 成功加载 {ok_count}/{len(state_dict)} 个参数") return ok_count @property def device(self): return next(self.parameters()).device @torch.no_grad() def sample( self, cond: float["b n d"] | float["b nw"], # noqa: F722 text: int["b nt"] | list[str], # noqa: F722 duration: int | int["b"], # noqa: F821 *, lens: int["b"] | None = None, # noqa: F821 steps=32, cfg_strength=1.0, sway_sampling_coef=None, seed: int | None = None, max_duration=4096, vocoder: Callable[[float["b d n"]], float["b nw"]] | None = None, # noqa: F722 no_ref_audio=False, duplicate_test=False, t_inter=0.1, edit_mask=None, use_acc_grl = True, use_prosody_encoder = True, ref_ratio = 1, ): self.eval() # raw wave -> mel, keep a copy for prosody encoder if available raw_audio = None if cond.ndim == 2: raw_audio = cond.clone() # (B, nw) cond = self.mel_spec(cond) cond = cond.permute(0, 2, 1) assert cond.shape[-1] == self.num_channels cond = cond.to(next(self.parameters()).dtype) cond_mean = cond.mean(dim=1, keepdim=True) batch, cond_seq_len, device = *cond.shape[:2], cond.device if not exists(lens): lens = torch.full((batch,), cond_seq_len, device=device, dtype=torch.long) # optional global prosody conditioning at inference (one embedding per sample) prosody_mel_cond = None prosody_text_cond = None prosody_embeds = None if self.prosody_encoder is not None and raw_audio is not None and use_prosody_encoder: embeds = [] for b in range(batch): audio_b = raw_audio[b].unsqueeze(0) # (1, nw) src_sr = self.mel_spec.target_sample_rate if src_sr != 16_000: audio_16k = torchaudio.functional.resample( audio_b, src_sr, 16_000 ).squeeze(0) else: audio_16k = audio_b.squeeze(0) fbank = extract_fbank_16k(audio_16k) fbank = fbank.unsqueeze(0).to(device=device, dtype=cond.dtype) emb = self.prosody_encoder(fbank, padding_mask=None)[0] # (512,) embeds.append(emb) prosody_embeds = torch.stack(embeds, dim=0) # (B, 512) # broadcast along mel and text prosody_mel_cond = prosody_embeds[:, None, :].expand(-1, cond_seq_len, -1) if use_acc_grl: # rand_mel = clip_and_shuffle(cond.permute(0, 2, 1).squeeze(0), cond.shape[1]) # rand_mel = rand_mel.unsqueeze(0).permute(0, 2, 1) # assert rand_mel.shape == cond.shape, f"Shape diff: rand_mel.shape: {rand_mel.shape}, cond.shape: {cond.shape}" # cond_grl = grad_reverse(rand_mel, lambda_=1.0) if ref_ratio < 1: rand_mel = clip_and_shuffle(cond.permute(0, 2, 1).squeeze(0), cond.shape[1], ratio=ref_ratio) rand_mel = rand_mel.unsqueeze(0).permute(0, 2, 1) assert rand_mel.shape == cond.shape, f"Shape diff: rand_mel.shape: {rand_mel.shape}, cond.shape: {cond.shape}" cond_grl = grad_reverse(rand_mel, lambda_=1.0) else: cond_grl = grad_reverse(cond, lambda_=1.0) # print("cond:", cond.shape, cond.mean(), cond.max(), cond.min(), "rand_mel:", rand_mel.mean(), rand_mel.max(), rand_mel.min(), "cond_grl:", cond_grl.mean(), cond_grl.max(), cond_grl.min()) # text if isinstance(text, list): if exists(self.vocab_char_map): text = list_str_to_idx(text, self.vocab_char_map).to(device) else: text = list_str_to_tensor(text).to(device) assert text.shape[0] == batch # duration cond_mask = lens_to_mask(lens) if edit_mask is not None: cond_mask = cond_mask & edit_mask if isinstance(duration, int): duration = torch.full((batch,), duration, device=device, dtype=torch.long) duration = torch.maximum( torch.maximum((text != -1).sum(dim=-1), lens) + 1, duration ) # duration at least text/audio prompt length plus one token, so something is generated # clamp and convert max_duration to python int for padding ops duration = duration.clamp(max=max_duration) max_duration = int(duration.amax().item()) # duplicate test corner for inner time step oberservation if duplicate_test: test_cond = F.pad(cond, (0, 0, cond_seq_len, max_duration - 2 * cond_seq_len), value=0.0) cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value=0.0) if prosody_mel_cond is not None: prosody_mel_cond = F.pad( prosody_mel_cond, (0, 0, 0, max_duration - cond_seq_len), value=0.0 ) prosody_mel_proj = self.prosody_to_mel(prosody_mel_cond) cond = cond + prosody_mel_proj if no_ref_audio: random_cond = torch.randn_like(cond) * 0.1 + cond_mean random_cond = random_cond / random_cond.mean(dim=1, keepdim=True) * cond_mean print("cond:", cond.mean(), cond.max(), cond.min(), "random_cond:", random_cond.mean(), random_cond.max(), random_cond.min(), "mean_cond:", cond_mean.shape) cond = random_cond cond_mask = F.pad(cond_mask, (0, max_duration - cond_mask.shape[-1]), value=False) cond_mask = cond_mask.unsqueeze(-1) if use_acc_grl: cond_grl = F.pad(cond_grl, (0, 0, 0, max_duration - cond_seq_len), value=0.0) step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond)) # allow direct control (cut cond audio) with lens passed in if batch > 1: mask = lens_to_mask(duration) else: # save memory and speed up, as single inference need no mask currently mask = None # neural ode def compute_sway_max(steps: int, t_start: float = 0.0, dtype=torch.float32, min_ratio: float | None = None, safety_factor: float = 0.5) -> float: """ Compute a safe upper bound for sway_sampling_coef given steps and t_start. - steps: number of ODE steps - t_start: start time in [0,1) - dtype: torch dtype (for machine eps) - min_ratio: smallest distinguishable dt^p (if None, use conservative default) - safety_factor: scale down the theoretical maximum to be safe """ assert 0.0 <= t_start < 1.0 dt = (1.0 - t_start) / max(1, steps) eps = torch.finfo(dtype).eps if min_ratio is None: # conservative default: ~100 * eps (float32 -> ~1e-5) min_ratio = max(1e-9, 1e2 * float(eps)) if dt >= 0.9: p_max = 1.0 + 10.0 else: # solve dt^p >= min_ratio => p <= log(min_ratio)/log(dt) p_max = math.log(min_ratio) / math.log(dt) sway_max = max(0.0, p_max - 1.0) sway_max = sway_max * float(safety_factor) return torch.tensor(sway_max, device=device, dtype=dtype) # prepare text-side prosody conditioning if embeddings available if prosody_embeds is not None: text_len = text.shape[1] prosody_text_cond = prosody_embeds[:, None, :].expand(-1, text_len, -1) else: prosody_text_cond = None def fn(t, x): # at each step, conditioning is fixed # if use_spk_enc: # mix_cond = t * cond + (1-t) * spk_emb # step_cond = torch.where(cond_mask, mix_cond, torch.zeros_like(mix_cond)) if use_acc_grl: step_cond = torch.where(cond_mask, cond_grl, torch.zeros_like(cond_grl)) else: step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond)) # predict flow pred = self.transformer( x=x, cond=step_cond, text=text, time=t, mask=mask, drop_audio_cond=False, drop_text=False, cache=True, prosody_text=prosody_text_cond, ) if cfg_strength < 1e-5: return pred null_pred = self.transformer( x=x, cond=step_cond, text=text, time=t, mask=mask, drop_audio_cond=True, drop_text=True, cache=True, prosody_text=prosody_text_cond, ) # cfg_t = cfg_strength * torch.cos(0.5 * torch.pi * t) # cfg_t = cfg_strength * (1 - t) cfg_t = cfg_strength * ((1 - t) ** 2) # print("t:", t, "cfg_t:", cfg_t) res = pred + (pred - null_pred) * cfg_t # print("t:", t.item(), "\tres:", res.shape, res.mean().item(), res.max().item(), res.min().item(), "\tpred:", pred.mean().item(), pred.max().item(), pred.min().item(), "\tnull_pred:", null_pred.mean().item(), null_pred.max().item(), null_pred.min().item(), "\tcfg_t:", cfg_t.item()) res = res.clamp(-20, 20) return res # noise input # to make sure batch inference result is same with different batch size, and for sure single inference # still some difference maybe due to convolutional layers y0 = [] for dur in duration: if exists(seed): torch.manual_seed(seed) y0.append(torch.randn(dur, self.num_channels, device=self.device, dtype=step_cond.dtype)) y0 = pad_sequence(y0, padding_value=0, batch_first=True) t_start = 0 # duplicate test corner for inner time step oberservation if duplicate_test: t_start = t_inter y0 = (1 - t_start) * y0 + t_start * test_cond steps = int(steps * (1 - t_start)) t = torch.linspace(t_start, 1, int(steps + 1), device=self.device, dtype=step_cond.dtype) sway_max = compute_sway_max(steps, t_start=t_start, dtype=step_cond.dtype, min_ratio=1e-9, safety_factor=0.7) if sway_sampling_coef is not None: sway_sampling_coef = min(sway_max, sway_sampling_coef) # t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t) t = t ** (1 + sway_sampling_coef) else: t = t ** (1 + sway_max) # print("t:",t, "sway_max:", sway_max, "sway_sampling_coef:", sway_sampling_coef) trajectory = odeint(fn, y0, t, **self.odeint_kwargs) self.transformer.clear_cache() sampled = trajectory[-1] out = sampled out = torch.where(cond_mask, cond, out) # out生成的部分，或者说pad补0的部分，单独计算mean, 然后和cond的mean做对齐（乘以系数，两个的均值要差不多） if no_ref_audio: out_mean = out[:,cond_seq_len:,:].mean(dim=1, keepdim=True) out[:,cond_seq_len:,:] = out[:,cond_seq_len:,:] - (out_mean - cond_mean) # print("out_mean:", out_mean.shape, out_mean.mean(), "cond_mean:", cond_mean.shape, cond_mean.mean(), "out:", out[:,cond_seq_len:,:].shape, out[:,cond_seq_len:,:].mean().item(), out[:,cond_seq_len:,:].max().item(), out[:,cond_seq_len:,:].min().item()) if exists(vocoder): out = out.permute(0, 2, 1) out = vocoder(out) # print("out:", out.shape, "trajectory:", trajectory.shape) return out, trajectory def info_nce_speaker(self, e_gt: torch.Tensor, e_pred: torch.Tensor, temperature: float = 0.1): """ InfoNCE loss for speaker encoder training. 同一条样本的 e_gt 与 e_pred 互为正例，其余均为负例。 Args: temperature: 温度缩放 τ Returns: loss: 标量 tensor，可 backward """ B = e_gt.size(0) # 2. L2 归一化 e_gt = F.normalize(e_gt, dim=1) e_pred = F.normalize(e_pred, dim=1) # 3. 计算 B×B 相似度矩阵（pred 对 gt） logits = torch.einsum('bd,cd->bc', e_pred, e_gt) / temperature # [B, B] # 4. 正例标签正好是对角线 labels = torch.arange(B, device=logits.device) # 5. InfoNCE = cross-entropy over in-batch negatives loss = F.cross_entropy(logits, labels) return loss def forward_old( self, batchs: Dict[str, torch.Tensor], # inp: float["b n d"] | float["b nw"], # mel or raw wave # noqa: F722 # text: int["b nt"] | list[str], # noqa: F722 *, # lens: int["b"] | None = None, # noqa: F821 noise_scheduler: str | None = None, ): inp = batchs["mel"].permute(0, 2, 1) lens = batchs["mel_lengths"] rand_mel = batchs["rand_mel"].permute(0, 2, 1) text = batchs["text"] target_text_lengths = torch.tensor([len(x) for x in text], device=inp.device) langs = batchs["langs"] # print("inp:", inp.shape, "rand_mel:", rand_mel.shape, "lens:", lens, "target_text_lengths:", target_text_lengths, "langs:", langs) # handle raw wave if inp.ndim == 2: inp = self.mel_spec(inp) inp = inp.permute(0, 2, 1) assert inp.shape[-1] == self.num_channels batch, seq_len, dtype, device, _σ1 = *inp.shape[:2], inp.dtype, self.device, self.sigma # print("inp_shape:", inp.shape, inp.max(), inp.min(), "dtype:", dtype, "device:", device, "σ1:", _σ1) # handle text as string if isinstance(text, list): if exists(self.vocab_char_map): text = list_str_to_idx(text, self.vocab_char_map).to(device) else: text = list_str_to_tensor(text).to(device) assert text.shape[0] == batch # lens and mask if not exists(lens): lens = torch.full((batch,), seq_len, device=device) mask = lens_to_mask(lens, length=seq_len) # useless here, as collate_fn will pad to max length in batch # get a random span to mask out for training conditionally frac_lengths = torch.zeros((batch,), device=self.device).float().uniform_(*self.frac_lengths_mask) rand_span_mask = mask_from_frac_lengths(lens, frac_lengths) if exists(mask): rand_span_mask &= mask # mel is x1 x1 = inp # x0 is gaussian noise x0 = torch.randn_like(x1) # time step time = torch.rand((batch,), dtype=dtype, device=self.device) # TODO. noise_scheduler # sample xt (φ_t(x) in the paper) t = time.unsqueeze(-1).unsqueeze(-1) φ = (1 - t) * x0 + t * x1 flow = x1 - x0 # cond = torch.where(rand_span_mask[..., None], torch.zeros_like(rand_mel), rand_mel) cond = torch.where(rand_span_mask[..., None], torch.zeros_like(x1), x1) # print("seq_len:", seq_len, "lens:", lens) if self.use_spk_enc: # 50%的概率使用spk_emb spk_emb = self.speaker_encoder(rand_mel, lens) # global_emb: [batch, 1, dim] -> 重复扩展到 [batch, seq_len, dim] spk_emb = spk_emb.unsqueeze(1).expand_as(x1) # print("spk_emb_shape:", spk_emb.shape) # 应用mask操作 cond = torch.where(rand_span_mask[..., None], torch.zeros_like(spk_emb), spk_emb) # assert cond.shape[0] == batch, "speaker encoder batch size mismatch" # print("x1.shape:", x1.shape, "cond_shape:", cond.shape) # 给一个随机数，把spk_emb * 随机数，再加上原来的cond *（1-随机数） rand_num = torch.rand((batch, 1, 1), dtype=dtype, device=self.device) cond = cond * rand_num + spk_emb * (1 - rand_num) cond_grl = grad_reverse(cond, lambda_=1.0) # print("inp_shape:", inp.shape, "rand_span_mask:", rand_span_mask.shape) # # # transformer and cfg training with a drop rate # drop_audio_cond = random() < self.audio_drop_prob # p_drop in voicebox paper # drop_text_cond = random() < self.text_drop_prob # p_drop in voicebox paper drop_audio_cond = random() < self.audio_drop_prob # p_drop in voicebox paper if random() < self.text_drop_prob: # p_uncond in voicebox paper drop_audio_cond = True drop_text_cond = True else: drop_text_cond = False # print("drop_audio_cond:", drop_audio_cond, "drop_text_cond:", drop_text_cond) # if want rigorously mask out padding, record in collate_fn in dataset.py, and pass in here # adding mask will use more memory, thus also need to adjust batchsampler with scaled down threshold for long sequences pred = self.transformer(x=φ, cond=cond_grl, text=text, time=time, drop_audio_cond=drop_audio_cond, drop_text=drop_text_cond) # flow matching loss pred_clamp = pred.float().clamp(-20, 20) loss = F.mse_loss(pred_clamp, flow, reduction="none") loss = loss[rand_span_mask] # [N] # # # 1. 全局截断：>2 或 NaN → 0（全局） # print("mse loss shape:", loss.shape, "loss max:", loss.max(), "loss min:", loss.min(), target_text_lengths[0]) # # 2. 统计非NaN值的百分比 # valid_mask = ~torch.isnan(loss) # total_count = loss.numel() # 总元素数量（所有维度） # valid_count = valid_mask.sum().item() # 非NaN元素数量 # valid_percentage = (valid_count / total_count) * 100 # print(f"mse loss: total_count: {total_count}", f"valid_count: {valid_count}", f"valid_percentage: {valid_percentage:.2f}%") # valid_loss = loss[~torch.isnan(loss)] loss = torch.where(torch.isnan(loss) | (loss > 300.0), 300.0, loss) loss = loss.mean() # loss = torch.tanh(torch.log1p(loss.mean())) # 对数缩放 # if len(valid_loss) > 0: # clipped_loss = torch.clamp(valid_loss, max=150) # loss = torch.tanh(torch.log1p(clipped_loss.mean())) # 对数缩放 # else: # loss = torch.tensor(0.0, device=pred.device) accent_logits = self.accent_classifier(cond_grl) accent_logits_mean = accent_logits.mean(dim=1) lang_labels = langs.to(accent_logits.device).long() # print("langs:", lang_labels, "accent_logits:", accent_logits.shape, "accent_logits_mean:", accent_logits_mean.shape) accent_loss = self.accent_criterion(accent_logits_mean, lang_labels) # guard against NaN / Inf in accent_loss if not torch.isfinite(accent_loss): accent_loss = torch.zeros_like(accent_loss, device=accent_loss.device) # accent_loss = torch.zeros_like(loss, device=loss.device, requires_grad=True) loss += 0.1 * accent_loss valid_indices = torch.where(time > 0.5)[0] # print("torch.where(time > 0.5):", valid_indices, torch.where(time > 0.5)) if valid_indices.size(0) > 2: # 动态选择符合条件的sample selected_gt = inp[valid_indices] selected_pred = pred[valid_indices] selected_text = text[valid_indices] selected_lens = lens[valid_indices] selected_target_lengths = target_text_lengths[valid_indices] # print("pred:", selected_pred.shape, "valid_indices:", valid_indices, "lens:", selected_lens, "target_lengths:", selected_target_lengths) if self.use_spk_enc and valid_indices.size(0) > 2: # speaker encoder loss e_gt = self.speaker_encoder(selected_gt, selected_lens) e_pred = self.speaker_encoder(selected_pred, selected_lens) spk_loss = self.info_nce_speaker(e_gt, e_pred) if not torch.isnan(spk_loss).any(): # and spk_loss.item() > 1e-6: loss = loss + spk_loss * 10.0 else: spk_loss = torch.zeros_like(loss, device=loss.device, requires_grad=False) else: spk_loss = torch.zeros_like(loss, device=loss.device, requires_grad=False) # print("spk_loss:", spk_loss) # ctc loss if self.use_ctc_loss and valid_indices.size(0) > 2: # 如果t大于0.5 则计算ctc loss ctc_loss = self.ctc( decoder_outputs=selected_pred, target_phones=selected_text, decoder_lengths=selected_lens, target_lengths=selected_target_lengths, ) # print("loss:", loss, "ctc_loss:", ctc_loss, "time: ", time.shape, time[valid_indices].mean()) # 如果ctc loss没有nan，才加上ctc loss if not torch.isnan(ctc_loss).any() and ctc_loss.item() > 1e-6: # ctc_scaled = torch.tanh(torch.log1p(ctc_loss)) ctc_scaled = ctc_loss loss = loss + 0.1 * ctc_scaled else: ctc_scaled = torch.zeros_like(loss, device=loss.device, requires_grad=False) # print("loss:", loss, "ctc_scaled:", ctc_scaled) else: ctc_scaled = torch.zeros_like(loss, device=loss.device, requires_grad=False) # 在计算完 total loss 之前 total_loss = loss # base flow loss + others you added # note: we intentionally do NOT add 0.0 * pred.sum() etc. here, to avoid # propagating NaNs from intermediate tensors into the loss scalar. return total_loss, ctc_scaled, accent_loss, len(valid_indices), cond, pred # accent_loss, def forward(self, batchs: Dict[str, torch.Tensor], *, noise_scheduler: str | None = None): """ Simplified forward version for accent-invariant flow matching. Removes speaker encoder and CTC parts, keeps accent GRL. """ inp = batchs["mel"].permute(0, 2, 1) # [B, T_mel, D] lens = batchs["mel_lengths"] text = batchs["text"] langs = batchs["langs"] audio_16k_list = batchs.get("audio_16k", None) prosody_idx_list = batchs.get("prosody_idx", None) # # ---- 4. 随机截取并打乱 segment ---- # rand_mel = [clip_and_shuffle(spec, spec.shape[-1]) for spec in batchs["mel"]] # padded_rand_mel = [] # for spec in rand_mel: # padding = (0, batchs["mel"].shape[-1] - spec.size(-1)) # padded_spec = F.pad(spec, padding, value=0) # padded_rand_mel.append(padded_spec) # rand_mel = torch.stack(padded_rand_mel).permute(0, 2, 1) # assert rand_mel.shape == inp.shape, f"shape diff: rand_mel.shape: {rand_mel.shape}, inp.shape: {inp.shape}" if inp.ndim == 2: inp = self.mel_spec(inp).permute(0, 2, 1) assert inp.shape[-1] == self.num_channels batch, seq_len, dtype, device = *inp.shape[:2], inp.dtype, self.device # --- handle text if isinstance(text, list): if exists(self.vocab_char_map): text = list_str_to_idx(text, self.vocab_char_map).to(device) else: text = list_str_to_tensor(text).to(device) assert text.shape[0] == batch # print("text:", batchs["text"][0], text.shape, text[0], batchs["text_lengths"][0]) # --- prosody conditioning (compute embeddings per sub-utterance) prosody_mel_cond = None prosody_text_cond = None if ( self.prosody_encoder is not None and audio_16k_list is not None and prosody_idx_list is not None ): # prepare zero tensors for each sample T_mel = seq_len T_text = text.shape[1] prosody_mel_cond = torch.zeros(batch, T_mel, 512, device=device, dtype=dtype) prosody_text_cond = torch.zeros(batch, T_text, 512, device=device, dtype=dtype) # collect all segments, run encoder per segment seg_embeds: list[Tensor] = [] seg_meta: list[tuple[int, int, int, int, int, int]] = [] for b in range(batch): audio_b = audio_16k_list[b] idx_list = prosody_idx_list[b] if audio_b is None or idx_list is None: continue audio_b = audio_b.to(device=device, dtype=dtype) for seg in idx_list: text_start, text_end, mel_start, mel_end, audio_start, audio_end = seg # clamp audio indices audio_start = max(0, min(audio_start, audio_b.shape[0] - 1)) audio_end = max(audio_start + 1, min(audio_end, audio_b.shape[0])) audio_seg = audio_b[audio_start:audio_end] if audio_seg.numel() == 0: continue fbank = extract_fbank_16k(audio_seg) # (T_fbank, 80) fbank = fbank.unsqueeze(0).to(device=device, dtype=dtype) # (1, T_fbank, 80) with torch.no_grad(): emb = self.prosody_encoder(fbank, padding_mask=None)[0] # (512,) seg_embeds.append(emb) seg_meta.append( (b, text_start, text_end, mel_start, mel_end) ) if seg_embeds: seg_embeds_tensor = torch.stack(seg_embeds, dim=0) # (N_seg, 512) # scatter embeddings back to per-sample tensors for emb, meta in zip(seg_embeds_tensor, seg_meta): b, ts, te, ms, me = meta emb_exp = emb.to(device=device, dtype=dtype) prosody_mel_cond[b, ms:me, :] = emb_exp prosody_text_cond[b, ts:te, :] = emb_exp # dropout on prosody conditioning prosody_mel_cond = self.prosody_dropout(prosody_mel_cond) prosody_text_cond = self.prosody_dropout(prosody_text_cond) # --- mask & random span mask = lens_to_mask(lens, length=seq_len) frac_lengths = torch.zeros((batch,), device=device).float().uniform_(*self.frac_lengths_mask) rand_span_mask = mask_from_frac_lengths(lens, frac_lengths) if exists(mask): rand_span_mask &= mask # --- flow setup x1 = inp x0 = torch.randn_like(x1) time = torch.rand((batch,), dtype=dtype, device=device) t = time[:, None, None] φ = (1 - t) * x0 + t * x1 flow = x1 - x0 # --- conditional input (masked mel) + optional prosody cond = torch.where(rand_span_mask[..., None], torch.zeros_like(x1), x1) # x1 # rand_mel if prosody_mel_cond is not None: prosody_mel_proj = self.prosody_to_mel(prosody_mel_cond) # (B, T_mel, num_channels) # if needed, pad/crop to seq_len if prosody_mel_proj.size(1) < seq_len: pad_len = seq_len - prosody_mel_proj.size(1) prosody_mel_proj = F.pad(prosody_mel_proj, (0, 0, 0, pad_len)) elif prosody_mel_proj.size(1) > seq_len: prosody_mel_proj = prosody_mel_proj[:, :seq_len, :] cond = cond + prosody_mel_proj # --- Gradient reversal: encourage accent-invariant cond cond_grl = grad_reverse(cond, lambda_=1.0) # # --- random drop condition for CFG-like robustness # drop_audio_cond = random() < self.audio_drop_prob # drop_text_cond = random() < self.text_drop_prob if not drop_audio_cond else True # safe per-batch random (tensor) rand_for_drop = torch.rand(1, device=device) drop_audio_cond = (rand_for_drop.item() < self.audio_drop_prob) rand_for_text = torch.rand(1, device=device) drop_text_cond = (rand_for_text.item() < self.text_drop_prob) # --- main prediction pred = self.transformer( x=φ, cond=cond_grl, text=text, time=time, drop_audio_cond=drop_audio_cond, drop_text=drop_text_cond, prosody_text=prosody_text_cond, ) # === FLOW LOSS (robust mask-weighted) === pred_clamp = pred.float().clamp(-20, 20) per_elem_loss = F.mse_loss(pred_clamp, flow, reduction="none") # [B, T, D] mask_exp = rand_span_mask.unsqueeze(-1).to(dtype=per_elem_loss.dtype) # [B, T, 1] masked_loss = per_elem_loss * mask_exp # zeros where mask False # total selected scalar (frames * dim) n_selected = mask_exp.sum() * per_elem_loss.size(-1) # scalar denom = torch.clamp(n_selected, min=1.0) loss_sum = masked_loss.sum() loss = loss_sum / denom # numeric safety loss = torch.where(torch.isnan(loss) | (loss > 300.0), torch.tensor(300.0, device=loss.device, dtype=loss.dtype), loss) # === ACCENT LOSS === accent_logits = self.accent_classifier(cond_grl) # pool across time -> [B, C] accent_logits_mean = accent_logits.mean(dim=1) lang_labels = langs.to(accent_logits_mean.device).long() accent_loss = self.accent_criterion(accent_logits_mean, lang_labels) # guard against NaN / Inf in accent_loss if not torch.isfinite(accent_loss): accent_loss = torch.zeros_like(accent_loss, device=accent_loss.device) base_loss = loss + 0.1 * accent_loss # === OPTIONAL CTC LOSS (robust, only on valid samples) === ctc_scaled = torch.tensor(0.0, device=device, dtype=dtype) if getattr(self, "use_ctc_loss", False) and getattr(self, "ctc", None) is not None: # select samples with larger t for CTC supervision (similar to forward_old) valid_indices = torch.where(time > 0.5)[0] if valid_indices.size(0) > 2: selected_pred = pred[valid_indices] selected_text = text[valid_indices] selected_lens = lens[valid_indices] # text was tokenized from list_str_to_idx, where padding is -1 selected_target_lengths = (selected_text != -1).sum(dim=-1) ctc_loss = self.ctc( decoder_outputs=selected_pred, target_phones=selected_text, decoder_lengths=selected_lens, target_lengths=selected_target_lengths, ) if torch.isfinite(ctc_loss) and ctc_loss.item() > 1e-6: ctc_scaled = ctc_loss base_loss = base_loss + 0.1 * ctc_scaled total_loss = base_loss # note: we intentionally do NOT add 0.0 * pred.sum() etc. here, to avoid # propagating NaNs from intermediate tensors into the loss scalar. return total_loss, accent_loss, ctc_scaled, cond, pred