Fix app mounting for Hugging Face Spaces: Export app at module level, add GPU wake-up retry logic, and fix indentation

app.py

@@ -166,9 +166,27 @@ def load_model(modality='T1'):
     
     if torch.cuda.is_available():
         try:
-            torch.cuda.empty_cache()
-            DEVICE = torch.device('cuda')
-            print(f"✓ Using device: {DEVICE}")
+            import time
+            max_retries = 3
+            retry_delay = 2
+            
+            for attempt in range(max_retries):
+                try:
+                    torch.cuda.empty_cache()
+                    test_tensor = torch.zeros(1).cuda()
+                    del test_tensor
+                    torch.cuda.synchronize()
+                    DEVICE = torch.device('cuda')
+                    print(f"✓ Using device: {DEVICE}")
+                    break
+                except RuntimeError as e:
+                    if "CUDA" in str(e) and attempt < max_retries - 1:
+                        print(f"⚠ GPU wake-up attempt {attempt + 1}/{max_retries}: {e}")
+                        print(f"⚠ Waiting {retry_delay}s for GPU to wake up...")
+                        time.sleep(retry_delay)
+                        retry_delay *= 2
+                    else:
+                        raise
         except Exception as e:
             print(f"⚠ CUDA available but failed to initialize: {e}. Falling back to CPU.")
             DEVICE = torch.device('cpu')
@@ -968,15 +986,14 @@ def create_interface():
     
     return demo
 
-if __name__ == "__main__":
-    print("=" * 60)
-    print("Initializing SRMA-Mamba Liver Segmentation Space...")
-    print("=" * 60)
-    
-    try:
-        demo = create_interface()
-        print("✓ Gradio interface created successfully")
-    except Exception as e:
+print("=" * 60)
+print("Initializing SRMA-Mamba Liver Segmentation Space...")
+print("=" * 60)
+
+try:
+    demo = create_interface()
+    print("✓ Gradio interface created successfully")
+except Exception as e:
         print(f"✗ Failed to create interface: {e}")
         import traceback
         tb = traceback.format_exc()
@@ -1007,145 +1024,157 @@ if __name__ == "__main__":
             print(f"✗ Even error interface failed: {e2}")
             raise
     
-    print("ℹ Models will be loaded on first inference")
-    print("=" * 60)
-    
-    from starlette.responses import JSONResponse
-    from starlette.requests import Request
+print("ℹ Models will be loaded on first inference")
+print("=" * 60)
+
+from starlette.responses import JSONResponse
+from starlette.requests import Request
+
+async def handle_segment(request: Request):
+    if request.method == "OPTIONS":
+        return JSONResponse({}, headers={
+            "Access-Control-Allow-Origin": "*",
+            "Access-Control-Allow-Methods": "POST, OPTIONS",
+            "Access-Control-Allow-Headers": "*"
+        })
     
-    async def handle_segment(request: Request):
-        if request.method == "OPTIONS":
-            return JSONResponse({}, headers={
-                "Access-Control-Allow-Origin": "*",
-                "Access-Control-Allow-Methods": "POST, OPTIONS",
-                "Access-Control-Allow-Headers": "*"
-            })
+    try:
+        form = await request.form()
+        file = form.get("file")
+        modality = form.get("modality", "T1")
+        slice_idx = form.get("slice_idx")
+        
+        if not file:
+            return JSONResponse({"error": "No file provided"}, status_code=400, headers={"Access-Control-Allow-Origin": "*"})
+        
+        if modality not in ['T1', 'T2']:
+            return JSONResponse({"error": "Modality must be 'T1' or 'T2'"}, status_code=400, headers={"Access-Control-Allow-Origin": "*"})
+        
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.nii.gz') as tmp_file:
+            content = await file.read()
+            tmp_file.write(content)
+            tmp_path = tmp_file.name
         
         try:
-            form = await request.form()
-            file = form.get("file")
-            modality = form.get("modality", "T1")
-            slice_idx = form.get("slice_idx")
-            
-            if not file:
-                return JSONResponse({"error": "No file provided"}, status_code=400, headers={"Access-Control-Allow-Origin": "*"})
-            
-            if modality not in ['T1', 'T2']:
-                return JSONResponse({"error": "Modality must be 'T1' or 'T2'"}, status_code=400, headers={"Access-Control-Allow-Origin": "*"})
-            
-            with tempfile.NamedTemporaryFile(delete=False, suffix='.nii.gz') as tmp_file:
-                content = await file.read()
-                tmp_file.write(content)
-                tmp_path = tmp_file.name
-            
-            try:
-                result = predict_volume_api(tmp_path, modality, int(slice_idx) if slice_idx else None)
-                if not result["success"]:
-                    error_detail = result.get("error", "Unknown error")
-                    print(f"API prediction failed: {error_detail}")
-                    return JSONResponse({
-                        "success": False,
-                        "error": error_detail,
-                        "error_type": "prediction_failed"
-                    }, status_code=500, headers={"Access-Control-Allow-Origin": "*"})
-                
-                with open(result["segmentation_path"], "rb") as seg_file:
-                    seg_data = seg_file.read()
-                    seg_base64 = base64.b64encode(seg_data).decode('utf-8')
-                
-                result["segmentation_file"] = f"data:application/octet-stream;base64,{seg_base64}"
-                os.unlink(tmp_path)
-                os.unlink(result["segmentation_path"])
-                
-                return JSONResponse(result, headers={"Access-Control-Allow-Origin": "*"})
-            except Exception as e:
-                error_msg = str(e)
-                import traceback
-                tb = traceback.format_exc()
-                print(f"Exception in handle_segment: {error_msg}")
-                print(f"Traceback: {tb}")
-                if os.path.exists(tmp_path):
-                    os.unlink(tmp_path)
+            result = predict_volume_api(tmp_path, modality, int(slice_idx) if slice_idx else None)
+            if not result["success"]:
+                error_detail = result.get("error", "Unknown error")
+                print(f"API prediction failed: {error_detail}")
                 return JSONResponse({
                     "success": False,
-                    "error": error_msg,
-                    "error_type": "exception",
-                    "traceback": tb if "traceback" in str(type(e)) else None
+                    "error": error_detail,
+                    "error_type": "prediction_failed"
                 }, status_code=500, headers={"Access-Control-Allow-Origin": "*"})
+            
+            with open(result["segmentation_path"], "rb") as seg_file:
+                seg_data = seg_file.read()
+                seg_base64 = base64.b64encode(seg_data).decode('utf-8')
+            
+            result["segmentation_file"] = f"data:application/octet-stream;base64,{seg_base64}"
+            os.unlink(tmp_path)
+            os.unlink(result["segmentation_path"])
+            
+            return JSONResponse(result, headers={"Access-Control-Allow-Origin": "*"})
         except Exception as e:
             error_msg = str(e)
             import traceback
             tb = traceback.format_exc()
-            print(f"Exception in handle_segment (outer): {error_msg}")
+            print(f"Exception in handle_segment: {error_msg}")
             print(f"Traceback: {tb}")
+            if os.path.exists(tmp_path):
+                os.unlink(tmp_path)
             return JSONResponse({
                 "success": False,
                 "error": error_msg,
-                "error_type": "outer_exception"
+                "error_type": "exception",
+                "traceback": tb if "traceback" in str(type(e)) else None
             }, status_code=500, headers={"Access-Control-Allow-Origin": "*"})
-    
-    async def handle_health(request: Request):
+    except Exception as e:
+        error_msg = str(e)
+        import traceback
+        tb = traceback.format_exc()
+        print(f"Exception in handle_segment (outer): {error_msg}")
+        print(f"Traceback: {tb}")
         return JSONResponse({
-            "status": "healthy",
-            "device": str(DEVICE) if DEVICE else "not initialized",
-            "model_t1_loaded": MODEL_T1 is not None,
-            "model_t2_loaded": MODEL_T2 is not None
-        }, headers={"Access-Control-Allow-Origin": "*"})
-    
-    import starlette.routing
-    
-    def add_api_routes():
+            "success": False,
+            "error": error_msg,
+            "error_type": "outer_exception"
+        }, status_code=500, headers={"Access-Control-Allow-Origin": "*"})
+
+async def handle_health(request: Request):
+    return JSONResponse({
+        "status": "healthy",
+        "device": str(DEVICE) if DEVICE else "not initialized",
+        "model_t1_loaded": MODEL_T1 is not None,
+        "model_t2_loaded": MODEL_T2 is not None
+    }, headers={"Access-Control-Allow-Origin": "*"})
+
+import starlette.routing
+
+def add_api_routes():
+    try:
+        demo.app.mount("/api", api_app)
+        print("✓ FastAPI app mounted at /api")
+    except Exception as mount_error:
+        print(f"Mounting failed: {mount_error}, trying direct routes...")
         try:
-            demo.app.mount("/api", api_app)
-            print("✓ FastAPI app mounted at /api")
-        except Exception as mount_error:
-            print(f"Mounting failed: {mount_error}, trying direct routes...")
+            existing_routes = [str(r) for r in demo.app.routes]
+            print(f"Existing routes: {existing_routes[:5]}...")
+            
+            routes_to_add = [
+                starlette.routing.Route("/api/segment", handle_segment, methods=["POST", "OPTIONS"]),
+                starlette.routing.Route("/api/health", handle_health, methods=["GET", "OPTIONS"])
+            ]
+            
+            if hasattr(demo.app, 'router') and hasattr(demo.app.router, 'routes'):
+                demo.app.router.routes.extend(routes_to_add)
+                print("✓ API routes added to router.routes")
+            else:
+                demo.app.routes.extend(routes_to_add)
+                print("✓ API routes added to app.routes")
+        except Exception as route_error:
+            print(f"Router routes failed: {route_error}, trying app.routes...")
             try:
-                existing_routes = [str(r) for r in demo.app.routes]
-                print(f"Existing routes: {existing_routes[:5]}...")
-                
-                routes_to_add = [
+                demo.app.routes.extend([
                     starlette.routing.Route("/api/segment", handle_segment, methods=["POST", "OPTIONS"]),
                     starlette.routing.Route("/api/health", handle_health, methods=["GET", "OPTIONS"])
-                ]
-                
-                if hasattr(demo.app, 'router') and hasattr(demo.app.router, 'routes'):
-                    demo.app.router.routes.extend(routes_to_add)
-                    print("✓ API routes added to router.routes")
-                else:
-                    demo.app.routes.extend(routes_to_add)
-                    print("✓ API routes added to app.routes")
-            except Exception as route_error:
-                print(f"Router routes failed: {route_error}, trying app.routes...")
-                try:
-                    demo.app.routes.extend([
-                        starlette.routing.Route("/api/segment", handle_segment, methods=["POST", "OPTIONS"]),
-                        starlette.routing.Route("/api/health", handle_health, methods=["GET", "OPTIONS"])
-                    ])
-                    print("✓ API routes added to app.routes")
-                except Exception as e2:
-                    print(f"All route addition methods failed: {e2}")
-                    import traceback
-                    traceback.print_exc()
-    
-    try:
-        add_api_routes()
-        print("✓ API routes configured")
-    except Exception as route_err:
-        print(f"⚠ API route configuration failed (non-critical): {route_err}")
-        import traceback
-        traceback.print_exc()
-    
-    try:
-        print("🚀 Launching Gradio interface...")
-        demo.launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            share=False,
-            show_error=True
-        )
-    except Exception as launch_err:
-        print(f"✗ Failed to launch interface: {launch_err}")
-        import traceback
-        traceback.print_exc()
-        raise
+                ])
+                print("✓ API routes added to app.routes")
+            except Exception as e2:
+                print(f"All route addition methods failed: {e2}")
+                import traceback
+                traceback.print_exc()
+
+try:
+    add_api_routes()
+    print("✓ API routes configured")
+except Exception as route_err:
+    print(f"⚠ API route configuration failed (non-critical): {route_err}")
+    import traceback
+    traceback.print_exc()
+
+print("✓ App initialization complete")
+print("=" * 60)
+
+is_spaces = bool(os.getenv("SPACE_ID") or os.getenv("SYSTEM") == "spaces" or os.getenv("HF_SPACE"))
+
+app = demo
+
+if __name__ == "__main__":
+    if not is_spaces:
+        print("🚀 Running locally - launching with demo.launch()")
+        try:
+            demo.launch(
+                server_name="0.0.0.0",
+                server_port=7860,
+                share=False,
+                show_error=True
+            )
+        except Exception as launch_err:
+            print(f"✗ Failed to launch interface: {launch_err}")
+            import traceback
+            traceback.print_exc()
+            raise
+    else:
+        print("🌐 Running on Hugging Face Spaces - app exported for Spaces to serve")
+        print("✓ App ready for Spaces to serve")