model.py

import torch
from transformers import AutoProcessor, AutoModelForVision2Seq

MODEL_NAME = "llava-hf/llava-interleave-qwen-0.5b-hf"

def load_model():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32

    processor = AutoProcessor.from_pretrained(MODEL_NAME)

    model = AutoModelForVision2Seq.from_pretrained(
        MODEL_NAME,
        torch_dtype=dtype,
        device_map="auto" if device == "cuda" else None,
    )

    model.eval()
    return processor, model, device


# Singleton
processor, model, device = load_model()