Update app.py

app.py

@@ -117,192 +117,20 @@ def health():
     mongo.admin.command("ping")
     return HealthResponse(status="ok", db=db.name, model="Qwen/Qwen-Image-Edit")
 
-@app.post("/generate")
-async def generate(
-    prompt: str = Form(...),
-    image1: UploadFile = File(...),
-    image2: Optional[UploadFile] = File(None),
-    user_id: Optional[str] = Form(None),
-    category_id: Optional[str] = Form(None),
-    user=Depends(verify_firebase_token)
-):
-    start_time = time.time()
-
-    # -------------------------
-    # 1. VALIDATE & READ IMAGES
-    # -------------------------
-    try:
-        img1_bytes = await image1.read()
-        pil_img1 = prepare_image(img1_bytes)
-        input1_id = fs.put(
-            img1_bytes,
-            filename=image1.filename,
-            contentType=image1.content_type,
-            metadata={"role": "input"}
-        )
-    except Exception as e:
-        raise HTTPException(400, f"Failed to read first image: {e}")
-
-    img2_bytes = None
-    input2_id = None
-    pil_img2 = None
-
-    if image2:
-        try:
-            img2_bytes = await image2.read()
-            pil_img2 = prepare_image(img2_bytes)
-            input2_id = fs.put(
-                img2_bytes,
-                filename=image2.filename,
-                contentType=image2.content_type,
-                metadata={"role": "input"}
-            )
-        except Exception as e:
-            raise HTTPException(400, f"Failed to read second image: {e}")
-
-    # -------------------------
-    # 2. COMBINE 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
-
-    # -------------------------
-    # 3. OPTIONAL CATEGORY CLICK LOGIC
-    # -------------------------
-    if user_id and category_id:
-        try:
-            admin_client = MongoClient(os.getenv("ADMIN_MONGODB_URI"))
-            admin_db = admin_client["adminPanel"]
-
-            categories_col = admin_db.categories
-            media_clicks_col = admin_db.media_clicks
-
-            # Validate user_oid & category_oid
-            user_oid = ObjectId(user_id)
-            category_oid = ObjectId(category_id)
-
-            # Check category exists
-            category_doc = categories_col.find_one({"_id": category_oid})
-            if not category_doc:
-                raise HTTPException(400, f"Invalid category_id: {category_id}")
-
-            now = datetime.utcnow()
-
-            # Try updating existing category click
-            update_res = media_clicks_col.update_one(
-                {"userId": user_oid, "categories.categoryId": category_oid},
-                {
-                    "$set": {
-                        "updatedAt": now,
-                        "categories.$.lastClickedAt": now
-                    },
-                    "$inc": {"categories.$.click_count": 1}
-                }
-            )
-
-            # If user or category does not exist → create/push new
-            if update_res.matched_count == 0:
-                media_clicks_col.update_one(
-                    {"userId": user_oid},
-                    {
-                        "$setOnInsert": {"createdAt": now},
-                        "$set": {"updatedAt": now},
-                        "$push": {
-                            "categories": {
-                                "categoryId": category_oid,
-                                "click_count": 1,
-                                "lastClickedAt": now
-                            }
-                        }
-                    },
-                    upsert=True
-                )
-
-        except Exception as e:
-            print("CATEGORY_LOG_ERROR:", e)
-            # but do NOT stop the image generation
-
-    # -------------------------
-    # 4. HF INFERENCE
-    # -------------------------
-    try:
-        pil_output = hf_client.image_to_image(
-            image=combined_img,
-            prompt=prompt,
-            model="Qwen/Qwen-Image-Edit"
-        )
-    except Exception as e:
-        response_time_ms = round((time.time() - start_time) * 1000)
-        logs_collection.insert_one({
-            "timestamp": datetime.utcnow(),
-            "status": "failure",
-            "input1_id": str(input1_id),
-            "input2_id": str(input2_id) if input2_id else None,
-            "prompt": prompt,
-            "user_email": user.get("email"),
-            "error": str(e),
-            "response_time_ms": response_time_ms
-        })
-        raise HTTPException(500, f"Inference failed: {e}")
-
-    # -------------------------
-    # 5. SAVE OUTPUT IMAGE
-    # -------------------------
-    out_buf = io.BytesIO()
-    pil_output.save(out_buf, format="PNG")
-    out_bytes = out_buf.getvalue()
-
-    out_id = fs.put(
-        out_bytes,
-        filename=f"result_{input1_id}.png",
-        contentType="image/png",
-        metadata={
-            "role": "output",
-            "prompt": prompt,
-            "input1_id": str(input1_id),
-            "input2_id": str(input2_id) if input2_id else None,
-            "user_email": user.get("email"),
-        }
-    )
-
-    response_time_ms = round((time.time() - start_time) * 1000)
-
-    # -------------------------
-    # 6. LOG SUCCESS
-    # -------------------------
-    logs_collection.insert_one({
-        "timestamp": datetime.utcnow(),
-        "status": "success",
-        "input1_id": str(input1_id),
-        "input2_id": str(input2_id) if input2_id else None,
-        "output_id": str(out_id),
-        "prompt": prompt,
-        "user_email": user.get("email"),
-        "response_time_ms": response_time_ms
-    })
-
-    return JSONResponse({
-        "output_image_id": str(out_id),
-        "user": user.get("email"),
-        "response_time_ms": response_time_ms
-    })
-
 # @app.post("/generate")
 # async def generate(
 #     prompt: str = Form(...),
 #     image1: UploadFile = File(...),
 #     image2: Optional[UploadFile] = File(None),
+#     user_id: Optional[str] = Form(None),
+#     category_id: Optional[str] = Form(None),
 #     user=Depends(verify_firebase_token)
 # ):
 #     start_time = time.time()
 
-#     # PREPARE IMAGE(S)
+#     # -------------------------
+#     # 1. VALIDATE & READ IMAGES
+#     # -------------------------
 #     try:
 #         img1_bytes = await image1.read()
 #         pil_img1 = prepare_image(img1_bytes)
@@ -313,7 +141,7 @@ async def generate(
 #             metadata={"role": "input"}
 #         )
 #     except Exception as e:
-#         raise HTTPException(status_code=400, detail=f"Failed to read first image: {e}")
+#         raise HTTPException(400, f"Failed to read first image: {e}")
 
 #     img2_bytes = None
 #     input2_id = None
@@ -330,9 +158,11 @@ async def generate(
 #                 metadata={"role": "input"}
 #             )
 #         except Exception as e:
-#             raise HTTPException(status_code=400, detail=f"Failed to read second image: {e}")
+#             raise HTTPException(400, f"Failed to read second image: {e}")
 
-#     # COMBINE IF NEEDED
+#     # -------------------------
+#     # 2. COMBINE IF NEEDED
+#     # -------------------------
 #     if pil_img2:
 #         total_width = pil_img1.width + pil_img2.width
 #         max_height = max(pil_img1.height, pil_img2.height)
@@ -342,7 +172,65 @@ async def generate(
 #     else:
 #         combined_img = pil_img1
 
-#     # INFERENCE
+#     # -------------------------
+#     # 3. OPTIONAL CATEGORY CLICK LOGIC
+#     # -------------------------
+#     if user_id and category_id:
+#         try:
+#             admin_client = MongoClient(os.getenv("ADMIN_MONGODB_URI"))
+#             admin_db = admin_client["adminPanel"]
+
+#             categories_col = admin_db.categories
+#             media_clicks_col = admin_db.media_clicks
+
+#             # Validate user_oid & category_oid
+#             user_oid = ObjectId(user_id)
+#             category_oid = ObjectId(category_id)
+
+#             # Check category exists
+#             category_doc = categories_col.find_one({"_id": category_oid})
+#             if not category_doc:
+#                 raise HTTPException(400, f"Invalid category_id: {category_id}")
+
+#             now = datetime.utcnow()
+
+#             # Try updating existing category click
+#             update_res = media_clicks_col.update_one(
+#                 {"userId": user_oid, "categories.categoryId": category_oid},
+#                 {
+#                     "$set": {
+#                         "updatedAt": now,
+#                         "categories.$.lastClickedAt": now
+#                     },
+#                     "$inc": {"categories.$.click_count": 1}
+#                 }
+#             )
+
+#             # If user or category does not exist → create/push new
+#             if update_res.matched_count == 0:
+#                 media_clicks_col.update_one(
+#                     {"userId": user_oid},
+#                     {
+#                         "$setOnInsert": {"createdAt": now},
+#                         "$set": {"updatedAt": now},
+#                         "$push": {
+#                             "categories": {
+#                                 "categoryId": category_oid,
+#                                 "click_count": 1,
+#                                 "lastClickedAt": now
+#                             }
+#                         }
+#                     },
+#                     upsert=True
+#                 )
+
+#         except Exception as e:
+#             print("CATEGORY_LOG_ERROR:", e)
+#             # but do NOT stop the image generation
+
+#     # -------------------------
+#     # 4. HF INFERENCE
+#     # -------------------------
 #     try:
 #         pil_output = hf_client.image_to_image(
 #             image=combined_img,
@@ -350,7 +238,6 @@ async def generate(
 #             model="Qwen/Qwen-Image-Edit"
 #         )
 #     except Exception as e:
-#         # LOG FAILURE
 #         response_time_ms = round((time.time() - start_time) * 1000)
 #         logs_collection.insert_one({
 #             "timestamp": datetime.utcnow(),
@@ -362,9 +249,11 @@ async def generate(
 #             "error": str(e),
 #             "response_time_ms": response_time_ms
 #         })
-#         raise HTTPException(status_code=500, detail=f"Inference failed: {e}")
+#         raise HTTPException(500, f"Inference failed: {e}")
 
-#     # SAVE OUTPUT IMAGE
+#     # -------------------------
+#     # 5. SAVE OUTPUT IMAGE
+#     # -------------------------
 #     out_buf = io.BytesIO()
 #     pil_output.save(out_buf, format="PNG")
 #     out_bytes = out_buf.getvalue()
@@ -382,19 +271,20 @@ async def generate(
 #         }
 #     )
 
-#     # FINAL RESPONSE TIME
 #     response_time_ms = round((time.time() - start_time) * 1000)
 
-#     # LOG SUCCESS - THIS WAS MISSING THE STATUS AND RESPONSE_TIME_MS FIELDS!
+#     # -------------------------
+#     # 6. LOG SUCCESS
+#     # -------------------------
 #     logs_collection.insert_one({
 #         "timestamp": datetime.utcnow(),
-#         "status": "success",  # ✅ FIXED: Added status field
+#         "status": "success",
 #         "input1_id": str(input1_id),
 #         "input2_id": str(input2_id) if input2_id else None,
 #         "output_id": str(out_id),
 #         "prompt": prompt,
-#         "user_email": user.get("email"),  # ✅ FIXED: Changed from "user" to "user_email"
-#         "response_time_ms": response_time_ms  # ✅ FIXED: Added response_time_ms field
+#         "user_email": user.get("email"),
+#         "response_time_ms": response_time_ms
 #     })
 
 #     return JSONResponse({
@@ -403,6 +293,116 @@ async def generate(
 #         "response_time_ms": response_time_ms
 #     })
 
+@app.post("/generate")
+async def generate(
+    prompt: str = Form(...),
+    image1: UploadFile = File(...),
+    image2: Optional[UploadFile] = File(None),
+    user=Depends(verify_firebase_token)
+):
+    start_time = time.time()
+
+    # PREPARE IMAGE(S)
+    try:
+        img1_bytes = await image1.read()
+        pil_img1 = prepare_image(img1_bytes)
+        input1_id = fs.put(
+            img1_bytes,
+            filename=image1.filename,
+            contentType=image1.content_type,
+            metadata={"role": "input"}
+        )
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Failed to read first image: {e}")
+
+    img2_bytes = None
+    input2_id = None
+    pil_img2 = None
+
+    if image2:
+        try:
+            img2_bytes = await image2.read()
+            pil_img2 = prepare_image(img2_bytes)
+            input2_id = fs.put(
+                img2_bytes,
+                filename=image2.filename,
+                contentType=image2.content_type,
+                metadata={"role": "input"}
+            )
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=f"Failed to read second image: {e}")
+
+    # COMBINE 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
+
+    # INFERENCE
+    try:
+        pil_output = hf_client.image_to_image(
+            image=combined_img,
+            prompt=prompt,
+            model="Qwen/Qwen-Image-Edit"
+        )
+    except Exception as e:
+        # LOG FAILURE
+        response_time_ms = round((time.time() - start_time) * 1000)
+        logs_collection.insert_one({
+            "timestamp": datetime.utcnow(),
+            "status": "failure",
+            "input1_id": str(input1_id),
+            "input2_id": str(input2_id) if input2_id else None,
+            "prompt": prompt,
+            "user_email": user.get("email"),
+            "error": str(e),
+            "response_time_ms": response_time_ms
+        })
+        raise HTTPException(status_code=500, detail=f"Inference failed: {e}")
+
+    # SAVE OUTPUT IMAGE
+    out_buf = io.BytesIO()
+    pil_output.save(out_buf, format="PNG")
+    out_bytes = out_buf.getvalue()
+
+    out_id = fs.put(
+        out_bytes,
+        filename=f"result_{input1_id}.png",
+        contentType="image/png",
+        metadata={
+            "role": "output",
+            "prompt": prompt,
+            "input1_id": str(input1_id),
+            "input2_id": str(input2_id) if input2_id else None,
+            "user_email": user.get("email"),
+        }
+    )
+
+    # FINAL RESPONSE TIME
+    response_time_ms = round((time.time() - start_time) * 1000)
+
+    # LOG SUCCESS - THIS WAS MISSING THE STATUS AND RESPONSE_TIME_MS FIELDS!
+    logs_collection.insert_one({
+        "timestamp": datetime.utcnow(),
+        "status": "success",  # ✅ FIXED: Added status field
+        "input1_id": str(input1_id),
+        "input2_id": str(input2_id) if input2_id else None,
+        "output_id": str(out_id),
+        "prompt": prompt,
+        "user_email": user.get("email"),  # ✅ FIXED: Changed from "user" to "user_email"
+        "response_time_ms": response_time_ms  # ✅ FIXED: Added response_time_ms field
+    })
+
+    return JSONResponse({
+        "output_image_id": str(out_id),
+        "user": user.get("email"),
+        "response_time_ms": response_time_ms
+    })
+
 
 
 @app.get("/image/{image_id}")