Update app.py

app.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import ViTForImageClassification, ViTFeatureExtractor
+from transformers import ViTForImageClassification, ViTFeatureExtractor, ViTConfig
 import gradio as gr
 from PIL import Image
 import os
@@ -26,21 +26,20 @@ else:
 id2label = {str(i): label for i, label in enumerate(labels)}
 label2id = {label: str(i) for i, label in enumerate(labels)}
 
-# Load the model with custom label mapping
-logging.info("Loading the model with custom label mapping")
-model = ViTForImageClassification.from_pretrained(
-    "google/vit-base-patch16-224-in21k",
-    num_labels=len(labels),
-    id2label=id2label,
-    label2id=label2id,
-    ignore_mismatched_sizes=True
-)
+# Create a configuration for the model
+config = ViTConfig.from_pretrained("google/vit-base-patch16-224-in21k")
+config.num_labels = len(labels)
+config.id2label = id2label
+config.label2id = label2id
+
+# Initialize the model with the configuration
+model = ViTForImageClassification(config)
 
 try:
-    # Load the state dict
+    # Load the state dict of the fine-tuned model
     state_dict = torch.load(model_path, map_location=torch.device('cpu'))
     model.load_state_dict(state_dict)
-    logging.info("Model loaded successfully")
+    logging.info("Fine-tuned model loaded successfully")
 except Exception as e:
     logging.error(f"Error loading model: {str(e)}")
     raise