models.py

from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchtune
from huggingface_hub import PyTorchModelHubMixin
from torchtune.models import llama3_2


def llama3_2_1B() -> torchtune.modules.transformer.TransformerDecoder:
    return llama3_2.llama3_2(
        vocab_size=128_256,
        num_layers=16,
        num_heads=32,
        num_kv_heads=8,
        embed_dim=2048,
        max_seq_len=2048,
        intermediate_dim=8192,
        attn_dropout=0.0,
        norm_eps=1e-5,
        rope_base=500_000,
        scale_factor=32,
    )


def llama3_2_100M() -> torchtune.modules.transformer.TransformerDecoder:
    return llama3_2.llama3_2(
        vocab_size=128_256,
        num_layers=4,
        num_heads=8,
        num_kv_heads=2,
        embed_dim=1024,
        max_seq_len=2048,
        intermediate_dim=8192,
        attn_dropout=0.0,
        norm_eps=1e-5,
        rope_base=500_000,
        scale_factor=32,
    )


FLAVORS = {
    "llama-1B": llama3_2_1B,
    "llama-100M": llama3_2_100M,
}


def _prepare_transformer(model):
    embed_dim = model.tok_embeddings.embedding_dim
    model.tok_embeddings = nn.Identity()
    model.output = nn.Identity()
    return model, embed_dim


def _create_causal_mask(seq_len: int, device: torch.device):
    return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device))


def _index_causal_mask(mask: torch.Tensor, input_pos: torch.Tensor):
    return mask[input_pos, :]


def sample_topk_topp(
    logits: torch.Tensor,
    topk: int,
    top_p: float,
    temperature: float,
) -> torch.Tensor:
    """
    Apply top-k, then nucleus (top-p), then sample.
    Returns a tensor of shape (batch_size, 1).
    """
    # scale + softmax
    scaled = logits / temperature
    probs = F.softmax(scaled, dim=-1)

    # --- top-k ---
    if topk < probs.size(-1):
        topk_vals, topk_idx = torch.topk(probs, topk, dim=-1)
        mask_k = torch.zeros_like(probs)
        mask_k.scatter_(-1, topk_idx, topk_vals)
        probs = mask_k

    # --- top-p (nucleus) ---
    sorted_probs, sorted_idx = torch.sort(probs, descending=True, dim=-1)
    cumulative = torch.cumsum(sorted_probs, dim=-1)
    keep = cumulative <= top_p
    keep[..., 0] = True  # always keep highest-prob

    # cast mask to same dtype as sorted_probs
    keep = keep.to(sorted_probs.dtype)

    # build final probabilities in correct dtype
    probs_final = torch.zeros_like(probs)
    src = sorted_probs * keep  # same dtype
    probs_final.scatter_(-1, sorted_idx, src)

    # renormalize
    probs_final = probs_final / probs_final.sum(dim=-1, keepdim=True)

    # sample once per batch, keep that extra dim
    return torch.multinomial(probs_final, num_samples=1)  # shape (batch,1)


@dataclass
class ModelArgs:
    backbone_flavor: str
    decoder_flavor: str
    text_vocab_size: int
    audio_vocab_size: int
    audio_num_codebooks: int


class Model(
    nn.Module,
    PyTorchModelHubMixin,
    repo_url="https://github.com/SesameAILabs/csm",
    pipeline_tag="text-to-speech",
    license="apache-2.0",
):
    def __init__(self, config: ModelArgs):
        super().__init__()
        self.config = config

        # Text+audio backbone
        self.backbone, backbone_dim = _prepare_transformer(FLAVORS[config.backbone_flavor]())
        # Audio decoder
        self.decoder, decoder_dim = _prepare_transformer(FLAVORS[config.decoder_flavor]())

        self.text_embeddings = nn.Embedding(config.text_vocab_size, backbone_dim)
        self.audio_embeddings = nn.Embedding(
            config.audio_vocab_size * config.audio_num_codebooks, backbone_dim
        )

        self.projection = nn.Linear(backbone_dim, decoder_dim, bias=False)
        self.codebook0_head = nn.Linear(backbone_dim, config.audio_vocab_size, bias=False)
        self.audio_head = nn.Parameter(
            torch.empty(config.audio_num_codebooks - 1, decoder_dim, config.audio_vocab_size)
        )

    def setup_caches(self, max_batch_size: int) -> None:
        dtype = next(self.parameters()).dtype
        device = next(self.parameters()).device

        with device:
            self.backbone.setup_caches(max_batch_size, dtype)
            self.decoder.setup_caches(
                max_batch_size, dtype, decoder_max_seq_len=self.config.audio_num_codebooks
            )

        self.register_buffer(
            "backbone_causal_mask", _create_causal_mask(self.backbone.max_seq_len, device)
        )
        self.register_buffer(
            "decoder_causal_mask",
            _create_causal_mask(self.config.audio_num_codebooks, device),
        )

    @torch.inference_mode()
    def generate_frame(
        self,
        tokens: torch.Tensor,           # (batch, seq, codebooks+1)
        tokens_mask: torch.Tensor,      # (batch, seq, codebooks+1)
        input_pos: torch.Tensor,        # (batch, seq)
        temperature: float,
        topk: int,
        top_p: float,
    ) -> torch.Tensor:
        dtype = next(self.parameters()).dtype

        # Backbone pass
        bb_mask = _index_causal_mask(self.backbone_causal_mask, input_pos)
        embeds = self._embed_tokens(tokens)
        h = self.backbone(
            (embeds * tokens_mask.unsqueeze(-1)).sum(dim=2),
            input_pos=input_pos,
            mask=bb_mask,
        ).to(dtype=dtype)

        # Last hidden state
        last_h = h[:, -1, :]               # (batch, hidden)
        last_h_unsq = last_h.unsqueeze(1)  # (batch,1,hidden)

        # --- codebook 0 ---
        c0_logits = self.codebook0_head(last_h)  # (batch, vocab)
        c0_sample = sample_topk_topp(c0_logits, topk, top_p, temperature)  # (batch,1)
        c0_embed = self._embed_audio(0, c0_sample.squeeze(-1)).unsqueeze(1)  # (batch,1,hidden)

        # Prepare decoder input
        curr_h = torch.cat([last_h_unsq, c0_embed], dim=1)  # (batch,2,hidden)
        curr_sample = c0_sample.clone()                     # (batch,1)
        curr_pos = torch.arange(0, curr_h.size(1)).unsqueeze(0).to(tokens.device).long()

        # --- remaining codebooks ---
        self.decoder.reset_caches()
        for i in range(1, self.config.audio_num_codebooks):
            dec_mask = _index_causal_mask(self.decoder_causal_mask, curr_pos)
            dec_h = self.decoder(self.projection(curr_h), input_pos=curr_pos, mask=dec_mask).to(dtype=dtype)

            ci_logits = torch.mm(dec_h[:, -1, :], self.audio_head[i - 1])
            ci_sample = sample_topk_topp(ci_logits, topk, top_p, temperature)  # (batch,1)
            ci_embed = self._embed_audio(i, ci_sample.squeeze(-1)).unsqueeze(1)  # (batch,1,hidden)

            curr_h = ci_embed
            curr_sample = torch.cat([curr_sample, ci_sample], dim=1)  # (batch,i+1)
            curr_pos = curr_pos[:, -1:] + 1

        return curr_sample  # (batch, audio_num_codebooks)

    def reset_caches(self):
        self.backbone.reset_caches()
        self.decoder.reset_caches()

    def _embed_audio(self, codebook: int, tokens: torch.Tensor) -> torch.Tensor:
        ids = tokens + codebook * self.config.audio_vocab_size
        return self.audio_embeddings(ids)

    def _embed_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
        text_ids = tokens[:, :, -1]
        text_emb = self.text_embeddings(text_ids).unsqueeze(-2)
        audio_ids = tokens[:, :, :-1] + (
            self.config.audio_vocab_size * torch.arange(self.config.audio_num_codebooks, device=tokens.device)
        )
        audio_emb = (
            self.audio_embeddings(audio_ids.reshape(-1))
            .reshape(tokens.size(0), tokens.size(1), self.config.audio_num_codebooks, -1)
        )
        return torch.cat([audio_emb, text_emb], dim=2)