Add Miso TTS 8B ZeroGPU demo

README.md

@@ -1,13 +1,17 @@
 ---
-title: MisoTTS
-emoji: 👀
-colorFrom: pink
-colorTo: blue
+title: Miso TTS 8B
+emoji: 🍲
+colorFrom: yellow
+colorTo: red
 sdk: gradio
-sdk_version: 6.16.0
-python_version: '3.12'
+sdk_version: 5.20.0
 app_file: app.py
 pinned: false
+license: other
+short_description: 8B CSM-style text-to-speech with voice cloning
+python_version: "3.10"
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Demo for [MisoLabs/MisoTTS](https://huggingface.co/MisoLabs/MisoTTS), an 8B Sesame CSM-style
+text-to-speech model. Inference code from [MisoLabsAI/MisoTTS](https://github.com/MisoLabsAI/MisoTTS).

app.py

@@ -0,0 +1,99 @@
+import os
+
+os.environ.setdefault("NO_TORCH_COMPILE", "1")
+
+import spaces
+import numpy as np
+import torch
+import torchaudio
+import gradio as gr
+
+from generator import Segment, load_miso_8b
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+generator = load_miso_8b(device=device, model_path_or_repo_id="MisoLabs/MisoTTS")
+SAMPLE_RATE = generator.sample_rate
+
+MAX_INPUT_CHARS = 1000
+
+DESCRIPTION = """
+# Miso TTS 8B
+
+Text-to-speech with the [MisoLabs/MisoTTS](https://huggingface.co/MisoLabs/MisoTTS) model — an
+8B [Sesame CSM](https://github.com/SesameAILabs/csm)-style model that generates Mimi audio codes
+from text, with optional voice continuation from a reference clip.
+
+Provide a reference audio + its transcript to clone a voice, or leave them empty for a default voice.
+Outputs carry an imperceptible watermark identifying the audio as AI-generated.
+"""
+
+
+def _resample_to_model(audio: torch.Tensor, sr: int) -> torch.Tensor:
+    audio = audio.mean(dim=0) if audio.ndim > 1 else audio
+    if sr != SAMPLE_RATE:
+        audio = torchaudio.functional.resample(audio, orig_freq=sr, new_freq=SAMPLE_RATE)
+    return audio
+
+
+@spaces.GPU(duration=120)
+def synthesize(text, ref_audio_path, ref_text, speaker_id, max_length_ms, temperature, topk):
+    text = (text or "").strip()
+    if not text:
+        raise gr.Error("Please enter some text to synthesize.")
+    if len(text) > MAX_INPUT_CHARS:
+        raise gr.Error(f"Text too long (>{MAX_INPUT_CHARS} characters).")
+
+    context = []
+    if ref_audio_path:
+        if not (ref_text or "").strip():
+            raise gr.Error("Please provide the transcript of the reference audio.")
+        wav, sr = torchaudio.load(ref_audio_path)
+        wav = _resample_to_model(wav, sr).to(device)
+        context = [Segment(speaker=int(speaker_id), text=ref_text.strip(), audio=wav)]
+
+    audio = generator.generate(
+        text=text,
+        speaker=int(speaker_id),
+        context=context,
+        max_audio_length_ms=float(max_length_ms),
+        temperature=float(temperature),
+        topk=int(topk),
+    )
+
+    audio_np = (audio * 32768).clamp(-32768, 32767).to(torch.int16).cpu().numpy()
+    return SAMPLE_RATE, audio_np
+
+
+with gr.Blocks(title="Miso TTS 8B") as demo:
+    gr.Markdown(DESCRIPTION)
+    with gr.Row():
+        with gr.Column():
+            text = gr.Textbox(
+                label="Text to synthesize",
+                placeholder="Hello from Miso.",
+                lines=3,
+                value="Hello from Miso. This is an eight billion parameter text to speech model.",
+            )
+            with gr.Accordion("Voice cloning (optional)", open=False):
+                ref_audio = gr.Audio(label="Reference audio", type="filepath")
+                ref_text = gr.Textbox(
+                    label="Reference transcript",
+                    placeholder="The exact words spoken in the reference audio.",
+                    lines=2,
+                )
+            with gr.Accordion("Advanced", open=False):
+                speaker_id = gr.Slider(0, 1, value=0, step=1, label="Speaker ID")
+                max_length = gr.Slider(2000, 30000, value=10000, step=1000, label="Max audio length (ms)")
+                temperature = gr.Slider(0.1, 1.5, value=0.9, step=0.05, label="Temperature")
+                topk = gr.Slider(1, 100, value=50, step=1, label="Top-k")
+            run = gr.Button("Generate", variant="primary")
+        with gr.Column():
+            out = gr.Audio(label="Generated speech")
+
+    run.click(
+        synthesize,
+        inputs=[text, ref_audio, ref_text, speaker_id, max_length, temperature, topk],
+        outputs=[out],
+    )
+
+demo.queue().launch()

generator.py

@@ -0,0 +1,238 @@
+from dataclasses import dataclass
+import os
+from typing import List, Optional, Tuple
+
+os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "60")
+os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "60")
+
+import torch
+import torchaudio
+from huggingface_hub import hf_hub_download
+from models import MISO_TTS_8B_CONFIG, Model, ModelArgs
+from moshi_compat import patch_bitsandbytes_import_for_unquantized_layers
+from moshi.models import loaders
+from tokenizers.processors import TemplateProcessing
+from transformers import AutoTokenizer
+from watermarking import MISO_TTS_WATERMARK, load_watermarker, watermark
+
+DEFAULT_MISO_TTS_REPO_ID = "MisoLabs/MisoTTS"
+patch_bitsandbytes_import_for_unquantized_layers()
+
+
+@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 = "NousResearch/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()
+        self._frame_size = self._model.config.audio_num_codebooks + 1
+
+        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(self._model.config.audio_num_codebooks)
+        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.lstrip()}")
+        text_frame = torch.zeros(len(text_tokens), self._frame_size).long()
+        text_frame_mask = torch.zeros(len(text_tokens), self._frame_size).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]:
+        assert audio.ndim == 1, "Audio must be single channel"
+
+        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), self._frame_size).long().to(self.device)
+        audio_frame_mask = torch.zeros(audio_tokens.size(1), self._frame_size).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, audio_num_codebooks + 1), (seq_len, audio_num_codebooks + 1)
+        """
+        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_generation_len = 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)
+
+        max_seq_len = 2048
+        max_context_len = max_seq_len - max_generation_len
+        if curr_tokens.size(1) >= max_context_len:
+            raise ValueError(
+                f"Inputs too long, must be below max_seq_len - max_generation_len: {max_context_len}"
+            )
+
+        for _ in range(max_generation_len):
+            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.
+        # If using Miso TTS in another application, use your own private key and keep it secret.
+        audio, wm_sample_rate = watermark(self._watermarker, audio, self.sample_rate, MISO_TTS_WATERMARK)
+        audio = torchaudio.functional.resample(audio, orig_freq=wm_sample_rate, new_freq=self.sample_rate)
+
+        return audio
+
+
+def _state_dict_from_checkpoint(checkpoint: object) -> dict[str, torch.Tensor]:
+    if not isinstance(checkpoint, dict):
+        raise TypeError(f"Expected checkpoint dict, got {type(checkpoint).__name__}")
+
+    for key in ("state_dict", "model_state_dict", "model"):
+        value = checkpoint.get(key)
+        if isinstance(value, dict):
+            checkpoint = value
+            break
+
+    state_dict = {}
+    for key, value in checkpoint.items():
+        if torch.is_tensor(value):
+            state_dict[key.removeprefix("module.")] = value
+    if not state_dict:
+        raise ValueError("Checkpoint did not contain any tensor state_dict entries")
+    return state_dict
+
+
+def _load_model(
+    model_path_or_repo_id: str,
+    config: ModelArgs,
+    device: str,
+    dtype: torch.dtype,
+) -> Model:
+    if os.path.isfile(model_path_or_repo_id):
+        model_file = model_path_or_repo_id
+    elif os.path.isdir(model_path_or_repo_id):
+        model_file = os.path.join(model_path_or_repo_id, "model.safetensors")
+    else:
+        model_file = hf_hub_download(repo_id=model_path_or_repo_id, filename="model.safetensors")
+
+    if os.path.isfile(model_file):
+        model = Model(config)
+        if model_file.endswith(".safetensors"):
+            try:
+                from safetensors.torch import load_file
+            except ImportError as exc:
+                raise ImportError("Install safetensors to load .safetensors checkpoint files") from exc
+
+            state_dict = load_file(model_file, device="cpu")
+        else:
+            checkpoint = torch.load(model_file, map_location="cpu")
+            state_dict = _state_dict_from_checkpoint(checkpoint)
+        model.load_state_dict(state_dict)
+    else:
+        raise FileNotFoundError(f"Could not resolve model checkpoint: {model_path_or_repo_id}")
+
+    model.to(device=device, dtype=dtype)
+    model.eval()
+    return model
+
+
+def load_miso_8b(
+    device: str = "cuda",
+    model_path_or_repo_id: Optional[str] = None,
+    dtype: torch.dtype = torch.bfloat16,
+) -> Generator:
+    source = model_path_or_repo_id or os.environ.get("MISO_TTS_8B_MODEL", DEFAULT_MISO_TTS_REPO_ID)
+    model = _load_model(source, MISO_TTS_8B_CONFIG, device=device, dtype=dtype)
+    return Generator(model)

models.py

@@ -0,0 +1,290 @@
+from dataclasses import dataclass
+import contextlib
+import io
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import PyTorchModelHubMixin
+
+with contextlib.redirect_stdout(io.StringIO()) as _torchtune_stdout:
+    from torchtune.models import llama3_2
+
+_torchtune_import_output = _torchtune_stdout.getvalue()
+if _torchtune_import_output.strip() != "import error: No module named 'triton'":
+    print(_torchtune_import_output, end="")
+
+
+def llama3_2_8B():
+    return llama3_2.llama3_2(
+        vocab_size=128_256,
+        num_layers=32,
+        num_heads=32,
+        num_kv_heads=8,
+        embed_dim=4096,
+        max_seq_len=2048,
+        intermediate_dim=14_336,
+        attn_dropout=0.1,
+        norm_eps=1e-5,
+        rope_base=500_000,
+        scale_factor=32,
+    )
+
+
+def llama3_2_300M():
+    return llama3_2.llama3_2(
+        vocab_size=128_256,
+        num_layers=8,
+        num_heads=24,
+        num_kv_heads=6,
+        embed_dim=1536,
+        max_seq_len=2048,
+        intermediate_dim=6912,
+        attn_dropout=0.1,
+        norm_eps=1e-5,
+        rope_base=500_000,
+        scale_factor=32,
+    )
+
+
+FLAVORS = {
+    "llama-8B": llama3_2_8B,
+    "llama-300M": llama3_2_300M,
+}
+
+
+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
+
+
+def _masked_cross_entropy(logits, targets, mask, vocab_size):
+    losses = F.cross_entropy(logits.reshape(-1, vocab_size), targets.reshape(-1), reduction="none")
+    weights = mask.reshape(-1).to(losses.dtype)
+    total = weights.sum().clamp_min(1.0)
+    return (losses * weights).sum() / total, total
+
+
+@dataclass
+class ModelArgs:
+    backbone_flavor: str
+    decoder_flavor: str
+    text_vocab_size: int
+    audio_vocab_size: int
+    audio_num_codebooks: int
+
+
+MISO_TTS_8B_CONFIG = ModelArgs(
+    backbone_flavor="llama-8B",
+    decoder_flavor="llama-300M",
+    text_vocab_size=128_256,
+    audio_vocab_size=2051,
+    audio_num_codebooks=32,
+)
+
+
+class Model(
+    nn.Module,
+    PyTorchModelHubMixin,
+    pipeline_tag="text-to-speech",
+    license="other",
+):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+
+        self.backbone, backbone_dim = _prepare_transformer(FLAVORS[config.backbone_flavor]())
+        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:
+        """Setup KV caches and return a causal mask."""
+        dtype = next(self.parameters()).dtype
+        device = next(self.parameters()).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))
+
+    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.config.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 forward(
+        self,
+        tokens: torch.Tensor,
+        tokens_mask: torch.Tensor,
+        targets: torch.Tensor,
+        targets_mask: torch.Tensor,
+        decoder_idx: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        dtype = next(self.parameters()).dtype
+        b, s, nc_plus_1 = tokens.size()
+        num_codebooks = nc_plus_1 - 1
+        _, s_amortized = decoder_idx.size()
+
+        targets_c0 = targets[:, :, 0]
+        am_idx = decoder_idx.view(b, s_amortized, 1).expand(b, s_amortized, num_codebooks - 1)
+        targets_c1_plus = torch.gather(targets[:, :, 1:], dim=1, index=am_idx)
+        valid_c0 = targets_mask[:, :, 0].bool()
+        valid_c1_plus = torch.gather(targets_mask[:, :, 1:].bool(), dim=1, index=am_idx)
+
+        embeds = self._embed_tokens(tokens)
+        masked_embeds = embeds * tokens_mask.unsqueeze(-1)
+        h = masked_embeds.sum(dim=2)
+        h = self.backbone(h).to(dtype=dtype)
+
+        c0_logits = self.codebook0_head(h)
+        h = h.unsqueeze(2)
+
+        target_frame = torch.cat([targets, torch.zeros(b, s, 1, device=h.device, dtype=targets.dtype)], dim=2)
+        target_embeds = self._embed_tokens(target_frame)
+        decoder_input = torch.cat([h, target_embeds[:, :, :-2, :]], dim=2)
+
+        idx = decoder_idx.view(b, s_amortized, 1, 1).expand(
+            b,
+            s_amortized,
+            num_codebooks,
+            decoder_input.size(-1),
+        )
+        decoder_input_amortized = torch.gather(decoder_input, dim=1, index=idx)
+        decoder_h = self.decoder(
+            self.projection(decoder_input_amortized).view(b * s_amortized, num_codebooks, -1).to(dtype=dtype)
+        )
+        decoder_h = decoder_h.view(b, s_amortized, num_codebooks, -1)
+
+        logits_c1_plus = torch.einsum(
+            "bsid,idv->bsiv",
+            decoder_h[:, :, 1:, :],
+            self.audio_head,
+        )
+
+        c0_loss, c0_weight = _masked_cross_entropy(
+            c0_logits,
+            targets_c0,
+            valid_c0,
+            self.config.audio_vocab_size,
+        )
+        c1_plus_loss, c1_weight = _masked_cross_entropy(
+            logits_c1_plus,
+            targets_c1_plus,
+            valid_c1_plus,
+            self.config.audio_vocab_size,
+        )
+        loss = (c0_loss * c0_weight + c1_plus_loss * c1_weight) / (c0_weight + c1_weight).clamp_min(1.0)
+        return c0_logits, logits_c1_plus, c0_loss, c1_plus_loss, loss
+
+    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.config.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.config.audio_vocab_size * torch.arange(self.config.audio_num_codebooks, device=tokens.device)
+        )
+        audio_embeds = self.audio_embeddings(audio_tokens.view(-1)).reshape(
+            tokens.size(0), tokens.size(1), self.config.audio_num_codebooks, -1
+        )
+
+        return torch.cat([audio_embeds, text_embeds], dim=-2)

moshi_compat.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from moshi.utils import quantize
+
+
+def patch_bitsandbytes_import_for_unquantized_layers() -> None:
+    original_linear = quantize.linear
+    original_multi_linear = quantize.multi_linear
+
+    def linear(module: nn.Module, x: torch.Tensor, name: str = "weight") -> torch.Tensor:
+        if quantize.is_quantized(module, name):
+            return original_linear(module, x, name)
+        return nn.functional.linear(x, getattr(module, name))
+
+    def multi_linear(
+        num_steps: int,
+        schedule: list[int] | None,
+        module: nn.Module,
+        x: torch.Tensor,
+        offset: int,
+        name: str = "weight",
+    ) -> torch.Tensor:
+        if quantize.is_quantized(module, name):
+            return original_multi_linear(num_steps, schedule, module, x, offset, name)
+
+        weight = getattr(module, name)
+        num_linear = num_steps if schedule is None else max(schedule) + 1
+        weight = weight.view(num_linear, -1, weight.shape[-1])
+
+        outputs = []
+        for t in range(x.shape[1]):
+            linear_index = t + offset
+            if schedule is not None:
+                linear_index = schedule[linear_index]
+            outputs.append(nn.functional.linear(x[:, t], weight[linear_index]))
+        return torch.stack(outputs, 1)
+
+    quantize.linear = linear
+    quantize.multi_linear = multi_linear

requirements.txt

@@ -0,0 +1,12 @@
+torch==2.8.0
+torchaudio==2.8.0
+tokenizers==0.21.0
+transformers==4.49.0
+huggingface_hub==0.28.1
+moshi==0.2.12
+torchtune==0.4.0
+torchao==0.9.0
+silentcipher @ git+https://github.com/SesameAILabs/silentcipher@master
+pydantic==2.10.6
+safetensors
+bitsandbytes; sys_platform == "linux"

watermarking.py

@@ -0,0 +1,87 @@
+import argparse
+
+import silentcipher
+import torch
+import torchaudio
+
+# Warning: When using MisoTTS in another application, you must set this key
+# and keep the watermark key secret.
+MISO_TTS_WATERMARK = [0, 0, 0, 0, 0]
+
+
+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:
+    try:
+        model = silentcipher.get_model(
+            model_type="44.1k",
+            device=device,
+        )
+    except Exception as exc:
+        raise RuntimeError(
+            "Failed to load the SilentCipher watermarking model. Miso TTS weights may "
+            "already be downloaded successfully; this is a separate Hugging Face Hub "
+            "download from sony/silentcipher. Retry the command to resume the cache, or "
+            "pre-download it with `uv run huggingface-cli download sony/silentcipher`."
+        ) from exc
+    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_miso_tts_watermarked = result["messages"][0] == watermark_key
+    else:
+        is_miso_tts_watermarked = False
+
+    return is_watermarked and is_miso_tts_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, MISO_TTS_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()