Initial release: ComfyUI-MisoTTS (modernized CSM 8B)

Modernized MisoTTS integration for ComfyUI with no torchtune/moshi: - vendored plain-torch Llama backbone (csm_llama), parity-verified Δ=0 vs torchtune - transformers.MimiModel codec (bit-identical codes to moshi), drops moshi/bnb/sphn - low-memory loader: streams 32GB fp32 checkpoint to GPU in bf16 (~18GB VRAM) - nodes: Model Loader, Generate (audiobook chunking + voice anchoring), EPUB Loader - pin-free requirements; runs on modern torch / Blackwell GPUs Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-06 23:37:54 +02:00
commit f7a6f7790d
13 changed files with 1110 additions and 0 deletions
@@ -0,0 +1,157 @@
+"""
+Modernized MisoTTS models.py — identical Model logic, torchtune removed.
+
+The ONLY change vs the upstream models.py is the transformer source:
+    upstream:  from torchtune.models import llama3_2  (drags torch==2.4, deprecated)
+    here:      import csm_llama                        (plain torch, parity-verified Δ=0)
+
+Everything else (generate_frame, _embed_*, setup_caches, sampling) is byte-for-byte
+the upstream logic, so the published checkpoint loads with identical keys.
+"""
+from dataclasses import dataclass
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+
+from . import csm_llama
+
+
+def llama3_2_8B():
+    return csm_llama.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 csm_llama.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):
+    # csm_llama decoders are already "prepared" (no tok_embeddings/output params).
+    return model, 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 _multinomial_sample_one_no_sync(probs):
+    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("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)
+    return _multinomial_sample_one_no_sync(probs)
+
+
+@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):
+    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:
+        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))
+
+    def generate_frame(self, tokens, tokens_mask, input_pos, temperature, topk):
+        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)
+
+        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 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)