Florence-metere1 / handler.py
VishalD1234's picture
Update handler.py
976dad6 verified
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
class VisionModelHandler:
def __init__(self, model_name_or_path="https://huggingface.co/VishalD1234/Florence-metere1"):
self.model_name_or_path = model_name_or_path
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Add the trust_remote_code parameter
self.model = AutoModelForImageClassification.from_pretrained(self.model_name_or_path, trust_remote_code=True)
self.feature_extractor = AutoFeatureExtractor.from_pretrained(self.model_name_or_path, trust_remote_code=True)
self.model.to(self.device)
self.model.eval()
# Rest of the code stays the same...