Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import ViTImageProcessor, ViTForImageClassification
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 from PIL import Image
 import numpy as np
 import logging
@@ -8,11 +8,11 @@ import logging
 # Configure logging
 logging.basicConfig(level=logging.DEBUG)
 
-# Load the pre-trained model and image processor
-model_name = "google/vit-base-patch16-224"
+# Load the pre-trained model and feature extractor
+model_name = "google/vit-base-patch16-224-in21k"
 logging.info("Loading image processor and model...")
-image_processor = ViTImageProcessor.from_pretrained(model_name)
-model = ViTForImageClassification.from_pretrained(model_name)
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+model = AutoModelForImageClassification.from_pretrained(model_name)
 
 # Define the prediction function
 def predict(image):
@@ -31,7 +31,7 @@ def predict(image):
         logging.debug("Image converted successfully.")
 
         logging.info("Processing image...")
-        inputs = image_processor(images=image, return_tensors="pt")
+        inputs = feature_extractor(images=image, return_tensors="pt")
         outputs = model(**inputs)
         logits = outputs.logits
         predicted_class_idx = logits.argmax(-1).item()