src/model_loader.py

"""Shared Granite Vision model loader.

Single global cache shared by all inference modules to avoid loading
the model twice and exhausting the memory limit.
"""

import os
import platform
from typing import Any

_processor: Any = None
_model: Any = None

_mlx_processor: Any = None
_mlx_model: Any = None
_mlx_tokenizer: Any = None


def use_api_mode() -> bool:
    """Check whether to use a remote OpenAI-compatible API for inference.

    Returns True when both MODEL_API_URL and RITS_API_KEY environment
    variables are set and non-empty.
    """
    return bool(os.environ.get("MODEL_API_URL") and os.environ.get("RITS_API_KEY"))


def use_mlx_mode() -> bool:
    """Check whether to use MLX for inference.

    Returns True when running on Apple Silicon (arm64 macOS) and
    mlx-vlm is installed. Takes priority over PyTorch MPS.
    """
    if platform.system() != "Darwin" or platform.machine() != "arm64":
        return False
    try:
        import mlx_vlm  # noqa: F401
        return True
    except ImportError:
        return False


def _is_zerogpu() -> bool:
    """Detect HuggingFace ZeroGPU Spaces."""
    return bool(os.environ.get("SPACE_ID")) and bool(os.environ.get("ZERO_GPU"))


def get_model_name() -> str:
    """Return the model identifier to send to the API."""
    return os.environ.get("MODEL_NAME", "ibm-granite/granite-vision-4.1-4b")


def get_mlx_model_name() -> str:
    """Return the MLX model identifier."""
    return os.environ.get("MLX_MODEL_NAME", "mlx-community/granite-4.0-3b-vision-5bit")


def load_mlx_model() -> tuple[Any, Any]:
    """Lazy-load the MLX Granite Vision model and processor.

    Also loads a tokenizer from the original model for chat template formatting,
    since the mlx-community model repo does not include a chat template.

    Returns:
        Tuple of (model, processor), or (None, None) if loading fails.
    """
    global _mlx_processor, _mlx_model, _mlx_tokenizer  # noqa: PLW0603

    if _mlx_processor is not None and _mlx_model is not None:
        return _mlx_model, _mlx_processor

    try:
        from mlx_vlm import load
        from transformers import AutoTokenizer

        model_id = get_mlx_model_name()
        _mlx_model, _mlx_processor = load(model_id)

        # Load tokenizer from the original model for chat template support
        _mlx_tokenizer = AutoTokenizer.from_pretrained(
            "ibm-granite/granite-vision-4.1-4b", trust_remote_code=True
        )

        print(f"Loaded {model_id} on MLX")
        return _mlx_model, _mlx_processor

    except ImportError:
        print("mlx-vlm not available")
        return None, None
    except Exception as e:  # noqa: BLE001
        import traceback

        print(f"MLX model load error: {e}")
        traceback.print_exc()
        return None, None


def get_mlx_tokenizer() -> Any:
    """Return the cached MLX tokenizer (loaded alongside the MLX model)."""
    return _mlx_tokenizer


def load_model() -> tuple[Any, Any]:
    """Lazy-load the Granite Vision model and processor.

    Returns:
        Tuple of (processor, model), or (None, None) if loading fails.
    """
    global _processor, _model  # noqa: PLW0603

    if _processor is not None and _model is not None:
        return _processor, _model

    try:
        import torch
        from transformers import AutoModelForImageTextToText, AutoProcessor

        model_id = get_model_name()
        if torch.cuda.is_available():
            device = "cuda"
        elif torch.backends.mps.is_available():
            device = "mps"
        else:
            device = "cpu"

        _processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True, backend="torchvision")

        _model = AutoModelForImageTextToText.from_pretrained(
            model_id,
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
        ).eval()

        if hasattr(_model, "merge_lora_adapters"):
            _model = _model.merge_lora_adapters()

        _model = _model.to(device)

        print(f"Loaded {model_id} on {device}")
        return _processor, _model

    except ImportError:
        print("Transformers not available, using stub")
        return None, None
    except Exception as e:  # noqa: BLE001
        import traceback

        print(f"Model load error: {e}")
        traceback.print_exc()
        return None, None