File size: 3,088 Bytes
427cb66
9fca7c3
456bcc4
0ed24f7
9fca7c3
 
 
 
 
0ed24f7
9fca7c3
 
427cb66
279af2f
427cb66
456bcc4
 
 
9fca7c3
 
 
279af2f
427cb66
b8c2f1b
 
 
 
9fca7c3
 
 
a73b3bc
456bcc4
9fca7c3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e9cafa4
9fca7c3
 
427cb66
9fca7c3
 
0ed24f7
9fca7c3
 
 
 
0ed24f7
9fca7c3
 
427cb66
9fca7c3
 
279af2f
9fca7c3
 
 
 
279af2f
9fca7c3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a73b3bc
 
9fca7c3
c7d6b3a
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
# -*- coding: UTF-8 -*-
import gradio as gr
import cv2
import numpy as np
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager
from insightface.app import FaceAnalysis
from insightface.model_zoo.face_swapper import FaceSwapper
from pathlib import Path
import requests
import os
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize FastAPI
app = FastAPI(title="Face Swap API", description="Gradio-based Face Swap App")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

# Health Check Endpoint
@app.get("/health")
async def health_check():
    return {"status": "healthy"}

# Download Model Function
MODEL_PATH = Path("models/inswapper_128.onnx")
MODEL_URL = "https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx"

def download_model():
    if not MODEL_PATH.exists():
        logger.info("Downloading model...")
        os.makedirs("models", exist_ok=True)
        response = requests.get(MODEL_URL, stream=True, timeout=30)
        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.")

download_model()

# Initialize Face Analyzer & Swapper
face_analyzer = FaceAnalysis(name="buffalo_l")
face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
face_swapper = FaceSwapper(MODEL_PATH)

# Face Swap Function
def swap_faces(source_img, target_img):
    source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
    target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)

    source_faces = face_analyzer.get(source_img)
    target_faces = face_analyzer.get(target_img)

    if not source_faces or not target_faces:
        return "Error: No faces detected in one or both images."
    if len(source_faces) > 1 or len(target_faces) > 1:
        return "Error: Multiple faces detected; only one face per image is supported."

    result_img = face_swapper.get(target_img, target_faces[0], source_faces[0], paste_back=True)
    return cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)

# Gradio UI
def gradio_face_swap(source_img, target_img):
    try:
        result = swap_faces(source_img, target_img)
        if isinstance(result, str):
            return result  # Error message
        return result
    except Exception as e:
        return f"Error: {str(e)}"

gr_interface = gr.Interface(
    fn=gradio_face_swap,
    inputs=[
        gr.Image(label="Upload Source Face"),
        gr.Image(label="Upload Target Image")
    ],
    outputs=gr.Image(label="Swapped Face"),
    title="AI Face Swap",
    description="Upload two images: one with the source face and another where you want the face swapped.",
    live=True
)

# Launch Gradio
app = gr.mount_gradio_app(app, gr_interface, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)