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)