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>
This commit is contained in:
@@ -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)
|
||||
Reference in New Issue
Block a user