Update app.py

app.py

@@ -14,13 +14,12 @@ import insightface
 from insightface.app import FaceAnalysis
 from huggingface_hub import hf_hub_download
 
-from fastapi import FastAPI, UploadFile, File, HTTPException, Response, Depends, Security, Query
+from fastapi import FastAPI, UploadFile, File, HTTPException, Response, Depends, Security, Form
 from fastapi.responses import RedirectResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
-from fastapi.concurrency import run_in_threadpool
-from pydantic import BaseModel
 from motor.motor_asyncio import AsyncIOMotorClient
-
+from bson import ObjectId
+import requests
 import uvicorn
 import gradio as gr
 from gradio import mount_gradio_app
@@ -28,36 +27,41 @@ from gradio import mount_gradio_app
 # DigitalOcean Spaces
 import boto3
 from botocore.client import Config
-from io import BytesIO
-from typing import Optional 
+from typing import Optional
+
 # --------------------- Logging ---------------------
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# --------------------- Paths -----------------------
+# --------------------- Secrets & Paths ---------------------
 REPO_ID = "HariLogicgo/face_swap_models"
-BASE_DIR = "./workspace"
 MODELS_DIR = "./models"
-
 os.makedirs(MODELS_DIR, exist_ok=True)
 
-# --------------------- Secrets ---------------------
-HF_TOKEN = os.getenv("HF_TOKEN")  # Hugging Face private repo token
-API_SECRET_TOKEN = os.getenv("API_SECRET_TOKEN")  # Bearer token for API
-# --------------------- DigitalOcean Spaces Credentials ---------------------
-DO_SPACES_REGION = os.getenv("DO_SPACES_REGION", "blr1")  # Default region = Bangalore
-DO_SPACES_ENDPOINT = os.getenv("DO_SPACES_ENDPOINT", f"https://{DO_SPACES_REGION}.digitaloceanspaces.com")
-DO_SPACES_KEY = os.getenv("DO_SPACES_KEY")  # Your Access Key
-DO_SPACES_SECRET = os.getenv("DO_SPACES_SECRET")  # Your Secret Key
-DO_SPACES_BUCKET = os.getenv("DO_SPACES_BUCKET")  # Bucket Name
-ADMIN_MONGO_URL = os.getenv("ADMIN_MONGO_URL") #MONGODB Storage for usage count
+HF_TOKEN = os.getenv("HF_TOKEN")
+API_SECRET_TOKEN = os.getenv("API_SECRET_TOKEN")
+
+DO_SPACES_REGION = os.getenv("DO_SPACES_REGION", "blr1")
+DO_SPACES_ENDPOINT = f"https://{DO_SPACES_REGION}.digitaloceanspaces.com"
+DO_SPACES_KEY = os.getenv("DO_SPACES_KEY")
+DO_SPACES_SECRET = os.getenv("DO_SPACES_SECRET")
+DO_SPACES_BUCKET = os.getenv("DO_SPACES_BUCKET")
+
+# NEW admin DB
+ADMIN_MONGO_URL = os.getenv("ADMIN_MONGO_URL")
 admin_client = AsyncIOMotorClient(ADMIN_MONGO_URL)
 admin_db = admin_client.adminPanel
 subcategories_col = admin_db.subcategories
 media_clicks_col = admin_db.media_clicks
+
+# OLD logs DB
+MONGODB_URL = os.getenv("MONGODB_URL")
+client = None
+database = None
+
 # --------------------- Download Models ---------------------
 def download_models():
-    logger.info("Downloading models from private HF repo...")
+    logger.info("Downloading models...")
     inswapper_path = hf_hub_download(
         repo_id=REPO_ID,
         filename="models/inswapper_128.onnx",
@@ -66,24 +70,20 @@ def download_models():
         token=HF_TOKEN
     )
 
-    buffalo_files = [
-        "1k3d68.onnx",
-        "2d106det.onnx",
-        "genderage.onnx",
-        "det_10g.onnx",
-        "w600k_r50.onnx"
-    ]
+    buffalo_files = ["1k3d68.onnx", "2d106det.onnx", "genderage.onnx", "det_10g.onnx", "w600k_r50.onnx"]
     for f in buffalo_files:
         hf_hub_download(
             repo_id=REPO_ID,
-            filename=f"models/buffalo_l/{f}",
+            filename=f"models/buffalo_l/" + f,
             repo_type="model",
             local_dir=MODELS_DIR,
             token=HF_TOKEN
         )
-    logger.info("Models downloaded successfully")
+
+    logger.info("Models downloaded.")
     return inswapper_path
 
+
 inswapper_path = download_models()
 
 # --------------------- Face Analysis + Swapper ---------------------
@@ -105,11 +105,6 @@ def ensure_codeformer():
 
 ensure_codeformer()
 
-# --------------------- MongoDB (for API logs only) ---------------------
-MONGODB_URL = os.getenv("MONGODB_URL")
-
-client = None
-database = None
 
 # --------------------- FastAPI ---------------------
 fastapi_app = FastAPI()
@@ -259,106 +254,63 @@ async def face_swap_api(
     source: UploadFile = File(...),
     target_category_id: str = Form(None),
     new_category_id: str = Form(None),
-    user_id: Optional[str] = Form(None),
-    credentials: HTTPAuthorizationCredentials = Security(security)
+    user_id: Optional[str] = Form(None)
 ):
     start_time = datetime.utcnow()
 
     try:
-        # ---------- VALIDATION ----------
+        # ============ VALIDATION ============
         if target_category_id and new_category_id:
-            raise HTTPException(
-                status_code=400,
-                detail="Provide only one of new_category_id or target_category_id, not both."
-            )
+            raise HTTPException(400, "Provide only one of new_category_id or target_category_id.")
 
         if not target_category_id and not new_category_id:
-            raise HTTPException(
-                status_code=400,
-                detail="Either new_category_id or target_category_id is required."
-            )
+            raise HTTPException(400, "Either new_category_id or target_category_id is required.")
 
-        # ---------- READ SOURCE ----------
+        # Read source image
         src_bytes = await source.read()
-
-        # Save source to Spaces
         src_key = f"faceswap/source/{uuid.uuid4().hex}_{source.filename}"
-        upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
+        upload_to_spaces(src_bytes, src_key)
 
-        # ==========================================================
-        # CASE 1: NEW_CATEGORY_ID → Fetch from MongoDB
-        # ==========================================================
+        # ============ CASE 1: MongoDB new_category_id ============
         if new_category_id:
-            from bson import ObjectId
-
-            doc = await subcategories_col.find_one({
-                "asset_images._id": ObjectId(new_category_id)
-            })
+            doc = await subcategories_col.find_one({"asset_images._id": ObjectId(new_category_id)})
 
             if not doc:
-                raise HTTPException(status_code=404, detail="Asset image not found in database")
+                raise HTTPException(404, "Asset image not found.")
 
-            # Find the asset image inside array
-            asset = None
-            for img in doc["asset_images"]:
-                if str(img["_id"]) == str(new_category_id):
-                    asset = img
-                    break
+            # find correct asset image
+            asset = next((img for img in doc["asset_images"] if str(img["_id"]) == new_category_id), None)
 
             if not asset:
-                raise HTTPException(status_code=404, detail="Asset image URL not found")
+                raise HTTPException(404, "Asset image URL not found.")
 
             target_url = asset["url"]
-            category_id = str(doc["categoryId"])
+            category_oid = doc["categoryId"]  # REAL ObjectId
 
-            # ---------- media_clicks update only if user_id exists ----------
-            # ---------- media_clicks update only if user_id exists ----------
+            # ---- MEDIA_CLICKS ONLY IF user_id EXISTS ----
             if user_id:
                 try:
-                    logger.info("======= MEDIA_CLICKS DEBUG LOG =======")
-                    logger.info(f"user_id received: {user_id}")
-                    logger.info(f"category_id resolved from subcategory: {category_id}")
-                    logger.info(f"ADMIN_MONGO_URL: {ADMIN_MONGO_URL}")
-                    
-                    try:
-                        user_oid = ObjectId(user_id)
-                        cat_oid = ObjectId(category_id)
-                    except Exception as conv_err:
-                        logger.error(f"❌ ObjectId conversion failed: {conv_err}")
-                        raise HTTPException(status_code=400, detail=f"Invalid user_id or category_id: {conv_err}")
-                    
+                    user_id = user_id.strip()
+                    user_oid = ObjectId(user_id)
                     now = datetime.utcnow()
-                    
+
                     update_result = await media_clicks_col.update_one(
+                        {"userId": user_oid, "categories.categoryId": category_oid},
                         {
-                            "userId": user_oid,
-                            "categories.categoryId": cat_oid
-                        },
-                        {
-                            "$set": {
-                                "updatedAt": now,
-                                "categories.$.lastClickedAt": now
-                            },
-                            "$inc": {
-                                "categories.$.click_count": 1
-                            }
+                            "$set": {"updatedAt": now, "categories.$.lastClickedAt": now},
+                            "$inc": {"categories.$.click_count": 1}
                         }
                     )
-                    
-                    logger.info(f"Existing category update matched_count: {update_result.matched_count}")
-                    logger.info(f"Existing category update modified_count: {update_result.modified_count}")
-                    
+
                     if update_result.matched_count == 0:
-                        logger.info("Category not found → pushing new category entry")
-                    
-                        insert_result = await media_clicks_col.update_one(
+                        await media_clicks_col.update_one(
                             {"userId": user_oid},
                             {
                                 "$setOnInsert": {"createdAt": now},
                                 "$set": {"updatedAt": now},
                                 "$push": {
                                     "categories": {
-                                        "categoryId": cat_oid,
+                                        "categoryId": category_oid,
                                         "click_count": 1,
                                         "lastClickedAt": now
                                     }
@@ -366,151 +318,62 @@ async def face_swap_api(
                             },
                             upsert=True
                         )
-                    
-                        logger.info(f"Insert/upsert raw_result: {insert_result.raw_result}")
-                    
-                    logger.info("======= END MEDIA_CLICKS LOG =======")
-
-            
-                except Exception as media_err:
-                    print("❌ ERROR while saving media_clicks:", str(media_err))
-        # DO NOT break the face-swap flow
-
-            # if user_id:
-            #     user_oid = ObjectId(user_id)
-            #     cat_oid = ObjectId(category_id)
-
-            #     now = datetime.utcnow()
-
-            #     # Try to update existing category entry
-            #     update_result = await media_clicks_col.update_one(
-            #         {
-            #             "userId": user_oid,
-            #             "categories.categoryId": cat_oid
-            #         },
-            #         {
-            #             "$set": {
-            #                 "updatedAt": now,
-            #                 "categories.$.lastClickedAt": now
-            #             },
-            #             "$inc": {
-            #                 "categories.$.click_count": 1
-            #             }
-            #         }
-            #     )
-
-            #     # If category did not exist, push new category object
-            #     if update_result.matched_count == 0:
-            #         await media_clicks_col.update_one(
-            #             {"userId": user_oid},
-            #             {
-            #                 "$setOnInsert": {"createdAt": now},
-            #                 "$set": {"updatedAt": now},
-            #                 "$push": {
-            #                     "categories": {
-            #                         "categoryId": cat_oid,
-            #                         "click_count": 1,
-            #                         "lastClickedAt": now
-            #                     }
-            #                 }
-            #             },
-            #             upsert=True
-            #         )
-
-        # ==========================================================
-        # CASE 2: EXISTING TARGET_CATEGORY_ID (DigitalOcean)
-        # ==========================================================
+
+                except Exception as err:
+                    logger.error(f"MEDIA_CLICK ERROR: {err}")
+                    # Do NOT break face-swap
+
+        # ============ CASE 2: DigitalOcean target_category_id ============
         if target_category_id:
             client = get_spaces_client()
             base_prefix = "faceswap/target/"
             resp = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=base_prefix, Delimiter="/")
-            categories = [prefix["Prefix"].split("/")[2] for prefix in resp.get("CommonPrefixes", [])]
 
+            categories = [p["Prefix"].split("/")[2] for p in resp.get("CommonPrefixes", [])]
             target_url = None
+
             for category in categories:
                 original_prefix = f"faceswap/target/{category}/original/"
-                thumb_prefix = f"faceswap/target/{category}/thumb/"
-
-                original_objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=original_prefix)
-                thumb_objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=thumb_prefix)
+                original_files = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=original_prefix)
 
-                original_files = [obj["Key"].split("/")[-1] for obj in original_objects.get("Contents", []) if obj["Key"].endswith(".png")]
-                thumb_files = [obj["Key"].split("/")[-1] for obj in thumb_objects.get("Contents", []) if obj["Key"].endswith(".png")]
-
-                for idx, filename in enumerate(sorted(original_files), start=1):
+                for idx, obj in enumerate(sorted(original_files.get("Contents", [])), start=1):
                     cid = f"{category.lower()}image_{idx}"
-                    if filename in thumb_files and cid == target_category_id:
-                        target_url = f"https://{DO_SPACES_BUCKET}.blr1.digitaloceanspaces.com/{original_prefix}{filename}"
+                    if cid == target_category_id:
+                        filename = obj["Key"].split("/")[-1]
+                        target_url = f"{DO_SPACES_ENDPOINT}/{DO_SPACES_BUCKET}/{original_prefix}{filename}"
                         break
+
                 if target_url:
                     break
 
             if not target_url:
-                raise HTTPException(status_code=404, detail="Target categoryId not found")
-
-        # ==========================================================
-        # DOWNLOAD TARGET IMAGE
-        # ==========================================================
-        resp = requests.get(target_url)
-        if resp.status_code != 200:
-            raise HTTPException(status_code=404, detail="Target image could not be fetched")
-
-        tgt_bytes = resp.content
-
-        # Decode images
-        src_array = np.frombuffer(src_bytes, np.uint8)
-        tgt_array = np.frombuffer(tgt_bytes, np.uint8)
+                raise HTTPException(404, "Target categoryId not found.")
 
-        src_bgr = cv2.imdecode(src_array, cv2.IMREAD_COLOR)
-        tgt_bgr = cv2.imdecode(tgt_array, cv2.IMREAD_COLOR)
+        # ============ DOWNLOAD TARGET IMAGE ============
+        target_bytes = requests.get(target_url).content
+        tgt_bgr = cv2.imdecode(np.frombuffer(target_bytes, np.uint8), cv2.IMREAD_COLOR)
+        src_bgr = cv2.imdecode(np.frombuffer(src_bytes, np.uint8), cv2.IMREAD_COLOR)
 
         if src_bgr is None or tgt_bgr is None:
-            raise HTTPException(status_code=400, detail="Invalid image data")
+            raise HTTPException(400, "Invalid image data")
 
+        # ============ FACE SWAP ============
         src_rgb = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2RGB)
         tgt_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
 
-        # ==========================================================
-        # FACE SWAP
-        # ==========================================================
         final_img, final_path, err = face_swap_and_enhance(src_rgb, tgt_rgb)
         if err:
-            raise HTTPException(status_code=500, detail=err)
+            raise HTTPException(500, err)
 
-        # Upload output to Spaces
-        with open(final_path, "rb") as f:
-            result_bytes = f.read()
-
-        result_key = f"faceswap/result/{uuid.uuid4().hex}_enhanced.png"
-        result_url = upload_to_spaces(result_bytes, result_key, content_type="image/png")
-
-        # Log success
-        end_time = datetime.utcnow()
-        response_time_ms = (end_time - start_time).total_seconds() * 1000
-
-        if database is not None:
-            await database.api_logs.insert_one({
-                "endpoint": "/face-swap",
-                "status": "success",
-                "response_time_ms": response_time_ms,
-                "timestamp": end_time
-            })
+        result_bytes = open(final_path, "rb").read()
+        result_key = f"faceswap/result/{uuid.uuid4().hex}.png"
+        result_url = upload_to_spaces(result_bytes, result_key)
 
         return {"result_key": result_key, "result_url": result_url}
 
     except Exception as e:
-        # Log failure
-        end_time = datetime.utcnow()
-        response_time_ms = (end_time - start_time).total_seconds() * 1000
-        if database is not None:
-            await database.api_logs.insert_one({
-                "endpoint": "/face-swap",
-                "status": "fail",
-                "response_time_ms": response_time_ms,
-                "timestamp": end_time,
-                "error": str(e)
-            })
-        raise HTTPException(status_code=500, detail=f"Face swap failed: {str(e)}")
+        raise HTTPException(500, f"Face swap failed: {str(e)}")
+
 
 # @fastapi_app.post("/face-swap", dependencies=[Depends(verify_token)])
 # async def face_swap_api(