freyatts/model.py

"""FreyaDiT: a non-autoregressive flow-matching DiT for Turkish TTS.

Latent frames self-attend with rotary position embeddings and cross-attend to
character-level text features refined by a small ConvNeXt stack. Trained with
an optimal-transport rectified-flow objective on frozen VoxCPM2 VAE latents.
"""

import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange



# spread of 2.3 Hz, against 14.9 Hz for the previous default of 0.
LEYLA_SEED = 9


def rope_freqs(dim, length, theta=10000.0, device=None):
    """Build rotary embedding angles of shape [length, dim]."""
    inv = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device).float() / dim))
    t = torch.arange(length, device=device).float()
    freqs = torch.outer(t, inv)
    return torch.cat([freqs, freqs], dim=-1)


def apply_rope(x, cos, sin):
    """Apply rotary position embedding to a [B, H, N, D] tensor."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    rotated = torch.cat([-x2, x1], dim=-1)
    return x * cos + rotated * sin


def additive_mask(mask, dtype):
    """Turn a bool key-padding mask [B, N] into an additive float mask [B, 1, 1, N].

    SDPA's memory-efficient backend wants an additive mask rather than bool.
    """
    if mask is None:
        return None
    out = torch.zeros(mask.shape[0], 1, 1, mask.shape[1], device=mask.device, dtype=dtype)
    return out.masked_fill(~mask[:, None, None, :], float("-inf"))


def modulate(x, shift, scale):
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)


def timestep_embed(t, dim, max_period=10000):
    """Sinusoidal embedding of a scalar diffusion time t in [0, 1]."""
    half = dim // 2
    freqs = torch.exp(-math.log(max_period) * torch.arange(half, device=t.device) / half)
    angles = t[:, None] * freqs[None]
    return torch.cat([angles.cos(), angles.sin()], dim=-1)


class SelfAttn(nn.Module):
    """Multi-head self-attention over latent frames, with RoPE."""

    def __init__(self, d, heads):
        super().__init__()
        self.h = heads
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.o = nn.Linear(d, d, bias=False)

    def forward(self, x, cos, sin, mask=None):
        q, k, v = [rearrange(t, "b n (h d) -> b h n d", h=self.h) for t in self.qkv(x).chunk(3, -1)]
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=additive_mask(mask, q.dtype))
        return self.o(rearrange(out, "b h n d -> b n (h d)"))


class CrossAttn(nn.Module):
    """Frames attend to character features. No positional encoding on the text side;
    alignment is learned through attention."""

    def __init__(self, d, heads):
        super().__init__()
        self.h = heads
        self.q = nn.Linear(d, d, bias=False)
        self.kv = nn.Linear(d, 2 * d, bias=False)
        self.o = nn.Linear(d, d, bias=False)

    def forward(self, x, ctx, ctx_mask=None):
        q = rearrange(self.q(x), "b n (h d) -> b h n d", h=self.h)
        k, v = [rearrange(t, "b n (h d) -> b h n d", h=self.h) for t in self.kv(ctx).chunk(2, -1)]
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=additive_mask(ctx_mask, q.dtype))
        return self.o(rearrange(out, "b h n d -> b n (h d)"))


class SwiGLU(nn.Module):
    def __init__(self, d, ff):
        super().__init__()
        self.w1 = nn.Linear(d, ff, bias=False)
        self.w2 = nn.Linear(d, ff, bias=False)
        self.w3 = nn.Linear(ff, d, bias=False)

    def forward(self, x):
        return self.w3(F.silu(self.w1(x)) * self.w2(x))


class Block(nn.Module):
    """DiT block: self-attention, cross-attention, SwiGLU FFN.

    Each of the three sub-layers is gated by adaLN-zero, so a single timestep
    conditioning vector produces nine modulation signals (shift/scale/gate x3).
    Gates start at zero so the block is initialized to identity.
    """

    def __init__(self, d, heads, ff):
        super().__init__()
        self.n1 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
        self.sa = SelfAttn(d, heads)
        self.nx = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
        self.xa = CrossAttn(d, heads)
        self.n2 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
        self.ff = SwiGLU(d, ff)
        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(d, 9 * d))
        nn.init.zeros_(self.ada[-1].weight)
        nn.init.zeros_(self.ada[-1].bias)

    def forward(self, x, c, ctx, cos, sin, fmask=None, cmask=None):
        s1, b1, g1, sx, bx, gx, s2, b2, g2 = self.ada(c).chunk(9, -1)
        x = x + g1.unsqueeze(1) * self.sa(modulate(self.n1(x), b1, s1), cos, sin, fmask)
        x = x + gx.unsqueeze(1) * self.xa(modulate(self.nx(x), bx, sx), ctx, cmask)
        x = x + g2.unsqueeze(1) * self.ff(modulate(self.n2(x), b2, s2))
        return x


class ConvNeXt1d(nn.Module):
    """1-D ConvNeXt block used to refine character embeddings before cross-attention."""

    def __init__(self, d, mult=2):
        super().__init__()
        self.dw = nn.Conv1d(d, d, 7, padding=3, groups=d)
        self.n = nn.LayerNorm(d)
        self.p1 = nn.Linear(d, d * mult)
        self.p2 = nn.Linear(d * mult, d)

    def forward(self, x):
        residual = x
        x = self.dw(x.transpose(1, 2)).transpose(1, 2)
        x = self.n(x)
        return residual + self.p2(F.gelu(self.p1(x)))


class FreyaDiT(nn.Module):
    """FreyaTTS acoustic model.

    Predicts the rectified-flow velocity field over 64-dim VoxCPM2 latent
    frames, conditioned on character ids. Also carries a small duration head
    that regresses log frame count from mean-pooled text features.
    """

    def __init__(self, vocab, feat=64, d=768, depth=22, heads=12, ff=2048, text_conv=4, fill_id=0):
        super().__init__()
        self.feat = feat
        self.d = d
        self.heads = heads
        self.char_emb = nn.Embedding(vocab, d)
        self.text_conv = nn.ModuleList([ConvNeXt1d(d) for _ in range(text_conv)])
        self.x_proj = nn.Linear(feat, d)
        self.t_mlp = nn.Sequential(nn.Linear(d, d), nn.SiLU(), nn.Linear(d, d))
        self.blocks = nn.ModuleList([Block(d, heads, ff) for _ in range(depth)])
        self.nf = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
        self.ada_f = nn.Sequential(nn.SiLU(), nn.Linear(d, 2 * d))
        nn.init.zeros_(self.ada_f[-1].weight)
        nn.init.zeros_(self.ada_f[-1].bias)
        self.out = nn.Linear(d, feat)
        nn.init.zeros_(self.out.weight)
        nn.init.zeros_(self.out.bias)
        self.dur = nn.Sequential(nn.Linear(d, d), nn.SiLU(), nn.Linear(d, 1))

    def text_encode(self, text_ids):
        """Embed character ids and refine them with the ConvNeXt stack."""
        x = self.char_emb(text_ids)
        for conv in self.text_conv:
            x = conv(x)
        return x

    def forward(self, x_t, t, ctx, fmask=None, cmask=None):
        """Predict the velocity field at noisy latents x_t and time t."""
        B, T, _ = x_t.shape
        head_dim = self.d // self.heads
        base = rope_freqs(head_dim, T, device=x_t.device)
        cos = torch.cos(base)[None, None]
        sin = torch.sin(base)[None, None]

        h = self.x_proj(x_t)
        c = self.t_mlp(timestep_embed(t, self.d))
        for blk in self.blocks:
            h = blk(h, c, ctx, cos, sin, fmask, cmask)

        s, g = self.ada_f(c).chunk(2, -1)
        return self.out(modulate(self.nf(h), s, g))

    def cfm_loss(self, x1, text_ids, fmask=None, cmask=None):
        """Conditional flow-matching loss on clean latents x1 (masked mean if fmask given)."""
        B = x1.shape[0]
        x0 = torch.randn_like(x1)
        t = torch.rand(B, device=x1.device)
        xt = (1 - t[:, None, None]) * x0 + t[:, None, None] * x1
        ctx = self.text_encode(text_ids)
        v = self(xt, t, ctx, fmask, cmask)
        target = x1 - x0
        if fmask is None:
            return ((v - target) ** 2).mean()
        m = fmask[..., None].float()
        return (((v - target) ** 2) * m).sum() / (m.sum() * self.feat + 1e-6)

    def dur_loss(self, text_ids, logT, cmask=None):
        """MSE loss of the duration head against log frame counts."""
        te = self.text_encode(text_ids)
        if cmask is not None:
            pooled = (te * cmask[..., None].float()).sum(1) / (cmask.sum(1, keepdim=True) + 1e-6)
        else:
            pooled = te.mean(1)
        pred = self.dur(pooled).squeeze(-1)
        return ((pred - logT) ** 2).mean(), pred

    @torch.no_grad()
    def sample(self, text_ids, T, steps=32, cmask=None, seed=LEYLA_SEED):
        """Integrate the ODE from noise to latents with a fixed-step Euler solver."""
        B = text_ids.shape[0]
        ctx = self.text_encode(text_ids)
        device = text_ids.device
        if seed is None:
            x = torch.randn(B, T, self.feat, device=device)
        else:
            gen = torch.Generator(device=device).manual_seed(int(seed))
            x = torch.randn(B, T, self.feat, device=device, generator=gen)
        for i in range(steps):
            t = torch.full((B,), i / steps, device=x.device)
            x = x + self(x, t, ctx, None, cmask) / steps
        return x