anyaccomp/fmt_model.py

import torch
import torch.nn as nn
from anyaccomp.llama_nar import DiffLlamaConcat


class FlowMatchingTransformerConcat(nn.Module):
    def __init__(
        self,
        vocab_size=1024,
        mel_dim=100,
        hidden_size=1024,
        num_layers=12,
        num_heads=16,
        cfg_scale=0.2,
        use_cond_code=False,
        cond_codebook_size=1024,
        cond_dim=1024,
        cond_scale_factor=1,
        sigma=1e-5,
        time_scheduler="linear",
        cfg=None,
    ):
        super().__init__()
        self.cfg = cfg

        mel_dim = (
            cfg.mel_dim if cfg is not None and hasattr(cfg, "mel_dim") else mel_dim
        )
        hidden_size = (
            cfg.hidden_size
            if cfg is not None and hasattr(cfg, "hidden_size")
            else hidden_size
        )
        num_layers = (
            cfg.num_layers
            if cfg is not None and hasattr(cfg, "num_layers")
            else num_layers
        )
        num_heads = (
            cfg.num_heads
            if cfg is not None and hasattr(cfg, "num_heads")
            else num_heads
        )
        cfg_scale = (
            cfg.cfg_scale
            if cfg is not None and hasattr(cfg, "cfg_scale")
            else cfg_scale
        )
        use_cond_code = (
            cfg.use_cond_code
            if cfg is not None and hasattr(cfg, "use_cond_code")
            else use_cond_code
        )
        cond_codebook_size = (
            cfg.cond_codebook_size
            if cfg is not None and hasattr(cfg, "cond_codebook_size")
            else cond_codebook_size
        )
        cond_dim = (
            cfg.cond_dim if cfg is not None and hasattr(cfg, "cond_dim") else cond_dim
        )
        time_scheduler = (
            cfg.time_scheduler
            if cfg is not None and hasattr(cfg, "time_scheduler")
            else time_scheduler
        )
        sigma = cfg.sigma if cfg is not None and hasattr(cfg, "sigma") else sigma
        cond_scale_factor = (
            cfg.cond_scale_factor
            if cfg is not None and hasattr(cfg, "cond_scale_factor")
            else 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_scale = cfg_scale
        self.use_cond_code = use_cond_code
        self.cond_codebook_size = cond_codebook_size
        self.cond_dim = cond_dim
        self.time_scheduler = time_scheduler
        self.sigma = sigma
        self.cond_scale_factor = cond_scale_factor

        self.vocab_size = (
            cfg.vocab_size
            if cfg is not None and hasattr(cfg, "vocab_size")
            else vocab_size
        )
        self.vocal_mel_proj = (
            nn.Linear(self.cfg.cond_code_dim, self.hidden_size)
            if not self.use_cond_code
            else nn.Sequential(
                nn.Embedding(
                    self.vocab_size, self.mel_dim
                ),  # [batch] -> [batch, mel_dim]
                nn.Linear(
                    self.mel_dim, self.hidden_size
                ),  # [batch, mel_dim] -> [batch, hidden_size]
            )
        )

        self.diff_estimator = DiffLlamaConcat(
            mel_dim=self.mel_dim,
            hidden_size=self.hidden_size,
            num_heads=self.num_heads,
            num_layers=self.num_layers,
            flash_attention=hasattr(cfg, "flash_attention") and cfg.flash_attention,
        )

        if hasattr(cfg, "repa_loss") and cfg.repa_loss.enable:
            repa_dim = (
                cfg.repa_loss.repa_dim
                if hasattr(cfg.repa_loss, "repa_dim")
                else self.hidden_size
            )
            self.repa_proj = nn.Sequential(
                nn.Linear(self.hidden_size, self.hidden_size),
                nn.SiLU(),
                nn.Linear(self.hidden_size, self.hidden_size),
                nn.SiLU(),
                nn.Linear(self.hidden_size, repa_dim),
            )

        self.reset_parameters()

    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)

    # removed forward_diffusion, loss_t, compute_loss, forward bcs they weren't needed for inference.

    @torch.no_grad()
    def reverse_diffusion(
        self,
        vocal_mel=None,
        prompt=None,
        right_prompt=None,
        x_mask=None,
        prompt_mask=None,
        right_prompt_mask=None,
        target_len=None,
        n_timesteps=10,
        cfg=1.0,
        rescale_cfg=0.75,
    ):
        h = 1.0 / n_timesteps
        prompt_len = prompt.shape[1] if prompt is not None else 0
        right_prompt_len = right_prompt.shape[1] if right_prompt is not None else 0
        # print(prompt_len, right_prompt_len)
        if vocal_mel is not None:
            target_len = vocal_mel.shape[1]
        elif target_len is None:
            target_len = 1000  # hardcode 50Hz 20s
        else:
            raise ValueError
        full_len = target_len
        target_len = target_len - prompt_len - right_prompt_len

        cond = self.vocal_mel_proj(vocal_mel)

        if x_mask is None:
            x_mask = torch.ones(cond.shape[0], target_len).to(cond.device)
        if prompt_mask is None and prompt is not None:
            prompt_mask = torch.ones(cond.shape[0], prompt_len).to(cond.device)
        if right_prompt_mask is None and right_prompt is not None:
            right_prompt_mask = torch.ones(cond.shape[0], right_prompt_len).to(
                cond.device
            )

        if prompt is not None and right_prompt is not None:
            xt_mask = torch.cat([prompt_mask, x_mask, right_prompt_mask], dim=1)
        elif prompt is not None and right_prompt is None:
            xt_mask = torch.cat([prompt_mask, x_mask], dim=1)
        elif prompt is None and right_prompt is not None:
            xt_mask = torch.cat([x_mask, right_prompt_mask], dim=1)
        else:
            xt_mask = x_mask

        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):
            if prompt is not None and right_prompt is not None:
                xt_input = torch.cat([prompt, xt, right_prompt], dim=1)
            elif prompt is not None and right_prompt is None:
                xt_input = torch.cat([prompt, xt], dim=1)
            elif prompt is None and right_prompt is not None:
                xt_input = torch.cat([xt, right_prompt], dim=1)
            else:
                xt_input = xt
            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, xt_mask, cond)
            flow_pred = flow_pred[:, prompt_len : prompt_len + target_len, :]
            # cfg

            if cfg > 0:
                uncond_flow_pred = self.diff_estimator(
                    xt_input, t, xt_mask, torch.zeros_like(cond)
                )
                uncond_flow_pred = uncond_flow_pred[
                    :, prompt_len : prompt_len + target_len, :
                ]
                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