Spaces:
Running on Zero
Running on Zero
File size: 2,474 Bytes
49574d5 | 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 | """Shared Granite Vision model loader.
Split into two stages for ZeroGPU compatibility:
- load_processor(): CPU-only, safe to call at startup or outside @spaces.GPU
- load_model(): requires CUDA, must only be called inside a @spaces.GPU context
The processor and model are cached globally so they are loaded at most once.
"""
import os
from typing import Any
_processor: Any = None
_model: Any = None
MODEL_ID = "ibm-granite/granite-vision-4.1-4b"
def load_processor() -> Any:
"""Load (or return cached) processor. CPU-only — safe outside @spaces.GPU.
Returns:
AutoProcessor instance, or None if loading fails.
"""
global _processor # noqa: PLW0603
if _processor is not None:
return _processor
try:
from transformers import AutoProcessor
token = os.environ.get("HF_TOKEN")
_processor = AutoProcessor.from_pretrained(
MODEL_ID, trust_remote_code=True, token=token, use_fast=True
)
print(f"Processor loaded for {MODEL_ID}")
return _processor
except Exception as e: # noqa: BLE001
import traceback
print(f"Processor load error: {e}")
traceback.print_exc()
return None
def load_model() -> tuple[Any, Any]:
"""Load (or return cached) model to CUDA. Must be called inside @spaces.GPU.
Returns:
Tuple of (processor, model), or (None, None) if loading fails.
"""
global _model # noqa: PLW0603
processor = load_processor()
if processor is None:
return None, None
if _model is not None:
return processor, _model
try:
import torch
from transformers import AutoModelForImageTextToText
token = os.environ.get("HF_TOKEN")
# Load on CPU first to avoid caching_allocator_warmup triggering
# torch._C._cuda_init() before ZeroGPU can intercept it.
_model = AutoModelForImageTextToText.from_pretrained(
MODEL_ID,
trust_remote_code=True,
dtype=torch.bfloat16,
token=token,
).eval()
_model = _model.to("cuda")
if hasattr(_model, "merge_lora_adapters"):
_model = _model.merge_lora_adapters()
print(f"Model loaded: {MODEL_ID} on cuda")
return processor, _model
except Exception as e: # noqa: BLE001
import traceback
print(f"Model load error: {e}")
traceback.print_exc()
return processor, None
|