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