Update app.py

app.py

@@ -1,162 +1,159 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException
-from fastapi.responses import Response
-from fastapi.middleware.cors import CORSMiddleware
-import cv2
+# -*- coding:UTF-8 -*-
+#!/usr/bin/env python
 import numpy as np
-import face_recognition
+import roop.globals
+from roop.core import start, decode_execution_providers
+from roop.processors.frame.core import get_frame_processors_modules
+from roop.utilities import normalize_output_path
+from roop.face_analyser import get_many_faces
 import os
+import requests
+import cv2
+from pathlib import Path
+from PIL import Image
+from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks
+from fastapi.responses import FileResponse
+from fastapi.middleware.cors import CORSMiddleware
 import shutil
 import logging
+import time
 
 app = FastAPI()
 
-# Enable CORS to allow Framer frontend to connect
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],  # Update with your Framer domain in production
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-
-def get_face_data(image):
-    """Detect face and landmarks in an image."""
-    face_locations = face_recognition.face_locations(image)
-    if len(face_locations) == 0:
-        raise ValueError("No face detected in the image.")
-    if len(face_locations) > 1:
-        raise ValueError("Multiple faces detected; only one face per image is supported.")
-
-    face_location = face_locations[0]  # (top, right, bottom, left)
-    landmarks = face_recognition.face_landmarks(image, face_locations=[face_location])
-    if not landmarks:
-        raise ValueError("Could not detect face landmarks.")
-
-    return face_location, landmarks[0]
-
-def get_face_size(face_location):
-    """Calculate the width and height of the face bounding box."""
-    top, right, bottom, left = face_location
-    width = right - left
-    height = bottom - top
-    return width, height
-
-def resize_face_image(source_img, target_face_size, source_face_location):
-    """Resize the source image to match the target face size."""
-    source_width, source_height = get_face_size(source_face_location)
-    target_width, target_height = target_face_size
-
-    # Calculate scaling factor to match the target face size
-    scale_x = target_width / source_width
-    scale_y = target_height / source_height
-    scale = min(scale_x, scale_y)  # Use the smaller scale to avoid distortion
-
-    # Resize the source image
-    new_width = int(source_img.shape[1] * scale)
-    new_height = int(source_img.shape[0] * scale)
-    resized_source = cv2.resize(source_img, (new_width, new_height), interpolation=cv2.INTER_AREA)
-
-    return resized_source, scale
-
-def swap_faces(source_img, target_img):
-    """Perform face swapping with size preservation and seamless blending."""
-    # Convert images to RGB (face_recognition expects RGB)
-    source_rgb = cv2.cvtColor(source_img, cv2.COLOR_BGR2RGB)
-    target_rgb = cv2.cvtColor(target_img, cv2.COLOR_BGR2RGB)
-
-    # Detect faces and landmarks
-    source_face_location, source_landmarks = get_face_data(source_rgb)
-    target_face_location, target_landmarks = get_face_data(target_rgb)
-
-    # Calculate face sizes
-    target_face_size = get_face_size(target_face_location)
-
-    # Resize source image to match target face size
-    resized_source, scale = resize_face_image(source_img, target_face_size, source_face_location)
-
-    # Adjust source face location after resizing
-    source_top, source_right, source_bottom, source_left = source_face_location
-    adjusted_source_location = (
-        int(source_top * scale),
-        int(source_right * scale),
-        int(source_bottom * scale),
-        int(source_left * scale)
-    )
-
-    # Extract the source face region
-    source_face = resized_source[
-        adjusted_source_location[0]:adjusted_source_location[2],
-        adjusted_source_location[3]:adjusted_source_location[1]
-    ]
-
-    # Calculate the center of the target face
-    target_top, target_right, target_bottom, target_left = target_face_location
-    target_center_x = (target_left + target_right) // 2
-    target_center_y = (target_top + target_bottom) // 2
-
-    # Create a mask for the source face
-    mask = 255 * np.ones(source_face.shape, source_face.dtype)
-
-    # Perform seamless cloning
-    try:
-        result = cv2.seamlessClone(
-            source_face,
-            target_img,
-            mask,
-            (target_center_x, target_center_y),
-            cv2.NORMAL_CLONE
-        )
-    except Exception as e:
-        logging.error(f"Seamless cloning failed: {str(e)}")
-        raise ValueError("Failed to blend the faces seamlessly.")
-
-    return result
+article_text = """
+<div style="text-align: center;">
+    <p>Enjoying the tool? Buy me a coffee and get exclusive prompt guides!</p>
+    <p><i>Instantly unlock helpful tips for creating better prompts!</i></p>
+    <div style="display: flex; justify-content: center;">
+        <a href=\"https://piczify.lemonsqueezy.com/buy/0f5206fa-68e8-42f6-9ca8-4f80c587c83e\">
+            <img src=\"https://www.buymeacoffee.com/assets/img/custom_images/yellow_img.png\" 
+                 alt=\"Buy Me a Coffee\" 
+                 style=\"height: 40px; width: auto; box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2); border-radius: 10px;\">
+        </a>
+    </div>
+</div>
+"""
+
+def download_model():
+    model_dir = Path("models")
+    model_path = model_dir / "inswapper_128.onnx"
+    model_url = "https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx"
+
+    if not model_path.exists():
+        logger.info("Model not found. Downloading inswapper_128.onnx...")
+        model_dir.mkdir(exist_ok=True)
+        try:
+            response = requests.get(model_url, stream=True)
+            response.raise_for_status()
+            with open(model_path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            logger.info("Model downloaded successfully.")
+        except Exception as e:
+            logger.error(f"Failed to download model: {e}")
+            raise RuntimeError("Could not download inswapper_128.onnx. Please download it manually from 'https://huggingface.co/ezioruan/inswapper_128.onnx' and place it in the 'models' folder.")
+    else:
+        logger.info("Model already exists at: %s", model_path)
+
+download_model()
+
+def cleanup_files(paths):
+    time.sleep(2)  # Wait 2 seconds to ensure response is sent
+    for path in paths:
+        if os.path.exists(path):
+            os.remove(path)
+            logger.info("Cleaned up: %s", path)
 
 @app.post("/swap-face/")
-async def swap_face(
-    source_file: UploadFile = File(...),
-    target_file: UploadFile = File(...),
-    doFaceEnhancer: bool = True
-):
+async def swap_face(source_file: UploadFile = File(...), target_file: UploadFile = File(...), doFaceEnhancer: bool = True, background_tasks: BackgroundTasks = BackgroundTasks()):
     try:
         # Save uploaded files temporarily
-        source_path = f"temp_source_{source_file.filename}"
-        target_path = f"temp_target_{target_file.filename}"
+        source_path = "input.jpg"
+        target_path = "target.jpg"
         output_path = "output.jpg"
 
+        # Read and save source image
+        source_content = await source_file.read()
         with open(source_path, "wb") as f:
-            shutil.copyfileobj(source_file.file, f)
+            f.write(source_content)
+        source_image = Image.open(source_path)
+        source_image.save(source_path, quality=100, subsampling=0)
+
+        # Read and save target image, get its size
+        target_content = await target_file.read()
         with open(target_path, "wb") as f:
-            shutil.copyfileobj(target_file.file, f)
+            f.write(target_content)
+        target_image = Image.open(target_path)
+        target_size = target_image.size
+        target_image.save(target_path, quality=100, subsampling=0)
 
-        # Read images
+        # Load images as NumPy arrays for face detection
         source_img = cv2.imread(source_path)
         target_img = cv2.imread(target_path)
-
         if source_img is None or target_img is None:
-            raise ValueError("Failed to load one or both images.")
-
-        # Perform custom face swap
-        result_img = swap_faces(source_img, target_img)
+            raise HTTPException(status_code=400, detail="Failed to load one or both images.")
 
-        # Save the result
-        cv2.imwrite(output_path, result_img)
-
-        # Read the output image
-        with open(output_path, "rb") as f:
-            image_data = f.read()
-
-        # Clean up temporary files
-        for path in [source_path, target_path, output_path]:
-            if os.path.exists(path):
-                os.remove(path)
-
-        # Return the swapped image
-        return Response(content=image_data, media_type="image/jpeg")
+        # Debug: Check face detection
+        logger.info("Source faces detected: %s", get_many_faces(source_img))
+        logger.info("Target faces detected: %s", get_many_faces(target_img))
 
+        # Set up roop globals
+        roop.globals.source_path = source_path
+        roop.globals.target_path = target_path
+        roop.globals.output_path = normalize_output_path(
+            roop.globals.source_path, roop.globals.target_path, output_path
+        )
+        roop.globals.frame_processors = ["face_swapper", "face_enhancer"]
+        roop.globals.headless = True
+        roop.globals.keep_fps = True
+        roop.globals.keep_audio = True
+        roop.globals.keep_frames = False
+        roop.globals.many_faces = False
+        roop.globals.reference_face_position = 0
+        roop.globals.similar_face_distance = 0.6
+        roop.globals.video_encoder = "libx264"
+        roop.globals.video_quality = 18
+        roop.globals.max_memory = 16
+        roop.globals.execution_providers = decode_execution_providers(["cpu"])
+        roop.globals.execution_threads = 4
+
+        logger.info("Using frame processors: %s", roop.globals.frame_processors)
+
+        # Validate frame processors
+        for frame_processor in get_frame_processors_modules(
+            roop.globals.frame_processors
+        ):
+            if not frame_processor.pre_check():
+                raise HTTPException(status_code=500, detail=f"Pre-check failed for {frame_processor.__name__}")
+
+        # Perform face swap
+        start()
+        if os.path.exists(output_path):
+            output_image = Image.open(output_path)
+            output_image = output_image.resize(target_size, Image.Resampling.LANCZOS)
+            output_image.save(output_path, quality=100, subsampling=0)
+            logger.info("Returning FileResponse for %s", output_path)
+            background_tasks.add_task(cleanup_files, [source_path, target_path, output_path])
+            return FileResponse(output_path, media_type="image/jpeg", filename="output.jpg")
+        else:
+            raise HTTPException(status_code=500, detail="Face swap failed, output not generated")
     except Exception as e:
-        logging.error(f"Error in face swap: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Face swap failed: {str(e)}")
+        logger.error("Error in swap_face: %s", str(e))
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="127.0.0.1", port=8000)