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>

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+.ruff_cache/
+.pytest_cache/
+*.wav
+.venv/

README.md

@@ -0,0 +1,62 @@
+# ComfyUI-MisoTTS
+
+ComfyUI nodes for [**MisoTTS**](https://github.com/MisoLabsAI/MisoTTS) — an 8B text-to-speech model built on the Sesame **CSM** architecture (Llama-3.2 backbone + audio decoder, Mimi codec, voice cloning).
+
+This is a **modernized** integration: the upstream model pins `torch==2.4`, `torchtune`, and `moshi`, which won't run on recent GPUs (e.g. Blackwell / RTX 50-series) or alongside modern ComfyUI. This pack removes those constraints with **no change in output**.
+
+## What's different from upstream
+
+| Upstream | Here | Why |
+|---|---|---|
+| `torchtune` (deprecated, pins torch 2.4) | vendored plain-torch Llama (`misotts/csm_llama.py`) | runs on any modern torch; numerically **identical** (verified Δ = 0, same weights) |
+| `moshi` + `bitsandbytes` + `sphn` for Mimi | `transformers.MimiModel` | **bit-identical** audio codes; drops 4 heavy deps |
+| gated `meta-llama/Llama-3.2-1B` tokenizer | ungated `unsloth/Llama-3.2-1B` mirror | no HF gating; same tokenizer |
+| loads 32 GB fp32 into CPU RAM | streams weights straight to GPU in bf16 | ~18 GB VRAM, ~0 CPU RAM |
+| watermark on by default | not bundled | minimal deps (re-addable) |
+
+Result: a pin-free `requirements.txt` (just `transformers`, `safetensors`, `tokenizers`, `torchaudio`).
+
+## Requirements
+
+- A CUDA build of PyTorch matching your GPU (for RTX 50-series: `cu128`+, torch ≥ 2.7).
+- ~18 GB free VRAM for bf16 inference.
+- First run downloads the model (~32 GB) from `MisoLabs/MisoTTS`, the Mimi codec (`kyutai/mimi`), and the tokenizer.
+
+```bash
+cd ComfyUI/custom_nodes
+git clone https://github.com/ethanfel/ComfyUI-MisoTTS
+pip install -r ComfyUI-MisoTTS/requirements.txt
+```
+
+## Nodes
+
+- **MisoTTS Model Loader** — loads the model (device / dtype). The 32 GB checkpoint is streamed to the GPU in the chosen dtype.
+- **MisoTTS Generate** — text → speech. Handles long text (whole EPUB chapters) via sentence-aware chunking and keeps a consistent voice across chunks. Optional `ref_audio` + `ref_text` clone a specific voice.
+- **MisoTTS EPUB Loader** — extracts a chapter range from an `.epub` as plain text.
+
+## Audiobook / EPUB workflow
+
+```
+MisoTTS EPUB Loader ──text──▶ MisoTTS Generate ──audio──▶ Save Audio
+MisoTTS Model Loader ─model──▶
+        (optional) Load Audio ──ref_audio──▶
+```
+
+**Voice consistency.** CSM-style models pick a fresh voice on each independent call, so a naïve chapter-at-a-time loop drifts. `MisoTTS Generate` avoids this by feeding the previous chunk(s) back as context (`context_window`, default `1`). For a *specific* narrator voice, connect a `ref_audio` clip (a few seconds) plus its `ref_text` — it's anchored across every chunk. Set a fixed `seed` for reproducible narration.
+
+Key `Generate` parameters:
+
+- `chunk_chars` (300) — target characters per chunk; larger = fewer joins, more VRAM/time.
+- `max_chunk_seconds` (30) — cap on audio generated per chunk.
+- `context_window` (1) — prior chunks reused as context for voice consistency (0 = independent).
+- `silence_ms` (250) — gap inserted between chunks.
+- `temperature` (0.9) / `topk` (50) — sampling.
+
+## Notes
+
+- **Speed**: an 8B autoregressive model at 12.5 Hz × 32 codebooks is ~0.2× realtime in eager mode — fine for batch/audiobook rendering, not live. A `torch.compile` path is a planned optimization.
+- **Watermarking** is not applied. If you redistribute generated audio, consider the upstream project's guidance.
+
+## Credits
+
+Model: [MisoLabsAI/MisoTTS](https://github.com/MisoLabsAI/MisoTTS) (Sesame CSM architecture). Mimi codec: Kyutai. This repo only provides the ComfyUI integration and the torchtune/moshi-free runtime.

__init__.py

@@ -0,0 +1,15 @@
+from .nodes import MisoTTSModelLoader, MisoTTSGenerate, MisoTTSEpubLoader
+
+NODE_CLASS_MAPPINGS = {
+    "MisoTTSModelLoader": MisoTTSModelLoader,
+    "MisoTTSGenerate": MisoTTSGenerate,
+    "MisoTTSEpubLoader": MisoTTSEpubLoader,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "MisoTTSModelLoader": "MisoTTS Model Loader",
+    "MisoTTSGenerate": "MisoTTS Generate",
+    "MisoTTSEpubLoader": "MisoTTS EPUB Loader",
+}
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

misotts/__init__.py

@@ -0,0 +1,3 @@
+from .inference import Generator, Segment, load_miso_8b
+
+__all__ = ["Generator", "Segment", "load_miso_8b"]

misotts/csm_llama.py

@@ -0,0 +1,308 @@
+"""
+torchtune-free reimplementation of the Llama3.2 TransformerDecoder used by CSM/MisoTTS.
+
+Goal: drop-in replacement for `torchtune.models.llama3_2.llama3_2(...)` followed by
+`_prepare_transformer` (tok_embeddings -> Identity, output -> Identity). The module:
+
+  * produces an IDENTICAL state_dict key layout, so MisoTTS weights load unchanged
+  * computes numerically identical outputs (RoPE scaling, GQA, RMSNorm, SwiGLU, KV-cache)
+  * exposes the methods models.py relies on: setup_caches / caches_are_enabled /
+    reset_caches / max_seq_len, and forward(h, input_pos=, mask=) returning float32.
+
+All math is copied 1:1 from torchtune 0.4.0 so this can be validated by diffing.
+No torchtune / torchao / torch-pin dependency: plain torch.
+"""
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ----------------------------------------------------------------------------- RoPE
+class Llama3ScaledRoPE(nn.Module):
+    """Verbatim port of torchtune.models.llama3_1._position_embeddings.Llama3ScaledRoPE."""
+
+    def __init__(self, dim, max_seq_len=4096, base=10_000, scale_factor=8,
+                 low_freq_factor=1, high_freq_factor=4, old_context_len=8192):
+        super().__init__()
+        self.dim = dim
+        self.base = base
+        self.max_seq_len = max_seq_len
+        self.scale_factor = scale_factor
+        self.low_freq_factor = low_freq_factor
+        self.high_freq_factor = high_freq_factor
+        self.old_context_len = old_context_len
+        self.is_cache_built = False
+        self.rope_init()
+
+    def rope_init(self):
+        freqs = 1.0 / (self.base ** (torch.arange(0, self.dim, 2)[: (self.dim // 2)].float() / self.dim))
+        if freqs.is_meta:
+            return
+        theta = self.apply_scaling(freqs, self.scale_factor, self.low_freq_factor,
+                                   self.high_freq_factor, self.old_context_len)
+        self.register_buffer("theta", theta, persistent=False)
+        self.build_rope_cache(self.max_seq_len)
+        self.is_cache_built = True
+
+    def build_rope_cache(self, max_seq_len=4096):
+        seq_idx = torch.arange(max_seq_len, dtype=self.theta.dtype, device=self.theta.device)
+        idx_theta = torch.einsum("i, j -> ij", seq_idx, self.theta).float()
+        cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)
+        self.register_buffer("cache", cache, persistent=False)
+
+    def apply_scaling(self, freqs, scale_factor, low_freq_factor, high_freq_factor, old_context_len):
+        low_freq_wavelen = old_context_len / low_freq_factor
+        high_freq_wavelen = old_context_len / high_freq_factor
+        new_freqs = []
+        for freq in freqs:
+            wavelen = 2 * math.pi / freq
+            if wavelen < high_freq_wavelen:
+                new_freqs.append(freq)
+            elif wavelen > low_freq_wavelen:
+                new_freqs.append(freq / scale_factor)
+            else:
+                assert low_freq_wavelen != high_freq_wavelen
+                smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
+                new_freqs.append((1 - smooth) * freq / scale_factor + smooth * freq)
+        return torch.tensor(new_freqs, dtype=freqs.dtype, device=freqs.device)
+
+    def forward(self, x, *, input_pos=None):
+        if not self.is_cache_built:
+            raise RuntimeError("RoPE cache is not built. Please call rope_init() first.")
+        seq_len = x.size(1)
+        rope_cache = self.cache[:seq_len] if input_pos is None else self.cache[input_pos]
+        xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+        rope_cache = rope_cache.view(-1, xshaped.size(1), 1, xshaped.size(3), 2)
+        x_out = torch.stack([
+            xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
+            xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
+        ], -1)
+        x_out = x_out.flatten(3)
+        return x_out.type_as(x)
+
+
+# ----------------------------------------------------------------------------- RMSNorm
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        x_fp32 = x.float()
+        x_normed = (x_fp32 * torch.rsqrt(x_fp32.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x)
+        return x_normed * self.scale
+
+
+# ----------------------------------------------------------------------------- KV cache
+class KVCache(nn.Module):
+    def __init__(self, batch_size, max_seq_len, num_heads, head_dim, dtype):
+        super().__init__()
+        cache_shape = (batch_size, num_heads, max_seq_len, head_dim)
+        self.register_buffer("k_cache", torch.zeros(cache_shape, dtype=dtype), persistent=False)
+        self.register_buffer("v_cache", torch.zeros(cache_shape, dtype=dtype), persistent=False)
+        self.register_buffer("cache_pos", torch.arange(0, cache_shape[2]), persistent=False)
+        self.batch_size = batch_size
+
+    def reset(self):
+        self.k_cache.zero_()
+        self.v_cache.zero_()
+        self.cache_pos -= self.size
+
+    @property
+    def size(self):
+        return self.cache_pos[0].item()
+
+    def update(self, k_val, v_val):
+        bsz, _, seq_len, _ = k_val.shape
+        if bsz > self.k_cache.shape[0]:
+            raise ValueError("batch size larger than cache")
+        assert (self.cache_pos[0] + seq_len) <= self.k_cache.shape[2]
+        k_out = self.k_cache
+        v_out = self.v_cache
+        k_out[:, :, self.cache_pos[:seq_len]] = k_val
+        v_out[:, :, self.cache_pos[:seq_len]] = v_val
+        self.cache_pos += seq_len
+        return k_out, v_out
+
+
+# ----------------------------------------------------------------------------- Attention
+class MultiHeadAttention(nn.Module):
+    def __init__(self, *, embed_dim, num_heads, num_kv_heads, head_dim,
+                 pos_embeddings, max_seq_len=4096, attn_dropout=0.0):
+        super().__init__()
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.embed_dim = embed_dim
+        self.attn_dropout = attn_dropout
+        self.head_dim = head_dim
+        self.max_seq_len = max_seq_len
+        self.is_causal = True
+        self.kv_cache = None
+        self.q_proj = nn.Linear(embed_dim, num_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(embed_dim, num_kv_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(embed_dim, num_kv_heads * head_dim, bias=False)
+        self.output_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.pos_embeddings = pos_embeddings
+        self.cache_enabled = False
+
+    def setup_cache(self, batch_size, dtype, max_seq_len):
+        if self.kv_cache is not None:
+            return
+        self.kv_cache = KVCache(batch_size=batch_size, max_seq_len=max_seq_len,
+                                num_heads=self.num_heads, head_dim=self.head_dim, dtype=dtype)
+        self.cache_enabled = True
+
+    def reset_cache(self):
+        self.kv_cache.reset()
+
+    @staticmethod
+    def _sdpa(q, k, v, mask, dropout_p, is_causal):
+        if mask is not None:
+            mask = mask[:, None, :, :]
+        return F.scaled_dot_product_attention(q, k, v, attn_mask=mask,
+                                              dropout_p=dropout_p, is_causal=is_causal)
+
+    def forward(self, x, y=None, *, mask=None, input_pos=None):
+        b, s_x, _ = x.shape
+        s_y = y.shape[1] if y is not None else 0
+        q = self.q_proj(x)
+        q_per_kv = self.num_heads // self.num_kv_heads
+        q = q.view(b, s_x, self.num_kv_heads * q_per_kv, self.head_dim)
+        if self.pos_embeddings is not None:
+            q = self.pos_embeddings(q, input_pos=input_pos)
+        q = q.transpose(1, 2)
+
+        if y is None:
+            k = self.kv_cache.k_cache
+            v = self.kv_cache.v_cache
+        else:
+            k = self.k_proj(y)
+            v = self.v_proj(y)
+            k = k.view(b, s_y, -1, self.head_dim)
+            if self.pos_embeddings is not None:
+                k = self.pos_embeddings(k, input_pos=input_pos)
+            k = k.view(b, s_y, self.num_kv_heads, 1, self.head_dim)
+            v = v.view(b, s_y, self.num_kv_heads, 1, self.head_dim)
+            if self.num_heads != self.num_kv_heads:
+                k = k.expand(b, s_y, self.num_kv_heads, q_per_kv, self.head_dim)
+                v = v.expand(b, s_y, self.num_kv_heads, q_per_kv, self.head_dim)
+            k = k.reshape(b, s_y, -1, self.head_dim)
+            v = v.reshape(b, s_y, -1, self.head_dim)
+            k = k.transpose(1, 2)
+            v = v.transpose(1, 2)
+            if self.kv_cache is not None and self.cache_enabled:
+                k, v = self.kv_cache.update(k, v)
+
+        output = self._sdpa(q, k, v, mask=mask,
+                            dropout_p=self.attn_dropout if self.training else 0.0,
+                            is_causal=self.kv_cache is None and mask is None and self.is_causal)
+        output = output.transpose(1, 2).contiguous().view(b, s_x, -1)
+        return self.output_proj(output)
+
+
+# ----------------------------------------------------------------------------- MLP
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+        self.activation = nn.SiLU()
+
+    def forward(self, x):
+        h = self.activation(self.w1(x))
+        h = h * self.w3(x)
+        return self.w2(h)
+
+
+# ----------------------------------------------------------------------------- Layer
+class TransformerSelfAttentionLayer(nn.Module):
+    def __init__(self, attn, mlp, sa_norm, mlp_norm):
+        super().__init__()
+        self.attn = attn
+        self.mlp = mlp
+        self.sa_norm = sa_norm
+        self.mlp_norm = mlp_norm
+
+    def setup_caches(self, batch_size, dtype, decoder_max_seq_len):
+        self.attn.setup_cache(batch_size, dtype, max_seq_len=decoder_max_seq_len)
+
+    def caches_are_enabled(self):
+        return self.attn.cache_enabled
+
+    def reset_cache(self):
+        self.attn.reset_cache()
+
+    def forward(self, x, *, mask=None, input_pos=None):
+        h = self.sa_norm(x)
+        attn_out = self.attn(h, h, mask=mask, input_pos=input_pos)
+        h = attn_out + x
+        mlp_out = self.mlp(self.mlp_norm(h))
+        return h + mlp_out
+
+
+# ----------------------------------------------------------------------------- Decoder
+class TransformerDecoder(nn.Module):
+    """Self-attention-only Llama decoder operating on pre-computed embeddings (h).
+
+    Mirrors torchtune's TransformerDecoder after _prepare_transformer replaced
+    tok_embeddings and output with Identity: forward takes embeddings, returns float32.
+    """
+
+    def __init__(self, *, layers, norm, max_seq_len, num_heads, head_dim):
+        super().__init__()
+        self.layers = nn.ModuleList(layers)
+        self.norm = norm
+        self.max_seq_len = max_seq_len
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.embed_dim = num_heads * head_dim
+
+    def setup_caches(self, batch_size, dtype, *, decoder_max_seq_len=None):
+        max_seq_len = decoder_max_seq_len if decoder_max_seq_len is not None else self.max_seq_len
+        self.decoder_max_cache_seq_len = max_seq_len
+        for layer in self.layers:
+            layer.setup_caches(batch_size, dtype, decoder_max_seq_len=max_seq_len)
+
+    def caches_are_enabled(self):
+        return self.layers[0].caches_are_enabled()
+
+    def reset_caches(self):
+        for layer in self.layers:
+            layer.reset_cache()
+
+    def forward(self, h, *, mask=None, input_pos=None):
+        for layer in self.layers:
+            h = layer(h, mask=mask, input_pos=input_pos)
+        h = self.norm(h)
+        return h.float()
+
+
+# ----------------------------------------------------------------------------- builder
+def llama3_2(*, vocab_size, num_layers, num_heads, num_kv_heads, embed_dim,
+             max_seq_len, intermediate_dim, attn_dropout=0.0, norm_eps=1e-5,
+             rope_base=500_000, scale_factor=32):
+    """Matches torchtune.models.llama3_2.llama3_2(...) + _prepare_transformer.
+
+    vocab_size is accepted for signature parity but unused (tok_embeddings/output
+    are Identity in CSM, so they carry no weights).
+    """
+    head_dim = embed_dim // num_heads
+    rope = Llama3ScaledRoPE(dim=head_dim, max_seq_len=max_seq_len, base=rope_base, scale_factor=scale_factor)
+    layers = []
+    for _ in range(num_layers):
+        attn = MultiHeadAttention(embed_dim=embed_dim, num_heads=num_heads, num_kv_heads=num_kv_heads,
+                                  head_dim=head_dim, pos_embeddings=rope, max_seq_len=max_seq_len,
+                                  attn_dropout=attn_dropout)
+        mlp = FeedForward(dim=embed_dim, hidden_dim=intermediate_dim)
+        layers.append(TransformerSelfAttentionLayer(
+            attn=attn, mlp=mlp,
+            sa_norm=RMSNorm(embed_dim, eps=norm_eps),
+            mlp_norm=RMSNorm(embed_dim, eps=norm_eps),
+        ))
+    return TransformerDecoder(layers=layers, norm=RMSNorm(embed_dim, eps=norm_eps),
+                              max_seq_len=max_seq_len, num_heads=num_heads, head_dim=head_dim)

misotts/inference.py

@@ -0,0 +1,192 @@
+"""
+MisoTTS inference engine — modernized, dependency-light.
+
+  * audio codec: transformers.MimiModel (bit-identical to moshi's Mimi; no moshi/sphn)
+  * transformer: vendored csm_llama (no torchtune / torch pin)
+  * loader: streams the 32GB fp32 checkpoint straight to GPU in bf16 (peak ~18GB VRAM,
+    ~0 CPU RAM) — works on machines with little free system RAM
+  * watermarking: not bundled (upstream applies silentcipher by default; omitted here to
+    keep deps to transformers/safetensors/torchaudio/tokenizers)
+"""
+import os
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import torch
+from huggingface_hub import hf_hub_download
+from safetensors import safe_open
+from tokenizers.processors import TemplateProcessing
+from transformers import AutoTokenizer, MimiModel
+
+from .models import MISO_TTS_8B_CONFIG, Model, ModelArgs
+
+DEFAULT_MISO_TTS_REPO_ID = "MisoLabs/MisoTTS"
+DEFAULT_TOKENIZER = "unsloth/Llama-3.2-1B"   # ungated, byte-identical to meta-llama/Llama-3.2-1B
+DEFAULT_MIMI = "kyutai/mimi"
+
+
+@dataclass
+class Segment:
+    speaker: int
+    text: str
+    audio: torch.Tensor   # (num_samples,), sample_rate = 24_000
+
+
+def load_llama3_tokenizer(name: str = DEFAULT_TOKENIZER):
+    tokenizer = AutoTokenizer.from_pretrained(name)
+    bos, eos = tokenizer.bos_token, 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 MimiAdapter:
+    """moshi-Mimi API (encode/decode/set_num_codebooks/sample_rate) over transformers.MimiModel."""
+
+    def __init__(self, mimi_model: MimiModel, num_codebooks: int):
+        self.m = mimi_model
+        self.num_q = num_codebooks
+        self.sample_rate = mimi_model.config.sampling_rate
+
+    def set_num_codebooks(self, n: int):
+        self.num_q = n
+
+    @torch.inference_mode()
+    def encode(self, x: torch.Tensor) -> torch.Tensor:        # [B,1,T] -> [B,K,T]
+        return self.m.encode(x, num_quantizers=self.num_q).audio_codes
+
+    @torch.inference_mode()
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:    # [B,K,T] -> [B,1,T']
+        return self.m.decode(codes).audio_values
+
+
+class Generator:
+    def __init__(self, model: Model, tokenizer_name: str = DEFAULT_TOKENIZER):
+        self._model = model
+        self._model.setup_caches(1)
+        self._text_tokenizer = load_llama3_tokenizer(tokenizer_name)
+        self._frame_size = self._model.config.audio_num_codebooks + 1
+
+        device = next(model.parameters()).device
+        mimi = MimiModel.from_pretrained(DEFAULT_MIMI).to(device).eval()
+        self._audio_tokenizer = MimiAdapter(mimi, self._model.config.audio_num_codebooks)
+
+        self.sample_rate = self._audio_tokenizer.sample_rate
+        self.device = device
+
+    def _tokenize_text_segment(self, text: str, speaker: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        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
+        return text_frame.to(self.device), text_frame_mask.to(self.device)
+
+    def _tokenize_audio(self, audio: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        assert audio.ndim == 1, "Audio must be single channel"
+        audio = audio.to(self.device)
+        audio_tokens = self._audio_tokenizer.encode(audio.unsqueeze(0).unsqueeze(0))[0]
+        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
+        return audio_frame, audio_frame_mask
+
+    def _tokenize_segment(self, segment: Segment) -> Tuple[torch.Tensor, torch.Tensor]:
+        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:
+            st, sm = self._tokenize_segment(segment)
+            tokens.append(st)
+            tokens_mask.append(sm)
+        gt, gm = self._tokenize_text_segment(text, speaker)
+        tokens.append(gt)
+        tokens_mask.append(gm)
+
+        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_context_len = 2048 - max_generation_len
+        if curr_tokens.size(1) >= max_context_len:
+            raise ValueError(
+                f"Inputs too long ({curr_tokens.size(1)} frames), must be below "
+                f"{max_context_len}. Reduce context_window or chunk size."
+            )
+
+        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
+
+        if not samples:
+            raise RuntimeError("No audio frames generated (immediate EOS).")
+
+        return self._audio_tokenizer.decode(torch.stack(samples).permute(1, 2, 0)).squeeze(0).squeeze(0)
+
+
+def _resolve_checkpoint(model_path_or_repo_id: str) -> str:
+    if os.path.isfile(model_path_or_repo_id):
+        return model_path_or_repo_id
+    if os.path.isdir(model_path_or_repo_id):
+        return os.path.join(model_path_or_repo_id, "model.safetensors")
+    return hf_hub_download(repo_id=model_path_or_repo_id, filename="model.safetensors")
+
+
+def _load_model_lowmem(safetensors_path: str, config: ModelArgs, device: str, dtype: torch.dtype) -> Model:
+    """Build the bf16 model directly on GPU and stream weights from disk, casting fp32->bf16
+    per tensor. Avoids holding the 32GB fp32 checkpoint in CPU RAM."""
+    prev = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    try:
+        with torch.device(device):
+            model = Model(config)
+    finally:
+        torch.set_default_dtype(prev)
+
+    msd = model.state_dict()
+    loaded = set()
+    with safe_open(safetensors_path, framework="pt", device=device) as f:
+        ckpt_keys = set(f.keys())
+        for k in f.keys():
+            if k not in msd:
+                continue
+            msd[k].copy_(f.get_tensor(k))
+            loaded.add(k)
+    missing = set(msd.keys()) - loaded
+    unexpected = ckpt_keys - set(msd.keys())
+    if missing or unexpected:
+        raise RuntimeError(f"checkpoint key mismatch:\n  missing={sorted(missing)}\n  unexpected={sorted(unexpected)}")
+    model.eval()
+    return model
+
+
+def load_miso_8b(device: str = "cuda", model_path_or_repo_id: Optional[str] = None,
+                 dtype: torch.dtype = torch.bfloat16, tokenizer_name: str = DEFAULT_TOKENIZER) -> Generator:
+    source = model_path_or_repo_id or os.environ.get("MISO_TTS_8B_MODEL", DEFAULT_MISO_TTS_REPO_ID)
+    ckpt = _resolve_checkpoint(source)
+    model = _load_model_lowmem(ckpt, MISO_TTS_8B_CONFIG, device=device, dtype=dtype)
+    return Generator(model, tokenizer_name=tokenizer_name)

misotts/models.py

@@ -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)

nodes/__init__.py

@@ -0,0 +1,5 @@
+from .loader import MisoTTSModelLoader
+from .generator import MisoTTSGenerate
+from .epub_loader import MisoTTSEpubLoader
+
+__all__ = ["MisoTTSModelLoader", "MisoTTSGenerate", "MisoTTSEpubLoader"]

nodes/epub_loader.py

@@ -0,0 +1,96 @@
+import re
+import zipfile
+import xml.etree.ElementTree as ET
+
+from bs4 import BeautifulSoup
+
+_BLOCK_TAGS = {"p", "h1", "h2", "h3", "h4", "h5", "h6", "li", "div", "br", "tr"}
+
+
+def _local(tag):
+    return tag.split("}")[-1]
+
+
+def _extract_chapters(epub_path):
+    chapters = []
+    with zipfile.ZipFile(epub_path, "r") as zf:
+        container = ET.fromstring(zf.read("META-INF/container.xml"))
+        rootfile = next(el for el in container.iter() if _local(el.tag) == "rootfile")
+        opf_path = rootfile.attrib["full-path"]
+        opf_dir = opf_path.rsplit("/", 1)[0] + "/" if "/" in opf_path else ""
+
+        opf = ET.fromstring(zf.read(opf_path))
+        manifest = {
+            el.attrib["id"]: el.attrib["href"]
+            for el in opf.iter()
+            if _local(el.tag) == "item" and "xhtml" in el.attrib.get("media-type", "")
+        }
+        spine = [el.attrib["idref"] for el in opf.iter() if _local(el.tag) == "itemref"]
+
+        for idref in spine:
+            href = manifest.get(idref)
+            if href is None:
+                continue
+            xhtml = zf.read(opf_dir + href).decode("utf-8", errors="replace")
+            soup = BeautifulSoup(xhtml, "html.parser")
+            for tag in soup(["script", "style"]):
+                tag.decompose()
+            title = None
+            if soup.title and soup.title.string:
+                title = soup.title.string.strip()
+            if not title:
+                for hn in ["h1", "h2", "h3"]:
+                    tag = soup.find(hn)
+                    if tag:
+                        title = tag.get_text(strip=True)
+                        break
+            if soup.title:
+                soup.title.decompose()
+            for hn in ["h1", "h2", "h3"]:
+                for tag in soup.find_all(hn):
+                    tag.decompose()
+            for tag in soup.find_all(_BLOCK_TAGS):
+                tag.append(soup.new_string("\n\n"))
+            text = soup.get_text(separator="")
+            text = re.sub(r"[^\S\n]+", " ", text)
+            text = re.sub(r" *\n *", "\n", text)
+            text = re.sub(r"\n{3,}", "\n\n", text)
+            chapters.append({"title": title, "text": text.strip()})
+    return chapters
+
+
+class MisoTTSEpubLoader:
+    """Load an EPUB and emit a chapter range as text, ready for the MisoTTS Generate node."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "epub_path": ("STRING", {"default": "", "tooltip": "Absolute path to the .epub file."}),
+                "chapter_start": ("INT", {"default": 1, "min": 1, "max": 9999, "step": 1,
+                                          "tooltip": "First chapter (1-indexed). Clamped to valid range."}),
+                "chapter_end": ("INT", {"default": 1, "min": 1, "max": 9999, "step": 1,
+                                        "tooltip": "Last chapter (1-indexed, inclusive). Clamped automatically."}),
+            },
+        }
+
+    RETURN_TYPES = ("STRING", "STRING", "STRING")
+    RETURN_NAMES = ("text", "chapter_title", "chapter_list")
+    FUNCTION = "load_epub"
+    CATEGORY = "MisoTTS"
+
+    def load_epub(self, epub_path, chapter_start, chapter_end):
+        chapters = _extract_chapters(epub_path)
+        n = len(chapters)
+        if n == 0:
+            return ("", "", "")
+        start = max(1, min(chapter_start, n))
+        end = max(start, min(chapter_end, n))
+        chapter_list = "\n".join(
+            f"{i}. {ch['title'] if ch['title'] else f'Chapter {i}'}"
+            for i, ch in enumerate(chapters, 1)
+        )
+        first = chapters[start - 1]
+        chapter_title = first["title"] if first["title"] else f"Chapter {start}"
+        text = "\n\n---\n\n".join(ch["text"] for ch in chapters[start - 1: end])
+        return (text, chapter_title, chapter_list)

nodes/generator.py

@@ -0,0 +1,175 @@
+import re
+
+import torch
+import torchaudio
+
+from ..misotts import Segment
+
+
+# --------------------------------------------------------------------------- audio helpers
+def _audio_to_mono24k(audio_dict, sr_target=24000):
+    """ComfyUI AUDIO dict -> 1-D mono tensor at 24 kHz (Mimi's rate)."""
+    wav = audio_dict["waveform"]
+    sr = int(audio_dict["sample_rate"])
+    if wav.dim() == 3:
+        wav = wav[0]                      # (C, T)
+    if wav.shape[0] > 1:
+        wav = wav.mean(0, keepdim=True)   # mix to mono
+    if sr != sr_target:
+        wav = torchaudio.functional.resample(wav, sr, sr_target)
+    return wav.squeeze(0).contiguous().float()
+
+
+# --------------------------------------------------------------------------- text chunking
+def _split_sentences(text):
+    parts = re.split(r"(?<=[.!?…])\s+", text.strip())
+    return [p.strip() for p in parts if p.strip()]
+
+
+def _hard_split(s, max_chars):
+    """Break an over-long sentence on commas, then on words, so no chunk exceeds max_chars."""
+    out, cur = [], ""
+    for tok in re.split(r"(?<=,)\s+", s):
+        if cur and len(cur) + 1 + len(tok) > max_chars:
+            out.append(cur)
+            cur = tok
+        else:
+            cur = f"{cur} {tok}".strip()
+    if cur:
+        out.append(cur)
+    final = []
+    for c in out:
+        if len(c) <= max_chars:
+            final.append(c)
+            continue
+        cc = ""
+        for w in c.split():
+            if cc and len(cc) + 1 + len(w) > max_chars:
+                final.append(cc)
+                cc = w
+            else:
+                cc = f"{cc} {w}".strip()
+        if cc:
+            final.append(cc)
+    return final
+
+
+def _chunk_text(text, max_chars):
+    """Sentence-aware chunking. Respects paragraph breaks and EPUB '---' chapter markers,
+    packs whole sentences up to max_chars, and hard-splits any sentence longer than that."""
+    chunks = []
+    paragraphs = re.split(r"\n\s*\n|\n?-{3,}\n?", text)
+    for para in paragraphs:
+        para = para.strip()
+        if not para:
+            continue
+        cur = ""
+        for s in _split_sentences(para):
+            if len(s) > max_chars:
+                if cur:
+                    chunks.append(cur)
+                    cur = ""
+                chunks.extend(_hard_split(s, max_chars))
+                continue
+            if cur and len(cur) + 1 + len(s) > max_chars:
+                chunks.append(cur)
+                cur = s
+            else:
+                cur = f"{cur} {s}".strip()
+        if cur:
+            chunks.append(cur)
+    return chunks
+
+
+# --------------------------------------------------------------------------- node
+class MisoTTSGenerate:
+    """Generate speech from text. Handles arbitrarily long text (audiobooks/EPUB chapters)
+    by sentence-aware chunking, and keeps a consistent voice across chunks by feeding prior
+    audio (and an optional reference clip) back as context — CSM models otherwise drift."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MISOTTS_MODEL", {"tooltip": "Loaded by the MisoTTS Model Loader node."}),
+                "text": ("STRING", {"multiline": True, "default": "",
+                                     "tooltip": "Text to synthesize. Long text is chunked automatically."}),
+            },
+            "optional": {
+                "ref_audio": ("AUDIO", {
+                    "tooltip": "Optional reference clip to clone the voice from. Anchored across every chunk.",
+                }),
+                "ref_text": ("STRING", {"default": "",
+                                        "tooltip": "Transcript of ref_audio. Improves cloning quality."}),
+                "speaker": ("INT", {"default": 0, "min": 0, "max": 31,
+                                    "tooltip": "Speaker id. Keep fixed for a single narrator."}),
+                "temperature": ("FLOAT", {"default": 0.9, "min": 0.1, "max": 2.0, "step": 0.05,
+                                          "tooltip": "Sampling temperature. Lower = steadier, higher = more varied."}),
+                "topk": ("INT", {"default": 50, "min": 1, "max": 500,
+                                 "tooltip": "Top-k sampling cutoff."}),
+                "max_chunk_seconds": ("FLOAT", {"default": 30.0, "min": 5.0, "max": 90.0, "step": 1.0,
+                                                "tooltip": "Max audio length generated per text chunk."}),
+                "chunk_chars": ("INT", {"default": 300, "min": 50, "max": 2000, "step": 10,
+                                        "tooltip": "Target characters per chunk. Larger = fewer joins, more VRAM/time."}),
+                "context_window": ("INT", {"default": 1, "min": 0, "max": 4,
+                                           "tooltip": (
+                                               "How many previous chunks to feed back as context to keep the voice "
+                                               "consistent. 1 is a good default; 0 makes each chunk independent "
+                                               "(voice may drift). Higher = steadier but slower / more VRAM.")}),
+                "silence_ms": ("INT", {"default": 250, "min": 0, "max": 2000, "step": 10,
+                                       "tooltip": "Silence inserted between chunks."}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 2**32 - 1,
+                                 "tooltip": "0 = random each run. Set a fixed value for reproducible narration."}),
+            },
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    FUNCTION = "generate"
+    CATEGORY = "MisoTTS"
+
+    def generate(self, model, text, ref_audio=None, ref_text="", speaker=0, temperature=0.9,
+                 topk=50, max_chunk_seconds=30.0, chunk_chars=300, context_window=1,
+                 silence_ms=250, seed=0):
+        if seed != 0:
+            torch.manual_seed(seed)
+        text = (text or "").strip()
+        if not text:
+            raise ValueError("MisoTTS Generate: text is empty.")
+
+        chunks = _chunk_text(text, int(chunk_chars))
+        if not chunks:
+            raise ValueError("MisoTTS Generate: no text chunks produced.")
+
+        sr = int(model.sample_rate)
+        ms = float(max_chunk_seconds) * 1000.0
+
+        ref_seg = None
+        if ref_audio is not None:
+            ref_seg = Segment(speaker=int(speaker), text=(ref_text or "").strip(),
+                              audio=_audio_to_mono24k(ref_audio, sr))
+
+        gap = torch.zeros(int(sr * silence_ms / 1000.0)) if silence_ms > 0 else None
+        keep = max(int(context_window), 1)
+
+        history, pieces = [], []
+        for i, chunk in enumerate(chunks):
+            ctx = []
+            if ref_seg is not None:
+                ctx.append(ref_seg)
+            if context_window > 0:
+                ctx.extend(history[-context_window:])
+
+            audio = model.generate(text=chunk, speaker=int(speaker), context=ctx,
+                                   max_audio_length_ms=ms, temperature=float(temperature), topk=int(topk))
+            audio = audio.detach().to("cpu", torch.float32)
+
+            if i > 0 and gap is not None:
+                pieces.append(gap)
+            pieces.append(audio)
+
+            history.append(Segment(speaker=int(speaker), text=chunk, audio=audio))
+            history = history[-keep:]
+
+        waveform = torch.cat(pieces).unsqueeze(0).unsqueeze(0)   # (1, 1, T)
+        return ({"waveform": waveform, "sample_rate": sr},)

nodes/loader.py

@@ -0,0 +1,70 @@
+import os
+
+import torch
+
+_import_error = None
+try:
+    from ..misotts import load_miso_8b
+except Exception as e:  # pragma: no cover - surfaced to the user at node runtime
+    load_miso_8b = None
+    _import_error = e
+
+try:
+    import folder_paths
+    CACHE_DIR = os.path.join(folder_paths.models_dir, "misotts")
+except ImportError:
+    CACHE_DIR = os.path.join(os.path.expanduser("~"), ".cache", "misotts")
+
+DTYPE_MAP = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}
+
+
+class MisoTTSModelLoader:
+    """Load the MisoTTS 8B model (modernized: no torchtune/moshi).
+
+    The 32 GB fp32 checkpoint is streamed straight to the GPU in the chosen dtype, so
+    loading needs ~18 GB VRAM (bf16) and almost no system RAM.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "device": (["cuda:0", "cuda:1", "cpu"], {"default": "cuda:0"}),
+                "dtype": (["bfloat16", "float16", "float32"], {"default": "bfloat16"}),
+            },
+            "optional": {
+                "model_repo_or_path": ("STRING", {
+                    "default": "MisoLabs/MisoTTS",
+                    "tooltip": "HF repo id or a local path to a model.safetensors / model dir.",
+                }),
+                "tokenizer": ("STRING", {
+                    "default": "unsloth/Llama-3.2-1B",
+                    "tooltip": (
+                        "Llama-3.2 tokenizer source. Default is an ungated mirror byte-identical "
+                        "to meta-llama/Llama-3.2-1B. Change only if you know what you're doing."
+                    ),
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("MISOTTS_MODEL",)
+    RETURN_NAMES = ("model",)
+    FUNCTION = "load_model"
+    CATEGORY = "MisoTTS"
+
+    def load_model(self, device, dtype, model_repo_or_path="MisoLabs/MisoTTS", tokenizer="unsloth/Llama-3.2-1B"):
+        if load_miso_8b is None:
+            raise ImportError(
+                "MisoTTS engine failed to import. Ensure transformers, safetensors, tokenizers "
+                f"and torchaudio are installed.\nOriginal error: {_import_error}"
+            )
+        os.makedirs(CACHE_DIR, exist_ok=True)
+        os.environ.setdefault("HF_HOME", CACHE_DIR)
+        source = model_repo_or_path.strip() or "MisoLabs/MisoTTS"
+        gen = load_miso_8b(
+            device=device,
+            model_path_or_repo_id=source,
+            dtype=DTYPE_MAP[dtype],
+            tokenizer_name=tokenizer.strip() or "unsloth/Llama-3.2-1B",
+        )
+        return (gen,)

pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "comfyui-misotts"
+description = "ComfyUI nodes for MisoTTS (Sesame CSM 8B) — modernized off torchtune/moshi, with EPUB/audiobook chunking and voice cloning."
+version = "0.1.0"
+license = { text = "Apache-2.0" }
+dependencies = []
+
+[project.urls]
+Repository = "https://github.com/ethanfel/ComfyUI-MisoTTS"
+
+[tool.comfy]
+PublisherId = "ethanfel"
+DisplayName = "ComfyUI-MisoTTS"
+Icon = ""

requirements.txt

@@ -0,0 +1,6 @@
+transformers>=4.46.0
+safetensors
+tokenizers
+torchaudio
+huggingface_hub
+beautifulsoup4