models/dit_backbone.py

"""
Diffusion Transformer (DiT) Backbone for Flow Matching in ViuAI_TTS_200M.
Enhanced with:
  1. In-Context Reference Audio Conditioning (Voice Prompting / Zero-Shot Cloning)
  2. Classifier-Free Guidance (CFG) Support
  3. AdaLN-Zero (Adaptive Layer Normalization) Modulation
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional


class TimestepEmbedder(nn.Module):
    """
    Embeds scalar timesteps t into vector representations using sinusoidal embeddings + MLP.
    """
    def __init__(self, hidden_dim: int, frequency_embedding_size: int = 256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_dim, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim, bias=True),
        )
        self.frequency_embedding_size = frequency_embedding_size

    @staticmethod
    def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
        half = dim // 2
        freqs = torch.exp(
            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half
        )
        args = t[:, None].float() * freqs[None]
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        if dim % 2:
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
        return embedding

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
        return self.mlp(t_freq)


def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)


class DiTBlock(nn.Module):
    """
    A DiT block with adaptive layer norm (AdaLN-Zero) conditioning and Cross-Attention on text & reference audio.
    """
    def __init__(self, hidden_dim: int, num_heads: int, cond_dim: int, ffn_dim: int, dropout: float = 0.1):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.head_dim = hidden_dim // num_heads

        # 1. Self-Attention
        self.norm1 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
        self.qkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=False)
        self.attn_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)

        # 2. Cross-Attention on Text Conditioning & Reference Audio
        self.norm2 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
        self.q_cross = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.kv_cross = nn.Linear(cond_dim, hidden_dim * 2, bias=False)
        self.cross_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)

        # 3. Feed Forward Network
        self.norm3 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
        self.ffn = nn.Sequential(
            nn.Linear(hidden_dim, ffn_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ffn_dim, hidden_dim),
            nn.Dropout(dropout),
        )

        # AdaLN modulation MLP produces 9 parameters:
        # shift, scale, gate for: self-attn, cross-attn, and ffn
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(cond_dim, 9 * hidden_dim, bias=True)
        )

        # Zero-initialize the modulation weights (AdaLN-Zero) for ultra-stable training
        nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
        nn.init.constant_(self.adaLN_modulation[-1].bias, 0)

    def forward(
        self,
        x: torch.Tensor,
        c: torch.Tensor,
        conditioning: torch.Tensor,
        cond_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Args:
            x: Audio latents/mel sequence [B, T_audio, hidden_dim]
            c: Timestep conditioning vector [B, cond_dim]
            conditioning: Combined Text + Reference Voice Latents [B, T_cond, cond_dim]
            cond_mask: Optional condition padding mask [B, T_cond]
        """
        B, T, C = x.shape
        params = self.adaLN_modulation(c).chunk(9, dim=-1)
        shift_msa, scale_msa, gate_msa, shift_cross, scale_cross, gate_cross, shift_mlp, scale_mlp, gate_mlp = params

        # 1. Self Attention with Modulation
        norm_x = modulate(self.norm1(x), shift_msa, scale_msa)
        qkv = self.qkv(norm_x).reshape(B, T, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        attn = F.scaled_dot_product_attention(q, k, v)
        attn = attn.transpose(1, 2).reshape(B, T, C)
        x = x + gate_msa.unsqueeze(1) * self.attn_proj(attn)

        # 2. Cross Attention on Combined Text + Voice Prompt Conditioning
        norm_x = modulate(self.norm2(x), shift_cross, scale_cross)
        q_c = self.q_cross(norm_x).reshape(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        kv_c = self.kv_cross(conditioning).reshape(B, conditioning.shape[1], 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        k_c, v_c = kv_c[0], kv_c[1]

        cross_attn = F.scaled_dot_product_attention(
            q_c, k_c, v_c,
            attn_mask=cond_mask.unsqueeze(1).unsqueeze(2) if cond_mask is not None else None
        )
        cross_attn = cross_attn.transpose(1, 2).reshape(B, T, C)
        x = x + gate_cross.unsqueeze(1) * self.cross_proj(cross_attn)

        # 3. Feed Forward with Modulation
        norm_x = modulate(self.norm3(x), shift_mlp, scale_mlp)
        x = x + gate_mlp.unsqueeze(1) * self.ffn(norm_x)
        return x


class DiTBackbone(nn.Module):
    """
    Diffusion Transformer (DiT) Core Model (~148M Parameters in 9 layers).
    Enhanced for 2026 SOTA Voice Cloning and In-Context Audio Prompting.
    """
    def __init__(
        self,
        in_channels: int = 80,
        hidden_dim: int = 896,
        num_layers: int = 9,
        num_heads: int = 14,
        ffn_dim: int = 3584,
        cond_dim: int = 512,
        time_embed_dim: int = 512,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.hidden_dim = hidden_dim
        self.cond_dim = cond_dim

        # Input projection for audio mel / latents
        self.in_proj = nn.Linear(in_channels, hidden_dim)

        # Reference voice audio prompt encoder (projects 80-mel to cond_dim)
        self.ref_audio_proj = nn.Sequential(
            nn.Linear(in_channels, cond_dim),
            nn.SiLU(),
            nn.Linear(cond_dim, cond_dim),
        )

        # Null condition vector for Classifier-Free Guidance (CFG)
        self.null_cond = nn.Parameter(torch.randn(1, 1, cond_dim) * 0.02)

        # Timestep embedder
        self.time_embedder = TimestepEmbedder(time_embed_dim)
        self.time_proj = nn.Linear(time_embed_dim, cond_dim)

        # DiT Blocks
        self.blocks = nn.ModuleList([
            DiTBlock(hidden_dim, num_heads, cond_dim, ffn_dim, dropout)
            for _ in range(num_layers)
        ])

        # Final Layer (AdaLN-Zero + Linear Out)
        self.final_norm = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
        self.final_adaLN = nn.Sequential(
            nn.SiLU(),
            nn.Linear(cond_dim, 2 * hidden_dim, bias=True)
        )
        self.out_proj = nn.Linear(hidden_dim, in_channels)

        # Zero init final layer for ultra-stable flow matching training
        nn.init.constant_(self.final_adaLN[-1].weight, 0)
        nn.init.constant_(self.final_adaLN[-1].bias, 0)
        nn.init.constant_(self.out_proj.weight, 0)
        nn.init.constant_(self.out_proj.bias, 0)

    @staticmethod
    def get_sinusoidal_pos_embedding(seq_len: int, dim: int, device: torch.device) -> torch.Tensor:
        """Computes 1D sinusoidal positional embeddings for sequence length seq_len."""
        half_dim = dim // 2
        inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half_dim, device=device).float() / half_dim))
        t = torch.arange(seq_len, device=device, dtype=torch.float)
        freqs = torch.einsum("i,j->ij", t, inv_freq)
        emb = torch.cat((freqs.sin(), freqs.cos()), dim=-1)
        if dim % 2:
            emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
        return emb.unsqueeze(0)  # [1, seq_len, dim]

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        text_cond: torch.Tensor,
        ref_mel: Optional[torch.Tensor] = None,
        drop_cond: bool = False,
        cond_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Args:
            x_t: Noisy audio latents [B, T_audio, in_channels]
            t: Continuous timesteps [B] in [0, 1]
            text_cond: Text representation [B, T_text, cond_dim]
            ref_mel: Optional 3-5 sec human reference audio mel [B, T_ref, 80]
            drop_cond: If True, uses null condition for CFG sampling
            cond_mask: Optional boolean padding mask for text_cond [B, T_text]
        Returns:
            v_t: Predicted velocity field [B, T_audio, in_channels]
        """
        B = x_t.shape[0]

        # Time conditioning
        t_emb = self.time_proj(self.time_embedder(t))

        # Handle Classifier-Free Guidance (CFG) conditioning
        ref_latents = None
        if ref_mel is not None:
            ref_latents = self.ref_audio_proj(ref_mel)

        total_cond_len = text_cond.shape[1] + (ref_latents.shape[1] if ref_latents is not None else 0)

        if drop_cond:
            full_cond = self.null_cond.expand(B, total_cond_len, -1)
            full_mask = None
        else:
            if ref_latents is not None:
                full_cond = torch.cat([ref_latents, text_cond], dim=1)
                if cond_mask is not None:
                    ref_mask = torch.ones((B, ref_latents.shape[1]), dtype=torch.bool, device=cond_mask.device)
                    full_mask = torch.cat([ref_mask, cond_mask], dim=1)
                else:
                    full_mask = None
            else:
                full_cond = text_cond
                full_mask = cond_mask

        # Audio projection with 1D Temporal Positional Embeddings
        h = self.in_proj(x_t)
        pos_emb = self.get_sinusoidal_pos_embedding(h.shape[1], self.hidden_dim, h.device)
        h = h + pos_emb

        # Pass through DiT blocks
        for block in self.blocks:
            h = block(h, t_emb, full_cond, cond_mask=full_mask)

        # Final modulation and projection to target velocity
        shift, scale = self.final_adaLN(t_emb).chunk(2, dim=-1)
        h = modulate(self.final_norm(h), shift, scale)
        return self.out_proj(h)