audio_modeling_omni.py

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,
        )