Spaces:
Running on Zero
Running on Zero
| """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 where CUDA init must happen inside @spaces.GPU.""" | |
| 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 | |