Update app.py

app.py

@@ -23,7 +23,7 @@ import base64
 import cv2
 import matplotlib.pyplot as plt
 from peft import PeftModel
-from gradcam_xception import load_xception_model, generate_smoothgrad_visualizations_xception
+from gradcam_xception import generate_smoothgrad_visualizations_xception
 warnings.filterwarnings("ignore", category=UserWarning)
 
 # Define Xception transform function directly in app.py
@@ -192,14 +192,20 @@ with st.sidebar:
     if not st.session_state.xception_model_loaded:
         if st.button("📥 Load Xception Model", type="primary"):
             # Load Xception model
-            model, device = load_detection_model_xception()
-            if model is not None:
-                st.session_state.xception_model = model
-                st.session_state.device = device
-                st.session_state.xception_model_loaded = True
-                st.success("✅ Xception model loaded!")
-            else:
-                st.error("❌ Failed to load Xception model.")
+            try:
+                from gradcam_xception import load_xception_model
+                model = load_xception_model()
+                device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+                
+                if model is not None:
+                    st.session_state.xception_model = model
+                    st.session_state.device = device
+                    st.session_state.xception_model_loaded = True
+                    st.success("✅ Xception model loaded!")
+                else:
+                    st.error("❌ Failed to load Xception model.")
+            except Exception as e:
+                st.error(f"Error loading model: {str(e)}")
     else:
         st.success("✅ Xception model loaded")