app.py

import gradio as gr
import insightface
from insightface.app import FaceAnalysis
import torch
from diffusers import StableDiffusionInpaintPipeline, ControlNetModel
import numpy as np
import cv2
from PIL import Image

assert insightface.__version__ >= '0.7'

def prepare_app():
    # Initialize InsightFace for face swapping
    app = FaceAnalysis(name='buffalo_l')
    app.prepare(ctx_id=0, det_size=(640, 640))
    swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=True, download_zip=True)
    
    # Initialize Stable Diffusion Inpainting with ControlNet for hair editing
    controlnet = ControlNetModel.from_pretrained(
        "lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16
    )
    pipe = StableDiffusionInpaintPipeline.from_pretrained(
        "runwayml/stable-diffusion-inpainting",
        controlnet=controlnet,
        torch_dtype=torch.float16
    )
    pipe = pipe.to("cuda")
    pipe.enable_xformers_memory_efficient_attention()  # Optional, for memory optimization
    return app, swapper, pipe

def sort_faces(faces):
    return sorted(faces, key=lambda x: x.bbox[0])

def get_face(faces, face_id):
    try:
        if len(faces) < face_id or face_id < 1:
            raise gr.Error(f"The image includes only {len(faces)} faces, however, you asked for face {face_id}")
        return faces[face_id-1]
    except Exception as e:
        raise gr.Error(f"An error occurred: {str(e)}")

def create_hair_mask(image, faces, face_id):
    """Create a precise mask for the hair region, avoiding the face."""
    face = get_face(faces, face_id)
    bbox = face.bbox
    x1, y1, x2, y2 = map(int, bbox)
    
    # Create a mask for the hair region (above and around the face, excluding facial features)
    mask = np.zeros_like(image)
    hair_height = int((y2 - y1) * 1.8)  # Extend above face for hair
    hair_width = int((x2 - x1) * 1.5)  # Widen to cover side hair
    mask_y1 = max(0, y1 - hair_height)
    mask_x1 = max(0, x1 - (hair_width - (x2 - x1)) // 2)
    mask_x2 = min(image.shape[1], x2 + (hair_width - (x2 - x1)) // 2)
    mask[max(0, y1 - hair_height):y1, mask_x1:mask_x2] = 255  # Focus on hair region, exclude face
    
    # Dilate mask to ensure full hair coverage
    kernel = np.ones((10, 10), np.uint8)
    mask = cv2.dilate(mask, kernel, iterations=1)
    
    # Convert to grayscale for inpainting
    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
    return Image.fromarray(mask)

def apply_hair_transformation(image, hair_style, hair_length, pipe):
    """Apply hairstyle transformation using Stable Diffusion Inpainting with ControlNet."""
    # Convert image to PIL if it's a numpy array
    if isinstance(image, np.ndarray):
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    
    # Detect faces to locate hair region
    faces = sort_faces(app.get(np.array(image)))
    if not faces:
        raise gr.Error("No faces detected in the source image.")
    
    # Create hair mask
    mask = create_hair_mask(np.array(image), faces, 1)
    
    # Generate Canny edge map for ControlNet
    image_np = np.array(image)
    gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
    canny = cv2.Canny(gray, 100, 200)
    canny_image = Image.fromarray(canny)
    
    # Construct prompt for hairstyle
    prompt = f"{hair_style.lower()} {hair_length.lower()} hair, high detail, realistic, natural lighting"
    negative_prompt = "distorted, blurry, low quality, unnatural, extra limbs"
    
    # Perform inpainting
    try:
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=image,
            mask_image=mask,
            control_image=canny_image,
            num_inference_steps=30,
            controlnet_conditioning_scale=0.8,
            guidance_scale=7.5
        ).images[0]
    except Exception as e:
        raise gr.Error(f"Error during hair transformation: {str(e)}")
    
    return result

app, swapper, pipe = prepare_app()

def swap_faces(sourceImage, sourceFaceIndex, destinationImage, destinationFaceIndex, hair_style, hair_length):
    """Swaps the entire face from the source to the destination image with hairstyle modification."""
    # Validate input images
    if sourceImage is None or destinationImage is None:
        raise gr.Error("Please upload both source and destination images.")
    
    # Convert images to numpy arrays if they are PIL Images
    if isinstance(sourceImage, Image.Image):
        sourceImage = np.array(sourceImage)
    if isinstance(destinationImage, Image.Image):
        destinationImage = np.array(destinationImage)
    
    # Apply hair transformation to the source image
    transformed_source = apply_hair_transformation(sourceImage, hair_style, hair_length, pipe)
    
    # Detect and sort faces in the transformed source image
    faces = sort_faces(app.get(np.array(transformed_source)))
    source_face = get_face(faces, sourceFaceIndex)
    
    # Detect and sort faces in the destination image
    res_faces = sort_faces(app.get(destinationImage))
    res_face = get_face(res_faces, destinationFaceIndex)

    # Perform face swap with full face transfer
    try:
        result = swapper.get(destinationImage, res_face, source_face, paste_back=True)
    except Exception as e:
        raise gr.Error(f"Error during face swap: {str(e)}")
    
    # Ensure result is in RGB format
    if isinstance(result, np.ndarray):
        result = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
    
    return result

gr.Interface(
    swap_faces, 
    [
        gr.Image(label="Source Image (the image with the face that you want to use)"),
        gr.Number(precision=0, value=1, label='Source Face Position', info='In case there are multiple faces on the image, specify which should be used from the left, starting at 1'),
        gr.Image(label="Destination Image (the image with the face that you want to replace)"),
        gr.Number(precision=0, value=1, label='Destination Face Position', info='In case there are multiple faces on the image, specify which should be replaced from the left, starting at 1'),
        gr.Dropdown(choices=["Curly", "Wavy", "Straight"], label="Hair Style", info="Select the desired hair style for the swapped face"),
        gr.Dropdown(choices=["Long", "Short", "Medium"], label="Hair Length", info="Select the desired hair length for the swapped face")
    ],
    gr.Image(),
    examples=[
        ['./examples/rihanna.jpg', 1, './examples/margaret_thatcher.jpg', 3, "Curly", "Long"],
        ['./examples/game_of_thrones.jpg', 5, './examples/game_of_thrones.jpg', 4, "Wavy", "Medium"],
    ],
    theme=gr.themes.Base(),
    title="Face Swapper App with Hair Editing 🔄",
    description="🌀 This app swaps the entire face from the source to the destination image, preserving all facial features and dimensions, and allows hairstyle modification. <br>➡️ Upload a source image and a destination image, specify the face positions, and choose the desired hair style and length for the source face. <br>⚡️ Uses Stable Diffusion with ControlNet Inpainting for hair editing and InsightFace for face swapping. <br>💡 Note: Hair editing quality depends on the input image and mask accuracy. For advanced hair transfer, consider Stable-Hair (https://github.com/Xiaojiu-z/Stable-Hair). <br>💡 At [Dentro](https://dentroai.com), we help you discover, develop, and implement AI within your organisation! <br>📖 Original face swap model: [InsightFace](https://github.com/deepinsight/insightface). <br>❤️ Feel free to like or duplicate this space!",
).launch()