From f7a6f7790d5b3f8b5c22f9a631e14ff080200b2e Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sat, 6 Jun 2026 23:37:54 +0200 Subject: [PATCH] Initial release: ComfyUI-MisoTTS (modernized CSM 8B) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .gitignore | 7 + README.md | 62 +++++++++ __init__.py | 15 +++ misotts/__init__.py | 3 + misotts/csm_llama.py | 308 +++++++++++++++++++++++++++++++++++++++++++ misotts/inference.py | 192 +++++++++++++++++++++++++++ misotts/models.py | 157 ++++++++++++++++++++++ nodes/__init__.py | 5 + nodes/epub_loader.py | 96 ++++++++++++++ nodes/generator.py | 175 ++++++++++++++++++++++++ nodes/loader.py | 70 ++++++++++ pyproject.toml | 14 ++ requirements.txt | 6 + 13 files changed, 1110 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 __init__.py create mode 100644 misotts/__init__.py create mode 100644 misotts/csm_llama.py create mode 100644 misotts/inference.py create mode 100644 misotts/models.py create mode 100644 nodes/__init__.py create mode 100644 nodes/epub_loader.py create mode 100644 nodes/generator.py create mode 100644 nodes/loader.py create mode 100644 pyproject.toml create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..9bd65f8 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +.ruff_cache/ +.pytest_cache/ +*.wav +.venv/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..a78dce0 --- /dev/null +++ b/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. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e32b7f2 --- /dev/null +++ b/__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"] diff --git a/misotts/__init__.py b/misotts/__init__.py new file mode 100644 index 0000000..1f2d0a8 --- /dev/null +++ b/misotts/__init__.py @@ -0,0 +1,3 @@ +from .inference import Generator, Segment, load_miso_8b + +__all__ = ["Generator", "Segment", "load_miso_8b"] diff --git a/misotts/csm_llama.py b/misotts/csm_llama.py new file mode 100644 index 0000000..9cd0a89 --- /dev/null +++ b/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) diff --git a/misotts/inference.py b/misotts/inference.py new file mode 100644 index 0000000..2b76a56 --- /dev/null +++ b/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) diff --git a/misotts/models.py b/misotts/models.py new file mode 100644 index 0000000..323f7a6 --- /dev/null +++ b/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) diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..bb27592 --- /dev/null +++ b/nodes/__init__.py @@ -0,0 +1,5 @@ +from .loader import MisoTTSModelLoader +from .generator import MisoTTSGenerate +from .epub_loader import MisoTTSEpubLoader + +__all__ = ["MisoTTSModelLoader", "MisoTTSGenerate", "MisoTTSEpubLoader"] diff --git a/nodes/epub_loader.py b/nodes/epub_loader.py new file mode 100644 index 0000000..7c0e76b --- /dev/null +++ b/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) diff --git a/nodes/generator.py b/nodes/generator.py new file mode 100644 index 0000000..574dd52 --- /dev/null +++ b/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},) diff --git a/nodes/loader.py b/nodes/loader.py new file mode 100644 index 0000000..96d34ba --- /dev/null +++ b/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,) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..8c6ac32 --- /dev/null +++ b/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 = "" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..a4bdfb1 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +transformers>=4.46.0 +safetensors +tokenizers +torchaudio +huggingface_hub +beautifulsoup4