Update api/main.py

api/main.py

@@ -1,3 +1,4 @@
+
 import os
 import uuid
 import shutil
@@ -11,20 +12,7 @@ from pydantic import BaseModel
 from PIL import Image
 
 from src.core import process_inpaint
-from pymongo import MongoClient
-
-# Connect to MongoDB (URI must be in HF env variables)
-MONGO_URI = os.environ.get("MONGO_URI")
-mongo_client = None
-mongo_collection = None
 
-if MONGO_URI:
-    try:
-        mongo_client = MongoClient(MONGO_URI)
-        mongo_db = mongo_client["object-remove-logs"]
-        mongo_collection = mongo_db["api_logs"]
-    except Exception as e:
-        print("Mongo connection failed:", e)
 # Directories (use writable space on HF Spaces)
 BASE_DIR = os.environ.get("DATA_DIR", "/data")
 if not os.path.isdir(BASE_DIR):
@@ -150,7 +138,7 @@ def inpaint(req: InpaintRequest, _: None = Depends(bearer_auth)) -> Dict[str, st
         raise HTTPException(status_code=404, detail="mask_id not found")
 
     img_rgba = _load_rgba_image(file_store[req.image_id]["path"])
-    mask_img = Image.open(file_store[req.mask_id]["path"])
+    mask_img = Image.open(file_store[req.mask_id]["path"])  # may be RGB/gray/RGBA
     mask_rgba = _load_rgba_mask_from_image(mask_img)
 
     result = process_inpaint(np.array(img_rgba), mask_rgba)
@@ -159,15 +147,6 @@ def inpaint(req: InpaintRequest, _: None = Depends(bearer_auth)) -> Dict[str, st
     Image.fromarray(result).save(result_path)
 
     logs.append({"result": result_name, "timestamp": datetime.utcnow().isoformat()})
-
-    # Mongo logging
-    if mongo_collection:
-        mongo_collection.insert_one({
-            "endpoint": "/inpaint",
-            "result": result_name,
-            "timestamp": datetime.utcnow().isoformat()
-        })
-
     return {"result": result_name}
 
 
@@ -200,6 +179,7 @@ def inpaint_multipart(
     request: Request = None,
     _: None = Depends(bearer_auth),
 ) -> Dict[str, str]:
+    # Load in-memory
     img = Image.open(image.file).convert("RGBA")
     m = Image.open(mask.file)
     mask_rgba = _load_rgba_mask_from_image(m)
@@ -220,16 +200,6 @@ def inpaint_multipart(
     if url:
         entry["url"] = url
     logs.append(entry)
-
-    # Mongo logging
-    if mongo_collection:
-        mongo_collection.insert_one({
-            "endpoint": "/inpaint-multipart",
-            "result": result_name,
-            "url": url,
-            "timestamp": datetime.utcnow().isoformat()
-        })
-
     resp: Dict[str, str] = {"result": result_name}
     if url:
         resp["url"] = url
@@ -255,4 +225,4 @@ def view_result(filename: str):
 
 @app.get("/logs")
 def get_logs(_: None = Depends(bearer_auth)) -> JSONResponse:
-    return JSONResponse(content=logs)
+    return JSONResponse(content=logs)