| """Host (torch, CPU) side of the DiffRhythm 1.2 DiT port to one Tenstorrent P100a. |
| |
| Upstream (DiffRhythm model/dit.py DiT.forward) as used by CFM.sample, per call: |
| c = time_embed(t) + start_time_embed(start) + duration_time_embed(duration) [512] |
| text_embed = TextEmbedding(text) (embedding + sinusoid + 4 ConvNeXtV2 blocks; zeros in for drop_text) [S, 512] |
| h = proj(cat(x, cond, text_embed, style, c)) -> [S, 2048]; h = ConvPos(h) + h (two k31 groups-16 convs + Mish) |
| 16 x LlamaDecoderLayer (RMSNorm, 32 x 64 heads, RoPE theta 1e4, SwiGLU 8192; no mask, bidirectional); |
| after each of the first 8: h += silu(text_fusion_linears[i](text_embed)) |
| out = proj_out(LayerNorm(h) * (1 + scale) + shift), scale, shift = linear(silu(c)).chunk(2) |
| The CFG branch (drop_audio_cond + drop_text) zeroes cond and the text tokens and uses the negative style. |
| |
| Split: the timestep-invariant text embedding and the small time MLPs stay in the DiffRhythm process on the host |
| (upstream modules, exact); everything weight-heavy runs on the card. The input projection is split by input: |
| h = x @ Wx + [cond @ Wcond + text_embed @ Wtext] (per song and branch) + [style @ Ws + c @ Wc + b] (per call row) |
| RefDecoder runs this formulation in float32 so it can be checked against upstream without a device. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Dict |
|
|
| import torch |
|
|
| from .common import (NEG, TILE, apply_rope_adjacent, interleave_pairs_permutation, linear_to_mm, |
| permute_heads_rows, rms, round_up, swiglu_interleave) |
|
|
| SEQ_BUCKET = 256 |
| LATENT_RATE = 44100 / 2048 |
| MIN_FRAMES, MAX_FRAMES = int(96 * LATENT_RATE), int(285 * LATENT_RATE) |
|
|
|
|
| def seq_buckets(): |
| """Every padded length a DiffRhythm song can have (warm_cache.py compiles them ahead of time).""" |
| return sorted({round_up(n, SEQ_BUCKET) for n in range(MIN_FRAMES, MAX_FRAMES + 1)}) |
|
|
|
|
| @dataclass(frozen=True) |
| class Config: |
| dim: int = 2048 |
| depth: int = 16 |
| heads: int = 32 |
| head_dim: int = 64 |
| intermediate: int = 8192 |
| text_dim: int = 512 |
| cond_dim: int = 512 |
| mel_dim: int = 64 |
| fusion_layers: int = 8 |
| conv_kernel: int = 31 |
| conv_groups: int = 16 |
| rope_theta: float = 10000.0 |
| eps: float = 1e-6 |
|
|
|
|
| class Checkpoint: |
| """transformer.* of DiffRhythm's cfm_model.pt (fp16 on disk), read as float32.""" |
|
|
| def __init__(self, path): |
| sd = torch.load(str(path), weights_only=True, map_location="cpu")["model_state_dict"] |
| self._sd = {k[len("transformer."):]: v for k, v in sd.items() if k.startswith("transformer.")} |
| self.cfg = Config() |
|
|
| def get(self, name: str) -> torch.Tensor: |
| return self._sd[name].float() |
|
|
|
|
| @dataclass |
| class LayerWeights: |
| wqkv: torch.Tensor |
| wo: torch.Tensor |
| w_gateup: torch.Tensor |
| w_down: torch.Tensor |
| ln1: torch.Tensor |
| ln2: torch.Tensor |
|
|
|
|
| def layer_weights(ckpt: Checkpoint, i: int) -> LayerWeights: |
| c = ckpt.cfg |
| g = lambda k: ckpt.get(f"transformer_blocks.{i}.{k}") |
| perm = interleave_pairs_permutation(c.head_dim) |
| wq = permute_heads_rows(g("self_attn.q_proj.weight"), c.heads, c.head_dim, perm) |
| wk = permute_heads_rows(g("self_attn.k_proj.weight"), c.heads, c.head_dim, perm) |
| return LayerWeights( |
| wqkv=torch.cat([linear_to_mm(wq), linear_to_mm(wk), linear_to_mm(g("self_attn.v_proj.weight"))], dim=1), |
| wo=linear_to_mm(g("self_attn.o_proj.weight")), |
| w_gateup=swiglu_interleave(g("mlp.gate_proj.weight"), g("mlp.up_proj.weight")), |
| w_down=linear_to_mm(g("mlp.down_proj.weight")), |
| ln1=g("input_layernorm.weight"), ln2=g("post_attention_layernorm.weight")) |
|
|
|
|
| @dataclass |
| class InputWeights: |
| """input_embed.proj split by input (matmul layout [in, out]).""" |
|
|
| wx: torch.Tensor |
| wcond: torch.Tensor |
| wtext: torch.Tensor |
| wstyle: torch.Tensor |
| wc: torch.Tensor |
| b: torch.Tensor |
|
|
|
|
| def input_weights(ckpt: Checkpoint) -> InputWeights: |
| c = ckpt.cfg |
| w = linear_to_mm(ckpt.get("input_embed.proj.weight")) |
| edges = [0, c.mel_dim, 2 * c.mel_dim, 2 * c.mel_dim + c.text_dim, 2 * c.mel_dim + c.text_dim + c.cond_dim] |
| parts = [w[a:b] for a, b in zip(edges, edges[1:] + [w.shape[0]])] |
| return InputWeights(*[p.contiguous() for p in parts], ckpt.get("input_embed.proj.bias")) |
|
|
|
|
| @dataclass(frozen=True) |
| class Geometry: |
| frames: int |
|
|
| @property |
| def seq_pad(self) -> int: |
| return round_up(self.frames, SEQ_BUCKET) |
|
|
|
|
| def key_mask(geo: Geometry) -> torch.Tensor: |
| """Additive [S_pad, S_pad] mask excluding padded keys (upstream attends over exactly S frames).""" |
| valid = torch.arange(geo.seq_pad)[None, :] < geo.frames |
| return torch.where(valid, 0.0, NEG).expand(geo.seq_pad, -1).contiguous() |
|
|
|
|
| def row_mask(geo: Geometry) -> torch.Tensor: |
| """[S_pad, 1]: 1 on real frames, 0 on padding (the position convs must see zeros past the song, like upstream).""" |
| return (torch.arange(geo.seq_pad) < geo.frames).float()[:, None] |
|
|
|
|
| def rope_tables(seq_pad: int, cfg: Config, dtype=torch.float32): |
| inv = 1.0 / (cfg.rope_theta ** (torch.arange(0, cfg.head_dim, 2, dtype=torch.float64) / cfg.head_dim)) |
| ang = torch.outer(torch.arange(seq_pad, dtype=torch.float64), inv).float() |
| return torch.cos(ang).repeat_interleave(2, -1).to(dtype), torch.sin(ang).repeat_interleave(2, -1).to(dtype) |
|
|
|
|
| def pad_rows(t: torch.Tensor, rows: int) -> torch.Tensor: |
| out = torch.zeros(rows, t.shape[-1], dtype=t.dtype) |
| out[: t.shape[0]] = t |
| return out |
|
|
|
|
| @dataclass |
| class CallRows: |
| """What changes per call: the style/time part of the input projection and the output modulation.""" |
|
|
| input_row: torch.Tensor |
| out_scale: torch.Tensor |
| out_shift: torch.Tensor |
|
|
|
|
| def call_rows(ckpt_or_iw, norm_linear_w, norm_linear_b, style: torch.Tensor, c: torch.Tensor) -> CallRows: |
| """style [512], c [512] (time + start + duration embeddings, from upstream on the host).""" |
| iw = ckpt_or_iw |
| input_row = style @ iw.wstyle + c @ iw.wc + iw.b |
| emb = torch.nn.functional.silu(c) @ norm_linear_w.t() + norm_linear_b |
| scale, shift = emb.chunk(2, dim=-1) |
| return CallRows(input_row, scale, shift) |
|
|
|
|
| class RefDecoder: |
| """float32 CPU DiT in the TT formulation.""" |
|
|
| def __init__(self, ckpt: Checkpoint): |
| self.ckpt, self.cfg = ckpt, ckpt.cfg |
| c = self.cfg |
| self.layers = [layer_weights(ckpt, i) for i in range(c.depth)] |
| self.iw = input_weights(ckpt) |
| self.conv = [(ckpt.get(f"input_embed.conv_pos_embed.conv1d.{i}.weight"), |
| ckpt.get(f"input_embed.conv_pos_embed.conv1d.{i}.bias")) for i in (0, 2)] |
| self.fusion = [(linear_to_mm(ckpt.get(f"text_fusion_linears.{i}.0.weight")), |
| ckpt.get(f"text_fusion_linears.{i}.0.bias")) for i in range(c.fusion_layers)] |
| self.w_out, self.b_out = linear_to_mm(ckpt.get("proj_out.weight")), ckpt.get("proj_out.bias") |
| self.norm_w, self.norm_b = ckpt.get("norm_out.linear.weight"), ckpt.get("norm_out.linear.bias") |
|
|
| def rows(self, style: torch.Tensor, c: torch.Tensor) -> CallRows: |
| return call_rows(self.iw, self.norm_w, self.norm_b, style, c) |
|
|
| def branch(self, text_embed: torch.Tensor, cond, geo: Geometry) -> Dict[str, object]: |
| """Per song and CFG branch: [cond @ Wcond + text @ Wtext] and the 8 text-fusion residuals, padded.""" |
| t = pad_rows(text_embed.float(), geo.seq_pad) |
| base = t @ self.iw.wtext |
| if cond is not None: |
| base = base + pad_rows(cond.float(), geo.seq_pad) @ self.iw.wcond |
| fusion = [torch.nn.functional.silu(t @ w + b) for w, b in self.fusion] |
| return {"base": base, "fusion": fusion} |
|
|
| def conv_pos(self, h: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor: |
| """Two k31 convs with Mish. Upstream's sequence ends at the song, so each conv sees zeros past it: the rows |
| past the song are zeroed before every conv (the first conv's Mish output is not zero there).""" |
| c = self.cfg |
| y = h |
| for w, b in self.conv: |
| if mask is not None: |
| y = y * mask |
| y = torch.nn.functional.mish(torch.nn.functional.conv1d(y.t()[None], w, b, padding=c.conv_kernel // 2, |
| groups=c.conv_groups))[0].t() |
| return y |
|
|
| def forward(self, x_t: torch.Tensor, branch: Dict[str, object], rows: CallRows, geo: Geometry, taps: list = None): |
| """[S, 64] noisy latents -> [S, 64] prediction.""" |
| c = self.cfg |
| mask = row_mask(geo) |
| h = (pad_rows(x_t.float(), geo.seq_pad) @ self.iw.wx + branch["base"] + rows.input_row) * mask |
| h = self.conv_pos(h, mask) + h |
| cos, sin = rope_tables(geo.seq_pad, c) |
| kmask = key_mask(geo) |
| for i, w in enumerate(self.layers): |
| a = rms(h, w.ln1, c.eps) |
| q, k, v = (a @ w.wqkv).split(c.dim, -1) |
| q = apply_rope_adjacent(q.view(-1, c.heads, c.head_dim).transpose(0, 1), cos, sin) |
| k = apply_rope_adjacent(k.view(-1, c.heads, c.head_dim).transpose(0, 1), cos, sin) |
| v = v.view(-1, c.heads, c.head_dim).transpose(0, 1) |
| scores = q @ k.transpose(-1, -2) * c.head_dim ** -0.5 + kmask |
| att = (torch.softmax(scores, -1) @ v).transpose(0, 1).reshape(geo.seq_pad, -1) |
| h = h + att @ w.wo |
| m = rms(h, w.ln2, c.eps) |
| gu = (m @ w.w_gateup).view(geo.seq_pad, -1, 2, TILE) |
| h = h + (torch.nn.functional.silu(gu[:, :, 0]) * gu[:, :, 1]).reshape(geo.seq_pad, -1) @ w.w_down |
| if i < c.fusion_layers: |
| h = h + branch["fusion"][i] |
| if taps is not None: |
| taps.append(h.clone()) |
| h = torch.nn.functional.layer_norm(h, (c.dim,), eps=c.eps) * (1 + rows.out_scale) + rows.out_shift |
| return (h @ self.w_out + self.b_out)[: geo.frames] |
|
|
