swap.py

import cv2
import numpy as np
import insightface
from insightface.app import FaceAnalysis
from gfpgan import GFPGANer
import os
import torch
import warnings
import gradio as gr
import time
from datetime import datetime
import shutil
import traceback

# Suppress specific warnings
warnings.filterwarnings("ignore", category=UserWarning, module="gradio_client.documentation")
warnings.filterwarnings("ignore", category=FutureWarning)

# Paths (giữ nguyên như bạn cung cấp)
model_path = os.path.join("models", "inswapper_128.onnx")
gfpgan_path = os.path.join("gfpgan", "weights", "GFPGANv1.4.pth")
buffalo_l_path = os.path.join("models", "buffalo_l")
output_dir = "output"

# Initialize logging
log_messages = []

def log_message(message):
    """Append message to log with timestamp."""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_messages.append(f"[{timestamp}] {message}")
    print(f"[{timestamp}] {message}")  # Also print to console
    return "\n".join(log_messages)

def validate_paths():
    """Validate required file and directory paths."""
    log_message("Validating file paths...")
    for path in [model_path, gfpgan_path]:
        if not os.path.isfile(path):
            return False, f"Error: File not found at {path}"
    if not os.path.isdir(buffalo_l_path):
        return False, f"Error: buffalo_l directory not found at {buffalo_l_path}. Please download and extract buffalo_l.zip from https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_l.zip to {buffalo_l_path}"
    # Kiểm tra các file cần thiết trong buffalo_l
    required_files = ["1k3d68.onnx", "2d106det.onnx", "det_10g.onnx", "genderage.onnx", "w600k_r50.onnx"]
    if not all(os.path.exists(os.path.join(buffalo_l_path, f)) for f in required_files):
        return False, f"Error: buffalo_l directory at {buffalo_l_path} is incomplete. Please ensure it contains {', '.join(required_files)}"
    return True, "All paths validated successfully"

def initialize_face_analysis():
    """Initialize FaceAnalysis model."""
    providers = [
        ('CUDAExecutionProvider', {
            'device_id': 0,
            'gpu_mem_limit': 10 * 1024 * 1024 * 1024,
            'arena_extend_strategy': 'kNextPowerOfTwo',
            'cudnn_conv_algo_search': 'EXHAUSTIVE',
            'do_copy_in_default_stream': True,
        }),
        'CPUExecutionProvider',
    ]
    try:
        log_message("Initializing FaceAnalysis...")
        # Sử dụng root="models" để tìm đúng models\buffalo_l
        app = FaceAnalysis(name="buffalo_l", root=os.path.dirname(buffalo_l_path), providers=providers)
        app.prepare(ctx_id=0, det_size=(640, 640))
        log_message(f"PyTorch CUDA available: {torch.cuda.is_available()}")
        log_message("FaceAnalysis initialized successfully")
        return app, None
    except Exception as e:
        error_msg = f"Error initializing FaceAnalysis: {str(e)}\n{traceback.format_exc()}"
        return None, log_message(error_msg)

def load_and_detect_faces(app, source_img, target_img):
    """Load images and detect faces."""
    try:
        log_message("Loading and detecting faces...")
        if source_img is None or target_img is None:
            return None, None, "Error: Source or target image is None"
        
        source_img_np = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
        target_img_np = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
        
        source_faces = app.get(source_img_np)
        target_faces = app.get(target_img_np)
        
        log_message(f"Source image: {len(source_faces)} faces detected")
        log_message(f"Target image: {len(target_faces)} faces detected")
        
        if len(source_faces) == 0 or len(target_faces) == 0:
            return None, None, "Error: No faces detected in source or target image!"
        
        return source_faces, target_faces, None
    except Exception as e:
        error_msg = f"Error in load_and_detect_faces: {str(e)}\n{traceback.format_exc()}"
        return None, None, log_message(error_msg)

def select_source_face(source_faces):
    """Select the first source face."""
    try:
        log_message("Selecting source face...")
        source_face = source_faces[0]
        log_message("Using first detected source face")
        return source_face, None
    except Exception as e:
        error_msg = f"Error selecting source face: {str(e)}\n{traceback.format_exc()}"
        return None, log_message(error_msg)

def perform_face_swap(source_face, target_face, target_img):
    """Perform face swapping with edge smoothing."""
    try:
        log_message("Loading inswapper model...")
        swapper = insightface.model_zoo.get_model(model_path, providers=[
            ('CUDAExecutionProvider', {
                'device_id': 0,
                'gpu_mem_limit': 10 * 1024 * 1024 * 1024,
                'arena_extend_strategy': 'kNextPowerOfTwo',
                'cudnn_conv_algo_search': 'EXHAUSTIVE',
                'do_copy_in_default_stream': True,
            }),
            'CPUExecutionProvider',
        ])
        log_message("Inswapper model loaded successfully")
        
        target_img_np = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
        result = target_img_np.copy()
        result = swapper.get(result, target_face, source_face, paste_back=True)
        
        x, y, w, h = target_face.bbox.astype(int)
        mask = np.zeros(result.shape[:2], dtype=np.float32)
        cv2.rectangle(mask, (x, y), (x + w, y + h), 1.0, -1)
        mask = cv2.GaussianBlur(mask, (9, 9), 0)
        mask = np.stack([mask]*3, axis=-1)
        result = (result * mask + target_img_np * (1 - mask)).astype(np.uint8)
        
        log_message("Face swapping completed")
        return result, None
    except Exception as e:
        error_msg = f"Error during face swapping: {str(e)}\n{traceback.format_exc()}"
        return None, log_message(error_msg)

def enhance_with_gfpgan(result):
    """Enhance swapped image using GFPGAN without resizing."""
    try:
        log_message("Enhancing with GFPGAN...")
        enhancer = GFPGANer(
            model_path=gfpgan_path,
            upscale=1,
            arch='clean',
            channel_multiplier=2,
            device='cuda' if torch.cuda.is_available() else 'cpu',
            bg_upsampler=None
        )
        _, _, enhanced_result = enhancer.enhance(result, paste_back=True)
        output_path = os.path.join(output_dir, "output.jpg")
        cv2.imwrite(output_path, enhanced_result)
        log_message(f"Enhanced image saved to {output_path}")
        return output_path, None
    except Exception as e:
        error_msg = f"Error during GFPGAN enhancement: {str(e)}\n{traceback.format_exc()}"
        return None, log_message(error_msg)

def face_swap(source_img, target_img):
    """Main face swap function for Gradio."""
    global log_messages
    log_messages = []  # Reset logs
    start_time = time.time()
    
    try:
        log_message("Starting face swap process...")
        if os.path.exists(output_dir):
            shutil.rmtree(output_dir)
        os.makedirs(output_dir, exist_ok=True)
        
        valid, message = validate_paths()
        log_message(message)
        if not valid:
            return None, log_message("Path validation failed")
        
        app, error = initialize_face_analysis()
        if error:
            return None, log_message(error)
        
        source_faces, target_faces, error = load_and_detect_faces(app, source_img, target_img)
        if error:
            return None, log_message(error)
        
        source_face, error = select_source_face(source_faces)
        if error:
            return None, log_message(error)
        target_face = target_faces[0]
        log_message(f"Target face attributes: {target_face.__dict__}")
        
        result, error = perform_face_swap(source_face, target_face, target_img)
        if error:
            return None, log_message(error)
        
        output_path, error = enhance_with_gfpgan(result)
        if error:
            return None, log_message(error)
        
        log_message(f"Processing completed in {time.time() - start_time:.2f} seconds")
        return output_path, "\n".join(log_messages)
    except Exception as e:
        error_msg = f"Unexpected error in face_swap: {str(e)}\n{traceback.format_exc()}"
        return None, log_message(error_msg)

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Face Swap Application")
    gr.Markdown("Upload source and target images to swap faces. The first detected face in the source image will be used.")
    
    with gr.Row():
        with gr.Column():
            source_img = gr.Image(type="pil", label="Source Image")
            target_img = gr.Image(type="pil", label="Target Image")
            submit_btn = gr.Button("Swap Faces")
        with gr.Column():
            output = gr.Image(label="Final Output")
    
    logs = gr.Textbox(label="Logs", interactive=False, lines=10)
    
    submit_btn.click(
        fn=face_swap,
        inputs=[source_img, target_img],
        outputs=[output, logs],
        api_name="faceswap"  # Đảm bảo endpoint /face_swap
    )

if __name__ == "__main__":
    try:
        log_message("Launching Gradio interface...")
        demo.launch(
            share=True,
            debug=True,
            allowed_paths=["models", "gfpgan/weights", output_dir],
            server_name="0.0.0.0",
            server_port=7860
        )
    except Exception as e:
        log_message(f"Error launching Gradio: {str(e)}\n{traceback.format_exc()}")
        print(f"Error launching Gradio: {str(e)}\n{traceback.format_exc()}")