|
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") |
|
|
|
|
|
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() |
|
|
|
|
|
|