Update app.py

app.py

@@ -24,30 +24,53 @@ def resize_to_4x6(img):
 
 @app.post("/process")
 async def process_img(file: UploadFile = File(...)):
-    img_bytes = await file.read()
-    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
-    np_img = np.array(img)
+    try:
+        # Read image
+        img_bytes = await file.read()
+        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
 
-    faces = face_app.get(np_img)
-    if not faces:
-        return JSONResponse({"error": "No face detected"}, status_code=400)
+        # ↓↓↓ Fix 1: Resize large images to safe size ↓↓↓
+        max_dim = 1500
+        if max(img.size) > max_dim:
+            img.thumbnail((max_dim, max_dim))
 
-    face = faces[0]
-    aligned = Image.fromarray(face.normed)
+        np_img = np.array(img)
 
-    # Remove background
-    no_bg = remove(aligned)
+        # Face detection
+        faces = face_app.get(np_img)
+        if not faces:
+            return JSONResponse({"error": "No face detected"}, status_code=400)
 
-    # Add white background
-    white = Image.new("RGB", no_bg.size, (255, 255, 255))
-    white.paste(no_bg, mask=no_bg.split()[3])
+        face = faces[0]
 
-    # Resize
-    final_img = resize_to_4x6(white)
+        # ↓↓↓ Fix 2: Some models return None in normed ↓↓↓
+        if face.normed is None:
+            return JSONResponse({"error": "Face alignment failed"}, status_code=500)
 
-    # Output
-    buf = io.BytesIO()
-    final_img.save(buf, format="JPEG")
-    buf.seek(0)
+        aligned_np = face.normed
+        aligned = Image.fromarray(aligned_np)
 
-    return Response(buf.getvalue(), media_type="image/jpeg")
+        # ↓↓↓ Fix 3: rembg requires proper RGB format ↓↓↓
+        if aligned.mode != "RGB":
+            aligned = aligned.convert("RGB")
+
+        no_bg = remove(aligned)
+
+        # Create white background
+        white = Image.new("RGB", no_bg.size, (255, 255, 255))
+        white.paste(no_bg, mask=no_bg.split()[3])
+
+        # Resize to 4×6
+        final_img = resize_to_4x6(white)
+
+        # Output
+        buf = io.BytesIO()
+        final_img.save(buf, format="JPEG")
+        buf.seek(0)
+
+        return Response(buf.getvalue(), media_type="image/jpeg")
+
+    except Exception as e:
+        # Log error inside HF Space
+        print("PROCESS ERROR:", e)
+        return JSONResponse({"error": str(e)}, status_code=500)