Update app.py

app.py

@@ -38,14 +38,14 @@ try:
 except Exception as e:
     print(f"⚠️  Error loading model from best.pt: {e}")
     try:
-        # Fallback to checkpoints directory
-        model.load_state_dict(torch.load('checkpoints/best_multiclass.pt', map_location=device))
-        print("✅ Multi-class model loaded successfully from checkpoints/best_multiclass.pt!")
+        # Fallback to checkpoints directory (where best.pt is actually stored)
+        model.load_state_dict(torch.load('checkpoints/best.pt', map_location=device))
+        print("✅ Multi-class model loaded successfully from checkpoints/best.pt!")
     except Exception as e2:
         print(f"❌ CRITICAL ERROR: Could not load model from any location!")
         print(f"   - best.pt error: {e}")
-        print(f"   - checkpoints/best_multiclass.pt error: {e2}")
-        raise RuntimeError("Model file not found! Please ensure best.pt is uploaded to the Space.")
+        print(f"   - checkpoints/best.pt error: {e2}")
+        raise RuntimeError("Model file not found! Please ensure best.pt is uploaded to the Space (in /checkpoints or root).")
 
 model = model.to(device)
 model.eval()
@@ -622,4 +622,3 @@ if __name__ == "__main__":
         server_port=7860,
         show_error=True
     )
-