import torch, fire from typing import Optional import torch.distributed from torch.nn import functional as F from flash_attn import flash_attn_varlen_func from torch import nn import numpy as np import deepspeed from transformers.activations import ACT2FN from dataclasses import dataclass from transformers.modeling_outputs import ModelOutput try: from .vector_quantize import VectorQuantize except: from vector_quantize import VectorQuantize from .flow_matching import ( ConditionalDecoder, ConditionalCFM, ) import math import copy def sinusoids(length, channels, max_timescale=10000): """Returns sinusoids for positional embedding""" assert channels % 2 == 0 log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1) inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2)) scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :] return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1) def get_sequence_mask(inputs, inputs_length): if inputs.dim() == 3: bsz, tgt_len, _ = inputs.size() else: bsz, tgt_len = inputs_length.shape[0], torch.max(inputs_length) sequence_mask = torch.arange(0, tgt_len).to(inputs.device) sequence_mask = torch.lt(sequence_mask, inputs_length.reshape(bsz, 1)).view(bsz, tgt_len, 1) unpacking_index = torch.cumsum(sequence_mask.to(torch.int64).view(-1), dim=0) - 1 # 转成下标 return sequence_mask, unpacking_index def unpack_hidden_states(hidden_states, lengths): bsz = lengths.shape[0] sequence_mask, unpacking_index = get_sequence_mask(hidden_states, lengths) hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view( bsz, torch.max(lengths), hidden_states.shape[-1] ) hidden_states = torch.where( sequence_mask, hidden_states, 0 ) # 3d (bsz, max_input_len, d) return hidden_states class RMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): """ RMSNorm is equivalent to T5LayerNorm """ super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states): variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) # convert into half-precision if necessary if self.weight.dtype in [torch.float16, torch.bfloat16]: hidden_states = hidden_states.to(self.weight.dtype) return self.weight * hidden_states class OmniWhisperAttention(nn.Module): def __init__(self, embed_dim, num_heads, causal=False): super().__init__() self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False) self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True) self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True) self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True) self.causal = causal def forward(self, hidden_states: torch.Tensor, seq_len: torch.Tensor): bsz, _ = hidden_states.size() query_states = self.q_proj(hidden_states).view(bsz, self.num_heads, self.head_dim) key_states = self.k_proj(hidden_states).view(bsz, self.num_heads, self.head_dim) value_states = self.v_proj(hidden_states).view(bsz, self.num_heads, self.head_dim) cu_len = F.pad(torch.cumsum(seq_len, dim=0), (1, 0), "constant", 0).to(torch.int32) max_seqlen = torch.max(seq_len).to(torch.int32).detach() attn_output = flash_attn_varlen_func(query_states, key_states, value_states, cu_len, cu_len, max_seqlen, max_seqlen, causal=self.causal) # (bsz * qlen, nheads, headdim) attn_output = attn_output.reshape(bsz, self.embed_dim) attn_output = self.out_proj(attn_output) return attn_output class OmniWhisperTransformerLayer(nn.Module): def __init__( self, act, d_model, encoder_attention_heads, encoder_ffn_dim, causal, ln_type="LayerNorm", ): super().__init__() self.embed_dim = d_model self.self_attn = OmniWhisperAttention( self.embed_dim, encoder_attention_heads, causal ) if ln_type == "LayerNorm": self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) elif ln_type == "RMSNorm": self.self_attn_layer_norm = RMSNorm(self.embed_dim) else: raise ValueError(f"Unknown ln_type: {ln_type}") self.activation_fn = act self.fc1 = nn.Linear(self.embed_dim, encoder_ffn_dim) self.fc2 = nn.Linear(encoder_ffn_dim, self.embed_dim) if ln_type == "LayerNorm": self.final_layer_norm = nn.LayerNorm(self.embed_dim) elif ln_type == "RMSNorm": self.final_layer_norm = RMSNorm(self.embed_dim) else: raise ValueError(f"Unknown ln_type: {ln_type}") def forward( self, hidden_states: torch.Tensor, seq_len: torch.Tensor ) -> torch.Tensor: residual = hidden_states hidden_states = self.self_attn_layer_norm(hidden_states) hidden_states = self.self_attn(hidden_states, seq_len) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.final_layer_norm(hidden_states) hidden_states = self.activation_fn(self.fc1(hidden_states)) hidden_states = self.fc2(hidden_states) hidden_states = residual + hidden_states if ( hidden_states.dtype == torch.float16 or hidden_states.dtype == torch.bfloat16 ) and (torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()): clamp_value = torch.finfo(hidden_states.dtype).max - 1000 hidden_states = torch.clamp( hidden_states, min=-clamp_value, max=clamp_value ) return hidden_states class OmniAudioEncoder(nn.Module): def __init__(self, config): super().__init__() config._attn_implementation = 'flash_attention_2' # self.config = config self.max_source_positions = (config.max_audio_seconds * config.sampling_rate // config.hop_length) // config.stride_size self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0 self.conv1 = nn.Conv1d(config.num_mel_bins, config.d_model, kernel_size=config.kernel_size, padding=1) self.conv2 = nn.Conv1d(config.d_model, config.d_model, kernel_size=config.kernel_size, stride=config.stride_size, padding=1) self.register_buffer("positional_embedding", sinusoids(self.max_source_positions, config.d_model)) # 1500 * d self.layers = nn.ModuleList([OmniWhisperTransformerLayer( ACT2FN[config.activation_function], config.d_model, config.encoder_attention_heads, config.encoder_ffn_dim, False) for _ in range(config.encoder_layers)]) self.layer_norm = nn.LayerNorm(config.d_model) @torch.no_grad() def fake_input(self, device): input_features = torch.rand([2, self.config.num_mel_bins, 10], dtype=torch.float32, device=device) encoder_length = torch.ones([2], dtype=torch.int32, device=device) * 3 bridge_length = torch.ones([2], dtype=torch.int32, device=device) return input_features, encoder_length, bridge_length def forward( self, input_features, output_length, ): input_features = input_features.to(self.conv1.weight.dtype) inputs_embeds = nn.functional.gelu(self.conv1(input_features)) # (bs, channels, frames) inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds)) # (bs, channels, frames // 2) inputs_embeds = inputs_embeds.permute(0, 2, 1) # (bs, frams, channels) bsz, tgt_len, _ = inputs_embeds.size() if tgt_len < self.positional_embedding.shape[0]: current_positional_embedding = self.positional_embedding[:tgt_len] else: current_positional_embedding = self.positional_embedding hidden_states = (inputs_embeds.to(torch.float32) + current_positional_embedding).to(inputs_embeds.dtype) # packing hidden states attention_mask, unpacking_index = get_sequence_mask(hidden_states, output_length) hidden_states = torch.masked_select(hidden_states, attention_mask).view(torch.sum(output_length), self.config.d_model) for idx, encoder_layer in enumerate(self.layers): hidden_states = encoder_layer(hidden_states, output_length) hidden_states = self.layer_norm(hidden_states) # unpacking hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view(bsz, tgt_len, self.config.d_model) hidden_states = torch.where(attention_mask, hidden_states, 0) return hidden_states class CasualConvTranspose1d(nn.Module): # 反卷积 def __init__(self, in_channels, out_channels, kernel_size, stride): super().__init__() self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride) self.norm = nn.GroupNorm(1, out_channels) self.in_channels = in_channels self.out_channels = out_channels def forward(self, hidden_states, input_length, output_dim=None): kernel_size = self.conv.kernel_size[0] stride = self.conv.stride[0] bsz = input_length.shape[0] if output_dim is None: output_dim = hidden_states.dim() if hidden_states.dim() <= 2: # unpack sequence to 3d sequence_mask, unpacking_index = get_sequence_mask(hidden_states, input_length) hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view(bsz, torch.max(input_length), self.in_channels) hidden_states = torch.where(sequence_mask, hidden_states, 0) # 3d (bsz, max_input_len, d) hidden_states = hidden_states.transpose(2, 1) # (N, L, C) -> (N, C, L) hidden_states = self.conv(hidden_states) hidden_states = self.norm(hidden_states) hidden_states = hidden_states.transpose(2, 1) # (N, C, L) -> (N, L, C) casual_padding_right = max(0, kernel_size - stride) hidden_states = hidden_states[:, :hidden_states.shape[1] - casual_padding_right, :] output_length = (input_length - 1) * stride + kernel_size - casual_padding_right sequence_mask, _ = get_sequence_mask(hidden_states, output_length) if output_dim <= 2: hidden_states = torch.masked_select(hidden_states, sequence_mask).view(-1, self.out_channels) else: hidden_states = torch.where(sequence_mask, hidden_states, 0) hidden_states = hidden_states[:, :torch.max(output_length), :] # 截断到最大有效长度 return hidden_states, output_length class MelSpecRefineNet(nn.Module): """ # post net, coarse to refined mel-spectrogram frames # ref1: Autoregressive Speech Synthesis without Vector Quantization # ref2: CosyVoice length_regulator.py # ref3: Neural Speech Synthesis with Transformer Network https://github.com/soobinseo/Transformer-TTS/blob/master/network.py """ def __init__(self, encoder_config, vocoder_config): super().__init__() self.encoder_config = encoder_config self.vocoder_config = vocoder_config layers = nn.ModuleList([]) in_channels = self.vocoder_config.num_mel_bins for i, out_channels in enumerate(self.vocoder_config.channels[:-1]): module = nn.Conv1d(in_channels, out_channels, 5, 1, 2) # cosyvoice kernel=3, stride=1, pad=1 in_channels = out_channels norm = nn.GroupNorm(1, out_channels) act = nn.Mish() layers.extend([module, norm, act]) layers.append(nn.Conv1d(in_channels, self.vocoder_config.num_mel_bins, 1, 1)) # projector self.layers = nn.Sequential(*layers) def compute_output_length(self, input_length): output_length = input_length.to( torch.float32) * self.encoder_config.hop_length / self.encoder_config.sampling_rate output_length = output_length * self.vocoder_config.sampling_rate / self.vocoder_config.hop_length return output_length.to(torch.int64) def forward(self, coarse_mel, input_length, output_length=None): bsz, _, d = coarse_mel.shape assert (d == self.vocoder_config.num_mel_bins) if output_length is None or not self.training: output_length = self.compute_output_length(input_length) coarse_mel, default_dtype = coarse_mel[:, :torch.max(input_length), :], coarse_mel.dtype coarse_mel = F.interpolate(coarse_mel.to(torch.float32).transpose(1, 2).contiguous(), size=output_length.max(), mode='nearest').to(default_dtype) refined_mel = self.layers(coarse_mel).transpose(1, 2).contiguous() # (bs, t, d) coarse_mel = coarse_mel.transpose(1, 2) # (bs, max(output_length), d) refined_mel += coarse_mel # residual conntection sequence_mask, _ = get_sequence_mask(refined_mel, output_length) coarse_mel = torch.where(sequence_mask, coarse_mel, 0) refined_mel = torch.where(sequence_mask, refined_mel, 0) return refined_mel, coarse_mel, output_length @dataclass class OmniAudioDecoderOutput(ModelOutput): refined_mel: Optional[torch.FloatTensor] = None coarse_mel: Optional[torch.FloatTensor] = None mel_length: Optional[torch.Tensor] = None hidden_states_before_dconv2: Optional[torch.FloatTensor] = None output_length_before_dconv2: Optional[torch.Tensor] = None class OmniAudioDecoder(nn.Module): def __init__(self, config): super().__init__() self.config = config.audio_config self.vocoder_config = config.vocoder_config self.max_source_positions = self.config.max_audio_seconds * self.config.sampling_rate // self.config.hop_length self.dconv1 = CasualConvTranspose1d( self.config.d_model, self.config.d_model, self.config.decoder_kernel_size, self.config.avg_pooler, ) self.register_buffer("positional_embedding", sinusoids(self.max_source_positions, self.config.d_model)) # causal transformer layers self.layers = nn.ModuleList( [OmniWhisperTransformerLayer( ACT2FN[self.config.activation_function], self.config.d_model, self.config.decoder_attention_heads, self.config.decoder_ffn_dim, True # causal ) for _ in range(self.config.decoder_layers) ]) self.layer_norm = nn.LayerNorm(self.config.d_model) self.dconv2 = CasualConvTranspose1d( self.config.d_model, self.vocoder_config.num_mel_bins, self.config.decoder_kernel_size, self.config.decoder_stride_size ) self.post_net = MelSpecRefineNet(config.audio_config, config.vocoder_config) self.gradient_checkpointing = True @torch.no_grad() def fake_input(self, device): audio_embed = torch.rand([1, 10, self.config.d_model], dtype=torch.float32, device=device) input_length = torch.ones([1], dtype=torch.int32, device=device) * 10 mel_labels_length = self.post_net.compute_output_length(input_length) return audio_embed, input_length, None, mel_labels_length def forward(self, audio_embed, input_length, mel_labels=None, mel_labels_length=None, fake_input=False, ): if fake_input: audio_embed, input_length, mel_labels, mel_labels_length = self.fake_input(self.layer_norm.weight.device) assert (audio_embed.shape[-1] == self.config.d_model) audio_embed = audio_embed.to(self.layer_norm.weight) # device and type audio_embed, output_length = self.dconv1(audio_embed, input_length, output_dim=3) # (b, l*2, d_model) _, tgt_len, _ = audio_embed.size() if tgt_len < self.positional_embedding.shape[0]: current_positional_embedding = self.positional_embedding[:tgt_len] else: current_positional_embedding = self.positional_embedding hidden_states = (audio_embed.to(torch.float32) + current_positional_embedding).to(audio_embed.dtype) # packing hidden states attention_mask, _ = get_sequence_mask(hidden_states, output_length) hidden_states = torch.masked_select(hidden_states, attention_mask).view(torch.sum(output_length), self.config.d_model) for idx, encoder_layer in enumerate(self.layers): hidden_states = encoder_layer(hidden_states, output_length) hidden_states = self.layer_norm(hidden_states) hidden_states_before_dconv2 = hidden_states output_length_before_dconv2 = output_length coarse_mel, output_length = self.dconv2(hidden_states, output_length, output_dim=3) refined_mel, coarse_mel, mel_labels_length = self.post_net(coarse_mel, output_length, mel_labels_length) return OmniAudioDecoderOutput( refined_mel=refined_mel, coarse_mel=coarse_mel, mel_length=mel_labels_length, hidden_states_before_dconv2=hidden_states_before_dconv2, output_length_before_dconv2=output_length_before_dconv2, ) class OmniAudioVQBridgeTokenizer(nn.Module): def __init__(self, config): super().__init__() self.config = config.audio_config self.gradient_checkpointing = False self.intermediate_dim = self.config.d_model * self.config.avg_pooler self.gate_proj = nn.Conv1d(self.config.d_model, self.intermediate_dim, self.config.avg_pooler, self.config.avg_pooler, bias=False) self.up_proj = nn.Conv1d(self.config.d_model, self.intermediate_dim, self.config.avg_pooler, self.config.avg_pooler, bias=False) self.down_proj = nn.Linear(self.intermediate_dim, self.intermediate_dim, bias=False) self.act_fn = ACT2FN['silu'] self.layer_norm = nn.LayerNorm(self.intermediate_dim) self.proj_decoder = nn.Linear(self.intermediate_dim, self.config.d_model) self.vq_list = nn.ModuleList([]) for idx, codebook_size in enumerate(self.config.vq_config.codebook_sizes): vq_config = copy.deepcopy(self.config.vq_config) vq_config.dim = self.intermediate_dim vq_config.codebook_size = codebook_size self.vq_list.append(VectorQuantize(vq_config)) for vq_layer in self.vq_list: deepspeed.zero.register_external_parameter(self, vq_layer.codebook.embed) def rvq_op(self, inputs, output_length): def rvq_layer_op(vq_layer, residual_encoding, output_length): q_v_i, code_ids_i = vq_layer(residual_encoding, output_length) residual_encoding = residual_encoding.float() - q_v_i.float() residual_encoding = residual_encoding.to(inputs.dtype) return residual_encoding, code_ids_i cmt_loss, residual_encoding = 0, inputs code_ids_list = [] for i, vq_layer in enumerate(self.vq_list): residual_encoding, code_ids_i = rvq_layer_op(vq_layer, residual_encoding, output_length) code_ids_list.append(code_ids_i) return torch.stack(code_ids_list, -1) def forward(self, x, output_length): batch_size, _, _ = x.shape output_length = output_length.to(x.device) if x.shape[1] % self.config.avg_pooler != 0: x = F.pad(x, (0, 0, 0, self.config.avg_pooler - x.shape[1] % self.config.avg_pooler), "constant", 0) xt = x.permute(0, 2, 1) g = self.gate_proj(xt).permute(0, 2, 1) # (bs, sl//poolersizre+1, d*2) u = self.up_proj(xt).permute(0, 2, 1) x = x.reshape(batch_size, -1, self.intermediate_dim) # (bs, sl//poolersizre+1, d*2) c = self.down_proj(self.act_fn(g) * u) res = self.layer_norm(c + x) valid_mask, _ = get_sequence_mask(res, output_length) code_ids = self.rvq_op(res, output_length) code_ids = torch.masked_select(code_ids, valid_mask).reshape(-1, len(self.vq_list)) # (sum(valid_sequence_length), vq_num) return code_ids @torch.no_grad() def decode(self, code_ids): vq_num = code_ids.shape[-1] res = sum(self.vq_list[i].get_output_from_indices(code_ids[:, i]).float() for i in range(vq_num-1,-1,-1)).to(self.proj_decoder.weight) decoder_emb = self.proj_decoder(res.to(self.proj_decoder.weight)) return decoder_emb @torch.no_grad() def recover(self, code_ids): vq_num = code_ids.shape[-1] res = sum(self.vq_list[i].get_output_from_indices(code_ids[:, i]).float() for i in range(vq_num-1,-1,-1)).to(self.proj_decoder.weight) return res class FlowmatchingPrenet(nn.Module): def __init__( self, input_feat_dim, out_feat_dim, d_model, attention_heads, ffn_dim, nlayers, activation_function, max_source_positions, target_mel_length_scale_ratio, ): super().__init__() self.d_model = d_model self.target_mel_length_scale_ratio = target_mel_length_scale_ratio self.gradient_checkpointing = False self.register_buffer( "positional_embedding", sinusoids(max_source_positions, d_model) ) self.in_mlp = nn.Sequential( nn.Linear(input_feat_dim, d_model * 4), nn.SiLU(), nn.Linear(d_model * 4, d_model), ) self.transformer_layers = nn.ModuleList( [ OmniWhisperTransformerLayer( act=ACT2FN[activation_function], d_model=d_model, encoder_attention_heads=attention_heads, encoder_ffn_dim=ffn_dim, causal=True, # causal ln_type="RMSNorm", ) for _ in range(nlayers) ] ) self.final_norm = RMSNorm(self.d_model) self.out_proj = nn.Linear(d_model, out_feat_dim, bias=False) def compute_output_length(self, input_length): output_length = input_length.float() * self.target_mel_length_scale_ratio return output_length.to(torch.int64) def forward(self, input_feat, input_length, output_length=None): """ Args: input_feat: [B, T, input_feat_dim] input_length: [B] output_length: [B] """ if output_length is None or not self.training: output_length = self.compute_output_length(input_length) input_feat = input_feat[:, : input_length.max(), :] # [B, T, D] orig_dtype = input_feat.dtype input_feat = F.interpolate( input=input_feat.to(torch.float32).transpose(1, 2).contiguous(), size=output_length.max(), mode="nearest", ).to(orig_dtype) input_feat = input_feat.transpose(1, 2).contiguous() # [B, T, D] hidden_states = self.in_mlp(input_feat) # packing hidden states bsz, tgt_len, d_model = hidden_states.shape attention_mask, unpacking_index = get_sequence_mask( hidden_states, output_length ) hidden_states = torch.masked_select(hidden_states, attention_mask).view( torch.sum(output_length), self.d_model ) for idx, encoder_layer in enumerate(self.transformer_layers): hidden_states = encoder_layer(hidden_states, output_length) # unpacking hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view( bsz, tgt_len, d_model ) hidden_states = torch.where(attention_mask, hidden_states, 0) hidden_states = self.final_norm(hidden_states) output = self.out_proj(hidden_states) return output, output_length @dataclass class OmniAudioFlowMatchingDecoderOutput(ModelOutput): flow_matching_mel: Optional[torch.FloatTensor] = None flow_matching_mel_lengths: Optional[torch.FloatTensor] = None class OmniAudioFlowMatchingDecoder(nn.Module): def __init__(self, config): super().__init__() self.config = config.flow_matching_config self.in_channels = self.config.in_channels self.spk_emb_dim = self.config.spk_emb_dim self.diffusion_steps = self.config.diffusion_steps self.cal_mel_mae = self.config.cal_mel_mae self.forward_step = -1 self.prenet = FlowmatchingPrenet( input_feat_dim=self.config.prenet_in_dim, out_feat_dim=self.config.prenet_out_dim, d_model=self.config.prenet_d_model, attention_heads=self.config.prenet_attention_heads, ffn_dim=self.config.prenet_ffn_dim, nlayers=self.config.prenet_nlayers, activation_function=self.config.prenet_activation_function, max_source_positions=self.config.prenet_max_source_positions, target_mel_length_scale_ratio=self.config.prenet_target_mel_length_scale_ratio, ) self.conditional_decoder = ConditionalDecoder( in_channels=self.in_channels * 2 + self.spk_emb_dim, out_channels=self.in_channels, causal=True, channels=self.config.channels, dropout=self.config.dropout, attention_head_dim=self.config.attention_head_dim, n_blocks=self.config.n_blocks, num_mid_blocks=self.config.num_mid_blocks, num_heads=self.config.num_heads, act_fn=self.config.act_fn, ) self.cfm = ConditionalCFM( in_channels=self.in_channels, cfm_params=self.config.cfm_params, n_spks=0, spk_emb_dim=self.spk_emb_dim, ) def unpack_hidden_states(self, hidden_states, output_length): unpacked = unpack_hidden_states(hidden_states, output_length) return unpacked, output_length def forward( self, refined_mel, input_length, mel_labels=None, mel_labels_length=None ): """ :param refined_mel: [bs, max_input_len, mel_bin] :param input_length: [batch_size] :param refined_mel: [bs, mel_bin, max_input_len] :return: """ self.forward_step += 1 orig_dtype = refined_mel.dtype prenet_mae_metric = torch.tensor(0.0).to(refined_mel.device) prenet_regression_loss = torch.tensor(0.0).to(refined_mel.device) if self.prenet is not None: refined_mel = refined_mel[:, : torch.max(input_length), :] if mel_labels_length is None: mel_labels_length = self.prenet.compute_output_length(input_length) refined_mel, input_length = self.prenet( refined_mel, input_length, mel_labels_length ) float_dtype = refined_mel.dtype refined_mel = refined_mel.float() input_length = input_length.long() refined_mel = refined_mel[:, : torch.max(input_length), :] sequence_mask, unpacking_index = get_sequence_mask(refined_mel, input_length) refined_mel = refined_mel.transpose(1, 2) # (bs, mel_bin, max_input_len) sequence_mask = sequence_mask.transpose(2, 1) # (bs, 1, sl) fm_mel = self.cfm.forward( estimator=self.conditional_decoder, mu=refined_mel.to(float_dtype), mask=sequence_mask.float(), n_timesteps=self.diffusion_steps, ) return OmniAudioFlowMatchingDecoderOutput( flow_matching_mel=fm_mel.transpose(1, 2), flow_matching_mel_lengths=mel_labels_length, )