app.py

import os
import time
import shutil
import argparse
import concurrent.futures

import cv2
import numpy as np
import torch
import gradio as gr
import spaces
import onnxruntime
from moviepy.editor import VideoFileClip
from tqdm import tqdm

from face_swapper import Inswapper, paste_to_whole
from face_analyser import detect_conditions, get_analysed_data, swap_options_list
from face_parsing import init_parsing_model, get_parsed_mask, mask_regions, mask_regions_to_list
from face_enhancer import get_available_enhancer_names, load_face_enhancer_model, cv2_interpolations
from utils import trim_video, open_directory, split_list_by_lengths, merge_img_sequence_from_ref, create_image_grid

parser = argparse.ArgumentParser(description="Free Face Swapper")
parser.add_argument("--out_dir", default=os.getcwd())
parser.add_argument("--batch_size", default=32)
parser.add_argument("--cuda", action="store_true", default=False)
parser.add_argument("--colab", action="store_true", default=False)
args, _ = parser.parse_known_args()

USE_COLAB = args.colab
DEF_OUTPUT_PATH = args.out_dir
BATCH_SIZE = int(args.batch_size)

WORKSPACE = None
OUTPUT_FILE = None
PREVIEW = None
STREAMER = None

DETECT_CONDITION = "best detection"
DETECT_SIZE = 640
DETECT_THRESH = 0.6
NUM_OF_SRC_SPECIFIC = 10

MASK_INCLUDE = ["Skin", "R-Eyebrow", "L-Eyebrow", "L-Eye", "R-Eye", "Nose", "Mouth", "L-Lip", "U-Lip"]
MASK_SOFT_KERNEL = 17
MASK_SOFT_ITERATIONS = 10
MASK_BLUR_AMOUNT = 0.1
MASK_ERODE_AMOUNT = 0.15

FACE_SWAPPER = None
FACE_ANALYSER = None
FACE_ENHANCER = None
FACE_PARSER = None

FACE_ENHANCER_LIST = ["NONE"]
FACE_ENHANCER_LIST.extend(get_available_enhancer_names())
FACE_ENHANCER_LIST.extend(cv2_interpolations)

PROVIDER = ["CPUExecutionProvider"]
device = "cpu"

def empty_cache():
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

def load_face_analyser_model(name="buffalo_l"):
    global FACE_ANALYSER
    if FACE_ANALYSER is None:
        FACE_ANALYSER = insightface.app.FaceAnalysis(name=name, providers=PROVIDER)
        FACE_ANALYSER.prepare(ctx_id=0, det_size=(DETECT_SIZE, DETECT_SIZE), det_thresh=DETECT_THRESH)

def load_face_swapper_model(path="./assets/pretrained_models/inswapper_128.onnx"):
    global FACE_SWAPPER
    if FACE_SWAPPER is None:
        FACE_SWAPPER = Inswapper(model_file=path, batch_size=1, providers=PROVIDER)

def load_face_parser_model(path="./assets/pretrained_models/79999_iter.pth"):
    global FACE_PARSER
    if FACE_PARSER is None:
        FACE_PARSER = init_parsing_model(path, device=device)

load_face_analyser_model()
load_face_swapper_model()

@spaces.GPU
def process(
    input_type,
    image_path,
    video_path,
    directory_path,
    source_path,
    output_path,
    output_name,
    keep_output_sequence,
    condition,
    age,
    distance,
    face_enhancer_name,
    enable_face_parser,
    mask_includes,
    mask_soft_kernel,
    mask_soft_iterations,
    blur_amount,
    erode_amount,
    face_scale,
    enable_laplacian_blend,
    crop_top,
    crop_bott,
    crop_left,
    crop_right,
    *specifics,
):
    global WORKSPACE, OUTPUT_FILE, PREVIEW, FACE_ENHANCER, FACE_PARSER
    start_time = time.time()

    def finish_text():
        mins, secs = divmod(time.time() - start_time, 60)
        return f"✔️ Completed in {int(mins)} min {int(secs)} sec."

    includes = mask_regions_to_list(mask_includes)
    specifics = list(specifics)
    half = len(specifics) // 2
    sources = specifics[:half]
    specifics = specifics[half:]

    if crop_top > crop_bott:
        crop_top, crop_bott = crop_bott, crop_top
    if crop_left > crop_right:
        crop_left, crop_right = crop_right, crop_left
    crop_mask = (crop_top, 511 - crop_bott, crop_left, 511 - crop_right)

    if face_enhancer_name != "NONE":
        FACE_ENHANCER = load_face_enhancer_model(name=face_enhancer_name, device=device)
    else:
        FACE_ENHANCER = None

    if enable_face_parser:
        load_face_parser_model()

    def swap_process(image_sequence):
        source_data = (source_path, age) if condition != "Specific Face" else ((sources, specifics), distance)

        analysed_targets, analysed_sources, whole_frame_list, num_faces_per_frame = get_analysed_data(
            FACE_ANALYSER,
            image_sequence,
            source_data,
            swap_condition=condition,
            detect_condition=DETECT_CONDITION,
            scale=face_scale,
        )

        preds = []
        matrs = []
        for batch_pred, batch_matr in FACE_SWAPPER.batch_forward(whole_frame_list, analysed_targets, analysed_sources):
            preds.extend(batch_pred)
            matrs.extend(batch_matr)
            empty_cache()

        if face_enhancer_name != "NONE" and FACE_ENHANCER is not None:
            enhancer_model, enhancer_model_runner = FACE_ENHANCER
            for idx, pred in enumerate(preds):
                preds[idx] = cv2.resize(enhancer_model_runner(pred, enhancer_model), (512, 512))
            empty_cache()

        if enable_face_parser:
            masks = []
            for batch_mask in get_parsed_mask(
                FACE_PARSER,
                preds,
                classes=includes,
                device=device,
                batch_size=BATCH_SIZE,
                softness=int(mask_soft_iterations),
            ):
                masks.append(batch_mask)
                empty_cache()
            masks = np.concatenate(masks, axis=0) if len(masks) >= 1 else masks
        else:
            masks = [None] * len(preds)

        split_preds = split_list_by_lengths(preds, num_faces_per_frame)
        split_matrs = split_list_by_lengths(matrs, num_faces_per_frame)
        split_masks = split_list_by_lengths(masks, num_faces_per_frame)

        def post_process(frame_idx, frame_img):
            whole_img = cv2.imread(frame_img)
            blend_method = "laplacian" if enable_laplacian_blend else "linear"
            for p, m, mask in zip(split_preds[frame_idx], split_matrs[frame_idx], split_masks[frame_idx]):
                p = cv2.resize(p, (512, 512))
                mask = cv2.resize(mask, (512, 512)) if mask is not None else None
                m /= 0.25
                whole_img = paste_to_whole(
                    p,
                    whole_img,
                    m,
                    mask=mask,
                    crop_mask=crop_mask,
                    blend_method=blend_method,
                    blur_amount=blur_amount,
                    erode_amount=erode_amount,
                )
            cv2.imwrite(frame_img, whole_img)

        with concurrent.futures.ThreadPoolExecutor() as executor:
            futures = [executor.submit(post_process, idx, frame_img) for idx, frame_img in enumerate(image_sequence)]
            for _ in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="Pasting back"):
                pass

    if input_type == "Image":
        output_file = os.path.join(output_path, output_name + ".png")
        cv2.imwrite(output_file, cv2.imread(image_path))
        swap_process([output_file])
        OUTPUT_FILE = output_file
        WORKSPACE = output_path
        PREVIEW = cv2.imread(output_file)[:, :, ::-1]
        return finish_text(), gr.update(visible=True, value=PREVIEW), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)

    if input_type == "Video":
        temp_path = os.path.join(output_path, output_name, "sequence")
        os.makedirs(temp_path, exist_ok=True)

        image_sequence = []
        cap = cv2.VideoCapture(video_path)
        curr_idx = 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            frame_path = os.path.join(temp_path, f"frame_{curr_idx}.jpg")
            cv2.imwrite(frame_path, frame)
            image_sequence.append(frame_path)
            curr_idx += 1
        cap.release()

        swap_process(image_sequence)

        output_video_path = os.path.join(output_path, output_name + ".mp4")
        merge_img_sequence_from_ref(video_path, image_sequence, output_video_path)

        if os.path.exists(temp_path) and not keep_output_sequence:
            shutil.rmtree(temp_path)

        WORKSPACE = output_path
        OUTPUT_FILE = output_video_path
        return finish_text(), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True, value=OUTPUT_FILE)

    return "Unsupported input type", gr.update(), gr.update(), gr.update(), gr.update()

def stop_running():
    global STREAMER
    if hasattr(STREAMER, "stop"):
        STREAMER.stop()
        STREAMER = None
    return "Cancelled"

css = "footer{display:none !important}"

with gr.Blocks(css=css) as interface:
    gr.Markdown("# 🗿 Free Face Swapper")

    with gr.Row():
        with gr.Column(scale=0.4):
            swap_option = gr.Dropdown(swap_options_list, value=swap_options_list[0], interactive=True, show_label=False)
            age = gr.Number(value=25, interactive=True, visible=False)
            detect_condition_dropdown = gr.Dropdown(detect_conditions, label="Condition", value=DETECT_CONDITION)
            detection_size = gr.Number(label="Detection Size", value=DETECT_SIZE)
            detection_threshold = gr.Number(label="Detection Threshold", value=DETECT_THRESH)
            output_directory = gr.Text(label="Output Directory", value=DEF_OUTPUT_PATH)
            output_name = gr.Text(label="Output Name", value="Result")
            keep_output_sequence = gr.Checkbox(label="Keep output sequence", value=False)
            face_scale = gr.Slider(label="Face Scale", minimum=0, maximum=2, value=1)
            face_enhancer_name = gr.Dropdown(FACE_ENHANCER_LIST, label="Face Enhancer", value="NONE")
            enable_face_parser_mask = gr.Checkbox(label="Enable Face Parsing", value=False)
            mask_include = gr.Dropdown(mask_regions.keys(), value=MASK_INCLUDE, multiselect=True, label="Include")
            mask_soft_kernel = gr.Number(label="Soft Erode Kernel", value=MASK_SOFT_KERNEL, visible=False)
            mask_soft_iterations = gr.Number(label="Soft Erode Iterations", value=MASK_SOFT_ITERATIONS)
            crop_top = gr.Slider(label="Top", minimum=0, maximum=511, value=0, step=1)
            crop_bott = gr.Slider(label="Bottom", minimum=0, maximum=511, value=511, step=1)
            crop_left = gr.Slider(label="Left", minimum=0, maximum=511, value=0, step=1)
            crop_right = gr.Slider(label="Right", minimum=0, maximum=511, value=511, step=1)
            erode_amount = gr.Slider(label="Mask Erode", minimum=0, maximum=1, value=MASK_ERODE_AMOUNT, step=0.05)
            blur_amount = gr.Slider(label="Mask Blur", minimum=0, maximum=1, value=MASK_BLUR_AMOUNT, step=0.05)
            enable_laplacian_blend = gr.Checkbox(label="Laplacian Blending", value=True)
            source_image_input = gr.Image(label="Source face", type="filepath", interactive=True)
            input_type = gr.Radio(["Image", "Video"], label="Target Type", value="Image")
            image_input = gr.Image(label="Target Image", interactive=True, type="filepath")
            video_input = gr.Video(label="Target Video", interactive=True)

        with gr.Column(scale=0.6):
            info = gr.Markdown(value="...")
            swap_button = gr.Button("✨ Swap", variant="primary")
            cancel_button = gr.Button("⛔ Cancel")
            preview_image = gr.Image(label="Output", interactive=False)
            preview_video = gr.Video(label="Output", interactive=False, visible=False)
            output_directory_button = gr.Button("📂", interactive=False, visible=False)
            output_video_button = gr.Button("🎬", interactive=False, visible=False)

    src_specific_inputs = []
    for i in range(NUM_OF_SRC_SPECIFIC):
        exec(f"src{i+1} = gr.Image(interactive=True, type='numpy', label='Source Face {i+1}')")
        exec(f"trg{i+1} = gr.Image(interactive=True, type='numpy', label='Specific Face {i+1}')")
    exec("src_specific_inputs = (" + ",".join([f"src{i+1}" for i in range(NUM_OF_SRC_SPECIFIC)] + [f\"trg{i+1}\" for i in range(NUM_OF_SRC_SPECIFIC)]) + ")")

    swap_inputs = [
        input_type, image_input, video_input, gr.Text(), source_image_input, output_directory, output_name,
        keep_output_sequence, swap_option, age, gr.Number(value=0.6), face_enhancer_name,
        enable_face_parser_mask, mask_include, mask_soft_kernel, mask_soft_iterations,
        blur_amount, erode_amount, face_scale, enable_laplacian_blend,
        crop_top, crop_bott, crop_left, crop_right, *src_specific_inputs
    ]

    swap_button.click(fn=process, inputs=swap_inputs, outputs=[info, preview_image, output_directory_button, output_video_button, preview_video], show_progress=True)
    cancel_button.click(fn=stop_running, inputs=None, outputs=[info])

if __name__ == "__main__":
    interface.queue().launch(server_name="0.0.0.0", share=False)