# https://github.com/open-mmlab/Amphion/blob/main/models/svc/flow_matching_transformer/fmt_model.py # Copyright (c) 2023 Amphion. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. import torch import numpy as np import torch.nn as nn import math from .llama import DiffLlama import torch.nn.functional as F class FlowMatchingTransformer(nn.Module): def __init__( self, mel_dim=100, hidden_size=1024, num_layers=12, num_heads=16, cfg_drop_prob=0.2, use_embedding=True, cond_codebook_size=1024, cond_scale_factor=1, sigma=1e-5, time_scheduler="linear", cfg=None, ): super().__init__() self.cfg = cfg if cfg is not None: mel_dim = getattr(cfg, "mel_dim", mel_dim) hidden_size = getattr(cfg, "hidden_size", hidden_size) num_layers = getattr(cfg, "num_layers", num_layers) num_heads = getattr(cfg, "num_heads", num_heads) cfg_drop_prob = getattr(cfg, "cfg_drop_prob", cfg_drop_prob) cond_codebook_size = getattr(cfg, "cond_codebook_size", cond_codebook_size) time_scheduler = getattr(cfg, "time_scheduler", time_scheduler) sigma = getattr(cfg, "sigma", sigma) cond_scale_factor = getattr(cfg, "cond_scale_factor", cond_scale_factor) self.mel_dim = mel_dim self.hidden_size = hidden_size self.num_layers = num_layers self.num_heads = num_heads self.cfg_drop_prob = cfg_drop_prob self.cond_codebook_size = cond_codebook_size self.time_scheduler = time_scheduler self.sigma = sigma self.cond_scale_factor = cond_scale_factor if use_embedding: self.cond_emb = nn.Embedding(cond_codebook_size, self.hidden_size) else: self.cond_emb = nn.Linear(cond_codebook_size, self.hidden_size) if cond_scale_factor != 1: self.do_resampling = True assert np.log2(cond_scale_factor).is_integer() up_layers = [] for _ in range(int(np.log2(cond_scale_factor))): up_layers.extend( [ nn.ConvTranspose1d( hidden_size, hidden_size, kernel_size=4, stride=2, padding=1 ), nn.GELU(), ] ) self.resampling_layers = nn.Sequential(*up_layers) else: self.do_resampling = False ### REPA: Use the Wav2Vec2Bert features to align. ### self.use_repa = "repa" in cfg self.repa_layer_index = None if self.use_repa: self.repa_layer_index = cfg.repa.layer_index self.repa_mlp_layer = nn.Sequential( nn.Linear(hidden_size, hidden_size * 4), nn.SiLU(), nn.Linear(hidden_size * 4, cfg.repa.output_dim), ) ### CTC: Use the ASR loss ### self.use_ctc = "ctc" in cfg self.ctc_layer_index = None if self.use_ctc: self.ctc_layer_index = cfg.ctc.layer_index self.ctc_mlp_layer = nn.Sequential( nn.Linear(hidden_size, hidden_size * 4), nn.SiLU(), nn.Linear(hidden_size * 4, cfg.ctc.output_dim), ) self.reset_parameters() self.diff_estimator = DiffLlama( mel_dim=mel_dim, hidden_size=hidden_size, num_heads=num_heads, num_layers=num_layers, ) self.sigma = sigma @torch.no_grad() def forward_diffusion(self, x, t, is_prompt=None): """ x: (B, T, mel_dim) t: (B,) """ new_t = t t = t.unsqueeze(-1).unsqueeze(-1) z = torch.randn( x.shape, dtype=x.dtype, device=x.device, requires_grad=False ) # (B, T, mel_dim) # get prompt len if torch.rand(1) <= self.cfg_drop_prob: prompt_len = torch.zeros(x.shape[0]).to(x) is_prompt = torch.zeros_like(x[:, :, 0]) else: if is_prompt is None: prompt_len = torch.randint( min(x.shape[1] // 4, 5), int(x.shape[1] * 0.4), (x.shape[0],) ).to( x.device ) # (B,) # get is_prompt is_prompt = torch.zeros_like(x[:, :, 0]) # (B, T) col_indices = ( torch.arange(is_prompt.shape[1]) .repeat(is_prompt.shape[0], 1) .to(prompt_len) ) # (B, T) is_prompt[col_indices < prompt_len.unsqueeze(1)] = 1 # (B, T) 1 if prompt else: prompt_len = is_prompt.sum(dim=1) # (B,) mask = torch.ones_like(x[:, :, 0]) # mask if 1, not mask if 0 mask[is_prompt.bool()] = 0 mask = mask[:, :, None] # flow matching: xt = (1 - (1 - sigma) * t) * x0 + t * x; where x0 ~ N(0, 1), x is a sample # flow gt: x - (1 - sigma) * x0 = x - (1 - sigma) * noise xt = ((1 - (1 - self.sigma) * t) * z + t * x) * mask + x * (1 - mask) return xt, z, new_t, prompt_len, mask def loss_t( self, x, x_mask, t, cond=None, is_prompt=None ): xt, z, new_t, prompt_len, mask = self.forward_diffusion(x, t, is_prompt) noise = z # drop all condition for cfg, so if prompt_len is 0, we also drop cond if cond is not None: cond = cond * torch.where( prompt_len > 0, torch.ones_like(prompt_len), torch.zeros_like(prompt_len), ).to(cond.device).unsqueeze(-1).unsqueeze(-1) dit_output = self.diff_estimator(xt, new_t, cond, x_mask, return_dict=True) flow_pred = dit_output["output"] # (B, T, mel_dim) # final mask used for loss calculation final_mask = mask * x_mask[..., None] # (B, T, 1) results = {"output": (noise, x, flow_pred, final_mask, prompt_len)} if self.use_repa: repa_hidden_states = dit_output["hidden_states"][ self.repa_layer_index ] # (B, T, hidden_size) repa_pred = self.repa_mlp_layer(repa_hidden_states) # (B, T, repa_dim) results["repa"] = repa_pred if self.use_ctc: ctc_hidden_states = dit_output["hidden_states"][ self.ctc_layer_index ] # (B, T, hidden_size) ctc_pred = self.ctc_mlp_layer(ctc_hidden_states) # (B, T, ctc_dim) results["ctc"] = ctc_pred return results def compute_loss(self, x, x_mask, cond=None, is_prompt=None): # x0: (B, T, num_quantizer) # x_mask: (B, T) mask is 0 for padding t = torch.rand(x.shape[0], device=x.device, requires_grad=False) t = torch.clamp(t, 1e-5, 1.0) # from CosyVoice: considering the generation process at the beginning is harder than follows, we involve a cosine scheduler for the timestep t if self.time_scheduler == "cos": t = 1 - torch.cos(t * math.pi * 0.5) else: pass return self.loss_t(x, x_mask, t, cond, is_prompt) def reset_parameters(self): def _reset_parameters(m): if isinstance(m, nn.MultiheadAttention): if m._qkv_same_embed_dim: nn.init.normal_(m.in_proj_weight, std=0.02) else: nn.init.normal_(m.q_proj_weight, std=0.02) nn.init.normal_(m.k_proj_weight, std=0.02) nn.init.normal_(m.v_proj_weight, std=0.02) if m.in_proj_bias is not None: nn.init.constant_(m.in_proj_bias, 0.0) nn.init.constant_(m.out_proj.bias, 0.0) if m.bias_k is not None: nn.init.xavier_normal_(m.bias_k) if m.bias_v is not None: nn.init.xavier_normal_(m.bias_v) elif ( isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d) or isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) ): m.weight.data.normal_(0.0, 0.02) elif isinstance(m, nn.Linear): m.weight.data.normal_(mean=0.0, std=0.02) if m.bias is not None: m.bias.data.zero_() elif isinstance(m, nn.Embedding): m.weight.data.normal_(mean=0.0, std=0.02) if m.padding_idx is not None: m.weight.data[m.padding_idx].zero_() self.apply(_reset_parameters) @torch.no_grad() def reverse_diffusion( self, cond, prompt, x_mask=None, prompt_mask=None, n_timesteps=10, cfg=1.0, rescale_cfg=0.75, ): h = 1.0 / n_timesteps prompt_len = prompt.shape[1] target_len = cond.shape[1] - prompt_len if x_mask == None: x_mask = torch.ones(cond.shape[0], target_len).to(cond.device) # (B, T) if prompt_mask == None: prompt_mask = torch.ones(cond.shape[0], prompt_len).to( cond.device ) # (B, prompt_len) xt_mask = torch.cat([prompt_mask, x_mask], dim=1) z = torch.randn( (cond.shape[0], target_len, self.mel_dim), dtype=cond.dtype, device=cond.device, requires_grad=False, ) xt = z # t from 0 to 1: x0 = z ~ N(0, 1) for i in range(n_timesteps): xt_input = torch.cat([prompt, xt], dim=1) t = (0 + (i + 0.5) * h) * torch.ones( z.shape[0], dtype=z.dtype, device=z.device ) flow_pred = self.diff_estimator(xt_input, t, cond, xt_mask) flow_pred = flow_pred[:, prompt_len:, :] # cfg if cfg > 0: uncond_flow_pred = self.diff_estimator( xt, t, torch.zeros_like(cond)[:, : xt.shape[1], :], x_mask ) pos_flow_pred_std = flow_pred.std() flow_pred_cfg = flow_pred + cfg * (flow_pred - uncond_flow_pred) rescale_flow_pred = ( flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std() ) flow_pred = ( rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg ) dxt = flow_pred * h xt = xt + dxt return xt @torch.no_grad() def reverse_diffusion_v2( self, cond, prompt, x_mask=None, prompt_mask=None, n_timesteps=10, cfg=1.0, rescale_cfg=0.75, ): h = 1.0 / n_timesteps prompt_len = prompt.shape[1] target_len = cond.shape[1] - prompt_len * 2 if x_mask == None: x_mask = torch.ones(cond.shape[0], target_len).to(cond.device) # (B, T) if prompt_mask == None: prompt_mask = torch.ones(cond.shape[0], prompt_len).to( cond.device ) # (B, prompt_len) xt_mask = torch.cat([prompt_mask, x_mask, prompt_mask], dim=1) z = torch.randn( (cond.shape[0], target_len, self.mel_dim), dtype=cond.dtype, device=cond.device, requires_grad=False, ) xt = z # t from 0 to 1: x0 = z ~ N(0, 1) for i in range(n_timesteps): xt_input = torch.cat([prompt, xt, prompt], dim=1) t = (0 + (i + 0.5) * h) * torch.ones( z.shape[0], dtype=z.dtype, device=z.device ) flow_pred = self.diff_estimator(xt_input, t, cond, xt_mask) flow_pred = flow_pred[:, prompt_len:-prompt_len, :] # cfg if cfg > 0: uncond_flow_pred = self.diff_estimator( xt, t, torch.zeros_like(cond)[:, : xt.shape[1], :], x_mask ) pos_flow_pred_std = flow_pred.std() flow_pred_cfg = flow_pred + cfg * (flow_pred - uncond_flow_pred) rescale_flow_pred = ( flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std() ) flow_pred = ( rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg ) dxt = flow_pred * h xt = xt + dxt return xt def forward(self, x, x_mask, cond_code, is_prompt=None): """ Args: x: (B, T, mel_dim) x_mask: (B, T) cond_code: (B, T), Note that cond_code might be not at 50Hz! """ T = x.shape[1] cond = self.cond_emb(cond_code) # (B, T, hidden_size) if self.do_resampling: # Align to the frame rate of Mels cond = self.resampling_layers(cond.transpose(1, 2)).transpose(1, 2) # print("cond_code: {}, after resampling: {}".format(cond_code.shape, cond.shape)) if cond.shape[1] >= T: # Check time dimension cond = cond[:, :T, :] else: padding_frames = T - cond.shape[1] last_frame = cond[:, -1:, :] padding = last_frame.repeat(1, padding_frames, 1) cond = torch.cat([cond, padding], dim=1) return self.compute_loss(x, x_mask, cond, is_prompt) if __name__ == "__main__": model_cfg = { "mel_dim": 128, "hidden_size": 256, "num_layers": 8, "num_heads": 8, "cfg_drop_prob": 0.2, "use_embedding": False, "cond_codebook_size": 256, "cond_scale_factor": 1, "sigma": 1e-5, "time_scheduler": "cos", } device = "cuda" x = torch.randn(2, 100, 128).to(device) x_mask = torch.ones(2, 100).to(device) # cond_code = torch.randint(0, 16384, (2, 25)).to(device) cond_code = torch.randn(2, 100, 256).to(device) model = FlowMatchingTransformer(cfg=model_cfg, **model_cfg).to(device) outputs = model(x, x_mask, cond_code) print(outputs) noise, x, flow_pred, final_mask, prompt_len = outputs["output"] final_mask = final_mask.squeeze(-1) flow_gt = x - (1 - 1e-5) * noise # [B, n_frames, D] diff_loss = F.l1_loss( flow_pred, flow_gt, reduction="none" ).float() * final_mask.unsqueeze(-1) diff_loss = torch.mean(diff_loss, dim=2).sum() / final_mask.sum() print("diff_loss:", diff_loss.item()) diffusion_cond = torch.randn(2, 150, 256).to(device) diffusion_cond_emb = model.cond_emb(diffusion_cond) diffusion_prompt = torch.randn(2, 50, 128).to(device) n_timesteps = 32 generated = model.reverse_diffusion( diffusion_cond_emb, diffusion_prompt, n_timesteps=n_timesteps ) print("generated:", generated.shape)