Vāgdhenu achitecture transformer

Vāgdhenu (https://huggingface.co/prathoshap/vagdhenu) is a Sanskrit chant TTS model. Its backbone is IndicF5 / F5-TTS: a flow-matching Diffusion Transformer (DiT) that does mel-spectrogram infilling.

Key specs (production backbone)

Sanskrit is routed through Kannada script. Prosody comes almost entirely from the reference audio (half-reference rule) + a voice-steering fine-tune. Vocoder is fine-tuned BigVGAN-v2.

Below is a complete, self-contained, educational PyTorch implementation of the DiT backbone used by Vāgdhenu / F5-TTS Base. It is deliberately simplified for clarity while preserving the real architecture (AdaLN-Zero, RoPE, ConvNeXt text path, concatenated noised+cond+text input, flow-matching style).

pythonCopy
1"""
2Educational DiT backbone matching Vāgdhenu / F5-TTS Base
3(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
4
5Based on:
6  - F5-TTS (Chen et al., 2024)  https://github.com/SWivid/F5-TTS
7  - IndicF5 (AI4Bharat)
8  - DiT (Peebles & Xie, 2023) with AdaLN-Zero
9"""
10
11import math
12import torch
13import torch.nn as nn
14import torch.nn.functional as F
15from typing import Optional, Tuple
16
17
18# ---------------------------------------------------------------------------
19# 1. Rotary Position Embedding (RoPE)
20# ---------------------------------------------------------------------------
21class RotaryEmbedding(nn.Module):
22    def __init__(self, dim: int, base: float = 10000.0):
23        super().__init__()
24        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
25        self.register_buffer("inv_freq", inv_freq, persistent=False)
26
27    def forward(self, seq_len: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
28        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
29        freqs = torch.einsum("i,j->ij", t, self.inv_freq)          # (T, dim/2)
30        emb = torch.cat((freqs, freqs), dim=-1)                    # (T, dim)
31        return emb.cos()[None, None, :, :], emb.sin()[None, None, :, :]  # (1,1,T,dim)
32
33
34def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
35    """x: (B, H, T, D)"""
36    x1, x2 = x.chunk(2, dim=-1)
37    return torch.cat((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1)
38
39
40# ---------------------------------------------------------------------------
41# 2. Timestep Embedding (sinusoidal + MLP)
42# ---------------------------------------------------------------------------
43class TimestepEmbedding(nn.Module):
44    def __init__(self, dim: int):
45        super().__init__()
46        self.mlp = nn.Sequential(
47            nn.Linear(dim, dim * 4),
48            nn.SiLU(),
49            nn.Linear(dim * 4, dim),
50        )
51
52    def forward(self, t: torch.Tensor) -> torch.Tensor:
53        """t: (B,) continuous time in [0,1]"""
54        half = self.mlp[0].in_features // 2
55        freqs = torch.exp(
56            -math.log(10000) * torch.arange(half, device=t.device) / half
57        )
58        args = t[:, None] * freqs[None]
59        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
60        return self.mlp(emb)                                       # (B, dim)
61
62
63# ---------------------------------------------------------------------------
64# 3. Adaptive LayerNorm (AdaLN-Zero) – core of DiT
65# ---------------------------------------------------------------------------
66class AdaLayerNormZero(nn.Module):
67    """Produces scale, shift, gate for attention and FFN (6 parameters)."""
68    def __init__(self, dim: int):
69        super().__init__()
70        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
71        self.linear = nn.Linear(dim, dim * 6)
72
73    def forward(self, x: torch.Tensor, emb: torch.Tensor):
74        # emb comes from timestep embedding
75        emb = self.linear(F.silu(emb))
76        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1)
77        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
78        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
79
80
81class AdaLayerNorm_Final(nn.Module):
82    def __init__(self, dim: int):
83        super().__init__()
84        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
85        self.linear = nn.Linear(dim, dim * 2)
86
87    def forward(self, x: torch.Tensor, emb: torch.Tensor):
88        emb = self.linear(F.silu(emb))
89        scale, shift = emb.chunk(2, dim=1)
90        return self.norm(x) * (1 + scale[:, None]) + shift[:, None]
91
92
93# ---------------------------------------------------------------------------
94# 4. Multi-Head Self-Attention with RoPE
95# ---------------------------------------------------------------------------
96class Attention(nn.Module):
97    def __init__(self, dim: int, heads: int = 16, dim_head: int = 64, dropout: float = 0.1):
98        super().__init__()
99        self.heads = heads
100        self.dim_head = dim_head
101        inner = heads * dim_head
102        self.to_qkv = nn.Linear(dim, inner * 3, bias=False)
103        self.to_out = nn.Sequential(nn.Linear(inner, dim), nn.Dropout(dropout))
104
105    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
106                mask: Optional[torch.Tensor] = None) -> torch.Tensor:
107        B, T, _ = x.shape
108        qkv = self.to_qkv(x).chunk(3, dim=-1)
109        q, k, v = map(lambda t: t.view(B, T, self.heads, self.dim_head).transpose(1, 2), qkv)
110
111        q = apply_rotary(q, cos, sin)
112        k = apply_rotary(k, cos, sin)
113
114        # scaled dot-product
115        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
116        out = out.transpose(1, 2).reshape(B, T, -1)
117        return self.to_out(out)
118
119
120# ---------------------------------------------------------------------------
121# 5. Feed-Forward (GEGLU style)
122# ---------------------------------------------------------------------------
123class FeedForward(nn.Module):
124    def __init__(self, dim: int, mult: int = 2, dropout: float = 0.1):
125        super().__init__()
126        inner = int(dim * mult * 2 / 3)          # GEGLU convention
127        self.net = nn.Sequential(
128            nn.Linear(dim, inner * 2),
129            nn.GELU(),                           # approximate GEGLU
130            nn.Dropout(dropout),
131            nn.Linear(inner, dim),
132            nn.Dropout(dropout),
133        )
134
135    def forward(self, x):
136        return self.net(x)
137
138
139# ---------------------------------------------------------------------------
140# 6. DiT Block (AdaLN-Zero + Attention + FFN)
141# ---------------------------------------------------------------------------
142class DiTBlock(nn.Module):
143    def __init__(self, dim: int, heads: int = 16, dim_head: int = 64,
144                 ff_mult: int = 2, dropout: float = 0.1):
145        super().__init__()
146        self.attn_norm = AdaLayerNormZero(dim)
147        self.attn = Attention(dim, heads, dim_head, dropout)
148        self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
149        self.ff = FeedForward(dim, ff_mult, dropout)
150
151    def forward(self, x: torch.Tensor, t_emb: torch.Tensor,
152                cos: torch.Tensor, sin: torch.Tensor,
153                mask: Optional[torch.Tensor] = None) -> torch.Tensor:
154        # Attention path
155        norm_x, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, t_emb)
156        attn_out = self.attn(norm_x, cos, sin, mask)
157        x = x + gate_msa.unsqueeze(1) * attn_out
158
159        # FFN path
160        norm_x = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
161        ff_out = self.ff(norm_x)
162        x = x + gate_mlp.unsqueeze(1) * ff_out
163        return x
164
165
166# ---------------------------------------------------------------------------
167# 7. Text Embedding + ConvNeXt V2 (the “alignment helper” of F5-TTS)
168# ---------------------------------------------------------------------------
169class ConvNeXtV2Block(nn.Module):
170    def __init__(self, dim: int, intermediate_dim: int):
171        super().__init__()
172        self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim)
173        self.norm = nn.LayerNorm(dim, eps=1e-6)
174        self.pwconv1 = nn.Linear(dim, intermediate_dim)
175        self.act = nn.GELU()
176        self.pwconv2 = nn.Linear(intermediate_dim, dim)
177        self.gamma = nn.Parameter(torch.ones(dim) * 1e-6)   # GRN-style scale
178
179    def forward(self, x: torch.Tensor) -> torch.Tensor:
180        # x: (B, T, C)
181        residual = x
182        x = x.transpose(1, 2)
183        x = self.dwconv(x).transpose(1, 2)
184        x = self.norm(x)
185        x = self.pwconv1(x)
186        x = self.act(x)
187        x = self.pwconv2(x)
188        return residual + self.gamma * x
189
190
191class TextEmbedding(nn.Module):
192    def __init__(self, text_num_embeds: int, text_dim: int, conv_layers: int = 4):
193        super().__init__()
194        self.embed = nn.Embedding(text_num_embeds + 1, text_dim)   # +1 for filler
195        self.conv_layers = nn.ModuleList([
196            ConvNeXtV2Block(text_dim, text_dim * 2) for _ in range(conv_layers)
197        ])
198
199    def forward(self, text: torch.Tensor) -> torch.Tensor:
200        x = self.embed(text)
201        for block in self.conv_layers:
202            x = block(x)
203        return x
204
205
206# ---------------------------------------------------------------------------
207# 8. Input Embedding (noised mel + masked cond + text → dim)
208# ---------------------------------------------------------------------------
209class InputEmbedding(nn.Module):
210    def __init__(self, mel_dim: int, text_dim: int, dim: int):
211        super().__init__()
212        self.proj = nn.Linear(mel_dim * 2 + text_dim, dim)
213        self.conv_pos = nn.Conv1d(dim, dim, kernel_size=31, padding=15, groups=dim)
214
215    def forward(self, x: torch.Tensor, cond: torch.Tensor, text_emb: torch.Tensor) -> torch.Tensor:
216        # x, cond: (B, T, mel_dim)   text_emb: (B, T, text_dim)
217        x = torch.cat([x, cond, text_emb], dim=-1)
218        x = self.proj(x)
219        x = x + self.conv_pos(x.transpose(1, 2)).transpose(1, 2)
220        return x
221
222
223# ---------------------------------------------------------------------------
224# 9. Full DiT (the heart of Vāgdhenu / F5-TTS)
225# ---------------------------------------------------------------------------
226class DiT(nn.Module):
227    """
228    Diffusion Transformer backbone used by Vāgdhenu.
229
230    Config matching production:
231        dim=1024, depth=22, heads=16, ff_mult=2,
232        text_dim=512, conv_layers=4, mel_dim=100
233    """
234    def __init__(
235        self,
236        dim: int = 1024,
237        depth: int = 22,
238        heads: int = 16,
239        dim_head: int = 64,
240        ff_mult: int = 2,
241        mel_dim: int = 100,
242        text_num_embeds: int = 256,          # vocab size (IndicF5 uses its own)
243        text_dim: int = 512,
244        conv_layers: int = 4,
245        dropout: float = 0.1,
246    ):
247        super().__init__()
248        self.dim = dim
249        self.time_embed = TimestepEmbedding(dim)
250        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers)
251        self.input_embed = InputEmbedding(mel_dim, text_dim, dim)
252        self.rotary = RotaryEmbedding(dim_head)
253
254        self.blocks = nn.ModuleList([
255            DiTBlock(dim, heads, dim_head, ff_mult, dropout)
256            for _ in range(depth)
257        ])
258        self.norm_out = AdaLayerNorm_Final(dim)
259        self.proj_out = nn.Linear(dim, mel_dim)
260
261        self._init_weights()
262
263    def _init_weights(self):
264        # AdaLN-Zero: final linear of each AdaLN starts at zero
265        for block in self.blocks:
266            nn.init.zeros_(block.attn_norm.linear.weight)
267            nn.init.zeros_(block.attn_norm.linear.bias)
268        nn.init.zeros_(self.norm_out.linear.weight)
269        nn.init.zeros_(self.norm_out.linear.bias)
270        nn.init.zeros_(self.proj_out.weight)
271        nn.init.zeros_(self.proj_out.bias)
272
273    def forward(
274        self,
275        x: torch.Tensor,          # (B, T, mel_dim)  noised mel
276        cond: torch.Tensor,       # (B, T, mel_dim)  masked reference mel
277        text: torch.Tensor,       # (B, T)           token ids (padded with filler)
278        time: torch.Tensor,       # (B,)             continuous t ∈ [0,1]
279        mask: Optional[torch.Tensor] = None,
280    ) -> torch.Tensor:
281        """
282        Returns predicted velocity / noise  (B, T, mel_dim)
283        used by the flow-matching ODE solver.
284        """
285        t_emb = self.time_embed(time)                       # (B, dim)
286        text_emb = self.text_embed(text)                    # (B, T, text_dim)
287        x = self.input_embed(x, cond, text_emb)             # (B, T, dim)
288
289        cos, sin = self.rotary(x.shape[1], x.device)
290
291        for block in self.blocks:
292            x = block(x, t_emb, cos, sin, mask)
293
294        x = self.norm_out(x, t_emb)
295        return self.proj_out(x)
296
297
298# ---------------------------------------------------------------------------
299# 10. Minimal usage example (for learning)
300# ---------------------------------------------------------------------------
301if __name__ == "__main__":
302    device = "cuda" if torch.cuda.is_available() else "cpu"
303
304    model = DiT(
305        dim=1024,
306        depth=22,
307        heads=16,
308        ff_mult=2,
309        text_dim=512,
310        conv_layers=4,
311        mel_dim=100,
312        text_num_embeds=256,
313    ).to(device)
314
315    B, T = 2, 256                       # batch, mel frames
316    x     = torch.randn(B, T, 100, device=device)   # noised mel
317    cond  = torch.randn(B, T, 100, device=device)   # masked reference
318    text  = torch.randint(0, 256, (B, T), device=device)
319    time  = torch.rand(B, device=device)
320
321    with torch.no_grad():
322        pred = model(x, cond, text, time)
323    print("Output shape:", pred.shape)          # (2, 256, 100)
324    print("Parameters:  ", sum(p.numel() for p in model.parameters()) / 1e6, "M")

How the real system uses this DiT

1. Training – Optimal-Transport Conditional Flow Matching (OT-CFM) on the text-guided speech-infilling task.
2. Inference – Solve the ODE with a small number of function evaluations (NFE ≈ 32–64) + classifier-free guidance; the reference clip supplies voice + swara + pace.
3. Vocoder – Fine-tuned BigVGAN-v2 turns the predicted mel into 24 kHz audio.
4. Text frontend (Vāgdhenu-specific) – Devanagari → SLP1 → Kannada routing + sandhi + meter/gaṇa handling.

Learning tips

• Start by running the tiny example above and inspect intermediate shapes.
• Replace depth=22 with depth=4 and dim=256 while debugging.
• The most important ideas are AdaLN-Zero (timestep controls every residual path) and the concatenated noised+cond+text input (no separate duration model).
• Official code lives in the F5-TTS / IndicF5 repositories; Vāgdhenu only fine-tunes the weights and adds the Sanskrit frontend.

This implementation is intentionally pure-PyTorch and comment-heavy so you can understand every moving part of the transformer that powers Vāgdhenu.