code/tt_diffrhythm/diffrhythm_host.py

"""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  # frames padded to multiples of this: 96-285 s songs (2067-6136 frames) -> 16 shapes
LATENT_RATE = 44100 / 2048  # latent frames per second (get_lrc_token: int(seconds * 44100 / 2048))
MIN_FRAMES, MAX_FRAMES = int(96 * LATENT_RATE), int(285 * LATENT_RATE)  # the 96-285 s the port was exercised over


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  # [2048, 3 * 2048], q/k rows permuted to adjacent pairs
    wo: torch.Tensor
    w_gateup: torch.Tensor  # [2048, 2 * 8192] tile-interleaved
    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  # [64, 2048]
    wcond: torch.Tensor  # [64, 2048]
    wtext: torch.Tensor  # [512, 2048]
    wstyle: torch.Tensor  # [512, 2048]
    wc: torch.Tensor  # [512, 2048]
    b: torch.Tensor


def input_weights(ckpt: Checkpoint) -> InputWeights:
    c = ckpt.cfg
    w = linear_to_mm(ckpt.get("input_embed.proj.weight"))  # [1664, 2048]
    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  # [2048]
    out_scale: torch.Tensor  # [2048]
    out_shift: torch.Tensor  # [2048]


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]