Update app.py

app.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import ViTForImageClassification, ViTFeatureExtractor
+from transformers import ViTForImageClassification, ViTFeatureExtractor, AutoConfig
 import gradio as gr
 from PIL import Image
 import os
@@ -12,8 +12,10 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 labels = ['Leggings', 'Jogger', 'Palazzo', 'Cargo', 'Dresspants', 'Chinos']
 logging.info(f"Labels: {labels}")
 
-# Define paths to the model files (all in the same directory as app.py)
+# Define the directory containing the model files
 model_dir = "."  # Use current directory
+
+# Define paths to the specific model files
 model_path = os.path.join(model_dir, "model.safetensors")
 config_path = os.path.join(model_dir, "config.json")
 preprocessor_path = os.path.join(model_dir, "preprocessor_config.json")
@@ -26,11 +28,18 @@ for path in [model_path, config_path, preprocessor_path]:
     else:
         logging.info(f"Found file: {path}")
 
-# Load the model and feature extractor using the local directory
-model_id = "google/vit-base-patch16-224"
+# Load the configuration
+config = AutoConfig.from_pretrained(config_path)
+
+# Load the feature extractor
+feature_extractor = ViTFeatureExtractor.from_pretrained(preprocessor_path)
 
-feature_extractor = ViTFeatureExtractor.from_pretrained(model_id)
-model = ViTForImageClassification.from_pretrained(model_path)
+# Load the model using the specific paths
+model = ViTForImageClassification.from_pretrained(
+    pretrained_model_name_or_path=None,
+    config=config,
+    state_dict=torch.load(model_path, map_location=torch.device('cpu'))
+)
 
 # Ensure the model is in evaluation mode
 model.eval()