Update app.py

app.py

@@ -1,8 +1,24 @@
 import gradio as gr
+from PIL import Image
+from vit_model_test import CustomModel
 
+# Initialize the model
+model = CustomModel()
 
+def predict(image: Image.Image):
+    # Get predictions from the model
+    label, confidence = model.predict(image)
+    return f"Predicted label: {label}", f"Confidence: {confidence:.2f}%"
 
+# Define the Gradio interface with updated API
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),  # Updated for new Gradio API
+    outputs=[gr.Textbox(), gr.Textbox()],  # Updated for new Gradio API
+    title="Vision Transformer Model",    # Title of the Gradio interface
+    description="Upload an image to classify it using the Vision Transformer model."  # Description
+)
 
-
-demo = gr.Interface(...
-demo.launch()
+# Launch the Gradio interface
+if __name__ == "__main__":
+    demo.launch()