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)