feat: add detailed error handling and logging to get-hairswap endpoint

server.py

@@ -72,13 +72,19 @@ _model = None  # type: ignore[assignment]
 def get_model():
     global _model
     if _model is None:
-        LOGGER.info("Loading StableHair model ...")
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        dtype = torch.float16 if device == "cuda" else torch.float32
-        # Import here to defer importing diffusers/transformers until needed
-        from infer_full import StableHair  # noqa: WPS433
-        _model = StableHair(config="./configs/hair_transfer.yaml", device=device, weight_dtype=dtype)
-        LOGGER.info("Model loaded")
+        try:
+            LOGGER.info("Loading StableHair model ...")
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            dtype = torch.float16 if device == "cuda" else torch.float32
+            LOGGER.info(f"Using device: {device}, dtype: {dtype}")
+            
+            # Import here to defer importing diffusers/transformers until needed
+            from infer_full import StableHair  # noqa: WPS433
+            _model = StableHair(config="./configs/hair_transfer.yaml", device=device, weight_dtype=dtype)
+            LOGGER.info("Model loaded successfully")
+        except Exception as e:
+            LOGGER.error(f"Failed to load model: {str(e)}")
+            raise Exception(f"Model loading failed: {str(e)}")
     return _model
 
 
@@ -118,48 +124,76 @@ async def upload_image(image: UploadFile = File(...), _=Depends(verify_bearer)):
 
 @app.post("/get-hairswap")
 def get_hairswap(req: HairSwapRequest, _=Depends(verify_bearer)):
-    # Resolve file paths
-    def find_file(image_id: str) -> str:
-        for name in os.listdir(UPLOAD_DIR):
-            if name.startswith(image_id):
-                return os.path.join(UPLOAD_DIR, name)
-        raise HTTPException(status_code=404, detail=f"Image id not found: {image_id}")
-
-    source_path = find_file(req.source_id)
-    reference_path = find_file(req.reference_id)
-
-    model = get_model()
-    # Prepare kwargs similar to infer_full
-    id_np, out_np, bald_np, ref_np = model.Hair_Transfer(
-        source_image=source_path,
-        reference_image=reference_path,
-        random_seed=-1,
-        step=30,
-        guidance_scale=req.guidance_scale,
-        scale=req.scale,
-        controlnet_conditioning_scale=req.controlnet_conditioning_scale,
-        size=512,
-    )
-
-    # Save result
-    result_id = str(uuid.uuid4())
-    out_img = Image.fromarray((out_np * 255.).astype(np.uint8))
-    filename = f"{result_id}.png"
-    out_path = os.path.join(RESULTS_DIR, filename)
-    out_img.save(out_path)
-    if results_col:
+    try:
+        # Resolve file paths
+        def find_file(image_id: str) -> str:
+            for name in os.listdir(UPLOAD_DIR):
+                if name.startswith(image_id):
+                    return os.path.join(UPLOAD_DIR, name)
+            raise HTTPException(status_code=404, detail=f"Image id not found: {image_id}")
+
+        source_path = find_file(req.source_id)
+        reference_path = find_file(req.reference_id)
+        
+        LOGGER.info(f"Found source: {source_path}, reference: {reference_path}")
+
+        # Load model with error handling
         try:
-            results_col.insert_one({
-                "_id": result_id,
-                "filename": filename,
-                "path": out_path,
-                "source_id": req.source_id,
-                "reference_id": req.reference_id,
-            })
-        except Exception:
-            pass
+            model = get_model()
+            LOGGER.info("Model loaded successfully")
+        except Exception as e:
+            LOGGER.error(f"Model loading failed: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"Model loading failed: {str(e)}")
 
-    return {"result": filename}
+        # Perform hair transfer with error handling
+        try:
+            LOGGER.info("Starting hair transfer...")
+            id_np, out_np, bald_np, ref_np = model.Hair_Transfer(
+                source_image=source_path,
+                reference_image=reference_path,
+                random_seed=-1,
+                step=30,
+                guidance_scale=req.guidance_scale,
+                scale=req.scale,
+                controlnet_conditioning_scale=req.controlnet_conditioning_scale,
+                size=512,
+            )
+            LOGGER.info("Hair transfer completed successfully")
+        except Exception as e:
+            LOGGER.error(f"Hair transfer failed: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"Hair transfer failed: {str(e)}")
+
+        # Save result
+        try:
+            result_id = str(uuid.uuid4())
+            out_img = Image.fromarray((out_np * 255.).astype(np.uint8))
+            filename = f"{result_id}.png"
+            out_path = os.path.join(RESULTS_DIR, filename)
+            out_img.save(out_path)
+            LOGGER.info(f"Result saved: {out_path}")
+            
+            if results_col:
+                try:
+                    results_col.insert_one({
+                        "_id": result_id,
+                        "filename": filename,
+                        "path": out_path,
+                        "source_id": req.source_id,
+                        "reference_id": req.reference_id,
+                    })
+                except Exception as e:
+                    LOGGER.warning(f"MongoDB save failed: {str(e)}")
+
+            return {"result": filename}
+        except Exception as e:
+            LOGGER.error(f"Result saving failed: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"Result saving failed: {str(e)}")
+            
+    except HTTPException:
+        raise
+    except Exception as e:
+        LOGGER.error(f"Unexpected error in get_hairswap: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")
 
 
 @app.get("/download/{filename}")