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...