File size: 744 Bytes
4183ac5
28421c3
4183ac5
 
 
 
28421c3
976dad6
 
 
28421c3
4183ac5
976dad6
4183ac5
976dad6
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
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...