filliptm-ComfyUI-FL-VoxtralTTS/modules/loader.py

"""Model download, weight loading, and caching for Voxtral TTS."""

from __future__ import annotations
import logging
import sys
import os
from pathlib import Path
from typing import Optional

import torch
from safetensors.torch import load_file

logger = logging.getLogger(__name__)

# Global model cache
_MODEL_CACHE = {}


def get_voxtral_models_dir() -> Path:
    """Get or create the VoxtralTTS model storage directory."""
    try:
        import folder_paths
        try:
            tts_paths = folder_paths.get_folder_paths("tts")
        except (KeyError, AttributeError):
            tts_paths = None
        if tts_paths:
            base = Path(tts_paths[0])
        else:
            base = Path(folder_paths.models_dir) / "tts"
    except ImportError:
        base = Path("models") / "tts"

    voxtral_dir = base / "VoxtralTTS"
    voxtral_dir.mkdir(parents=True, exist_ok=True)
    return voxtral_dir


def download_model(repo_id: str, local_dir: Path) -> Path:
    """Download model from HuggingFace Hub if not already present."""
    marker = local_dir / "consolidated.safetensors"
    if marker.exists():
        logger.info(f"Model already downloaded at {local_dir}")
        return local_dir

    logger.info(f"Downloading {repo_id} to {local_dir}...")
    from huggingface_hub import snapshot_download
    snapshot_download(
        repo_id=repo_id,
        local_dir=str(local_dir),
        local_dir_use_symlinks=False,
    )
    logger.info(f"Download complete: {local_dir}")
    return local_dir


def _remap_weights(state_dict: dict) -> dict:
    """Remap Mistral weight names to our module structure.

    Returns dict with keys:
        "backbone": {name: tensor}
        "acoustic_transformer": {name: tensor}
        "codec_decoder": {name: tensor}
        "audio_embeddings": {name: tensor}
    """
    parts = {
        "backbone": {},
        "acoustic_transformer": {},
        "codec_decoder": {},
        "audio_embeddings": {},
    }

    for key, value in state_dict.items():
        if key.startswith("acoustic_transformer."):
            new_key = key[len("acoustic_transformer."):]
            parts["acoustic_transformer"][new_key] = value
        elif key.startswith("audio_tokenizer."):
            new_key = key[len("audio_tokenizer."):]
            parts["codec_decoder"][new_key] = value
        elif key.startswith("mm_audio_embeddings.audio_codebook_embeddings."):
            new_key = key[len("mm_audio_embeddings.audio_codebook_embeddings."):]
            parts["audio_embeddings"][new_key] = value
        elif key.startswith("mm_audio_embeddings.tok_embeddings."):
            # Shared with backbone tok_embeddings
            new_key = key[len("mm_audio_embeddings."):]
            parts["backbone"][new_key] = value
        else:
            # Everything else is backbone (layers, norm, tok_embeddings, output)
            parts["backbone"][key] = value

    return parts


class VoxtralLoader:
    """Handles downloading, loading, and caching the Voxtral TTS pipeline."""

    @staticmethod
    def load(device: str = "auto", dtype: str = "bfloat16",
             force_reload: bool = False) -> "VoxtralTTSPipeline":
        """Load the full Voxtral TTS pipeline.

        Args:
            device: "auto", "cuda", "mps", or "cpu"
            dtype: "bfloat16", "float16", or "float32"
            force_reload: Force re-download and reload

        Returns:
            VoxtralTTSPipeline instance ready for generation
        """
        from .model_info import MODEL_REPO_ID, get_default_device

        # Resolve device
        if device == "auto":
            device = get_default_device()
        torch_device = torch.device(device)

        # Resolve dtype
        dtype_map = {
            "bfloat16": torch.bfloat16,
            "float16": torch.float16,
            "float32": torch.float32,
        }
        torch_dtype = dtype_map.get(dtype, torch.bfloat16)

        # MPS doesn't support bfloat16 well — fall back to float16
        if device == "mps" and torch_dtype == torch.bfloat16:
            logger.info("MPS device detected — using float16 instead of bfloat16")
            torch_dtype = torch.float16

        cache_key = f"{device}_{dtype}"
        if not force_reload and cache_key in _MODEL_CACHE:
            logger.info("Using cached Voxtral TTS pipeline")
            return _MODEL_CACHE[cache_key]

        # Download
        models_dir = get_voxtral_models_dir()
        model_dir = models_dir / "Voxtral-4B-TTS-2603"
        download_model(MODEL_REPO_ID, model_dir)

        # Add src to path for imports
        src_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "src")
        if src_dir not in sys.path:
            sys.path.insert(0, src_dir)

        from voxtral_tts.config import VoxtralConfig
        from voxtral_tts.backbone import MistralBackbone
        from voxtral_tts.acoustic_transformer import FlowMatchingAcousticTransformer
        from voxtral_tts.codec_decoder import VoxtralCodecDecoder
        from voxtral_tts.embeddings import MultiVocabEmbeddings
        from voxtral_tts.tokenizer import VoxtralTokenizer
        from voxtral_tts.pipeline import VoxtralTTSPipeline

        # Parse config
        config = VoxtralConfig.from_params_json(model_dir / "params.json")

        # Load weights
        logger.info("Loading model weights...")
        raw_weights = load_file(str(model_dir / "consolidated.safetensors"))
        parts = _remap_weights(raw_weights)
        del raw_weights  # Free memory

        # Build backbone
        logger.info("Building LLM backbone...")
        backbone = MistralBackbone(config.transformer)
        backbone.load_state_dict(parts["backbone"], strict=False)
        backbone.to(device=torch_device, dtype=torch_dtype)
        backbone.eval()

        # Build acoustic transformer
        logger.info("Building acoustic transformer...")
        acoustic_tf = FlowMatchingAcousticTransformer(config.acoustic_transformer)
        acoustic_tf.load_state_dict(parts["acoustic_transformer"], strict=False)
        acoustic_tf.to(device=torch_device, dtype=torch_dtype)
        acoustic_tf.eval()

        # Build codec decoder
        logger.info("Building codec decoder...")
        codec = VoxtralCodecDecoder(config.codec_decoder)
        codec.load_state_dict(parts["codec_decoder"], strict=False)
        codec.to(device=torch_device, dtype=torch_dtype)
        codec.eval()

        # Build audio embeddings — size from checkpoint
        logger.info("Building audio embeddings...")
        emb_weight = parts["audio_embeddings"].get("embeddings.weight")
        total_entries = emb_weight.shape[0] if emb_weight is not None else 9088
        audio_emb = MultiVocabEmbeddings.from_config(
            total_entries=total_entries, embedding_dim=config.transformer.dim
        )
        audio_emb.load_state_dict(parts["audio_embeddings"], strict=False)
        audio_emb.to(device=torch_device, dtype=torch_dtype)
        audio_emb.eval()

        del parts  # Free memory

        # Load tokenizer
        logger.info("Loading tokenizer...")
        tokenizer = VoxtralTokenizer.from_model_dir(model_dir)

        # Build pipeline
        pipeline = VoxtralTTSPipeline(
            backbone=backbone,
            acoustic_transformer=acoustic_tf,
            codec_decoder=codec,
            audio_embeddings=audio_emb,
            tokenizer=tokenizer,
            config=config,
            voice_embeddings_dir=model_dir / "voice_embedding",
            device=torch_device,
            dtype=torch_dtype,
        )

        _MODEL_CACHE[cache_key] = pipeline
        logger.info(f"Voxtral TTS pipeline ready on {device} ({dtype})")
        return pipeline

    @staticmethod
    def unload():
        """Clear the model cache and free memory."""
        _MODEL_CACHE.clear()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        logger.info("Voxtral TTS pipeline unloaded")