anyonehomep1mane
Code Changes
5aa6736
import torch
from transformers import AutoModel, AutoProcessor
from config.settings import MODEL_ID, TORCH_DTYPE, ATTN_IMPLEMENTATION, DEVICE
def load_model():
model = AutoModel.from_pretrained(
MODEL_ID,
torch_dtype=TORCH_DTYPE,
attn_implementation=ATTN_IMPLEMENTATION
)
model = model.to(DEVICE)
model.eval()
processor = AutoProcessor.from_pretrained(MODEL_ID)
return model, processor