Update app.py

app.py

@@ -1,22 +1,15 @@
 # -*- coding:UTF-8 -*-
-#!/usr/bin/env python
-import numpy as np
-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
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.responses import Response
+from fastapi.middleware.cors import CORSMiddleware
 import cv2
-from pathlib import Path
+import numpy as np
 from PIL import Image
-from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks
-from fastapi.responses import FileResponse
-from fastapi.middleware.cors import CORSMiddleware
+import os
 import shutil
 import logging
-import time
+import requests
+from pathlib import Path
 
 app = FastAPI()
 
@@ -27,26 +20,12 @@ logger = logging.getLogger(__name__)
 # Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Update with your Framer domain in production
+    allow_origins=["*"],  # Update with Framer domain in production
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-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"
@@ -64,96 +43,96 @@ def download_model():
             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)
+            raise RuntimeError("Could not download inswapper_128.onnx. Please check logs.")
 
+# Download model on startup
 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)
+def get_many_faces(image):
+    """Simplified face detection using insightface (placeholder)."""
+    from insightface.app import FaceAnalysis
+    app = FaceAnalysis(name="buffalo_l")
+    app.prepare(ctx_id=0, det_size=(640, 640))
+    faces = app.get(image)
+    return faces if faces else []
+
+def swap_faces(source_img, target_img):
+    """Perform face swapping using insightface and inswapper model."""
+    from insightface.utils import face_align
+    from insightface.model_zoo import face_swapper
+
+    # Initialize face analysis
+    face_analyzer = FaceAnalysis(name="buffalo_l")
+    face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
+
+    # Detect faces
+    source_faces = face_analyzer.get(source_img)
+    target_faces = face_analyzer.get(target_img)
+
+    if not source_faces or not target_faces:
+        raise ValueError("No faces detected in one or both images.")
+    if len(source_faces) > 1 or len(target_faces) > 1:
+        raise ValueError("Multiple faces detected; only one face per image is supported.")
+
+    source_face = source_faces[0]
+    target_face = target_faces[0]
+
+    # Load the face swapper model
+    model_path = Path("models/inswapper_128.onnx")
+    swapper = face_swapper.FaceSwapper(model_path)
+
+    # Perform face swap
+    result = swapper.get(target_img, target_face, source_face, paste_back=True)
+
+    # Resize to match target image size
+    target_pil = Image.fromarray(cv2.cvtColor(target_img, cv2.COLOR_BGR2RGB))
+    result_pil = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
+    result_pil = result_pil.resize(target_pil.size, Image.Resampling.LANCZOS)
+
+    return cv2.cvtColor(np.array(result_pil), cv2.COLOR_RGB2BGR)
 
 @app.post("/swap-face/")
-async def swap_face(source_file: UploadFile = File(...), target_file: UploadFile = File(...), doFaceEnhancer: bool = True, background_tasks: BackgroundTasks = BackgroundTasks()):
+async def swap_face(source_file: UploadFile = File(...), target_file: UploadFile = File(...), doFaceEnhancer: bool = True):
     try:
         # Save uploaded files temporarily
-        source_path = "input.jpg"
-        target_path = "target.jpg"
+        source_path = "temp_source.jpg"
+        target_path = "temp_target.jpg"
         output_path = "output.jpg"
 
         # Read and save source image
         source_content = await source_file.read()
         with open(source_path, "wb") as f:
             f.write(source_content)
-        source_image = Image.open(source_path)
-        source_image.save(source_path, quality=100, subsampling=0)
+        source_img = cv2.imread(source_path)
+        if source_img is None:
+            raise ValueError("Failed to load source image.")
 
-        # Read and save target image, get its size
+        # Read and save target image
         target_content = await target_file.read()
         with open(target_path, "wb") as 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)
-
-        # 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 HTTPException(status_code=400, detail="Failed to load one or both images.")
-
-        # 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__}")
+        if target_img is None:
+            raise ValueError("Failed to load target image.")
 
         # 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")
+        result_img = swap_faces(source_img, target_img)
+
+        # 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")
+
     except Exception as 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)
+        raise HTTPException(status_code=500, detail=str(e))