Update app.py

app.py

@@ -31,11 +31,10 @@ class BacterialMorphologyClassifier(nn.Module):
         x = self.fc(x)
         return x
 
-# Load the model and weights at app startup
+# Load the model and weights at the start of the app
 model = BacterialMorphologyClassifier()
-MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth"
-state_dict = torch.hub.load_state_dict_from_url(MODEL_PATH, map_location=torch.device('cpu'))
-model.load_state_dict(state_dict, strict=False)
+MODEL_PATH = 'https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth'  # Replace this with the local path if needed
+model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu')), strict=False)
 model.eval()
 
 # Move model to GPU if available
@@ -79,4 +78,4 @@ def predict():
         return jsonify({'error': str(e)})
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=5000, debug=True)
+    app.run(host='0.0.0.0', port=5000, debug=False)  # Set debug=False for production