Update app.py

app.py

@@ -1,15 +1,21 @@
 import gradio as gr
+from transformers import ViTForImageClassification, ViTProcessor
+
+# Load the model and processor
+model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")
+processor = ViTProcessor.from_pretrained("google/vit-base-patch16-224")
 
-# Define the function to predict using the model
 def predict_image(img):
-    model = gr.load("models/google/vit-base-patch16-224")
-    return model.predict(img)
+    inputs = processor(img, return_tensors="pt")
+    outputs = model(**inputs)
+    predictions = outputs.logits.argmax(-1)
+    return model.config.id2label[predictions.item()]
 
 # Create the interface
 iface = gr.Interface(
     fn=predict_image,
-    inputs=gr.Image(),
-    outputs="label",
+    inputs=gr.Image(shape=(224, 224)),
+    outputs="text",
     live=True,
     capture_session=True,
     title="Image recognition",