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,3 @@
|
||||
from .inference import Generator, Segment, load_miso_8b
|
||||
|
||||
__all__ = ["Generator", "Segment", "load_miso_8b"]
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user