| """Per-level teacher-forced depth graph (one ONNX graph per level 1..15). |
| |
| Each graph is the upstream DepthV2 computation for one level, with the same classes (_SelfBlock, _CrossBlock) and the |
| same prefix (_prefix) and last-position read-out (_run) as tts/model/depth_v2.py, mask None (the upstream kmask is |
| all True for named voices, so padding is never masked). It holds only the weights that level needs: its weight group |
| (4 self-attention blocks), the 4 shared cross-attention blocks, the code embeddings, the input projection and the |
| level's head. The 15 graphs are chained by the runtime; the group's weights are not shared between levels 9..15 here |
| (each level graph carries group 8), which costs bandwidth but keeps every part small. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from tts.model.depth_v2 import D, GROUP_OF_LEVEL, LAYERS, RANK, _CrossBlock, _SelfBlock |
| from tts.tokens import N_CODES, N_Q |
|
|
| Q_SERVE = 16 |
|
|
|
|
| class DepthLevel(nn.Module): |
| def __init__(self, level: int, n_codes: int = N_CODES, d_backbone: int = 2048) -> None: |
| super().__init__() |
| if not 1 <= level <= N_Q - 1: |
| raise ValueError(f"level {level} outside 1..{N_Q - 1}") |
| self.level = level |
| self.group = GROUP_OF_LEVEL[level] |
| self.n_codes = n_codes |
| self.in_proj = nn.Linear(d_backbone, D, bias=False) |
| self.code_lo = nn.Embedding((N_Q - 1) * n_codes, RANK) |
| self.code_up = nn.Linear(RANK, D, bias=False) |
| self.level_emb = nn.Embedding(N_Q, D) |
| self.trunk = nn.ModuleList(_SelfBlock() for _ in range(LAYERS)) |
| self.cross = nn.ModuleList(_CrossBlock() for _ in range(LAYERS)) |
| self.norm = nn.LayerNorm(D) |
| self.head = nn.Linear(D, n_codes, bias=False) |
|
|
| def forward(self, h: torch.Tensor, codes: torch.Tensor, mem: torch.Tensor) -> torch.Tensor: |
| """h (B, 2048) fp32 backbone hidden; codes (B, Q_SERVE) int64 with columns cb0..cb_{L-1} used; mem (B, 1, 1024). |
| Returns the level's logits (B, 2048).""" |
| L = self.level |
| ks = torch.arange(L, device=h.device) |
| flat = codes[:, :L] + ks * self.n_codes |
| e = self.code_up(self.code_lo(flat)) + self.level_emb(ks + 1)[None] |
| x = torch.cat([self.in_proj(h)[:, None] + self.level_emb.weight[:1], e], 1) |
| T = x.shape[1] |
| mask = torch.ones(T, T, dtype=torch.bool, device=x.device).tril()[None, None] |
| for li in range(LAYERS): |
| x = self.trunk[li](x, mask) |
| x = self.cross[li](x, mem, None) |
| return self.head(self.norm(x[:, -1])) |
|
|
|
|
| def level_group(level: int) -> int: |
| return GROUP_OF_LEVEL[level] |
|
|
