Update app.py

app.py

@@ -9,6 +9,9 @@ import cv2
 import traceback
 import gc
 import uuid
+import time
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
 from PIL import Image, ImageFilter, ImageEnhance
 from torchvision.transforms import functional as TF
 from scipy.ndimage import label
@@ -23,8 +26,6 @@ from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 import io
-import asyncio
-from concurrent.futures import ThreadPoolExecutor
 import logging
 
 logging.basicConfig(level=logging.INFO)
@@ -37,7 +38,7 @@ AGING_MODEL_PATH = "face_aging_model/best_unet_model.pth"
 BEARD_MODEL_PATH = "models/best_hair_117_epoch_v4.pt"
 GFPGAN_MODEL_PATH = "GFPGANv1.4.pth"
 
-SAFE_IMG_SIZE = 384
+SAFE_IMG_SIZE = 384               # used only for aging model
 SOURCE_AGE = 20
 TARGET_AGE = 80
 WRINKLE_STRENGTH = 0.42
@@ -52,7 +53,6 @@ GFPGAN_WEIGHT = 0.5
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 USE_FP16 = DEVICE.type == "cuda" and torch.cuda.is_available()
-
 logger.info(f"🚀 Device: {DEVICE}, FP16: {USE_FP16}")
 
 os.environ["HF_HOME"] = "/tmp/hf_cache"
@@ -199,7 +199,7 @@ def load_aging_model():
     return age_model
 
 # ================================================
-# 6. LOAD FACE PARSER & BEARD MODEL
+# 6. LOAD FACE PARSER & BEARD MODEL (optimized sizes)
 # ================================================
 def load_face_parser():
     global face_processor, face_parser
@@ -225,7 +225,7 @@ def load_beard_model():
     return beard_model
 
 # ================================================
-# 7. MASK FUNCTIONS (shortened for space - same as before)
+# 7. MASK FUNCTIONS (optimized: 256px inference)
 # ================================================
 def get_lips_mask(pil_image: Image.Image) -> np.ndarray:
     img_np = np.array(pil_image.resize((256, 256), Image.LANCZOS))
@@ -254,9 +254,10 @@ def exclude_lips_from_mask(beard_mask, pil_image):
 def get_beard_mask(pil_image: Image.Image) -> np.ndarray:
     temp = f"temp_{uuid.uuid4().hex[:8]}.jpg"
     try:
-        pil_image.resize((384,384), Image.LANCZOS).save(temp)
+        # OPTIMIZATION: use 256px instead of 384
+        pil_image.resize((256, 256), Image.LANCZOS).save(temp)
         model = load_beard_model()
-        res = model(temp, device=DEVICE.type, conf=0.3, iou=0.5, verbose=False, half=USE_FP16, imgsz=384)
+        res = model(temp, device=DEVICE.type, conf=0.3, iou=0.5, verbose=False, half=USE_FP16, imgsz=256)
         h, w = np.array(pil_image).shape[:2]
         mask = np.zeros((h,w), dtype=np.uint8)
         if res[0].masks is not None:
@@ -287,7 +288,8 @@ def clean_mask(mask, min_area=100):
 
 def get_hair_mask_segformer(pil_image: Image.Image) -> np.ndarray:
     processor, parser = load_face_parser()
-    img_r = pil_image.resize((384,384), Image.LANCZOS)
+    # OPTIMIZATION: 256px instead of 384
+    img_r = pil_image.resize((256, 256), Image.LANCZOS)
     inputs = processor(images=img_r, return_tensors="pt").to(DEVICE)
     if USE_FP16: inputs['pixel_values'] = inputs['pixel_values'].half()
     with torch.no_grad():
@@ -345,42 +347,83 @@ def process_masks_parallel(image):
         return h.result(), b.result()
 
 # ================================================
-# 8. MAIN PROCESSING (FIXED tensor dimension)
+# 8. FAST FALLBACK (when full pipeline times out)
+# ================================================
+def fast_aging_fallback(input_image: Image.Image) -> Image.Image:
+    """Very fast aging: no GFPGAN, no beard mask, simplified hair mask."""
+    logger.info("⚡ Using fast fallback aging")
+    orig = input_image.convert("RGB")
+    ow, oh = orig.size
+    target_size = min(256, max(ow, oh))
+    if target_size % 2 == 0:
+        target_size -= 1
+
+    img_resized = orig.resize((target_size, target_size), Image.LANCZOS)
+    rgb_tensor = TF.to_tensor(img_resized)
+    src_age = torch.full((1, target_size, target_size), SOURCE_AGE / 100.0)
+    tgt_age = torch.full((1, target_size, target_size), TARGET_AGE / 100.0)
+    cond_input = torch.cat([
+        rgb_tensor.unsqueeze(0),
+        src_age.unsqueeze(1),
+        tgt_age.unsqueeze(1)
+    ], dim=1).to(DEVICE)
+    if USE_FP16:
+        cond_input = cond_input.half()
+
+    with torch.no_grad():
+        raw_output = load_aging_model()(cond_input).squeeze(0)
+
+    alpha = WRINKLE_STRENGTH
+    blended = (1 - alpha) * rgb_tensor + alpha * raw_output
+    blended = blended.clamp(0, 1).float() if USE_FP16 else blended
+
+    aged = TF.to_pil_image(blended).resize((ow, oh), Image.LANCZOS)
+    aged = enhance_texture(aged)
+    aged = post_correct_aged(orig, aged)
+
+    # Simple luminance-based hair whitening (no segmentation)
+    gray = np.array(aged.convert('L')) / 255.0
+    hair_mask = (gray > 0.65).astype(np.float32)
+    hair_mask = cv2.GaussianBlur(hair_mask, (9,9), 3)
+    beard_mask = np.zeros_like(hair_mask)
+
+    final = apply_hair_and_beard_color(aged, hair_mask, beard_mask)
+    return final
+
+# ================================================
+# 9. MAIN PROCESSING (with optional time check)
 # ================================================
 def process_face_aging(input_image: Image.Image) -> Image.Image:
+    start_time = time.time()
     try:
         logger.info(f"→ Processing image: {input_image.size}")
         orig = input_image.convert("RGB")
         ow, oh = orig.size
 
+        # Aging model step
         target_size = min(SAFE_IMG_SIZE, max(ow, oh))
         if target_size % 2 == 1:
             target_size -= 1
 
         img_resized = orig.resize((target_size, target_size), Image.LANCZOS)
-        rgb_tensor = TF.to_tensor(img_resized)                    # shape: (3, H, W)
-
+        rgb_tensor = TF.to_tensor(img_resized)
         src_age = torch.full((1, target_size, target_size), SOURCE_AGE / 100.0)
         tgt_age = torch.full((1, target_size, target_size), TARGET_AGE / 100.0)
-
-        # FIXED: Make all 4D tensors
         cond_input = torch.cat([
-            rgb_tensor.unsqueeze(0),          # (1, 3, H, W)
-            src_age.unsqueeze(1),             # (1, 1, H, W)
-            tgt_age.unsqueeze(1)              # (1, 1, H, W)
-        ], dim=1).to(DEVICE)                  # Final: (1, 5, H, W)
-
+            rgb_tensor.unsqueeze(0),
+            src_age.unsqueeze(1),
+            tgt_age.unsqueeze(1)
+        ], dim=1).to(DEVICE)
         if USE_FP16:
             cond_input = cond_input.half()
 
         with torch.no_grad():
             aging_net = load_aging_model()
-            raw_output = aging_net(cond_input).squeeze(0)   # (3, H, W)
+            raw_output = aging_net(cond_input).squeeze(0)
 
         alpha = WRINKLE_STRENGTH
         blended = (1 - alpha) * rgb_tensor + alpha * raw_output
         blended = blended.clamp(0, 1)
-
         if USE_FP16:
             blended = blended.float()
 
@@ -388,14 +431,17 @@ def process_face_aging(input_image: Image.Image) -> Image.Image:
         final_aged = enhance_texture(final_aged)
         final_aged = post_correct_aged(orig, final_aged)
 
-        logger.info(" Generating masks...")
+        # Masks (parallel)
+        logger.info("🔄 Generating masks...")
         hair_mask, beard_mask = process_masks_parallel(final_aged)
 
-        logger.info(" Applying white hair & beard...")
+        logger.info("🎨 Applying white hair & beard...")
         final_img = apply_hair_and_beard_color(final_aged, hair_mask, beard_mask)
 
-        if USE_GFPGAN and (ow * oh) < 2000000:
-            logger.info(" Applying GFPGAN...")
+        # GFPGAN only if image is not too large and we have time left (> 2 sec)
+        elapsed = time.time() - start_time
+        if USE_GFPGAN and (ow * oh) < 1000000 and elapsed < 7.0:
+            logger.info("✨ Applying GFPGAN...")
             gfpgan = load_gfpgan()
             if gfpgan:
                 try:
@@ -406,8 +452,10 @@ def process_face_aging(input_image: Image.Image) -> Image.Image:
                     final_img = Image.fromarray(cv2.cvtColor(restored, cv2.COLOR_BGR2RGB))
                 except Exception as e:
                     logger.warning(f"GFPGAN skipped: {e}")
+        else:
+            logger.info("⏭️ Skipping GFPGAN (image too large or time low)")
 
-        logger.info("✅ Processing completed!")
+        logger.info(f"✅ Processing completed in {time.time()-start_time:.2f}s")
         gc.collect()
         return final_img
 
@@ -417,12 +465,17 @@ def process_face_aging(input_image: Image.Image) -> Image.Image:
         raise HTTPException(status_code=500, detail=f"Processing failed: {str(e)}")
 
 # ================================================
-# 9. FASTAPI
+# 10. FASTAPI WITH TIMEOUT
 # ================================================
 app = FastAPI(title="Face Aging API")
 
-app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True,
-                   allow_methods=["*"], allow_headers=["*"])
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 @app.on_event("startup")
 async def startup_event():
@@ -440,20 +493,23 @@ async def age_face(file: UploadFile = File(...)):
     if not file.content_type.startswith("image/"):
         raise HTTPException(400, "Only image files allowed")
     contents = await file.read()
+    input_image = Image.open(io.BytesIO(contents)).convert("RGB")
+
+    loop = asyncio.get_event_loop()
     try:
-        input_image = Image.open(io.BytesIO(contents)).convert("RGB")
-        loop = asyncio.get_event_loop()
-        result = await loop.run_in_executor(executor, process_face_aging, input_image)
-
-        buf = io.BytesIO()
-        result.save(buf, format="JPEG", quality=90, optimize=True)
-        buf.seek(0)
-        return StreamingResponse(buf, media_type="image/jpeg")
-    except Exception as e:
-        logger.error(f"Endpoint error: {e}")
-        raise HTTPException(500, f"Failed: {str(e)}")
-    finally:
-        gc.collect()
+        # 9.5 second timeout for full pipeline
+        result = await asyncio.wait_for(
+            loop.run_in_executor(executor, process_face_aging, input_image),
+            timeout=9.5
+        )
+    except asyncio.TimeoutError:
+        logger.warning("⏱️ Full processing timeout – using fast fallback")
+        result = await loop.run_in_executor(executor, fast_aging_fallback, input_image)
+
+    buf = io.BytesIO()
+    result.save(buf, format="JPEG", quality=90, optimize=True)
+    buf.seek(0)
+    return StreamingResponse(buf, media_type="image/jpeg")
 
 if __name__ == "__main__":
     import uvicorn