src/model_loader.py

"""Singleton model loader for GGUF models.

Keeps one model loaded at a time. Supports dual-mode deployment:
- HF Spaces (CPU Basic): CPU-only inference (n_gpu_layers=0, n_threads=2)
- Local: Full GPU offload via CUDA when libcudart.so.12 is present.
"""

import gc
import logging
import os
from pathlib import Path
from dataclasses import dataclass

from huggingface_hub import hf_hub_download

logger = logging.getLogger(__name__)

MODELS_DIR = Path(__file__).resolve().parent.parent / "models"

GARMENT_TYPES = frozenset({
    "shirt", "blouse", "t-shirt", "top", "tank-top",
    "sweater", "cardigan", "hoodie", "sweatshirt",
    "jacket", "coat", "blazer", "vest",
    "pants", "jeans", "trousers", "shorts", "skirt",
    "dress", "jumpsuit", "romper",
    "boots", "shoes", "sneakers", "sandals", "heels", "flats", "loafers",
    "hat", "cap", "beanie", "scarf", "gloves", "belt",
    "bag", "purse", "backpack", "clutch",
    "tie", "bow-tie", "watch", "sunglasses", "glasses",
    "socks", "stockings", "tights",
    "underwear", "bra", "swimsuit", "bikini",
})


@dataclass
class ModelConfig:
    repo_id: str
    model_file: str
    mmproj_file: str | None
    handler_type: str  # "mtmd", "qwen25vl", "text_only"
    n_ctx: int = 4096


VISION_MODEL = ModelConfig(
    repo_id="ggml-org/gemma-3-4b-it-GGUF",
    model_file="gemma-3-4b-it-Q4_K_M.gguf",
    mmproj_file="mmproj-model-f16.gguf",
    handler_type="mtmd",
    n_ctx=4096,
)

TEXT_MODEL = ModelConfig(
    repo_id="ggml-org/gemma-3-4b-it-GGUF",
    model_file="gemma-3-4b-it-Q4_K_M.gguf",
    mmproj_file=None,
    handler_type="text_only",
    n_ctx=4096,
)


class _ModelManager:
    """Singleton that keeps one Llama model loaded at a time."""

    def __init__(self):
        self._llm = None
        self._current_config: ModelConfig | None = None

    def _ensure_downloaded(self, config: ModelConfig) -> tuple[Path, Path | None]:
        """Download model files if not present. Returns (model_path, mmproj_path)."""
        model_dir = MODELS_DIR / config.repo_id.split("/")[-1]
        model_dir.mkdir(parents=True, exist_ok=True)

        model_path = model_dir / config.model_file
        if not model_path.exists():
            logger.info("Downloading %s from %s...", config.model_file, config.repo_id)
            hf_hub_download(
                repo_id=config.repo_id,
                filename=config.model_file,
                local_dir=model_dir,
            )

        mmproj_path = None
        if config.mmproj_file:
            mmproj_path = model_dir / config.mmproj_file
            if not mmproj_path.exists():
                logger.info("Downloading %s from %s...", config.mmproj_file, config.repo_id)
                hf_hub_download(
                    repo_id=config.repo_id,
                    filename=config.mmproj_file,
                    local_dir=model_dir,
                )

        return model_path, mmproj_path

    @staticmethod
    def _detect_gpu_layers() -> int:
        """Detect whether to offload layers to GPU.

        On HF Spaces (CPU Basic) there is no CUDA — always CPU.
        Locally, probe for libcudart.so.12 to confirm real CUDA.
        """
        if os.environ.get("SPACE_ID"):
            logger.info("Running on HF Spaces — using CPU inference")
            return 0
        if os.environ.get("CUDA_VISIBLE_DEVICES") == "":
            return 0
        try:
            import ctypes
            ctypes.CDLL("libcudart.so.12")
            return -1
        except OSError:
            logger.warning("CUDA runtime not found — running on CPU")
            return 0

    def _is_same_model(self, config: ModelConfig) -> bool:
        if self._current_config is None:
            return False
        return (
            self._current_config.repo_id == config.repo_id
            and self._current_config.model_file == config.model_file
            and self._current_config.mmproj_file == config.mmproj_file
        )

    def load(self, config: ModelConfig):
        """Load a model. Unloads current model first if different."""
        if self._is_same_model(config):
            logger.debug("Model already loaded: %s", config.model_file)
            return self._llm

        self.unload()

        model_path, mmproj_path = self._ensure_downloaded(config)

        from llama_cpp import Llama

        chat_handler = None
        if config.handler_type == "mtmd" and mmproj_path:
            from llama_cpp.llama_chat_format import MTMDChatHandler
            chat_handler = MTMDChatHandler(clip_model_path=str(mmproj_path))
        elif config.handler_type == "qwen25vl" and mmproj_path:
            from llama_cpp.llama_chat_format import Qwen25VLChatHandler
            chat_handler = Qwen25VLChatHandler(clip_model_path=str(mmproj_path))

        n_gpu_layers = self._detect_gpu_layers()
        n_threads = 2 if os.environ.get("SPACE_ID") else None
        logger.info(
            "Loading model: %s (handler: %s, gpu_layers: %s, threads: %s)",
            config.model_file, config.handler_type, n_gpu_layers, n_threads or "default",
        )

        llama_kwargs: dict = {
            "model_path": str(model_path),
            "chat_handler": chat_handler,
            "n_gpu_layers": n_gpu_layers,
            "n_ctx": config.n_ctx,
            "verbose": False,
        }
        if n_threads is not None:
            llama_kwargs["n_threads"] = n_threads

        self._llm = Llama(**llama_kwargs)
        self._current_config = config
        logger.info("Model loaded successfully")
        return self._llm

    def unload(self):
        """Free the current model from memory."""
        if self._llm is not None:
            logger.info("Unloading model: %s", self._current_config.model_file)
            del self._llm
            self._llm = None
            self._current_config = None
            gc.collect()

    def get_vision_model(self):
        """Load and return the vision model."""
        return self.load(VISION_MODEL)

    def get_text_model(self):
        """Load and return the text-only model (reuses same model without mmproj for now)."""
        return self.load(VISION_MODEL)

    @property
    def is_loaded(self) -> bool:
        return self._llm is not None


model_manager = _ModelManager()