release

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+prompts/conversational_a.wav filter=lfs diff=lfs merge=lfs -text
+prompts/conversational_b.wav filter=lfs diff=lfs merge=lfs -text
+prompts/read_speech_a.wav filter=lfs diff=lfs merge=lfs -text
+prompts/read_speech_b.wav filter=lfs diff=lfs merge=lfs -text
+prompts/read_speech_c.wav filter=lfs diff=lfs merge=lfs -text
+prompts/read_speech_d.wav filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,5 +1,5 @@
 ---
-title: CSM 1B
+title: Sesame CSM Space
 emoji: 🚀
 colorFrom: blue
 colorTo: blue
@@ -8,7 +8,36 @@ sdk_version: 5.20.0
 app_file: app.py
 pinned: false
 license: apache-2.0
-short_description: 1B variant of Sesame's Conversational Speech Model
+short_description: Generation using Sesame's Conversational Speech Model
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## CSM 1B
+
+**2025/03/13** - We are releasing the 1B CSM variant. Code is available on GitHub: [SesameAILabs/csm](https://github.com/SesameAILabs/csm). Checkpoint is [hosted on HuggingFace](https://huggingface.co/sesame/csm-1b).
+
+Try out the interactive demo of our fine-tuned version [sesame.com/voicedemo](https://www.sesame.com/voicedemo).
+
+Generate from the open-source base model [hosted on HuggingFace](https://huggingface.co/spaces/sesame/csm-1b).
+
+---
+
+CSM (Conversational Speech Model) is a speech generation model from [Sesame](sesame.com) that generates RVQ audio codes from text and audio inputs. A fine-tuned version of this model powers the interactive demo in our [technical blog post](https://www.sesame.com/research/crossing_the_uncanny_valley_of_voice).
+
+The model architecture employs a [Llama](https://www.llama.com/) backbone and a smaller audio decoder that produces [Mimi](https://huggingface.co/kyutai/mimi) audio codes.
+
+## Misuse and abuse ⚠️
+
+This project provides a high-quality speech generation model for research and educational purposes. While we encourage responsible and ethical use, we **explicitly prohibit** the following:
+
+- **Impersonation or Fraud**: Do not use this model to generate speech that mimics real individuals without their explicit consent.
+- **Misinformation or Deception**: Do not use this model to create deceptive or misleading content, such as fake news or fraudulent calls.
+- **Illegal or Harmful Activities**: Do not use this model for any illegal, harmful, or malicious purposes.
+
+By using this model, you agree to comply with all applicable laws and ethical guidelines. We are **not responsible** for any misuse, and we strongly condemn unethical applications of this technology.
+
+**Prompts**
+Conversational prompts are from the [EdAcc dataset](https://groups.inf.ed.ac.uk/edacc/)
+Read speech prompts are form the [LibriTTS-R dataset](https://google.github.io/df-conformer/librittsr/)
+
+**Authors**
+Johan Schalkwyk, Ankit Kumar, Dan Lyth, Sefik Emre Eskimez, Zack Hodari, Cinjon Resnick, Ramon Sanabria, Raven Jiang, and the Sesame team.

app.py

@@ -0,0 +1,232 @@
+import os
+
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+import torchaudio
+from generator import Segment, load_csm_1b
+from huggingface_hub import hf_hub_download, login
+from watermarking import watermark
+
+api_key = os.getenv("HF_TOKEN")
+gpu_timeout = int(os.getenv("GPU_TIMEOUT", 60))
+CSM_1B_HF_WATERMARK = list(map(int, os.getenv("WATERMARK_KEY").split(" ")))
+
+login(token=api_key)
+
+SPACE_INTRO_TEXT = """\
+# Sesame CSM 1B
+
+Generate from CSM 1B (Conversational Speech Model). 
+Code is available on GitHub: [SesameAILabs/csm](https://github.com/SesameAILabs/csm). 
+Checkpoint is [hosted on HuggingFace](https://huggingface.co/sesame/csm-1b).
+
+Try out the interactive demo of our fine-tuned model [sesame.com/voicedemo](https://www.sesame.com/voicedemo).
+
+The model has some capacity for non-English languages due to data contamination in the training 
+data, but it is likely not to perform well.
+
+---
+
+"""
+
+CONVO_INTRO_TEXT = """\
+## Conversation content
+
+Each line is an utterance in the conversation to generate. Speakers alternate between A and B, starting with speaker A.
+"""
+
+DEFAULT_CONVERSATION = """\
+Hey how are you doing.
+Pretty good, pretty good.
+I'm great, so happy to be speaking to you.
+Me too, this is some cool stuff huh?
+Yeah, I've been reading more about speech generation, and it really seems like context is important.
+Definitely.
+"""
+
+SPEAKER_PROMPTS = {
+    "conversational_a": {
+        "text": (
+            "like revising for an exam I'd have to try and like keep up the momentum because I'd "
+            "start really early I'd be like okay I'm gonna start revising now and then like "
+            "you're revising for ages and then I just like start losing steam I didn't do that "
+            "for the exam we had recently to be fair that was a more of a last minute scenario "
+            "but like yeah I'm trying to like yeah I noticed this yesterday that like Mondays I "
+            "sort of start the day with this not like a panic but like a"
+        ),
+        "audio": "prompts/conversational_a.wav",
+    },
+    "conversational_b": {
+        "text": (
+            "like a super Mario level. Like it's very like high detail. And like, once you get "
+            "into the park, it just like, everything looks like a computer game and they have all "
+            "these, like, you know, if, if there's like a, you know, like in a Mario game, they "
+            "will have like a question block. And if you like, you know, punch it, a coin will "
+            "come out. So like everyone, when they come into the park, they get like this little "
+            "bracelet and then you can go punching question blocks around."
+        ),
+        "audio": "prompts/conversational_b.wav",
+    },
+    "read_speech_a": {
+        "text": (
+            "And Lake turned round upon me, a little abruptly, his odd yellowish eyes, a little "
+            "like those of the sea eagle, and the ghost of his smile that flickered on his "
+            "singularly pale face, with a stern and insidious look, confronted me."
+        ),
+        "audio": "prompts/read_speech_a.wav",
+    },
+    "read_speech_b": {
+        "text": (
+            "He was such a big boy that he wore high boots and carried a jack knife. He gazed and "
+            "gazed at the cap, and could not keep from fingering the blue tassel."
+        ),
+        "audio": "prompts/read_speech_b.wav",
+    },
+    "read_speech_c": {
+        "text": (
+            "All passed so quickly, there was so much going on around him, the Tree quite forgot "
+            "to look to himself."
+        ),
+        "audio": "prompts/read_speech_c.wav",
+    },
+    "read_speech_d": {
+        "text": (
+            "Suddenly I was back in the old days Before you felt we ought to drift apart. It was "
+            "some trick-the way your eyebrows raise."
+        ),
+        "audio": "prompts/read_speech_d.wav",
+    },
+}
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_path = hf_hub_download(repo_id="sesame/csm-1b", filename="ckpt.pt")
+generator = load_csm_1b(model_path, device)
+
+
+@spaces.GPU(duration=gpu_timeout)
+def infer(
+    text_prompt_speaker_a,
+    text_prompt_speaker_b,
+    audio_prompt_speaker_a,
+    audio_prompt_speaker_b,
+    gen_conversation_input,
+) -> tuple[np.ndarray, int]:
+    audio_prompt_a = prepare_prompt(text_prompt_speaker_a, 0, audio_prompt_speaker_a)
+    audio_prompt_b = prepare_prompt(text_prompt_speaker_b, 1, audio_prompt_speaker_b)
+
+    prompt_segments: list[Segment] = [audio_prompt_a, audio_prompt_b]
+    generated_segments: list[Segment] = []
+
+    conversation_lines = [line.strip() for line in gen_conversation_input.strip().split("\n") if line.strip()]
+    for i, line in enumerate(conversation_lines):
+        # Alternating speakers A and B, starting with A
+        speaker_id = i % 2
+
+        audio_tensor = generator.generate(
+            text=line,
+            speaker=speaker_id,
+            context=prompt_segments + generated_segments,
+        )
+        generated_segments.append(Segment(text=line, speaker=speaker_id, audio=audio_tensor))
+
+    # Concatenate all generations and convert to 16-bit int format
+    audio_tensors = [segment.audio for segment in generated_segments]
+    audio_tensor = torch.cat(audio_tensors, dim=0)
+
+    # This applies an imperceptible watermark to identify audio as AI-generated.
+    # Watermarking ensures transparency, dissuades misuse, and enables traceability.
+    # Please be a responsible AI citizen and keep the watermarking in place.
+    # If using CSM 1B in another application, use your own private key and keep it secret.
+    audio_tensor, wm_sample_rate = watermark(
+        generator._watermarker, audio_tensor, generator.sample_rate, CSM_1B_HF_WATERMARK
+    )
+    audio_tensor = torchaudio.functional.resample(
+        audio_tensor, orig_freq=wm_sample_rate, new_freq=generator.sample_rate
+    )
+
+    audio_array = (audio_tensor * 32768).to(torch.int16).cpu().numpy()
+
+    return generator.sample_rate, audio_array
+
+
+def prepare_prompt(text: str, speaker: int, audio_path: str) -> Segment:
+    audio_tensor, _ = load_prompt_audio(audio_path)
+    return Segment(text=text, speaker=speaker, audio=audio_tensor)
+
+
+def load_prompt_audio(audio_path: str) -> torch.Tensor:
+    audio_tensor, sample_rate = torchaudio.load(audio_path)
+    audio_tensor = audio_tensor.squeeze(0)
+    if sample_rate != generator.sample_rate:
+        audio_tensor = torchaudio.functional.resample(
+            audio_tensor, orig_freq=sample_rate, new_freq=generator.sample_rate
+        )
+    return audio_tensor, generator.sample_rate
+
+
+def create_speaker_prompt_ui(speaker_name: str):
+    speaker_dropdown = gr.Dropdown(
+        choices=list(SPEAKER_PROMPTS.keys()), label="Select a predefined speaker", value=speaker_name
+    )
+    with gr.Accordion("Or add your own voice prompt", open=False):
+        text_prompt_speaker = gr.Textbox(label="Speaker prompt", lines=4, value=SPEAKER_PROMPTS[speaker_name]["text"])
+        audio_prompt_speaker = gr.Audio(
+            label="Speaker prompt", type="filepath", value=SPEAKER_PROMPTS[speaker_name]["audio"]
+        )
+
+    return speaker_dropdown, text_prompt_speaker, audio_prompt_speaker
+
+
+with gr.Blocks() as app:
+    gr.Markdown(SPACE_INTRO_TEXT)
+    gr.Markdown("## Voices")
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### Speaker A")
+            speaker_a_dropdown, text_prompt_speaker_a, audio_prompt_speaker_a = create_speaker_prompt_ui(
+                "conversational_a"
+            )
+
+        with gr.Column():
+            gr.Markdown("### Speaker B")
+            speaker_b_dropdown, text_prompt_speaker_b, audio_prompt_speaker_b = create_speaker_prompt_ui(
+                "conversational_b"
+            )
+
+    def update_audio(speaker):
+        if speaker in SPEAKER_PROMPTS:
+            return SPEAKER_PROMPTS[speaker]["audio"]
+        return None
+
+    def update_text(speaker):
+        if speaker in SPEAKER_PROMPTS:
+            return SPEAKER_PROMPTS[speaker]["text"]
+        return None
+
+    speaker_a_dropdown.change(fn=update_audio, inputs=[speaker_a_dropdown], outputs=[audio_prompt_speaker_a])
+    speaker_b_dropdown.change(fn=update_audio, inputs=[speaker_b_dropdown], outputs=[audio_prompt_speaker_b])
+
+    speaker_a_dropdown.change(fn=update_text, inputs=[speaker_a_dropdown], outputs=[text_prompt_speaker_a])
+    speaker_b_dropdown.change(fn=update_text, inputs=[speaker_b_dropdown], outputs=[text_prompt_speaker_b])
+
+    gr.Markdown(CONVO_INTRO_TEXT)
+
+    gen_conversation_input = gr.TextArea(label="conversation", lines=20, value=DEFAULT_CONVERSATION)
+    generate_btn = gr.Button("Generate conversation", variant="primary")
+    audio_output = gr.Audio(label="Synthesized audio")
+
+    generate_btn.click(
+        infer,
+        inputs=[
+            text_prompt_speaker_a,
+            text_prompt_speaker_b,
+            audio_prompt_speaker_a,
+            audio_prompt_speaker_b,
+            gen_conversation_input,
+        ],
+        outputs=[audio_output],
+    )
+
+app.launch(ssr_mode=True)

generator.py

@@ -0,0 +1,178 @@
+import os
+from dataclasses import dataclass
+from typing import List, Tuple
+
+import torch
+import torchaudio
+from huggingface_hub import hf_hub_download
+from models import Model, ModelArgs
+from moshi.models import loaders
+from tokenizers.processors import TemplateProcessing
+from transformers import AutoTokenizer
+from watermarking import load_watermarker, watermark
+
+CSM_1B_HF_WATERMARK = list(map(int, os.getenv("WATERMARK_KEY").split(" ")))
+
+
+@dataclass
+class Segment:
+    speaker: int
+    text: str
+    # (num_samples,), sample_rate = 24_000
+    audio: torch.Tensor
+
+
+def load_llama3_tokenizer():
+    """
+    https://github.com/huggingface/transformers/issues/22794#issuecomment-2092623992
+    """
+    tokenizer_name = "meta-llama/Llama-3.2-1B"
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    bos = tokenizer.bos_token
+    eos = tokenizer.eos_token
+    tokenizer._tokenizer.post_processor = TemplateProcessing(
+        single=f"{bos}:0 $A:0 {eos}:0",
+        pair=f"{bos}:0 $A:0 {eos}:0 {bos}:1 $B:1 {eos}:1",
+        special_tokens=[(f"{bos}", tokenizer.bos_token_id), (f"{eos}", tokenizer.eos_token_id)],
+    )
+
+    return tokenizer
+
+
+class Generator:
+    def __init__(
+        self,
+        model: Model,
+    ):
+        self._model = model
+        self._model.setup_caches(1)
+
+        self._text_tokenizer = load_llama3_tokenizer()
+
+        device = next(model.parameters()).device
+        mimi_weight = hf_hub_download(loaders.DEFAULT_REPO, loaders.MIMI_NAME)
+        mimi = loaders.get_mimi(mimi_weight, device=device)
+        mimi.set_num_codebooks(32)
+        self._audio_tokenizer = mimi
+
+        self._watermarker = load_watermarker(device=device)
+
+        self.sample_rate = mimi.sample_rate
+        self.device = device
+
+    def _tokenize_text_segment(self, text: str, speaker: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        frame_tokens = []
+        frame_masks = []
+
+        text_tokens = self._text_tokenizer.encode(f"[{speaker}]{text}")
+        text_frame = torch.zeros(len(text_tokens), 33).long()
+        text_frame_mask = torch.zeros(len(text_tokens), 33).bool()
+        text_frame[:, -1] = torch.tensor(text_tokens)
+        text_frame_mask[:, -1] = True
+
+        frame_tokens.append(text_frame.to(self.device))
+        frame_masks.append(text_frame_mask.to(self.device))
+
+        return torch.cat(frame_tokens, dim=0), torch.cat(frame_masks, dim=0)
+
+    def _tokenize_audio(self, audio: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        frame_tokens = []
+        frame_masks = []
+
+        # (K, T)
+        audio = audio.to(self.device)
+        audio_tokens = self._audio_tokenizer.encode(audio.unsqueeze(0).unsqueeze(0))[0]
+        # add EOS frame
+        eos_frame = torch.zeros(audio_tokens.size(0), 1).to(self.device)
+        audio_tokens = torch.cat([audio_tokens, eos_frame], dim=1)
+
+        audio_frame = torch.zeros(audio_tokens.size(1), 33).long().to(self.device)
+        audio_frame_mask = torch.zeros(audio_tokens.size(1), 33).bool().to(self.device)
+        audio_frame[:, :-1] = audio_tokens.transpose(0, 1)
+        audio_frame_mask[:, :-1] = True
+
+        frame_tokens.append(audio_frame)
+        frame_masks.append(audio_frame_mask)
+
+        return torch.cat(frame_tokens, dim=0), torch.cat(frame_masks, dim=0)
+
+    def _tokenize_segment(self, segment: Segment) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Returns:
+            (seq_len, 33), (seq_len, 33)
+        """
+        text_tokens, text_masks = self._tokenize_text_segment(segment.text, segment.speaker)
+        audio_tokens, audio_masks = self._tokenize_audio(segment.audio)
+
+        return torch.cat([text_tokens, audio_tokens], dim=0), torch.cat([text_masks, audio_masks], dim=0)
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        text: str,
+        speaker: int,
+        context: List[Segment],
+        max_audio_length_ms: float = 90_000,
+        temperature: float = 0.9,
+        topk: int = 50,
+    ) -> torch.Tensor:
+        self._model.reset_caches()
+
+        max_audio_frames = int(max_audio_length_ms / 80)
+        tokens, tokens_mask = [], []
+        for segment in context:
+            segment_tokens, segment_tokens_mask = self._tokenize_segment(segment)
+            tokens.append(segment_tokens)
+            tokens_mask.append(segment_tokens_mask)
+
+        gen_segment_tokens, gen_segment_tokens_mask = self._tokenize_text_segment(text, speaker)
+        tokens.append(gen_segment_tokens)
+        tokens_mask.append(gen_segment_tokens_mask)
+
+        prompt_tokens = torch.cat(tokens, dim=0).long().to(self.device)
+        prompt_tokens_mask = torch.cat(tokens_mask, dim=0).bool().to(self.device)
+
+        samples = []
+        curr_tokens = prompt_tokens.unsqueeze(0)
+        curr_tokens_mask = prompt_tokens_mask.unsqueeze(0)
+        curr_pos = torch.arange(0, prompt_tokens.size(0)).unsqueeze(0).long().to(self.device)
+
+        for _ in range(max_audio_frames):
+            sample = self._model.generate_frame(curr_tokens, curr_tokens_mask, curr_pos, temperature, topk)
+            if torch.all(sample == 0):
+                break  # eos
+
+            samples.append(sample)
+
+            curr_tokens = torch.cat([sample, torch.zeros(1, 1).long().to(self.device)], dim=1).unsqueeze(1)
+            curr_tokens_mask = torch.cat(
+                [torch.ones_like(sample).bool(), torch.zeros(1, 1).bool().to(self.device)], dim=1
+            ).unsqueeze(1)
+            curr_pos = curr_pos[:, -1:] + 1
+
+        audio = self._audio_tokenizer.decode(torch.stack(samples).permute(1, 2, 0)).squeeze(0).squeeze(0)
+
+        # This applies an imperceptible watermark to identify audio as AI-generated.
+        # Watermarking ensures transparency, dissuades misuse, and enables traceability.
+        # Please be a responsible AI citizen and keep the watermarking in place.
+        # If using CSM 1B in another application, use your own private key and keep it secret.
+        audio, wm_sample_rate = watermark(self._watermarker, audio, self.sample_rate, CSM_1B_HF_WATERMARK)
+        audio = torchaudio.functional.resample(audio, orig_freq=wm_sample_rate, new_freq=self.sample_rate)
+
+        return audio
+
+
+def load_csm_1b(ckpt_path: str = "ckpt.pt", device: str = "cuda") -> Generator:
+    model_args = ModelArgs(
+        backbone_flavor="llama-1B",
+        decoder_flavor="llama-100M",
+        text_vocab_size=128256,
+        audio_vocab_size=2051,
+        audio_num_codebooks=32,
+    )
+    model = Model(model_args).to(device=device, dtype=torch.bfloat16)
+    state_dict = torch.load(ckpt_path)
+    model.load_state_dict(state_dict)
+
+    generator = Generator(model)
+    return generator

models.py

@@ -0,0 +1,196 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torchtune
+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):
+    """
+    Args:
+        mask: (max_seq_len, max_seq_len)
+        input_pos: (batch_size, seq_len)
+
+    Returns:
+        (batch_size, seq_len, max_seq_len)
+    """
+    r = mask[input_pos, :]
+    return r
+
+
+def _multinomial_sample_one_no_sync(probs):  # Does multinomial sampling without a cuda synchronization
+    q = torch.empty_like(probs).exponential_(1)
+    return torch.argmax(probs / q, dim=-1, keepdim=True).to(dtype=torch.int)
+
+
+def sample_topk(logits: torch.Tensor, topk: int, temperature: float):
+    logits = logits / temperature
+
+    filter_value: float = -float("Inf")
+    indices_to_remove = logits < torch.topk(logits, topk)[0][..., -1, None]
+    scores_processed = logits.masked_fill(indices_to_remove, filter_value)
+    scores_processed = torch.nn.functional.log_softmax(scores_processed, dim=-1)
+    probs = torch.nn.functional.softmax(scores_processed, dim=-1)
+
+    sample_token = _multinomial_sample_one_no_sync(probs)
+    return sample_token
+
+
+@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):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+
+        self.backbone, backbone_dim = _prepare_transformer(FLAVORS[args.backbone_flavor]())
+        self.decoder, decoder_dim = _prepare_transformer(FLAVORS[args.decoder_flavor]())
+
+        self.text_embeddings = nn.Embedding(args.text_vocab_size, backbone_dim)
+        self.audio_embeddings = nn.Embedding(args.audio_vocab_size * args.audio_num_codebooks, backbone_dim)
+
+        self.projection = nn.Linear(backbone_dim, decoder_dim, bias=False)
+        self.codebook0_head = nn.Linear(backbone_dim, args.audio_vocab_size, bias=False)
+        self.audio_head = nn.Parameter(torch.empty(args.audio_num_codebooks - 1, decoder_dim, args.audio_vocab_size))
+
+    def setup_caches(self, max_batch_size: int) -> torch.Tensor:
+        """Setup KV caches and return a causal mask."""
+        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.args.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.args.audio_num_codebooks, device))
+
+    def generate_frame(
+        self,
+        tokens: torch.Tensor,
+        tokens_mask: torch.Tensor,
+        input_pos: torch.Tensor,
+        temperature: float,
+        topk: int,
+    ) -> torch.Tensor:
+        """
+        Args:
+            tokens: (batch_size, seq_len, audio_num_codebooks+1)
+            tokens_mask: (batch_size, seq_len, audio_num_codebooks+1)
+            input_pos: (batch_size, seq_len) positions for each token
+            mask: (batch_size, seq_len, max_seq_len
+
+        Returns:
+            (batch_size, audio_num_codebooks) sampled tokens
+        """
+        dtype = next(self.parameters()).dtype
+        b, s, _ = tokens.size()
+
+        assert self.backbone.caches_are_enabled(), "backbone caches are not enabled"
+        curr_backbone_mask = _index_causal_mask(self.backbone_causal_mask, input_pos)
+        embeds = self._embed_tokens(tokens)
+        masked_embeds = embeds * tokens_mask.unsqueeze(-1)
+        h = masked_embeds.sum(dim=2)
+        h = self.backbone(h, input_pos=input_pos, mask=curr_backbone_mask).to(dtype=dtype)
+
+        last_h = h[:, -1, :]
+        c0_logits = self.codebook0_head(last_h)
+        c0_sample = sample_topk(c0_logits, topk, temperature)
+        c0_embed = self._embed_audio(0, c0_sample)
+
+        curr_h = torch.cat([last_h.unsqueeze(1), c0_embed], dim=1)
+        curr_sample = c0_sample.clone()
+        curr_pos = torch.arange(0, curr_h.size(1), device=curr_h.device).unsqueeze(0).repeat(curr_h.size(0), 1)
+
+        # Decoder caches must be reset every frame.
+        self.decoder.reset_caches()
+        for i in range(1, self.args.audio_num_codebooks):
+            curr_decoder_mask = _index_causal_mask(self.decoder_causal_mask, curr_pos)
+            decoder_h = self.decoder(self.projection(curr_h), input_pos=curr_pos, mask=curr_decoder_mask).to(
+                dtype=dtype
+            )
+            ci_logits = torch.mm(decoder_h[:, -1, :], self.audio_head[i - 1])
+            ci_sample = sample_topk(ci_logits, topk, temperature)
+            ci_embed = self._embed_audio(i, ci_sample)
+
+            curr_h = ci_embed
+            curr_sample = torch.cat([curr_sample, ci_sample], dim=1)
+            curr_pos = curr_pos[:, -1:] + 1
+
+        return curr_sample
+
+    def reset_caches(self):
+        self.backbone.reset_caches()
+        self.decoder.reset_caches()
+
+    def _embed_audio(self, codebook: int, tokens: torch.Tensor) -> torch.Tensor:
+        return self.audio_embeddings(tokens + codebook * self.args.audio_vocab_size)
+
+    def _embed_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
+        text_embeds = self.text_embeddings(tokens[:, :, -1]).unsqueeze(-2)
+
+        audio_tokens = tokens[:, :, :-1] + (
+            self.args.audio_vocab_size * torch.arange(self.args.audio_num_codebooks, device=tokens.device)
+        )
+        audio_embeds = self.audio_embeddings(audio_tokens.view(-1)).reshape(
+            tokens.size(0), tokens.size(1), self.args.audio_num_codebooks, -1
+        )
+
+        return torch.cat([audio_embeds, text_embeds], dim=-2)

prompts/conversational_a.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:356648c1bc6c1da7883004557e9b21a2ef7d01682d8b9d02d6dcb950b348b04f
+size 2646044

prompts/conversational_b.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c247153011385d33aaeed193adfec562c32182e2facd30cc8cd0b3e820e94afb
+size 2646044

prompts/read_speech_a.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59480708f84c77ab2967d14d821c2ccade9d7761685d060575121f49a149005b
+size 831412

prompts/read_speech_b.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f582640265864499cbe6a8c687ea0f9e08e7fa41eeb2caa923d0a3bada55fcef
+size 576052

prompts/read_speech_c.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7da15ab3ee7f8bbc8abfce73ce65936a80a535ae4a86db2d9c4756caba69e9c3
+size 385964

prompts/read_speech_d.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:09cad0494f9d0038b0f0eb039f47d752c45e56d92679f96587e20f67b2c1b7d8
+size 435884

requirements.txt

@@ -0,0 +1,11 @@
+torch==2.4.0
+torchaudio==2.4.0
+tokenizers==0.21.0
+transformers==4.49.0
+huggingface_hub==0.28.1
+spaces==0.32.0
+gradio==5.20.1
+moshi==0.2.2
+torchtune==0.4.0
+torchao==0.9.0
+silentcipher @ git+https://github.com/SesameAILabs/silentcipher@master

watermarking.py

@@ -0,0 +1,78 @@
+import os
+import argparse
+
+import silentcipher
+import torch
+import torchaudio
+
+CSM_1B_HF_WATERMARK = list(map(int, os.getenv("WATERMARK_KEY").split(" ")))
+
+
+def cli_check_audio() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--audio_path", type=str, required=True)
+    args = parser.parse_args()
+
+    check_audio_from_file(args.audio_path)
+
+
+def load_watermarker(device: str = "cuda") -> silentcipher.server.Model:
+    model = silentcipher.get_model(
+        model_type="44.1k",
+        device=device,
+    )
+    return model
+
+
+@torch.inference_mode()
+def watermark(
+    watermarker: silentcipher.server.Model,
+    audio_array: torch.Tensor,
+    sample_rate: int,
+    watermark_key: list[int],
+) -> tuple[torch.Tensor, int]:
+    audio_array_44khz = torchaudio.functional.resample(audio_array, orig_freq=sample_rate, new_freq=44100)
+    encoded, _ = watermarker.encode_wav(audio_array_44khz, 44100, watermark_key, calc_sdr=False, message_sdr=36)
+
+    output_sample_rate = min(44100, sample_rate)
+    encoded = torchaudio.functional.resample(encoded, orig_freq=44100, new_freq=output_sample_rate)
+    return encoded, output_sample_rate
+
+
+@torch.inference_mode()
+def verify(
+    watermarker: silentcipher.server.Model,
+    watermarked_audio: torch.Tensor,
+    sample_rate: int,
+    watermark_key: list[int],
+) -> bool:
+    watermarked_audio_44khz = torchaudio.functional.resample(watermarked_audio, orig_freq=sample_rate, new_freq=44100)
+    result = watermarker.decode_wav(watermarked_audio_44khz, 44100, phase_shift_decoding=True)
+
+    is_watermarked = result["status"]
+    if is_watermarked:
+        is_csm_watermarked = result["messages"][0] == watermark_key
+    else:
+        is_csm_watermarked = False
+
+    return is_watermarked and is_csm_watermarked
+
+
+def check_audio_from_file(audio_path: str) -> None:
+    watermarker = load_watermarker(device="cuda")
+
+    audio_array, sample_rate = load_audio(audio_path)
+    is_watermarked = verify(watermarker, audio_array, sample_rate, CSM_1B_HF_WATERMARK)
+
+    outcome = "Watermarked" if is_watermarked else "Not watermarked"
+    print(f"{outcome}: {audio_path}")
+
+
+def load_audio(audio_path: str) -> tuple[torch.Tensor, int]:
+    audio_array, sample_rate = torchaudio.load(audio_path)
+    audio_array = audio_array.mean(dim=0)
+    return audio_array, int(sample_rate)
+
+
+if __name__ == "__main__":
+    cli_check_audio()