Update app.py

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 from PIL import Image
 import os
 import logging
+from safetensors.torch import load_file  # Import safetensors loading function
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -34,11 +35,12 @@ config = AutoConfig.from_pretrained(config_path)
 # Load the feature extractor
 feature_extractor = ViTFeatureExtractor.from_pretrained(preprocessor_path)
 
-# Load the model using the specific paths
+# Load the model using the safetensors file
+state_dict = load_file(model_path)  # Use safetensors to load the model weights
 model = ViTForImageClassification.from_pretrained(
     pretrained_model_name_or_path=None,
     config=config,
-    state_dict=torch.load(model_path, map_location=torch.device('cpu'))
+    state_dict=state_dict
 )
 
 # Ensure the model is in evaluation mode