images accesed through digital ocean

app.py

@@ -17,13 +17,17 @@ from fastapi import FastAPI, UploadFile, File, HTTPException, Response, Depends,
 from fastapi.responses import RedirectResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
-from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorGridFSBucket
-from bson import ObjectId
+from motor.motor_asyncio import AsyncIOMotorClient
 
 import uvicorn
 import gradio as gr
 from gradio import mount_gradio_app
 
+# DigitalOcean Spaces
+import boto3
+from botocore.client import Config
+from io import BytesIO
+from typing import Optional 
 # --------------------- Logging ---------------------
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -31,18 +35,21 @@ logger = logging.getLogger(__name__)
 # --------------------- Paths -----------------------
 REPO_ID = "HariLogicgo/face_swap_models"
 BASE_DIR = "./workspace"
-UPLOAD_DIR = os.path.join(BASE_DIR, "uploads")
-RESULT_DIR = os.path.join(BASE_DIR, "results")
 MODELS_DIR = "./models"
 
-os.makedirs(UPLOAD_DIR, exist_ok=True)
-os.makedirs(RESULT_DIR, exist_ok=True)
 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 config
+DO_SPACES_KEY = os.getenv("DO_SPACES_KEY")
+DO_SPACES_SECRET = os.getenv("DO_SPACES_SECRET")
+DO_SPACES_REGION = os.getenv("DO_SPACES_REGION")
+DO_SPACES_ENDPOINT = os.getenv("DO_SPACES_ENDPOINT", f"https://{DO_SPACES_REGION}.digitaloceanspaces.com")
+DO_SPACES_BUCKET = os.getenv("DO_SPACES_BUCKET")
+
 # --------------------- Download Models ---------------------
 def download_models():
     logger.info("Downloading models from private HF repo...")
@@ -93,24 +100,22 @@ def ensure_codeformer():
 
 ensure_codeformer()
 
-# --------------------- MongoDB ---------------------
+# --------------------- MongoDB (for API logs only) ---------------------
 MONGODB_URL = os.getenv("MONGODB_URL")
 
 client: AsyncIOMotorClient = None
 database = None
-fs_bucket: AsyncIOMotorGridFSBucket = None
 
 # --------------------- FastAPI ---------------------
 fastapi_app = FastAPI()
 
 @fastapi_app.on_event("startup")
 async def startup_db():
-    global client, database, fs_bucket
-    logger.info("Initializing MongoDB + GridFS...")
+    global client, database
+    logger.info("Initializing MongoDB for API logs...")
     client = AsyncIOMotorClient(MONGODB_URL)
     database = client.FaceSwap
-    fs_bucket = AsyncIOMotorGridFSBucket(database)
-    logger.info("MongoDB + GridFS initialized")
+    logger.info("MongoDB initialized for API logs")
 
 @fastapi_app.on_event("shutdown")
 async def shutdown_db():
@@ -130,7 +135,7 @@ def verify_token(credentials: HTTPAuthorizationCredentials = Security(security))
 # --------------------- Logging API Hits ---------------------
 async def log_faceswap_hit(token: str, status: str = "success"):
     global database
-    if database is None:  # ✅ FIXED: explicit None check
+    if database is None:
         return
     await database.api_logs.insert_one({
         "token": token,
@@ -142,13 +147,13 @@ async def log_faceswap_hit(token: str, status: str = "success"):
 # --------------------- Face Swap Pipeline ---------------------
 swap_lock = threading.Lock()
 
-def face_swap_and_enhance(src_img, tgt_img):
+def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
     try:
         with swap_lock:
-            shutil.rmtree(UPLOAD_DIR, ignore_errors=True)
-            shutil.rmtree(RESULT_DIR, ignore_errors=True)
-            os.makedirs(UPLOAD_DIR, exist_ok=True)
-            os.makedirs(RESULT_DIR, exist_ok=True)
+            # Use a temp dir for intermediate files
+            if os.path.exists(temp_dir):
+                shutil.rmtree(temp_dir)
+            os.makedirs(temp_dir, exist_ok=True)
 
             src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
             tgt_bgr = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
@@ -158,19 +163,19 @@ def face_swap_and_enhance(src_img, tgt_img):
             if not src_faces or not tgt_faces:
                 return None, None, "❌ Face not detected in one of the images"
 
-            swapped_path = os.path.join(UPLOAD_DIR, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
+            swapped_path = os.path.join(temp_dir, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
             swapped_bgr = swapper.get(tgt_bgr, tgt_faces[0], src_faces[0])
             if swapped_bgr is None:
                 return None, None, "❌ Face swap failed"
 
             cv2.imwrite(swapped_path, swapped_bgr)
 
-            cmd = f"python {CODEFORMER_PATH} -w 0.7 --input_path {swapped_path} --output_path {RESULT_DIR} --bg_upsampler realesrgan --face_upsample"
+            cmd = f"python {CODEFORMER_PATH} -w 0.7 --input_path {swapped_path} --output_path {temp_dir} --bg_upsampler realesrgan --face_upsample"
             result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
             if result.returncode != 0:
                 return None, None, f"❌ CodeFormer failed:\n{result.stderr}"
 
-            final_results_dir = os.path.join(RESULT_DIR, "final_results")
+            final_results_dir = os.path.join(temp_dir, "final_results")
             final_files = [f for f in os.listdir(final_results_dir) if f.endswith(".png")]
             if not final_files:
                 return None, None, "❌ No enhanced image found"
@@ -202,6 +207,29 @@ with gr.Blocks() as demo:
 
     btn.click(process, [src_input, tgt_input], [output_img, download, error_box])
 
+# --------------------- DigitalOcean Spaces Helper ---------------------
+def get_spaces_client():
+    session = boto3.session.Session()
+    client = session.client(
+        's3',
+        region_name=DO_SPACES_REGION,
+        endpoint_url=DO_SPACES_ENDPOINT,
+        aws_access_key_id=DO_SPACES_KEY,
+        aws_secret_access_key=DO_SPACES_SECRET,
+        config=Config(signature_version='s3v4')
+    )
+    return client
+
+def upload_to_spaces(file_bytes, key, content_type="image/png"):
+    client = get_spaces_client()
+    client.put_object(Bucket=DO_SPACES_BUCKET, Key=key, Body=file_bytes, ContentType=content_type, ACL='public-read')
+    return f"{DO_SPACES_ENDPOINT}/{DO_SPACES_BUCKET}/{key}"
+
+def download_from_spaces(key):
+    client = get_spaces_client()
+    obj = client.get_object(Bucket=DO_SPACES_BUCKET, Key=key)
+    return obj['Body'].read()
+
 # --------------------- API Endpoints ---------------------
 @fastapi_app.get("/")
 def root():
@@ -211,82 +239,66 @@ def root():
 async def health():
     return {"status": "healthy"}
 
-@fastapi_app.post("/source", dependencies=[Depends(verify_token)])
-async def upload_source(image: UploadFile = File(...)):
-    contents = await image.read()
-    file_id = await fs_bucket.upload_from_stream(image.filename, contents, metadata={"type": "source"})
-    return {"source_id": str(file_id)}
-
-@fastapi_app.get("/targets", dependencies=[Depends(verify_token)])
-async def list_targets():
-    files = []
-    async for file in database.fs.files.find({"metadata.type": "target", "metadata.predefined": True}):
-        files.append({
-            "id": str(file["_id"]),
-            "filename": file["filename"]
-        })
-    return {"targets": files}
-
-@fastapi_app.post("/target", dependencies=[Depends(verify_token)])
-async def upload_target(image: UploadFile = File(...)):
-    contents = await image.read()
-    file_id = await fs_bucket.upload_from_stream(image.filename, contents, metadata={"type": "target"})
-    return {"target_id": str(file_id)}
-
-class FaceSwapRequest(BaseModel):
-    source_id: str
-    target_id: str
-
-@fastapi_app.post("/faceswap", dependencies=[Depends(verify_token)])
-async def perform_faceswap(request: FaceSwapRequest, credentials: HTTPAuthorizationCredentials = Security(security)):
+@fastapi_app.post("/face-swap", dependencies=[Depends(verify_token)])
+async def face_swap_api(
+    source: UploadFile = File(...),
+    target: UploadFile = File(...),
+    credentials: HTTPAuthorizationCredentials = Security(security)
+):
     try:
-        src_stream = await fs_bucket.open_download_stream(ObjectId(request.source_id))
-        src_bytes = await src_stream.read()
-        tgt_stream = await fs_bucket.open_download_stream(ObjectId(request.target_id))
-        tgt_bytes = await tgt_stream.read()
-    except Exception:
-        await log_faceswap_hit(credentials.credentials, status="error")
-        raise HTTPException(status_code=404, detail="Source or Target not found")
+        # Read images
+        src_bytes = await source.read()
+        tgt_bytes = await target.read()
+
+        # Save to Spaces (source/target)
+        src_key = f"faceswap/source/{uuid.uuid4().hex}_{source.filename}"
+        tgt_key = f"faceswap/target/{uuid.uuid4().hex}_{target.filename}"
+        upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
+        upload_to_spaces(tgt_bytes, tgt_key, content_type=target.content_type)
+
+        # Decode for processing
+        src_array = np.frombuffer(src_bytes, np.uint8)
+        tgt_array = np.frombuffer(tgt_bytes, np.uint8)
+        src_bgr = cv2.imdecode(src_array, cv2.IMREAD_COLOR)
+        tgt_bgr = cv2.imdecode(tgt_array, cv2.IMREAD_COLOR)
+        if src_bgr is None or tgt_bgr is None:
+            await log_faceswap_hit(credentials.credentials, status="error")
+            raise HTTPException(status_code=400, detail="Invalid image data")
+
+        src_rgb = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2RGB)
+        tgt_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
+
+        # Run face swap pipeline
+        final_img, final_path, err = face_swap_and_enhance(src_rgb, tgt_rgb)
+        if err:
+            await log_faceswap_hit(credentials.credentials, status="error")
+            raise HTTPException(status_code=500, detail=err)
+
+        # Upload result 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
+        await log_faceswap_hit(credentials.credentials, status="success")
+
+        return {"result_key": result_key, "result_url": result_url}
 
-    src_array = np.frombuffer(src_bytes, np.uint8)
-    src_bgr = cv2.imdecode(src_array, cv2.IMREAD_COLOR)
-    tgt_array = np.frombuffer(tgt_bytes, np.uint8)
-    tgt_bgr = cv2.imdecode(tgt_array, cv2.IMREAD_COLOR)
-
-    if src_bgr is None or tgt_bgr is None:
-        await log_faceswap_hit(credentials.credentials, status="error")
-        raise HTTPException(status_code=400, detail="Invalid image data")
-
-    src_rgb = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2RGB)
-    tgt_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
-
-    final_img, final_path, err = face_swap_and_enhance(src_rgb, tgt_rgb)
-    if err:
+    except Exception as e:
         await log_faceswap_hit(credentials.credentials, status="error")
-        raise HTTPException(status_code=500, detail=err)
-
-    with open(final_path, "rb") as f:
-        final_bytes = f.read()
+        raise HTTPException(status_code=500, detail=f"Face swap failed: {str(e)}")
 
-    result_id = await fs_bucket.upload_from_stream("enhanced.png", final_bytes, metadata={"type": "result"})
-
-    # ✅ Log success
-    await log_faceswap_hit(credentials.credentials, status="success")
-
-    return {"result_id": str(result_id)}
-
-@fastapi_app.get("/download/{result_id}", dependencies=[Depends(verify_token)])
-async def download_result(result_id: str):
+@fastapi_app.get("/preview/{result_key:path}")
+async def preview_result(result_key: str):
     try:
-        stream = await fs_bucket.open_download_stream(ObjectId(result_id))
-        data = await stream.read()
+        img_bytes = download_from_spaces(result_key)
     except Exception:
         raise HTTPException(status_code=404, detail="Result not found")
-
     return Response(
-        content=data,
+        content=img_bytes,
         media_type="image/png",
-        headers={"Content-Disposition": "attachment; filename=result.png"}
+        headers={"Content-Disposition": "inline; filename=result.png"}
     )
 
 # --------------------- Mount Gradio ---------------------

requirements.txt

@@ -35,6 +35,6 @@ huggingface_hub
 fastapi
 uvicorn
 motor
-
+boto3
 # Optional but included in CodeFormer repo
 tb-nightly