Create app.py

app.py

@@ -0,0 +1,26 @@
+import gradio as gr
+import torch
+from transformers import ViTForImageClassification, ViTFeatureExtractor
+from PIL import Image
+
+# Load ViT 221k model for image classification (pre-trained on ImageNet21k)
+vit_221k_model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224-in21k")
+vit_221k_feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
+
+# Inference function for predicting with ViT 221k
+def predict(image):
+    # Preprocess the input image
+    inputs = vit_221k_feature_extractor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = vit_221k_model(**inputs)
+        predicted_class_idx = outputs.logits.argmax(-1).item()
+    
+    # Get the label corresponding to the prediction
+    vit_221k_label = vit_221k_model.config.id2label[predicted_class_idx]
+    return f"Prediction from ViT 221k Model: {vit_221k_label}"
+
+# Create Gradio Interface
+interface = gr.Interface(fn=predict, inputs="image", outputs="text", title="ViT 221k Image Classification")
+
+# Launch the interface
+interface.launch()