File size: 1,526 Bytes
8839278 | 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 | from __future__ import annotations
from typing import Any
import torch
from backend.config import (
MAX_PIXELS,
MIN_PIXELS,
MODEL_GLM,
MODEL_OVIS,
TEST_MODE,
)
ovis_processor: Any = None
ovis_model: Any = None
glm_processor: Any = None
glm_model: Any = None
def load_models() -> None:
"""Eager-load both OCR models onto CUDA (ZeroGPU contract)."""
global ovis_processor, ovis_model, glm_processor, glm_model
if TEST_MODE:
return
from transformers import (
AutoModelForImageTextToText,
AutoProcessor,
Qwen3_5ForConditionalGeneration,
)
ovis_processor = AutoProcessor.from_pretrained(
MODEL_OVIS,
min_pixels=MIN_PIXELS,
max_pixels=MAX_PIXELS,
)
ovis_model = Qwen3_5ForConditionalGeneration.from_pretrained(
MODEL_OVIS,
dtype=torch.bfloat16,
attn_implementation="sdpa",
).to("cuda")
ovis_model.eval()
glm_processor = AutoProcessor.from_pretrained(MODEL_GLM, trust_remote_code=True)
glm_model = AutoModelForImageTextToText.from_pretrained(
MODEL_GLM,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
).to("cuda")
glm_model.eval()
def models_loaded() -> dict[str, bool]:
if TEST_MODE:
return {"ovis": True, "glm": True, "test_mode": True}
return {
"ovis": ovis_processor is not None and ovis_model is not None,
"glm": glm_processor is not None and glm_model is not None,
"test_mode": False,
}
|