rag_server/applications/clip_embedding/clip_model.py

import os
import torch
from transformers import AutoModel, AutoConfig, CLIPProcessor

MODEL_NAME = "BAAI/EVA-CLIP-8B"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float8 if DEVICE == "cuda" else torch.float32
MAX_BATCH = int(os.getenv("MAX_BATCH", "32"))

TRUST_REMOTE_CODE = True

print(f"[model_config] Loading {MODEL_NAME} on {DEVICE} dtype={DTYPE} ...")

config = AutoConfig.from_pretrained(MODEL_NAME, trust_remote_code=TRUST_REMOTE_CODE)
model = AutoModel.from_pretrained(
    MODEL_NAME, config=config, trust_remote_code=TRUST_REMOTE_CODE
).to(dtype=DTYPE, device=DEVICE).eval()

processor = CLIPProcessor.from_pretrained(MODEL_NAME)

def get_model():
    return model, processor, DEVICE, DTYPE, MAX_BATCH