Update app.py

app.py

@@ -13,9 +13,10 @@ 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)
-model_path = "model.safetensors"
-config_path = "config.json"
-preprocessor_path = "preprocessor_config.json"
+model_dir = "."  # Use current directory
+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")
 
 # Check if all required files exist
 for path in [model_path, config_path, preprocessor_path]:
@@ -25,12 +26,9 @@ for path in [model_path, config_path, preprocessor_path]:
     else:
         logging.info(f"Found file: {path}")
 
-# Load the configuration and feature extractor
-config = ViTForImageClassification.from_pretrained(".", config=config_path)
-feature_extractor = ViTFeatureExtractor.from_pretrained(".")
-
-# Initialize the model with the configuration
-model = ViTForImageClassification.from_pretrained(".", config=config)
+# Load the model and feature extractor using the local directory
+feature_extractor = ViTFeatureExtractor.from_pretrained(model_dir)
+model = ViTForImageClassification.from_pretrained(model_dir)
 
 # Ensure the model is in evaluation mode
 model.eval()