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