ComfyUI-MisoTTS/misotts/inference.py

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