LogicGoInfotechSpaces commited on
Commit
903e397
·
verified ·
1 Parent(s): 9f5afa6

Create media_clicks_app.py

Browse files
Files changed (1) hide show
  1. media_clicks_app.py +1259 -0
media_clicks_app.py ADDED
@@ -0,0 +1,1259 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #####################FASTAPI___________________##############
2
+ import os
3
+ os.environ["OMP_NUM_THREADS"] = "1"
4
+ import shutil
5
+ import uuid
6
+ import cv2
7
+ import numpy as np
8
+ import threading
9
+ import asyncio
10
+ import subprocess
11
+ import logging
12
+ import tempfile
13
+ import sys
14
+ import time
15
+ from datetime import datetime,timedelta
16
+ import tempfile
17
+ import insightface
18
+ from insightface.app import FaceAnalysis
19
+ from huggingface_hub import hf_hub_download
20
+ from fastapi import FastAPI, UploadFile, File, HTTPException, Response, Depends, Security, Form
21
+ from fastapi.responses import RedirectResponse
22
+ from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
23
+ from motor.motor_asyncio import AsyncIOMotorClient
24
+ from bson import ObjectId
25
+ from bson.errors import InvalidId
26
+ import httpx
27
+ import uvicorn
28
+ from PIL import Image
29
+ import io
30
+ import requests
31
+ # DigitalOcean Spaces
32
+ import boto3
33
+ from botocore.client import Config
34
+ from typing import Optional
35
+
36
+ # --------------------- Logging ---------------------
37
+ logging.basicConfig(level=logging.INFO)
38
+ logger = logging.getLogger(__name__)
39
+
40
+ # --------------------- Secrets & Paths ---------------------
41
+ REPO_ID = "HariLogicgo/face_swap_models"
42
+ MODELS_DIR = "./models"
43
+ os.makedirs(MODELS_DIR, exist_ok=True)
44
+
45
+ HF_TOKEN = os.getenv("HF_TOKEN")
46
+ API_SECRET_TOKEN = os.getenv("API_SECRET_TOKEN")
47
+
48
+ DO_SPACES_REGION = os.getenv("DO_SPACES_REGION", "blr1")
49
+ DO_SPACES_ENDPOINT = f"https://{DO_SPACES_REGION}.digitaloceanspaces.com"
50
+ DO_SPACES_KEY = os.getenv("DO_SPACES_KEY")
51
+ DO_SPACES_SECRET = os.getenv("DO_SPACES_SECRET")
52
+ DO_SPACES_BUCKET = os.getenv("DO_SPACES_BUCKET")
53
+
54
+ # NEW admin DB (with error handling for missing env vars)
55
+ ADMIN_MONGO_URL = os.getenv("ADMIN_MONGO_URL")
56
+ admin_client = None
57
+ admin_db = None
58
+ subcategories_col = None
59
+ if ADMIN_MONGO_URL:
60
+ try:
61
+ admin_client = AsyncIOMotorClient(ADMIN_MONGO_URL)
62
+ admin_db = admin_client.adminPanel
63
+ subcategories_col = admin_db.subcategories
64
+ except Exception as e:
65
+ logger.warning(f"MongoDB admin connection failed (optional): {e}")
66
+
67
+ # Collage Maker DB (optional)
68
+ COLLAGE_MAKER_DB_URL = os.getenv("COLLAGE_MAKER_DB_URL")
69
+ collage_maker_client = None
70
+ collage_maker_db = None
71
+ collage_subcategories_col = None
72
+ if COLLAGE_MAKER_DB_URL:
73
+ try:
74
+ collage_maker_client = AsyncIOMotorClient(COLLAGE_MAKER_DB_URL)
75
+ collage_maker_db = collage_maker_client.adminPanel
76
+ collage_subcategories_col = collage_maker_db.subcategories
77
+ except Exception as e:
78
+ logger.warning(f"MongoDB collage-maker connection failed (optional): {e}")
79
+
80
+ # AI Enhancer DB (optional)
81
+
82
+ AI_ENHANCER_DB_URL = os.getenv("AI_ENHANCER_DB_URL")
83
+ ai_enhancer_client = None
84
+ ai_enhancer_db = None
85
+ ai_enhancer_subcategories_col = None
86
+
87
+ if AI_ENHANCER_DB_URL:
88
+ try:
89
+ ai_enhancer_client = AsyncIOMotorClient(AI_ENHANCER_DB_URL)
90
+ ai_enhancer_db = ai_enhancer_client.test # 🔴 test database
91
+ ai_enhancer_subcategories_col = ai_enhancer_db.subcategories
92
+ except Exception as e:
93
+ logger.warning(f"MongoDB ai-enhancer connection failed (optional): {e}")
94
+
95
+
96
+ # OLD logs DB
97
+ MONGODB_URL = os.getenv("MONGODB_URL_LOGS")
98
+ client = None
99
+ database = None
100
+
101
+ # --------------------- Download Models ---------------------
102
+ def download_models():
103
+ try:
104
+ logger.info("Downloading models...")
105
+ inswapper_path = hf_hub_download(
106
+ repo_id=REPO_ID,
107
+ filename="models/inswapper_128.onnx",
108
+ repo_type="model",
109
+ local_dir=MODELS_DIR,
110
+ token=HF_TOKEN
111
+ )
112
+
113
+ buffalo_files = ["1k3d68.onnx", "2d106det.onnx", "genderage.onnx", "det_10g.onnx", "w600k_r50.onnx"]
114
+ for f in buffalo_files:
115
+ hf_hub_download(
116
+ repo_id=REPO_ID,
117
+ filename=f"models/buffalo_l/" + f,
118
+ repo_type="model",
119
+ local_dir=MODELS_DIR,
120
+ token=HF_TOKEN
121
+ )
122
+
123
+ logger.info("Models downloaded successfully.")
124
+ return inswapper_path
125
+ except Exception as e:
126
+ logger.error(f"Model download failed: {e}")
127
+ raise
128
+
129
+ try:
130
+ inswapper_path = download_models()
131
+
132
+ # --------------------- Face Analysis + Swapper ---------------------
133
+ providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
134
+ face_analysis_app = FaceAnalysis(name="buffalo_l", root=MODELS_DIR, providers=providers)
135
+ face_analysis_app.prepare(ctx_id=0, det_size=(640, 640))
136
+ swapper = insightface.model_zoo.get_model(inswapper_path, providers=providers)
137
+ logger.info("Face analysis models loaded successfully")
138
+ except Exception as e:
139
+ logger.error(f"Failed to initialize face analysis models: {e}")
140
+ # Set defaults to prevent crash
141
+ inswapper_path = None
142
+ face_analysis_app = None
143
+ swapper = None
144
+
145
+ # --------------------- CodeFormer ---------------------
146
+ CODEFORMER_PATH = "CodeFormer/inference_codeformer.py"
147
+
148
+ def ensure_codeformer():
149
+ """
150
+ Ensure CodeFormer's local basicsr + facelib are importable and
151
+ pretrained weights are downloaded. No setup.py needed — we use
152
+ sys.path / PYTHONPATH instead.
153
+ """
154
+ try:
155
+ if not os.path.exists("CodeFormer"):
156
+ logger.info("CodeFormer not found, cloning repository...")
157
+ subprocess.run("git clone https://github.com/sczhou/CodeFormer.git", shell=True, check=True)
158
+ subprocess.run("pip install -r CodeFormer/requirements.txt", shell=True, check=False)
159
+
160
+ # Add CodeFormer root to sys.path so `import basicsr` and
161
+ # `import facelib` resolve to the local (compatible) versions
162
+ # instead of the broken PyPI basicsr==1.4.2.
163
+ codeformer_root = os.path.join(os.getcwd(), "CodeFormer")
164
+ if codeformer_root not in sys.path:
165
+ sys.path.insert(0, codeformer_root)
166
+ logger.info(f"Added {codeformer_root} to sys.path for local basicsr/facelib")
167
+
168
+ # NOTE: We do NOT need the PyPI 'realesrgan' package.
169
+ # Both in-process and subprocess paths use CodeFormer's local
170
+ # basicsr.utils.realesrgan_utils.RealESRGANer instead.
171
+ # Installing PyPI realesrgan at runtime would re-install the
172
+ # broken basicsr==1.4.2 and break everything.
173
+
174
+ # Download pretrained weights if not already present
175
+ if os.path.exists("CodeFormer"):
176
+ try:
177
+ subprocess.run("python CodeFormer/scripts/download_pretrained_models.py facelib", shell=True, check=False, timeout=300)
178
+ except (subprocess.TimeoutExpired, subprocess.CalledProcessError):
179
+ logger.warning("Failed to download facelib models (optional)")
180
+ try:
181
+ subprocess.run("python CodeFormer/scripts/download_pretrained_models.py CodeFormer", shell=True, check=False, timeout=300)
182
+ except (subprocess.TimeoutExpired, subprocess.CalledProcessError):
183
+ logger.warning("Failed to download CodeFormer models (optional)")
184
+ except Exception as e:
185
+ logger.error(f"CodeFormer setup failed: {e}")
186
+ logger.warning("Continuing without CodeFormer features...")
187
+
188
+ ensure_codeformer()
189
+
190
+ # --------------------- In-Process CodeFormer (No Subprocess!) ---------------------
191
+ # Load CodeFormer models ONCE at startup instead of spawning a new Python process per request.
192
+ # This eliminates 15-40s of model loading overhead per request.
193
+
194
+ codeformer_net = None
195
+ codeformer_upsampler = None
196
+ codeformer_face_helper = None
197
+ codeformer_device = None
198
+
199
+ def init_codeformer_in_process():
200
+ """Load CodeFormer models once into memory for fast per-request inference."""
201
+ global codeformer_net, codeformer_upsampler, codeformer_face_helper, codeformer_device
202
+ try:
203
+ import torch
204
+ from torchvision.transforms.functional import normalize as torch_normalize
205
+
206
+ # Add CodeFormer to Python path
207
+ codeformer_root = os.path.join(os.getcwd(), "CodeFormer")
208
+ if codeformer_root not in sys.path:
209
+ sys.path.insert(0, codeformer_root)
210
+
211
+ from basicsr.utils.registry import ARCH_REGISTRY
212
+ from basicsr.utils.download_util import load_file_from_url
213
+ from facelib.utils.face_restoration_helper import FaceRestoreHelper
214
+
215
+ codeformer_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
216
+ logger.info(f"Initializing CodeFormer on device: {codeformer_device}")
217
+
218
+ # 1) Load CodeFormer network
219
+ net = ARCH_REGISTRY.get('CodeFormer')(
220
+ dim_embd=512, codebook_size=1024, n_head=8, n_layers=9,
221
+ connect_list=['32', '64', '128', '256']
222
+ ).to(codeformer_device)
223
+
224
+ ckpt_path = load_file_from_url(
225
+ url='https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth',
226
+ model_dir='weights/CodeFormer', progress=True, file_name=None
227
+ )
228
+ checkpoint = torch.load(ckpt_path, map_location=codeformer_device)['params_ema']
229
+ net.load_state_dict(checkpoint)
230
+ net.eval()
231
+ codeformer_net = net
232
+
233
+ # 2) RealESRGAN upsampler — SKIPPED for face swap
234
+ # Background/face upsampling is the #1 bottleneck (~20s per image).
235
+ # For face swap we only need CodeFormer face restoration, not super-resolution.
236
+ # The upsampler is kept as None; we no longer download the 64MB model at startup.
237
+ codeformer_upsampler = None
238
+
239
+ # 3) Create FaceRestoreHelper (reused per request)
240
+ # NOTE: local CodeFormer uses "upscale_factor" (not "upscale")
241
+ # upscale_factor=1 → keep original resolution (no 2x upscale needed for face swap)
242
+ codeformer_face_helper = FaceRestoreHelper(
243
+ upscale_factor=1,
244
+ face_size=512,
245
+ crop_ratio=(1, 1),
246
+ det_model='retinaface_resnet50',
247
+ save_ext='png',
248
+ use_parse=True,
249
+ device=codeformer_device
250
+ )
251
+
252
+ logger.info("✅ CodeFormer models loaded in-process successfully!")
253
+ return True
254
+ except Exception as e:
255
+ logger.error(f"Failed to load CodeFormer in-process: {e}")
256
+ logger.warning("CodeFormer enhancement will be unavailable.")
257
+ return False
258
+
259
+ # Try to load CodeFormer models in-process
260
+ _codeformer_available = init_codeformer_in_process()
261
+ # --------------------- FastAPI ---------------------
262
+ fastapi_app = FastAPI()
263
+
264
+ @fastapi_app.on_event("startup")
265
+ async def startup_db():
266
+ global client, database
267
+ if MONGODB_URL:
268
+ try:
269
+ logger.info("Initializing MongoDB for API logs...")
270
+ client = AsyncIOMotorClient(MONGODB_URL)
271
+ database = client.logs
272
+ logger.info("MongoDB initialized for API logs")
273
+ except Exception as e:
274
+ logger.warning(f"MongoDB connection failed (optional): {e}")
275
+ client = None
276
+ database = None
277
+ else:
278
+ logger.warning("MONGODB_URL not set, skipping MongoDB initialization")
279
+
280
+ @fastapi_app.on_event("shutdown")
281
+ async def shutdown_db():
282
+ global client, admin_client, collage_maker_client
283
+ if client is not None:
284
+ client.close()
285
+ logger.info("MongoDB connection closed")
286
+ if admin_client is not None:
287
+ admin_client.close()
288
+ logger.info("Admin MongoDB connection closed")
289
+ if collage_maker_client is not None:
290
+ collage_maker_client.close()
291
+ logger.info("Collage Maker MongoDB connection closed")
292
+
293
+ # --------------------- Auth ---------------------
294
+ security = HTTPBearer()
295
+
296
+ def verify_token(credentials: HTTPAuthorizationCredentials = Security(security)):
297
+ if credentials.credentials != API_SECRET_TOKEN:
298
+ raise HTTPException(status_code=401, detail="Invalid or missing token")
299
+ return credentials.credentials
300
+
301
+ # --------------------- DB Selector ---------------------
302
+ def get_app_db_collections(appname: Optional[str] = None):
303
+ """
304
+ Returns subcategories_collection based on appname.
305
+ """
306
+
307
+ if appname:
308
+ app = appname.strip().lower()
309
+
310
+ if app == "collage-maker":
311
+ if collage_subcategories_col is not None:
312
+ return collage_subcategories_col
313
+ logger.warning("Collage-maker DB not configured, falling back to admin")
314
+
315
+ elif app == "ai-enhancer":
316
+ if ai_enhancer_subcategories_col is not None:
317
+ return ai_enhancer_subcategories_col
318
+ logger.warning("AI-Enhancer DB not configured, falling back to admin")
319
+
320
+ # default fallback
321
+ return subcategories_col
322
+
323
+
324
+
325
+ # --------------------- Logging API Hits ---------------------
326
+ async def log_faceswap_hit(token: str, status: str = "success"):
327
+ global database
328
+ if database is None:
329
+ return
330
+ await database.faceswap.insert_one({
331
+ "token": token,
332
+ "endpoint": "/faceswap",
333
+ "status": status,
334
+ "timestamp": datetime.utcnow()
335
+ })
336
+
337
+ # --------------------- Face Swap Pipeline ---------------------
338
+ swap_lock = threading.Lock()
339
+
340
+ def enhance_image_with_codeformer(rgb_img, temp_dir=None, w=0.7):
341
+ """
342
+ Enhance face image using CodeFormer.
343
+ Uses in-process models (fast) if available, falls back to subprocess (slow).
344
+ """
345
+ global codeformer_net, codeformer_upsampler, codeformer_face_helper, codeformer_device
346
+
347
+ t0 = time.time()
348
+
349
+ # ── FAST PATH: In-process CodeFormer (no subprocess!) ──
350
+ if codeformer_net is not None and codeformer_face_helper is not None:
351
+ import torch
352
+ from torchvision.transforms.functional import normalize as torch_normalize
353
+ from basicsr.utils import img2tensor, tensor2img
354
+ from facelib.utils.misc import is_gray
355
+
356
+ bgr_img = cv2.cvtColor(rgb_img, cv2.COLOR_RGB2BGR)
357
+
358
+ # Reset face helper state
359
+ codeformer_face_helper.clean_all()
360
+ codeformer_face_helper.read_image(bgr_img)
361
+
362
+ num_faces = codeformer_face_helper.get_face_landmarks_5(
363
+ only_center_face=False, resize=640, eye_dist_threshold=5
364
+ )
365
+ logger.info(f"[CodeFormer] Detected {num_faces} faces in {time.time()-t0:.2f}s")
366
+
367
+ codeformer_face_helper.align_warp_face()
368
+
369
+ # Enhance each cropped face with CodeFormer neural net
370
+ t_faces = time.time()
371
+ for idx, cropped_face in enumerate(codeformer_face_helper.cropped_faces):
372
+ cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
373
+ torch_normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
374
+ cropped_face_t = cropped_face_t.unsqueeze(0).to(codeformer_device)
375
+
376
+ try:
377
+ with torch.no_grad():
378
+ output = codeformer_net(cropped_face_t, w=w, adain=True)[0]
379
+ restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
380
+ del output
381
+ torch.cuda.empty_cache()
382
+ except Exception as e:
383
+ logger.warning(f"[CodeFormer] Face {idx} inference failed: {e}")
384
+ restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))
385
+
386
+ restored_face = restored_face.astype('uint8')
387
+ codeformer_face_helper.add_restored_face(restored_face, cropped_face)
388
+ logger.info(f"[CodeFormer] Face restoration ({num_faces} faces): {time.time()-t_faces:.2f}s")
389
+
390
+ # Paste restored faces back onto original image
391
+ # NOTE: We skip RealESRGAN background/face upsampling — it's the #1 bottleneck
392
+ # (~20s) and unnecessary for face swap. We only need CodeFormer face restoration.
393
+ t_paste = time.time()
394
+ codeformer_face_helper.get_inverse_affine(None)
395
+ restored_img = codeformer_face_helper.paste_faces_to_input_image(
396
+ upsample_img=None, draw_box=False
397
+ )
398
+ logger.info(f"[CodeFormer] Paste back: {time.time()-t_paste:.2f}s")
399
+
400
+ logger.info(f"[CodeFormer] In-process enhancement done in {time.time()-t0:.2f}s")
401
+ return cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
402
+
403
+ # ── SLOW FALLBACK: Subprocess CodeFormer (with timeout!) ──
404
+ logger.warning("[CodeFormer] In-process models unavailable, falling back to subprocess")
405
+ if temp_dir is None:
406
+ temp_dir = os.path.join(tempfile.gettempdir(), f"enhance_{uuid.uuid4().hex[:8]}")
407
+ os.makedirs(temp_dir, exist_ok=True)
408
+
409
+ input_path = os.path.join(temp_dir, "input.jpg")
410
+ cv2.imwrite(input_path, cv2.cvtColor(rgb_img, cv2.COLOR_RGB2BGR))
411
+
412
+ python_cmd = sys.executable if sys.executable else "python3"
413
+ cmd = (
414
+ f"{python_cmd} {CODEFORMER_PATH} "
415
+ f"-w {w} "
416
+ f"--input_path {input_path} "
417
+ f"--output_path {temp_dir} "
418
+ f"--bg_upsampler None "
419
+ f"--upscale 1"
420
+ )
421
+
422
+ result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=120)
423
+ if result.returncode != 0:
424
+ raise RuntimeError(result.stderr)
425
+
426
+ final_dir = os.path.join(temp_dir, "final_results")
427
+ files = [f for f in os.listdir(final_dir) if f.endswith(".png")]
428
+ if not files:
429
+ raise RuntimeError("No enhanced output")
430
+
431
+ final_path = os.path.join(final_dir, files[0])
432
+ enhanced = cv2.imread(final_path)
433
+ logger.info(f"[CodeFormer] Subprocess enhancement done in {time.time()-t0:.2f}s")
434
+ return cv2.cvtColor(enhanced, cv2.COLOR_BGR2RGB)
435
+
436
+ def multi_face_swap(src_img, tgt_img):
437
+ pipeline_start = time.time()
438
+ src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
439
+ tgt_bgr = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
440
+
441
+ t0 = time.time()
442
+ src_faces = face_analysis_app.get(src_bgr)
443
+ tgt_faces = face_analysis_app.get(tgt_bgr)
444
+ logger.info(f"[Pipeline] Multi-face detection: {time.time()-t0:.2f}s")
445
+
446
+ if not src_faces or not tgt_faces:
447
+ raise ValueError("No faces detected")
448
+
449
+ def face_sort_key(face):
450
+ x1, y1, x2, y2 = face.bbox
451
+ area = (x2 - x1) * (y2 - y1)
452
+ cx = (x1 + x2) / 2
453
+ return (-area, cx)
454
+
455
+ src_male = sorted([f for f in src_faces if f.gender == 1], key=face_sort_key)
456
+ src_female = sorted([f for f in src_faces if f.gender == 0], key=face_sort_key)
457
+ tgt_male = sorted([f for f in tgt_faces if f.gender == 1], key=face_sort_key)
458
+ tgt_female = sorted([f for f in tgt_faces if f.gender == 0], key=face_sort_key)
459
+
460
+ pairs = []
461
+ for s, t in zip(src_male, tgt_male):
462
+ pairs.append((s, t))
463
+ for s, t in zip(src_female, tgt_female):
464
+ pairs.append((s, t))
465
+
466
+ if not pairs:
467
+ src_faces = sorted(src_faces, key=face_sort_key)
468
+ tgt_faces = sorted(tgt_faces, key=face_sort_key)
469
+ pairs = list(zip(src_faces, tgt_faces))
470
+
471
+ t0 = time.time()
472
+ result_img = tgt_bgr.copy()
473
+ for src_face, _ in pairs:
474
+ if face_analysis_app is None:
475
+ raise ValueError("Face analysis models not initialized.")
476
+ current_faces = sorted(face_analysis_app.get(result_img), key=face_sort_key)
477
+ candidates = [f for f in current_faces if f.gender == src_face.gender] or current_faces
478
+ target_face = candidates[0]
479
+
480
+ if swapper is None:
481
+ raise ValueError("Face swap models not initialized.")
482
+ result_img = swapper.get(result_img, target_face, src_face, paste_back=True)
483
+ logger.info(f"[Pipeline] Multi-face swap ({len(pairs)} pairs): {time.time()-t0:.2f}s")
484
+
485
+ logger.info(f"[Pipeline] TOTAL multi_face_swap: {time.time()-pipeline_start:.2f}s")
486
+ return cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
487
+
488
+
489
+
490
+ def face_swap_and_enhance(src_img, tgt_img, temp_dir=None):
491
+ try:
492
+ with swap_lock:
493
+ pipeline_start = time.time()
494
+
495
+ if temp_dir is None:
496
+ temp_dir = os.path.join(tempfile.gettempdir(), f"faceswap_work_{uuid.uuid4().hex[:8]}")
497
+ if os.path.exists(temp_dir):
498
+ shutil.rmtree(temp_dir)
499
+ os.makedirs(temp_dir, exist_ok=True)
500
+
501
+ if face_analysis_app is None:
502
+ return None, None, "❌ Face analysis models not initialized."
503
+ if swapper is None:
504
+ return None, None, "❌ Face swap models not initialized."
505
+
506
+ src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
507
+ tgt_bgr = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
508
+
509
+ t0 = time.time()
510
+ src_faces = face_analysis_app.get(src_bgr)
511
+ tgt_faces = face_analysis_app.get(tgt_bgr)
512
+ logger.info(f"[Pipeline] Face detection: {time.time()-t0:.2f}s")
513
+
514
+ if not src_faces or not tgt_faces:
515
+ return None, None, "❌ Face not detected in one of the images"
516
+
517
+ t0 = time.time()
518
+ swapped_bgr = swapper.get(tgt_bgr, tgt_faces[0], src_faces[0])
519
+ logger.info(f"[Pipeline] Face swap: {time.time()-t0:.2f}s")
520
+
521
+ if swapped_bgr is None:
522
+ return None, None, "❌ Face swap failed"
523
+
524
+ # Use in-process CodeFormer enhancement (fast path)
525
+ t0 = time.time()
526
+ swapped_rgb = cv2.cvtColor(swapped_bgr, cv2.COLOR_BGR2RGB)
527
+ try:
528
+ enhanced_rgb = enhance_image_with_codeformer(swapped_rgb)
529
+ enhanced_bgr = cv2.cvtColor(enhanced_rgb, cv2.COLOR_RGB2BGR)
530
+ except Exception as e:
531
+ logger.error(f"[Pipeline] CodeFormer failed, using raw swap: {e}")
532
+ enhanced_bgr = swapped_bgr
533
+ logger.info(f"[Pipeline] Enhancement: {time.time()-t0:.2f}s")
534
+
535
+ final_path = os.path.join(temp_dir, f"result_{uuid.uuid4().hex[:8]}.png")
536
+ cv2.imwrite(final_path, enhanced_bgr)
537
+
538
+ final_img = cv2.cvtColor(enhanced_bgr, cv2.COLOR_BGR2RGB)
539
+
540
+ logger.info(f"[Pipeline] TOTAL face_swap_and_enhance: {time.time()-pipeline_start:.2f}s")
541
+ return final_img, final_path, ""
542
+
543
+ except Exception as e:
544
+ return None, None, f"❌ Error: {str(e)}"
545
+
546
+ def compress_image(
547
+ image_bytes: bytes,
548
+ max_size=(1280, 1280), # max width/height
549
+ quality=75 # JPEG quality (60–80 is ideal)
550
+ ) -> bytes:
551
+ """
552
+ Compress image by resizing and lowering quality.
553
+ Returns compressed image bytes.
554
+ """
555
+ img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
556
+
557
+ # Resize while maintaining aspect ratio
558
+ img.thumbnail(max_size, Image.LANCZOS)
559
+
560
+ output = io.BytesIO()
561
+ img.save(
562
+ output,
563
+ format="JPEG",
564
+ quality=quality,
565
+ optimize=True,
566
+ progressive=True
567
+ )
568
+
569
+ return output.getvalue()
570
+
571
+ # --------------------- DigitalOcean Spaces Helper ---------------------
572
+ def get_spaces_client():
573
+ session = boto3.session.Session()
574
+ client = session.client(
575
+ 's3',
576
+ region_name=DO_SPACES_REGION,
577
+ endpoint_url=DO_SPACES_ENDPOINT,
578
+ aws_access_key_id=DO_SPACES_KEY,
579
+ aws_secret_access_key=DO_SPACES_SECRET,
580
+ config=Config(signature_version='s3v4')
581
+ )
582
+ return client
583
+
584
+ def upload_to_spaces(file_bytes, key, content_type="image/png"):
585
+ client = get_spaces_client()
586
+ client.put_object(Bucket=DO_SPACES_BUCKET, Key=key, Body=file_bytes, ContentType=content_type, ACL='public-read')
587
+ return f"{DO_SPACES_ENDPOINT}/{DO_SPACES_BUCKET}/{key}"
588
+
589
+ def download_from_spaces(key):
590
+ client = get_spaces_client()
591
+ obj = client.get_object(Bucket=DO_SPACES_BUCKET, Key=key)
592
+ return obj['Body'].read()
593
+
594
+ def mandatory_enhancement(rgb_img):
595
+ """
596
+ Always runs CodeFormer on the final image.
597
+ Fail-safe: returns original if enhancement fails.
598
+ """
599
+ try:
600
+ return enhance_image_with_codeformer(rgb_img)
601
+ except Exception as e:
602
+ logger.error(f"CodeFormer failed, returning original: {e}")
603
+ return rgb_img
604
+
605
+ # --------------------- API Endpoints ---------------------
606
+ @fastapi_app.get("/")
607
+ async def root():
608
+ """Root endpoint"""
609
+ return {
610
+ "success": True,
611
+ "message": "FaceSwap API",
612
+ "data": {
613
+ "version": "1.0.0",
614
+ "Product Name":"Beauty Camera - GlowCam AI Studio",
615
+ "Released By" : "LogicGo Infotech"
616
+ }
617
+ }
618
+ @fastapi_app.get("/health")
619
+ async def health():
620
+ return {"status": "healthy"}
621
+
622
+ @fastapi_app.get("/test-admin-db")
623
+ async def test_admin_db():
624
+ try:
625
+ doc = await admin_db.list_collection_names()
626
+ return {"ok": True, "collections": doc}
627
+ except Exception as e:
628
+ return {"ok": False, "error": str(e), "url": ADMIN_MONGO_URL}
629
+
630
+ @fastapi_app.post("/face-swap", dependencies=[Depends(verify_token)])
631
+ async def face_swap_api(
632
+ source: UploadFile = File(...),
633
+ image2: Optional[UploadFile] = File(None),
634
+ target_category_id: str = Form(None),
635
+ new_category_id: str = Form(None),
636
+ user_id: Optional[str] = Form(None),
637
+ appname: Optional[str] = Form(None),
638
+ credentials: HTTPAuthorizationCredentials = Security(security)
639
+ ):
640
+ start_time = datetime.utcnow()
641
+
642
+ try:
643
+ # ------------------------------------------------------------------
644
+ # VALIDATION
645
+ # ------------------------------------------------------------------
646
+ # --------------------------------------------------------------
647
+ # BACKWARD COMPATIBILITY FOR OLD ANDROID VERSIONS
648
+ # --------------------------------------------------------------
649
+ if target_category_id == "":
650
+ target_category_id = None
651
+
652
+ if new_category_id == "":
653
+ new_category_id = None
654
+
655
+ if user_id == "":
656
+ user_id = None
657
+
658
+ subcategories_collection = get_app_db_collections(appname)
659
+
660
+ logger.info(f"[FaceSwap] Incoming request → target_category_id={target_category_id}, new_category_id={new_category_id}, user_id={user_id}")
661
+
662
+ if target_category_id and new_category_id:
663
+ raise HTTPException(400, "Provide only one of new_category_id or target_category_id.")
664
+
665
+ if not target_category_id and not new_category_id:
666
+ raise HTTPException(400, "Either new_category_id or target_category_id is required.")
667
+
668
+ # ------------------------------------------------------------------
669
+ # READ SOURCE IMAGE
670
+ # ------------------------------------------------------------------
671
+ src_bytes = await source.read()
672
+ src_key = f"faceswap/source/{uuid.uuid4().hex}_{source.filename}"
673
+ upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
674
+
675
+ # ------------------------------------------------------------------
676
+ # CASE 1 : new_category_id → MongoDB lookup
677
+ # ------------------------------------------------------------------
678
+ if new_category_id:
679
+
680
+ # doc = await subcategories_col.find_one({
681
+ # "asset_images._id": ObjectId(new_category_id)
682
+ # })
683
+ doc = await subcategories_collection.find_one({
684
+ "asset_images._id": ObjectId(new_category_id)
685
+ })
686
+
687
+
688
+ if not doc:
689
+ raise HTTPException(404, "Asset image not found in database")
690
+
691
+ # extract correct asset
692
+ asset = next(
693
+ (img for img in doc["asset_images"] if str(img["_id"]) == new_category_id),
694
+ None
695
+ )
696
+
697
+ if not asset:
698
+ raise HTTPException(404, "Asset image URL not found")
699
+
700
+ # correct URL
701
+ target_url = asset["url"]
702
+
703
+ # correct categoryId (ObjectId)
704
+ #category_oid = doc["categoryId"] # <-- DO NOT CONVERT TO STRING
705
+ subcategory_oid = doc["_id"]
706
+
707
+ # # ------------------------------------------------------------------
708
+ # # CASE 2 : target_category_id → DigitalOcean path (unchanged logic)
709
+ # # ------------------------------------------------------------------
710
+ if target_category_id:
711
+ client = get_spaces_client()
712
+ base_prefix = "faceswap/target/"
713
+ resp = client.list_objects_v2(
714
+ Bucket=DO_SPACES_BUCKET, Prefix=base_prefix, Delimiter="/"
715
+ )
716
+
717
+ # Extract categories from the CommonPrefixes
718
+ categories = [p["Prefix"].split("/")[2] for p in resp.get("CommonPrefixes", [])]
719
+
720
+ target_url = None
721
+
722
+ # --- FIX STARTS HERE ---
723
+ for category in categories:
724
+ original_prefix = f"faceswap/target/{category}/original/"
725
+ thumb_prefix = f"faceswap/target/{category}/thumb/" # Keep for file list check (optional but safe)
726
+
727
+ # List objects in original/
728
+ original_objects = client.list_objects_v2(
729
+ Bucket=DO_SPACES_BUCKET, Prefix=original_prefix
730
+ ).get("Contents", [])
731
+
732
+ # List objects in thumb/ (optional: for the old code's extra check)
733
+ thumb_objects = client.list_objects_v2(
734
+ Bucket=DO_SPACES_BUCKET, Prefix=thumb_prefix
735
+ ).get("Contents", [])
736
+
737
+ # Extract only the filenames and filter for .png
738
+ original_filenames = sorted([
739
+ obj["Key"].split("/")[-1] for obj in original_objects
740
+ if obj["Key"].split("/")[-1].endswith(".png")
741
+ ])
742
+ thumb_filenames = [
743
+ obj["Key"].split("/")[-1] for obj in thumb_objects
744
+ ]
745
+
746
+ # Replicate the old indexing logic based on sorted filenames
747
+ for idx, filename in enumerate(original_filenames, start=1):
748
+ cid = f"{category.lower()}image_{idx}"
749
+
750
+ # Optional: Replicate the thumb file check for 100% parity
751
+ # if filename in thumb_filenames and cid == target_category_id:
752
+ # Simpler check just on the ID, assuming thumb files are present
753
+ if cid == target_category_id:
754
+ # Construct the final target URL using the full prefix and the filename
755
+ target_url = f"{DO_SPACES_ENDPOINT}/{DO_SPACES_BUCKET}/{original_prefix}{filename}"
756
+ break
757
+
758
+ if target_url:
759
+ break
760
+ # --- FIX ENDS HERE ---
761
+
762
+ if not target_url:
763
+ raise HTTPException(404, "Target categoryId not found")
764
+ # # ------------------------------------------------------------------
765
+ # # DOWNLOAD TARGET IMAGE
766
+ # # ------------------------------------------------------------------
767
+ async with httpx.AsyncClient(timeout=30.0) as client:
768
+ response = await client.get(target_url)
769
+ response.raise_for_status()
770
+ tgt_bytes = response.content
771
+
772
+ src_bgr = cv2.imdecode(np.frombuffer(src_bytes, np.uint8), cv2.IMREAD_COLOR)
773
+ tgt_bgr = cv2.imdecode(np.frombuffer(tgt_bytes, np.uint8), cv2.IMREAD_COLOR)
774
+
775
+ if src_bgr is None or tgt_bgr is None:
776
+ raise HTTPException(400, "Invalid image data")
777
+
778
+ src_rgb = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2RGB)
779
+ tgt_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
780
+
781
+ # ------------------------------------------------------------------
782
+ # READ OPTIONAL IMAGE2
783
+ # ------------------------------------------------------------------
784
+ img2_rgb = None
785
+ if image2:
786
+ img2_bytes = await image2.read()
787
+ img2_bgr = cv2.imdecode(np.frombuffer(img2_bytes, np.uint8), cv2.IMREAD_COLOR)
788
+ if img2_bgr is not None:
789
+ img2_rgb = cv2.cvtColor(img2_bgr, cv2.COLOR_BGR2RGB)
790
+
791
+ # ------------------------------------------------------------------
792
+ # FACE SWAP EXECUTION (run in thread to not block event loop)
793
+ # ------------------------------------------------------------------
794
+ if img2_rgb is not None:
795
+ def _couple_swap():
796
+ pipeline_start = time.time()
797
+ src_images = [src_rgb, img2_rgb]
798
+
799
+ all_src_faces = []
800
+ t0 = time.time()
801
+ for img in src_images:
802
+ faces = face_analysis_app.get(cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
803
+ all_src_faces.extend(faces)
804
+
805
+ tgt_faces = face_analysis_app.get(cv2.cvtColor(tgt_rgb, cv2.COLOR_RGB2BGR))
806
+ logger.info(f"[Pipeline] Couple face detection: {time.time()-t0:.2f}s")
807
+
808
+ if not all_src_faces:
809
+ raise ValueError("No faces detected in source images")
810
+ if not tgt_faces:
811
+ raise ValueError("No faces detected in target image")
812
+
813
+ def face_sort_key(face):
814
+ x1, y1, x2, y2 = face.bbox
815
+ area = (x2 - x1) * (y2 - y1)
816
+ cx = (x1 + x2) / 2
817
+ return (-area, cx)
818
+
819
+ src_male = sorted([f for f in all_src_faces if f.gender == 1], key=face_sort_key)
820
+ src_female = sorted([f for f in all_src_faces if f.gender == 0], key=face_sort_key)
821
+ tgt_male = sorted([f for f in tgt_faces if f.gender == 1], key=face_sort_key)
822
+ tgt_female = sorted([f for f in tgt_faces if f.gender == 0], key=face_sort_key)
823
+
824
+ pairs = []
825
+ for s, t in zip(src_male, tgt_male):
826
+ pairs.append((s, t))
827
+ for s, t in zip(src_female, tgt_female):
828
+ pairs.append((s, t))
829
+
830
+ if not pairs:
831
+ src_all = sorted(all_src_faces, key=face_sort_key)
832
+ tgt_all = sorted(tgt_faces, key=face_sort_key)
833
+ pairs = list(zip(src_all, tgt_all))
834
+
835
+ t0 = time.time()
836
+ with swap_lock:
837
+ result_img = cv2.cvtColor(tgt_rgb, cv2.COLOR_RGB2BGR)
838
+ for src_face, _ in pairs:
839
+ current_faces = sorted(face_analysis_app.get(result_img), key=face_sort_key)
840
+ candidates = [f for f in current_faces if f.gender == src_face.gender] or current_faces
841
+ target_face = candidates[0]
842
+ result_img = swapper.get(result_img, target_face, src_face, paste_back=True)
843
+ logger.info(f"[Pipeline] Couple face swap: {time.time()-t0:.2f}s")
844
+
845
+ result_rgb_out = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
846
+
847
+ t0 = time.time()
848
+ enhanced_rgb = mandatory_enhancement(result_rgb_out)
849
+ logger.info(f"[Pipeline] Couple enhancement: {time.time()-t0:.2f}s")
850
+
851
+ enhanced_bgr = cv2.cvtColor(enhanced_rgb, cv2.COLOR_RGB2BGR)
852
+
853
+ temp_dir = tempfile.mkdtemp(prefix="faceswap_")
854
+ final_path = os.path.join(temp_dir, "result.png")
855
+ cv2.imwrite(final_path, enhanced_bgr)
856
+
857
+ with open(final_path, "rb") as f:
858
+ result_bytes = f.read()
859
+
860
+ logger.info(f"[Pipeline] TOTAL couple swap: {time.time()-pipeline_start:.2f}s")
861
+ return result_bytes
862
+
863
+ try:
864
+ result_bytes = await asyncio.to_thread(_couple_swap)
865
+ except ValueError as ve:
866
+ raise HTTPException(400, str(ve))
867
+
868
+ else:
869
+ # ----- SINGLE SOURCE SWAP (run in thread) -----
870
+ def _single_swap():
871
+ return face_swap_and_enhance(src_rgb, tgt_rgb)
872
+
873
+ final_img, final_path, err = await asyncio.to_thread(_single_swap)
874
+
875
+ if err:
876
+ raise HTTPException(500, err)
877
+
878
+ with open(final_path, "rb") as f:
879
+ result_bytes = f.read()
880
+
881
+ result_key = f"faceswap/result/{uuid.uuid4().hex}_enhanced.png"
882
+ result_url = upload_to_spaces(result_bytes, result_key)
883
+ # -------------------------------------------------
884
+ # COMPRESS IMAGE (2–3 MB target)
885
+ # -------------------------------------------------
886
+ compressed_bytes = compress_image(
887
+ image_bytes=result_bytes,
888
+ max_size=(1280, 1280),
889
+ quality=72
890
+ )
891
+
892
+ compressed_key = f"faceswap/result/{uuid.uuid4().hex}_enhanced_compressed.jpg"
893
+ compressed_url = upload_to_spaces(
894
+ compressed_bytes,
895
+ compressed_key,
896
+ content_type="image/jpeg"
897
+ )
898
+ end_time = datetime.utcnow()
899
+ response_time_ms = (end_time - start_time).total_seconds() * 1000
900
+
901
+ if database is not None:
902
+ log_entry = {
903
+ "endpoint": "/face-swap",
904
+ "status": "success",
905
+ "response_time_ms": response_time_ms,
906
+ "timestamp": end_time
907
+ }
908
+ if appname:
909
+ log_entry["appname"] = appname
910
+ await database.faceswap.insert_one(log_entry)
911
+
912
+
913
+ return {
914
+ "result_key": result_key,
915
+ "result_url": result_url,
916
+ "Compressed_Image_URL": compressed_url
917
+ }
918
+
919
+ except Exception as e:
920
+ end_time = datetime.utcnow()
921
+ response_time_ms = (end_time - start_time).total_seconds() * 1000
922
+
923
+ if database is not None:
924
+ log_entry = {
925
+ "endpoint": "/face-swap",
926
+ "status": "fail",
927
+ "response_time_ms": response_time_ms,
928
+ "timestamp": end_time,
929
+ "error": str(e)
930
+ }
931
+ if appname:
932
+ log_entry["appname"] = appname
933
+ await database.faceswap.insert_one(log_entry)
934
+
935
+ raise HTTPException(500, f"Face swap failed: {str(e)}")
936
+
937
+ @fastapi_app.get("/preview/{result_key:path}")
938
+ async def preview_result(result_key: str):
939
+ try:
940
+ img_bytes = download_from_spaces(result_key)
941
+ except Exception:
942
+ raise HTTPException(status_code=404, detail="Result not found")
943
+ return Response(
944
+ content=img_bytes,
945
+ media_type="image/png",
946
+ headers={"Content-Disposition": "inline; filename=result.png"}
947
+ )
948
+
949
+ @fastapi_app.post("/multi-face-swap", dependencies=[Depends(verify_token)])
950
+ async def multi_face_swap_api(
951
+ source_image: UploadFile = File(...),
952
+ target_image: UploadFile = File(...)
953
+ ):
954
+ start_time = datetime.utcnow()
955
+
956
+ try:
957
+ # -----------------------------
958
+ # Read images
959
+ # -----------------------------
960
+ src_bytes = await source_image.read()
961
+ tgt_bytes = await target_image.read()
962
+
963
+ src_bgr = cv2.imdecode(np.frombuffer(src_bytes, np.uint8), cv2.IMREAD_COLOR)
964
+ tgt_bgr = cv2.imdecode(np.frombuffer(tgt_bytes, np.uint8), cv2.IMREAD_COLOR)
965
+
966
+ if src_bgr is None or tgt_bgr is None:
967
+ raise HTTPException(400, "Invalid image data")
968
+
969
+ src_rgb = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2RGB)
970
+ tgt_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
971
+
972
+ # -----------------------------
973
+ # Multi-face swap (run in thread to not block event loop)
974
+ # -----------------------------
975
+ def _multi_swap_and_enhance():
976
+ swapped_rgb = multi_face_swap(src_rgb, tgt_rgb)
977
+ return mandatory_enhancement(swapped_rgb)
978
+
979
+ final_rgb = await asyncio.to_thread(_multi_swap_and_enhance)
980
+
981
+ final_bgr = cv2.cvtColor(final_rgb, cv2.COLOR_RGB2BGR)
982
+
983
+ # -----------------------------
984
+ # Save temp result
985
+ # -----------------------------
986
+ temp_dir = tempfile.mkdtemp(prefix="multi_faceswap_")
987
+ result_path = os.path.join(temp_dir, "result.png")
988
+ cv2.imwrite(result_path, final_bgr)
989
+
990
+ with open(result_path, "rb") as f:
991
+ result_bytes = f.read()
992
+
993
+ # -----------------------------
994
+ # Upload
995
+ # -----------------------------
996
+ result_key = f"faceswap/multi/{uuid.uuid4().hex}.png"
997
+ result_url = upload_to_spaces(
998
+ result_bytes,
999
+ result_key,
1000
+ content_type="image/png"
1001
+ )
1002
+
1003
+ return {
1004
+ "result_key": result_key,
1005
+ "result_url": result_url
1006
+ }
1007
+
1008
+ except Exception as e:
1009
+ raise HTTPException(status_code=500, detail=str(e))
1010
+
1011
+
1012
+ @fastapi_app.post("/face-swap-couple", dependencies=[Depends(verify_token)])
1013
+ async def face_swap_couple_api(
1014
+ image1: UploadFile = File(...),
1015
+ image2: Optional[UploadFile] = File(None),
1016
+ target_category_id: str = Form(None),
1017
+ new_category_id: str = Form(None),
1018
+ user_id: Optional[str] = Form(None),
1019
+ appname: Optional[str] = Form(None),
1020
+ credentials: HTTPAuthorizationCredentials = Security(security)
1021
+ ):
1022
+ """
1023
+ Production-ready face swap endpoint supporting:
1024
+ - Multiple source images (image1 + optional image2)
1025
+ - Gender-based pairing
1026
+ - Merged faces from multiple sources
1027
+ - Mandatory CodeFormer enhancement
1028
+ """
1029
+ start_time = datetime.utcnow()
1030
+
1031
+ try:
1032
+ # -----------------------------
1033
+ # Validate input
1034
+ # -----------------------------
1035
+ if target_category_id == "":
1036
+ target_category_id = None
1037
+ if new_category_id == "":
1038
+ new_category_id = None
1039
+ if user_id == "":
1040
+ user_id = None
1041
+
1042
+ subcategories_collection = get_app_db_collections(appname)
1043
+
1044
+ if target_category_id and new_category_id:
1045
+ raise HTTPException(400, "Provide only one of new_category_id or target_category_id.")
1046
+ if not target_category_id and not new_category_id:
1047
+ raise HTTPException(400, "Either new_category_id or target_category_id is required.")
1048
+
1049
+ logger.info(f"[FaceSwap] Incoming request → target_category_id={target_category_id}, new_category_id={new_category_id}, user_id={user_id}")
1050
+
1051
+ # -----------------------------
1052
+ # Read source images
1053
+ # -----------------------------
1054
+ src_images = []
1055
+ img1_bytes = await image1.read()
1056
+ src1 = cv2.imdecode(np.frombuffer(img1_bytes, np.uint8), cv2.IMREAD_COLOR)
1057
+ if src1 is None:
1058
+ raise HTTPException(400, "Invalid image1 data")
1059
+ src_images.append(cv2.cvtColor(src1, cv2.COLOR_BGR2RGB))
1060
+
1061
+ if image2:
1062
+ img2_bytes = await image2.read()
1063
+ src2 = cv2.imdecode(np.frombuffer(img2_bytes, np.uint8), cv2.IMREAD_COLOR)
1064
+ if src2 is not None:
1065
+ src_images.append(cv2.cvtColor(src2, cv2.COLOR_BGR2RGB))
1066
+
1067
+ # -----------------------------
1068
+ # Resolve target image
1069
+ # -----------------------------
1070
+ target_url = None
1071
+ if new_category_id:
1072
+ doc = await subcategories_collection.find_one({
1073
+ "asset_images._id": ObjectId(new_category_id)
1074
+ })
1075
+
1076
+ if not doc:
1077
+ raise HTTPException(404, "Asset image not found in database")
1078
+
1079
+ asset = next(
1080
+ (img for img in doc["asset_images"] if str(img["_id"]) == new_category_id),
1081
+ None
1082
+ )
1083
+
1084
+ if not asset:
1085
+ raise HTTPException(404, "Asset image URL not found")
1086
+
1087
+ target_url = asset["url"]
1088
+ subcategory_oid = doc["_id"]
1089
+
1090
+ if target_category_id:
1091
+ client = get_spaces_client()
1092
+ base_prefix = "faceswap/target/"
1093
+ resp = client.list_objects_v2(
1094
+ Bucket=DO_SPACES_BUCKET, Prefix=base_prefix, Delimiter="/"
1095
+ )
1096
+
1097
+ categories = [p["Prefix"].split("/")[2] for p in resp.get("CommonPrefixes", [])]
1098
+
1099
+ for category in categories:
1100
+ original_prefix = f"faceswap/target/{category}/original/"
1101
+ thumb_prefix = f"faceswap/target/{category}/thumb/"
1102
+
1103
+ original_objects = client.list_objects_v2(
1104
+ Bucket=DO_SPACES_BUCKET, Prefix=original_prefix
1105
+ ).get("Contents", [])
1106
+
1107
+ thumb_objects = client.list_objects_v2(
1108
+ Bucket=DO_SPACES_BUCKET, Prefix=thumb_prefix
1109
+ ).get("Contents", [])
1110
+
1111
+ original_filenames = sorted([
1112
+ obj["Key"].split("/")[-1] for obj in original_objects
1113
+ if obj["Key"].split("/")[-1].endswith(".png")
1114
+ ])
1115
+
1116
+ for idx, filename in enumerate(original_filenames, start=1):
1117
+ cid = f"{category.lower()}image_{idx}"
1118
+ if cid == target_category_id:
1119
+ target_url = f"{DO_SPACES_ENDPOINT}/{DO_SPACES_BUCKET}/{original_prefix}{filename}"
1120
+ break
1121
+
1122
+ if target_url:
1123
+ break
1124
+
1125
+ if not target_url:
1126
+ raise HTTPException(404, "Target categoryId not found")
1127
+
1128
+ async with httpx.AsyncClient(timeout=30.0) as client:
1129
+ response = await client.get(target_url)
1130
+ response.raise_for_status()
1131
+ tgt_bytes = response.content
1132
+
1133
+ tgt_bgr = cv2.imdecode(np.frombuffer(tgt_bytes, np.uint8), cv2.IMREAD_COLOR)
1134
+ if tgt_bgr is None:
1135
+ raise HTTPException(400, "Invalid target image data")
1136
+
1137
+ # -----------------------------
1138
+ # Couple face swap + enhance (run in thread)
1139
+ # -----------------------------
1140
+ def _couple_face_swap_and_enhance():
1141
+ pipeline_start = time.time()
1142
+
1143
+ all_src_faces = []
1144
+ t0 = time.time()
1145
+ for img in src_images:
1146
+ faces = face_analysis_app.get(cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
1147
+ all_src_faces.extend(faces)
1148
+
1149
+ tgt_faces = face_analysis_app.get(tgt_bgr)
1150
+ logger.info(f"[Pipeline] Couple-ep face detection: {time.time()-t0:.2f}s")
1151
+
1152
+ if not all_src_faces:
1153
+ raise ValueError("No faces detected in source images")
1154
+ if not tgt_faces:
1155
+ raise ValueError("No faces detected in target image")
1156
+
1157
+ def face_sort_key(face):
1158
+ x1, y1, x2, y2 = face.bbox
1159
+ area = (x2 - x1) * (y2 - y1)
1160
+ cx = (x1 + x2) / 2
1161
+ return (-area, cx)
1162
+
1163
+ src_male = sorted([f for f in all_src_faces if f.gender == 1], key=face_sort_key)
1164
+ src_female = sorted([f for f in all_src_faces if f.gender == 0], key=face_sort_key)
1165
+ tgt_male = sorted([f for f in tgt_faces if f.gender == 1], key=face_sort_key)
1166
+ tgt_female = sorted([f for f in tgt_faces if f.gender == 0], key=face_sort_key)
1167
+
1168
+ pairs = []
1169
+ for s, t in zip(src_male, tgt_male):
1170
+ pairs.append((s, t))
1171
+ for s, t in zip(src_female, tgt_female):
1172
+ pairs.append((s, t))
1173
+
1174
+ if not pairs:
1175
+ src_all = sorted(all_src_faces, key=face_sort_key)
1176
+ tgt_all = sorted(tgt_faces, key=face_sort_key)
1177
+ pairs = list(zip(src_all, tgt_all))
1178
+
1179
+ t0 = time.time()
1180
+ with swap_lock:
1181
+ result_img = tgt_bgr.copy()
1182
+ for src_face, _ in pairs:
1183
+ current_faces = sorted(face_analysis_app.get(result_img), key=face_sort_key)
1184
+ candidates = [f for f in current_faces if f.gender == src_face.gender] or current_faces
1185
+ target_face = candidates[0]
1186
+ result_img = swapper.get(result_img, target_face, src_face, paste_back=True)
1187
+ logger.info(f"[Pipeline] Couple-ep face swap: {time.time()-t0:.2f}s")
1188
+
1189
+ result_rgb = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
1190
+
1191
+ t0 = time.time()
1192
+ enhanced_rgb = mandatory_enhancement(result_rgb)
1193
+ logger.info(f"[Pipeline] Couple-ep enhancement: {time.time()-t0:.2f}s")
1194
+
1195
+ enhanced_bgr = cv2.cvtColor(enhanced_rgb, cv2.COLOR_RGB2BGR)
1196
+
1197
+ temp_dir = tempfile.mkdtemp(prefix="faceswap_")
1198
+ final_path = os.path.join(temp_dir, "result.png")
1199
+ cv2.imwrite(final_path, enhanced_bgr)
1200
+
1201
+ with open(final_path, "rb") as f:
1202
+ result_bytes = f.read()
1203
+
1204
+ logger.info(f"[Pipeline] TOTAL couple-ep swap: {time.time()-pipeline_start:.2f}s")
1205
+ return result_bytes
1206
+
1207
+ try:
1208
+ result_bytes = await asyncio.to_thread(_couple_face_swap_and_enhance)
1209
+ except ValueError as ve:
1210
+ raise HTTPException(400, str(ve))
1211
+
1212
+ result_key = f"faceswap/result/{uuid.uuid4().hex}_enhanced.png"
1213
+ result_url = upload_to_spaces(result_bytes, result_key)
1214
+
1215
+ compressed_bytes = compress_image(result_bytes, max_size=(1280, 1280), quality=72)
1216
+ compressed_key = f"faceswap/result/{uuid.uuid4().hex}_enhanced_compressed.jpg"
1217
+ compressed_url = upload_to_spaces(compressed_bytes, compressed_key, content_type="image/jpeg")
1218
+
1219
+ # -----------------------------
1220
+ # Log API usage
1221
+ # -----------------------------
1222
+ end_time = datetime.utcnow()
1223
+ response_time_ms = (end_time - start_time).total_seconds() * 1000
1224
+ if database is not None:
1225
+ log_entry = {
1226
+ "endpoint": "/face-swap-couple",
1227
+ "status": "success",
1228
+ "response_time_ms": response_time_ms,
1229
+ "timestamp": end_time
1230
+ }
1231
+ if appname:
1232
+ log_entry["appname"] = appname
1233
+ await database.faceswap.insert_one(log_entry)
1234
+
1235
+ return {
1236
+ "result_key": result_key,
1237
+ "result_url": result_url,
1238
+ "compressed_url": compressed_url
1239
+ }
1240
+
1241
+ except Exception as e:
1242
+ end_time = datetime.utcnow()
1243
+ response_time_ms = (end_time - start_time).total_seconds() * 1000
1244
+ if database is not None:
1245
+ log_entry = {
1246
+ "endpoint": "/face-swap-couple",
1247
+ "status": "fail",
1248
+ "response_time_ms": response_time_ms,
1249
+ "timestamp": end_time,
1250
+ "error": str(e)
1251
+ }
1252
+ if appname:
1253
+ log_entry["appname"] = appname
1254
+ await database.faceswap.insert_one(log_entry)
1255
+ raise HTTPException(500, f"Face swap failed: {str(e)}")
1256
+
1257
+
1258
+ if __name__ == "__main__":
1259
+ uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)