Update app.py

app.py

@@ -3,7 +3,6 @@ from torchvision import transforms
 from transformers import AutoModelForImageClassification
 import gradio as gr
 import torch
-from model import vit
 
 def predict(inp):
     inputs = data_transforms(inp)[None]
@@ -14,40 +13,14 @@ def predict(inp):
         confidences = {labels[i]: probs[0][i] for i in range(num_classes)}
     return confidences
 
-"""height=28
-width=28
-batch_size=128
-n_channels=3
-patch_size=14
-dim=384
-n_head=12
-feed_forward=1024
-num_blocks=8"""
-height=224
-batch_size=128
-width=224
-n_channels=3
-patch_size=16
-dim=256
-n_head=8
-feed_forward=512
-num_blocks=12
-num_classes=2
 data_transforms = transforms.Compose([
-	    transforms.Resize((height,width)),  # Resize the images to a specific size
+	    transforms.Resize((224,224)),  # Resize the images to a specific size
 	    transforms.ToTensor(),           # Convert images to tensors
 	    #transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # Normalize the image data
 	])
 
-model = vit(height,width,n_channels,patch_size,batch_size,dim,n_head,feed_forward,num_blocks,num_classes)# Load saved weights
-model.load_state_dict(
-    torch.load(f="vit_model.pt",
-               map_location=torch.device("cpu")) # load to CPU
-)
-print(model.state_dict())
-"""labels = [
-    'airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'
-]"""
+# Load model directly
+model = AutoModelForImageClassification.from_pretrained("Manu8/vit_cats-vs-dogs", trust_remote_code=True)
 labels = [
     'cat','dog'
 ]