Update app.py

app.py

@@ -45,8 +45,11 @@ hf_client = InferenceClient(token=HF_TOKEN)
 # ---------------------------------------------------------------------
 # MongoDB setup
 # ---------------------------------------------------------------------
-MONGODB_URI = "mongodb+srv://harilogicgo_db_user:EpWocJKRpXau8eBW@images.nh9bx9l.mongodb.net/?appName=images"
-DB_NAME = "polaroid_db"
+MONGODB_URI = os.getenv("MONGODB_URI")
+if not MONGODB_URI:
+    raise RuntimeError("MONGODB_URI not set in environment variables")
+
+DB_NAME = os.getenv("DB_NAME", "polaroid_db")
 
 mongo = MongoClient(MONGODB_URI)
 db = mongo[DB_NAME]
@@ -69,7 +72,7 @@ app.add_middleware(
 # Auth dependency
 # ---------------------------------------------------------------------
 async def verify_firebase_token(request: Request):
-    """Middleware-like dependency to verify Firebase JWT from Authorization header."""
+    """Verify Firebase JWT from Authorization header"""
     auth_header = request.headers.get("Authorization")
     if not auth_header or not auth_header.startswith("Bearer "):
         raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
@@ -96,8 +99,11 @@ class HealthResponse(BaseModel):
 @app.get("/health", response_model=HealthResponse)
 def health():
     """Public health check"""
-    mongo.admin.command("ping")
-    return HealthResponse(status="ok", db=db.name, model="Qwen/Qwen-Image-Edit")
+    try:
+        mongo.admin.command("ping")
+        return HealthResponse(status="ok", db=db.name, model="Qwen/Qwen-Image-Edit")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"DB ping failed: {e}")
 
 
 @app.post("/generate")
@@ -108,70 +114,77 @@ async def generate(
     user=Depends(verify_firebase_token)
 ):
     """Generate edited image using HF Inference (secured by Firebase Auth) and log only response time."""
-    start_time = datetime.utcnow()  # Start time for response time logging
+    start_time = datetime.utcnow()  # Start timer
+
     try:
+        # Read first image
         img1_bytes = await image1.read()
         pil_img1 = Image.open(io.BytesIO(img1_bytes)).convert("RGB")
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Failed to read first image: {e}")
 
-    pil_img2 = None
-    img2_bytes = None
-    if image2:
-        try:
+        pil_img2 = None
+        img2_bytes = None
+        if image2:
             img2_bytes = await image2.read()
             pil_img2 = Image.open(io.BytesIO(img2_bytes)).convert("RGB")
-        except Exception as e:
-            raise HTTPException(status_code=400, detail=f"Failed to read second image: {e}")
-
-    if pil_img2:
-        total_width = pil_img1.width + pil_img2.width
-        max_height = max(pil_img1.height, pil_img2.height)
-        combined_img = Image.new("RGB", (total_width, max_height))
-        combined_img.paste(pil_img1, (0, 0))
-        combined_img.paste(pil_img2, (pil_img1.width, 0))
-    else:
-        combined_img = pil_img1
 
-    try:
+        # Combine images if two provided
+        if pil_img2:
+            total_width = pil_img1.width + pil_img2.width
+            max_height = max(pil_img1.height, pil_img2.height)
+            combined_img = Image.new("RGB", (total_width, max_height))
+            combined_img.paste(pil_img1, (0, 0))
+            combined_img.paste(pil_img2, (pil_img1.width, 0))
+        else:
+            combined_img = pil_img1
+
+        # Run HF inference
         pil_output = hf_client.image_to_image(
             image=combined_img,
             prompt=prompt,
             model="Qwen/Qwen-Image-Edit"
         )
+
         out_buf = io.BytesIO()
         pil_output.save(out_buf, format="PNG")
         out_bytes = out_buf.getvalue()
-    except Exception as e:
-        traceback.print_exc()
-        raise HTTPException(status_code=500, detail=f"Inference failed: {e}")
-
-    out_id = fs.put(
-        out_bytes,
-        filename=f"result_{datetime.utcnow().timestamp()}.png",
-        contentType="image/png",
-        role="output",
-        user_email=user.get("email"),
-    )
 
-    # Calculate response time in milliseconds
-    response_time_ms = (datetime.utcnow() - start_time).total_seconds() * 1000
+        # Save output to GridFS
+        out_id = fs.put(
+            out_bytes,
+            filename=f"result_{datetime.utcnow().timestamp()}.png",
+            contentType="image/png",
+            user_email=user.get("email"),
+        )
 
-    # Log only response time
-    try:
-        logs_collection.insert_one({
-            "timestamp": datetime.utcnow(),
-            "user": user.get("email"),
-            "response_time_ms": response_time_ms
-        })
-    except Exception as e:
-        print("⚠️ Failed to log response time:", e)
+        # Calculate response time
+        end_time = datetime.utcnow()
+        response_time_ms = (end_time - start_time).total_seconds() * 1000
+
+        # Log response time only
+        try:
+            logs_collection.insert_one({
+                "timestamp": end_time,
+                "user": user.get("email"),
+                "response_time_ms": response_time_ms
+            })
+        except Exception as e:
+            print("⚠️ Failed to log response time:", e)
+
+        return JSONResponse({"output_image_id": str(out_id), "user": user.get("email")})
 
-    return JSONResponse({
-        "output_image_id": str(out_id),
-        "user": user.get("email"),
-        "response_time_ms": response_time_ms
-    })
+    except Exception as e:
+        traceback.print_exc()
+        end_time = datetime.utcnow()
+        response_time_ms = (end_time - start_time).total_seconds() * 1000
+        try:
+            logs_collection.insert_one({
+                "timestamp": end_time,
+                "user": user.get("email") if user else "unknown",
+                "response_time_ms": response_time_ms
+            })
+        except Exception:
+            pass
+        raise HTTPException(status_code=500, detail=f"Processing failed: {e}")
 
 
 @app.get("/image/{image_id}")
@@ -196,7 +209,6 @@ def get_image(image_id: str, download: Optional[bool] = False):
         headers=headers
     )
 
-
 # ---------------------------------------------------------------------
 # Run locally
 # ---------------------------------------------------------------------