Update app.py

app.py

@@ -4,7 +4,7 @@ import json
 import traceback
 from datetime import datetime
 from typing import Optional
-
+import time
 from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request, Depends
 from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
@@ -45,11 +45,8 @@ hf_client = InferenceClient(token=HF_TOKEN)
 # ---------------------------------------------------------------------
 # MongoDB setup
 # ---------------------------------------------------------------------
-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")
+MONGODB_URI = "mongodb+srv://harilogicgo_db_user:EpWocJKRpXau8eBW@images.nh9bx9l.mongodb.net/?appName=images"
+DB_NAME = "polaroid_db"
 
 mongo = MongoClient(MONGODB_URI)
 db = mongo[DB_NAME]
@@ -72,7 +69,7 @@ app.add_middleware(
 # Auth dependency
 # ---------------------------------------------------------------------
 async def verify_firebase_token(request: Request):
-    """Verify Firebase JWT from Authorization header"""
+    """Middleware-like dependency to 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")
@@ -99,11 +96,8 @@ class HealthResponse(BaseModel):
 @app.get("/health", response_model=HealthResponse)
 def health():
     """Public health check"""
-    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}")
+    mongo.admin.command("ping")
+    return HealthResponse(status="ok", db=db.name, model="Qwen/Qwen-Image-Edit")
 
 
 @app.post("/generate")
@@ -113,78 +107,87 @@ async def generate(
     image2: Optional[UploadFile] = File(None),
     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 timer
+    start_time = time.time()   # ⬅️ Start timing
 
     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:
+    pil_img2 = None
+    img2_bytes = None
+    if image2:
+        try:
             img2_bytes = await image2.read()
             pil_img2 = Image.open(io.BytesIO(img2_bytes)).convert("RGB")
-
-        # 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
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=f"Failed to read second image: {e}")
+
+    input1_id = fs.put(img1_bytes, filename=image1.filename, contentType=image1.content_type, role="input")
+    input2_id = None
+    if pil_img2:
+        input2_id = fs.put(img2_bytes, filename=image2.filename, contentType=image2.content_type, role="input")
+
+    # Combine images if needed
+    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
+
+    # HuggingFace inference
+    try:
         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}")
+
+    # Save output image
+    out_id = fs.put(
+        out_bytes,
+        filename=f"result_{input1_id}.png",
+        contentType="image/png",
+        prompt=prompt,
+        role="output",
+        input1_id=str(input1_id),
+        input2_id=str(input2_id) if input2_id else None,
+        user_email=user.get("email"),
+    )
 
-        # 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"),
-        )
-
-        # Calculate response time
-        end_time = datetime.utcnow()
-        response_time_ms = (end_time - start_time).total_seconds() * 1000
+    # -------------------------------------------
+    # LOGGING WITH RESPONSE TIME (NEW ADDITION)
+    # -------------------------------------------
+    end_time = time.time()
+    response_time_ms = round((end_time - start_time) * 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)
+    logs_collection.insert_one({
+        "timestamp": datetime.utcnow(),
+        "input1_id": str(input1_id),
+        "input2_id": str(input2_id) if input2_id else None,
+        "output_id": str(out_id),
+        "prompt": prompt,
+        "user": user.get("email"),
+        "response_time_ms": response_time_ms   # <-- LOGGED HERE
+    })
+    # -------------------------------------------
 
-        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}")
@@ -209,9 +212,11 @@ def get_image(image_id: str, download: Optional[bool] = False):
         headers=headers
     )
 
+
 # ---------------------------------------------------------------------
 # Run locally
 # ---------------------------------------------------------------------
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)
+