File size: 4,374 Bytes
082393b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | """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
|