merge

#### FACE_ENHANCER.PY CODE START ###

import os
import cv2
import torch
import gfpgan
from PIL import Image
from upscaler.RealESRGAN import RealESRGAN
from upscaler.codeformer import CodeFormerEnhancer

def gfpgan_runner(img, model):
    _, imgs, _ = model.enhance(img, paste_back=True, has_aligned=True)
    return imgs[0]


def realesrgan_runner(img, model):
    img = model.predict(img)
    return img


def codeformer_runner(img, model):
    img = model.enhance(img)
    return img


supported_enhancers = {
    "CodeFormer": ("./assets/pretrained_models/codeformer.onnx", codeformer_runner),
    "GFPGAN": ("./assets/pretrained_models/GFPGANv1.4.pth", gfpgan_runner),
    "REAL-ESRGAN 2x": ("./assets/pretrained_models/RealESRGAN_x2.pth", realesrgan_runner),
    "REAL-ESRGAN 4x": ("./assets/pretrained_models/RealESRGAN_x4.pth", realesrgan_runner),
    "REAL-ESRGAN 8x": ("./assets/pretrained_models/RealESRGAN_x8.pth", realesrgan_runner)
}

cv2_interpolations = ["LANCZOS4", "CUBIC", "NEAREST"]

def get_available_enhancer_names():
    available = []
    for name, data in supported_enhancers.items():
        path = os.path.join(os.path.abspath(os.path.dirname(__file__)), data[0])
        if os.path.exists(path):
            available.append(name)
    return available


def load_face_enhancer_model(name='GFPGAN', device="cpu"):
    assert name in get_available_enhancer_names() + cv2_interpolations, f"Face enhancer {name} unavailable."
    if name in supported_enhancers.keys():
        model_path, model_runner = supported_enhancers.get(name)
        model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
    if name == 'CodeFormer':
        model = CodeFormerEnhancer(model_path=model_path, device=device)
    elif name == 'GFPGAN':
        model = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=device)
    elif name == 'REAL-ESRGAN 2x':
        model = RealESRGAN(device, scale=2)
        model.load_weights(model_path, download=False)
    elif name == 'REAL-ESRGAN 4x':
        model = RealESRGAN(device, scale=4)
        model.load_weights(model_path, download=False)
    elif name == 'REAL-ESRGAN 8x':
        model = RealESRGAN(device, scale=8)
        model.load_weights(model_path, download=False)
    elif name == 'LANCZOS4':
        model = None
        model_runner = lambda img, _: cv2.resize(img, (512,512), interpolation=cv2.INTER_LANCZOS4)
    elif name == 'CUBIC':
        model = None
        model_runner = lambda img, _: cv2.resize(img, (512,512), interpolation=cv2.INTER_CUBIC)
    elif name == 'NEAREST':
        model = None
        model_runner = lambda img, _: cv2.resize(img, (512,512), interpolation=cv2.INTER_NEAREST)
    else:
        model = None
    return (model, model_runner)


#### FACE_EHNANCER.PY CODE END ###

#### FACE_SWAPPER.PY CODE START ###

import time
import torch
import onnx
import cv2
import onnxruntime
import numpy as np
from tqdm import tqdm
import torch.nn as nn
from onnx import numpy_helper
from skimage import transform as trans
import torchvision.transforms.functional as F
import torch.nn.functional as F
from utils import mask_crop, laplacian_blending


arcface_dst = np.array(
    [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366],
     [41.5493, 92.3655], [70.7299, 92.2041]],
    dtype=np.float32)

def estimate_norm(lmk, image_size=112, mode='arcface'):
    assert lmk.shape == (5, 2)
    assert image_size % 112 == 0 or image_size % 128 == 0
    if image_size % 112 == 0:
        ratio = float(image_size) / 112.0
        diff_x = 0
    else:
        ratio = float(image_size) / 128.0
        diff_x = 8.0 * ratio
    dst = arcface_dst * ratio
    dst[:, 0] += diff_x
    tform = trans.SimilarityTransform()
    tform.estimate(lmk, dst)
    M = tform.params[0:2, :]
    return M


def norm_crop2(img, landmark, image_size=112, mode='arcface'):
    M = estimate_norm(landmark, image_size, mode)
    warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0)
    return warped, M


class Inswapper():
    def __init__(self, model_file=None, batch_size=32, providers=['CPUExecutionProvider']):
        self.model_file = model_file
        self.batch_size = batch_size

        model = onnx.load(self.model_file)
        graph = model.graph
        self.emap = numpy_helper.to_array(graph.initializer[-1])

        self.session_options = onnxruntime.SessionOptions()
        self.session = onnxruntime.InferenceSession(self.model_file, sess_options=self.session_options, providers=providers)

    def forward(self, imgs, latents):
        preds = []
        for img, latent in zip(imgs, latents):
            img = img / 255
            pred = self.session.run(['output'], {'target': img, 'source': latent})[0]
            preds.append(pred)

    def get(self, imgs, target_faces, source_faces):
        imgs = list(imgs)

        preds = [None] * len(imgs)
        matrs = [None] * len(imgs)

        for idx, (img, target_face, source_face) in enumerate(zip(imgs, target_faces, source_faces)):
            matrix, blob, latent = self.prepare_data(img, target_face, source_face)
            pred = self.session.run(['output'], {'target': blob, 'source': latent})[0]
            pred = pred.transpose((0, 2, 3, 1))[0]
            pred = np.clip(255 * pred, 0, 255).astype(np.uint8)[:, :, ::-1]

            preds[idx] = pred
            matrs[idx] = matrix

        return (preds, matrs)

    def prepare_data(self, img, target_face, source_face):
        if isinstance(img, str):
            img = cv2.imread(img)

        aligned_img, matrix = norm_crop2(img, target_face.kps, 128)

        blob = cv2.dnn.blobFromImage(aligned_img, 1.0 / 255, (128, 128), (0., 0., 0.), swapRB=True)

        latent = source_face.normed_embedding.reshape((1, -1))
        latent = np.dot(latent, self.emap)
        latent /= np.linalg.norm(latent)

        return (matrix, blob, latent)

    def batch_forward(self, img_list, target_f_list, source_f_list):
        num_samples = len(img_list)
        num_batches = (num_samples + self.batch_size - 1) // self.batch_size

        for i in tqdm(range(num_batches), desc="Generating face"):
            start_idx = i * self.batch_size
            end_idx = min((i + 1) * self.batch_size, num_samples)

            batch_img = img_list[start_idx:end_idx]
            batch_target_f = target_f_list[start_idx:end_idx]
            batch_source_f = source_f_list[start_idx:end_idx]

            batch_pred, batch_matr = self.get(batch_img, batch_target_f, batch_source_f)

            yield batch_pred, batch_matr


def paste_to_whole(foreground, background, matrix, mask=None, crop_mask=(0,0,0,0), blur_amount=0.1, erode_amount = 0.15, blend_method='linear'):
    inv_matrix = cv2.invertAffineTransform(matrix)
    fg_shape = foreground.shape[:2]
    bg_shape = (background.shape[1], background.shape[0])
    foreground = cv2.warpAffine(foreground, inv_matrix, bg_shape, borderValue=0.0)

    if mask is None:
        mask = np.full(fg_shape, 1., dtype=np.float32)
        mask = mask_crop(mask, crop_mask)
        mask = cv2.warpAffine(mask, inv_matrix, bg_shape, borderValue=0.0)
    else:
        assert fg_shape == mask.shape[:2], "foreground & mask shape mismatch!"
        mask = mask_crop(mask, crop_mask).astype('float32')
        mask = cv2.warpAffine(mask, inv_matrix, (background.shape[1], background.shape[0]), borderValue=0.0)

    _mask = mask.copy()
    _mask[_mask > 0.05] = 1.
    non_zero_points = cv2.findNonZero(_mask)
    _, _, w, h = cv2.boundingRect(non_zero_points)
    mask_size = int(np.sqrt(w * h))

    if erode_amount > 0:
        kernel_size = max(int(mask_size * erode_amount), 1)
        structuring_element = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_size, kernel_size))
        mask = cv2.erode(mask, structuring_element)

    if blur_amount > 0:
        kernel_size = max(int(mask_size * blur_amount), 3)
        if kernel_size % 2 == 0:
            kernel_size += 1
        mask = cv2.GaussianBlur(mask, (kernel_size, kernel_size), 0)

    mask = np.tile(np.expand_dims(mask, axis=-1), (1, 1, 3))

    if blend_method == 'laplacian':
        composite_image = laplacian_blending(foreground, background, mask.clip(0,1), num_levels=4)
    else:
        composite_image = mask * foreground + (1 - mask) * background

    return composite_image.astype("uint8").clip(0, 255)

#### FACE_SWAPPER.PY CODE END ###


#### FACE_ANALYSER.PY CODE START ###

import os
import cv2
import numpy as np
from tqdm import tqdm
from utils import scale_bbox_from_center

detect_conditions = [
    "best detection",
    "left most",
    "right most",
    "top most",
    "bottom most",
    "middle",
    "biggest",
    "smallest",
]

swap_options_list = [
    "All Face",
    "Specific Face",
    "Age less than",
    "Age greater than",
    "All Male",
    "All Female",
    "Left Most",
    "Right Most",
    "Top Most",
    "Bottom Most",
    "Middle",
    "Biggest",
    "Smallest",
]

def get_single_face(faces, method="best detection"):
    total_faces = len(faces)
    if total_faces == 1:
        return faces[0]

    print(f"{total_faces} face detected. Using {method} face.")
    if method == "best detection":
        return sorted(faces, key=lambda face: face["det_score"])[-1]
    elif method == "left most":
        return sorted(faces, key=lambda face: face["bbox"][0])[0]
    elif method == "right most":
        return sorted(faces, key=lambda face: face["bbox"][0])[-1]
    elif method == "top most":
        return sorted(faces, key=lambda face: face["bbox"][1])[0]
    elif method == "bottom most":
        return sorted(faces, key=lambda face: face["bbox"][1])[-1]
    elif method == "middle":
        return sorted(faces, key=lambda face: (
                (face["bbox"][0] + face["bbox"][2]) / 2 - 0.5) ** 2 +
                ((face["bbox"][1] + face["bbox"][3]) / 2 - 0.5) ** 2)[len(faces) // 2]
    elif method == "biggest":
        return sorted(faces, key=lambda face: (face["bbox"][2] - face["bbox"][0]) * (face["bbox"][3] - face["bbox"][1]))[-1]
    elif method == "smallest":
        return sorted(faces, key=lambda face: (face["bbox"][2] - face["bbox"][0]) * (face["bbox"][3] - face["bbox"][1]))[0]


def analyse_face(image, model, return_single_face=True, detect_condition="best detection", scale=1.0):
    faces = model.get(image)
    if scale != 1: # landmark-scale
        for i, face in enumerate(faces):
            landmark = face['kps']
            center = np.mean(landmark, axis=0)
            landmark = center + (landmark - center) * scale
            faces[i]['kps'] = landmark

    if not return_single_face:
        return faces

    return get_single_face(faces, method=detect_condition)


def cosine_distance(a, b):
    a /= np.linalg.norm(a)
    b /= np.linalg.norm(b)
    return 1 - np.dot(a, b)


def get_analysed_data(face_analyser, image_sequence, source_data, swap_condition="All face", detect_condition="left most", scale=1.0):
    if swap_condition != "Specific Face":
        source_path, age = source_data
        source_image = cv2.imread(source_path)
        analysed_source = analyse_face(source_image, face_analyser, return_single_face=True, detect_condition=detect_condition, scale=scale)
    else:
        analysed_source_specifics = []
        source_specifics, threshold = source_data
        for source, specific in zip(*source_specifics):
            if source is None or specific is None:
                continue
            analysed_source = analyse_face(source, face_analyser, return_single_face=True, detect_condition=detect_condition, scale=scale)
            analysed_specific = analyse_face(specific, face_analyser, return_single_face=True, detect_condition=detect_condition, scale=scale)
            analysed_source_specifics.append([analysed_source, analysed_specific])

    analysed_target_list = []
    analysed_source_list = []
    whole_frame_eql_list = []
    num_faces_per_frame = []

    total_frames = len(image_sequence)
    curr_idx = 0
    for curr_idx, frame_path in tqdm(enumerate(image_sequence), total=total_frames, desc="Analysing face data"):
        frame = cv2.imread(frame_path)
        analysed_faces = analyse_face(frame, face_analyser, return_single_face=False, detect_condition=detect_condition, scale=scale)

        n_faces = 0
        for analysed_face in analysed_faces:
            if swap_condition == "All Face":
                analysed_target_list.append(analysed_face)
                analysed_source_list.append(analysed_source)
                whole_frame_eql_list.append(frame_path)
                n_faces += 1
            elif swap_condition == "Age less than" and analysed_face["age"] < age:
                analysed_target_list.append(analysed_face)
                analysed_source_list.append(analysed_source)
                whole_frame_eql_list.append(frame_path)
                n_faces += 1
            elif swap_condition == "Age greater than" and analysed_face["age"] > age:
                analysed_target_list.append(analysed_face)
                analysed_source_list.append(analysed_source)
                whole_frame_eql_list.append(frame_path)
                n_faces += 1
            elif swap_condition == "All Male" and analysed_face["gender"] == 1:
                analysed_target_list.append(analysed_face)
                analysed_source_list.append(analysed_source)
                whole_frame_eql_list.append(frame_path)
                n_faces += 1
            elif swap_condition == "All Female" and analysed_face["gender"] == 0:
                analysed_target_list.append(analysed_face)
                analysed_source_list.append(analysed_source)
                whole_frame_eql_list.append(frame_path)
                n_faces += 1
            elif swap_condition == "Specific Face":
                for analysed_source, analysed_specific in analysed_source_specifics:
                    distance = cosine_distance(analysed_specific["embedding"], analysed_face["embedding"])
                    if distance < threshold:
                        analysed_target_list.append(analysed_face)
                        analysed_source_list.append(analysed_source)
                        whole_frame_eql_list.append(frame_path)
                        n_faces += 1

        if swap_condition == "Left Most":
            analysed_face = get_single_face(analysed_faces, method="left most")
            analysed_target_list.append(analysed_face)
            analysed_source_list.append(analysed_source)
            whole_frame_eql_list.append(frame_path)
            n_faces += 1

        elif swap_condition == "Right Most":
            analysed_face = get_single_face(analysed_faces, method="right most")
            analysed_target_list.append(analysed_face)
            analysed_source_list.append(analysed_source)
            whole_frame_eql_list.append(frame_path)
            n_faces += 1

        elif swap_condition == "Top Most":
            analysed_face = get_single_face(analysed_faces, method="top most")
            analysed_target_list.append(analysed_face)
            analysed_source_list.append(analysed_source)
            whole_frame_eql_list.append(frame_path)
            n_faces += 1

        elif swap_condition == "Bottom Most":
            analysed_face = get_single_face(analysed_faces, method="bottom most")
            analysed_target_list.append(analysed_face)
            analysed_source_list.append(analysed_source)
            whole_frame_eql_list.append(frame_path)
            n_faces += 1

        elif swap_condition == "Middle":
            analysed_face = get_single_face(analysed_faces, method="middle")
            analysed_target_list.append(analysed_face)
            analysed_source_list.append(analysed_source)
            whole_frame_eql_list.append(frame_path)
            n_faces += 1

        elif swap_condition == "Biggest":
            analysed_face = get_single_face(analysed_faces, method="biggest")
            analysed_target_list.append(analysed_face)
            analysed_source_list.append(analysed_source)
            whole_frame_eql_list.append(frame_path)
            n_faces += 1

        elif swap_condition == "Smallest":
            analysed_face = get_single_face(analysed_faces, method="smallest")
            analysed_target_list.append(analysed_face)
            analysed_source_list.append(analysed_source)
            whole_frame_eql_list.append(frame_path)
            n_faces += 1

        num_faces_per_frame.append(n_faces)

    return analysed_target_list, analysed_source_list, whole_frame_eql_list, num_faces_per_frame


#### FACE_ANALYSER.PY CODE END ###

#### UTILS.PY CODE START ###


import os
import cv2
import time
import glob
import shutil
import platform
import datetime
import subprocess
import numpy as np
from threading import Thread
from moviepy.editor import VideoFileClip, ImageSequenceClip
from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip


logo_image = cv2.imread("./assets/images/logo.png", cv2.IMREAD_UNCHANGED)


quality_types = ["poor", "low", "medium", "high", "best"]


bitrate_quality_by_resolution = {
    240: {"poor": "300k", "low": "500k", "medium": "800k", "high": "1000k", "best": "1200k"},
    360: {"poor": "500k","low": "800k","medium": "1200k","high": "1500k","best": "2000k"},
    480: {"poor": "800k","low": "1200k","medium": "2000k","high": "2500k","best": "3000k"},
    720: {"poor": "1500k","low": "2500k","medium": "4000k","high": "5000k","best": "6000k"},
    1080: {"poor": "2500k","low": "4000k","medium": "6000k","high": "7000k","best": "8000k"},
    1440: {"poor": "4000k","low": "6000k","medium": "8000k","high": "10000k","best": "12000k"},
    2160: {"poor": "8000k","low": "10000k","medium": "12000k","high": "15000k","best": "20000k"}
}


crf_quality_by_resolution = {
    240: {"poor": 45, "low": 35, "medium": 28, "high": 23, "best": 20},
    360: {"poor": 35, "low": 28, "medium": 23, "high": 20, "best": 18},
    480: {"poor": 28, "low": 23, "medium": 20, "high": 18, "best": 16},
    720: {"poor": 23, "low": 20, "medium": 18, "high": 16, "best": 14},
    1080: {"poor": 20, "low": 18, "medium": 16, "high": 14, "best": 12},
    1440: {"poor": 18, "low": 16, "medium": 14, "high": 12, "best": 10},
    2160: {"poor": 16, "low": 14, "medium": 12, "high": 10, "best": 8}
}


def get_bitrate_for_resolution(resolution, quality):
    available_resolutions = list(bitrate_quality_by_resolution.keys())
    closest_resolution = min(available_resolutions, key=lambda x: abs(x - resolution))
    return bitrate_quality_by_resolution[closest_resolution][quality]


def get_crf_for_resolution(resolution, quality):
    available_resolutions = list(crf_quality_by_resolution.keys())
    closest_resolution = min(available_resolutions, key=lambda x: abs(x - resolution))
    return crf_quality_by_resolution[closest_resolution][quality]


def get_video_bitrate(video_file):
    ffprobe_cmd = ['ffprobe', '-v', 'error', '-select_streams', 'v:0', '-show_entries',
        'stream=bit_rate', '-of', 'default=noprint_wrappers=1:nokey=1', video_file]
    result = subprocess.run(ffprobe_cmd, stdout=subprocess.PIPE)
    kbps = max(int(result.stdout) // 1000, 10)
    return str(kbps) + 'k'


def trim_video(video_path, output_path, start_frame, stop_frame):
    video_name, _ = os.path.splitext(os.path.basename(video_path))
    trimmed_video_filename = video_name + "_trimmed" + ".mp4"
    temp_path = os.path.join(output_path, "trim")
    os.makedirs(temp_path, exist_ok=True)
    trimmed_video_file_path = os.path.join(temp_path, trimmed_video_filename)

    video = VideoFileClip(video_path, fps_source="fps")
    fps = video.fps
    start_time = start_frame / fps
    duration = (stop_frame - start_frame) / fps

    bitrate = get_bitrate_for_resolution(min(*video.size), "high")

    trimmed_video = video.subclip(start_time, start_time + duration)
    trimmed_video.write_videofile(
        trimmed_video_file_path, codec="libx264", audio_codec="aac", bitrate=bitrate,
    )
    trimmed_video.close()
    video.close()

    return trimmed_video_file_path


def open_directory(path=None):
    if path is None:
        return
    try:
        os.startfile(path)
    except:
        subprocess.Popen(["xdg-open", path])


class StreamerThread(object):
    def __init__(self, src=0):
        self.capture = cv2.VideoCapture(src)
        self.capture.set(cv2.CAP_PROP_BUFFERSIZE, 2)
        self.FPS = 1 / 30
        self.FPS_MS = int(self.FPS * 1000)
        self.thread = None
        self.stopped = False
        self.frame = None

    def start(self):
        self.thread = Thread(target=self.update, args=())
        self.thread.daemon = True
        self.thread.start()

    def stop(self):
        self.stopped = True
        self.thread.join()
        print("stopped")

    def update(self):
        while not self.stopped:
            if self.capture.isOpened():
                (self.status, self.frame) = self.capture.read()
            time.sleep(self.FPS)


class ProcessBar:
    def __init__(self, bar_length, total, before="⬛", after="🟨"):
        self.bar_length = bar_length
        self.total = total
        self.before = before
        self.after = after
        self.bar = [self.before] * bar_length
        self.start_time = time.time()

    def get(self, index):
        total = self.total
        elapsed_time = time.time() - self.start_time
        average_time_per_iteration = elapsed_time / (index + 1)
        remaining_iterations = total - (index + 1)
        estimated_remaining_time = remaining_iterations * average_time_per_iteration

        self.bar[int(index / total * self.bar_length)] = self.after
        info_text = f"({index+1}/{total}) {''.join(self.bar)} "
        info_text += f"(ETR: {int(estimated_remaining_time // 60)} min {int(estimated_remaining_time % 60)} sec)"
        return info_text


def add_logo_to_image(img, logo=logo_image):
    logo_size = int(img.shape[1] * 0.1)
    logo = cv2.resize(logo, (logo_size, logo_size))
    if logo.shape[2] == 4:
        alpha = logo[:, :, 3]
    else:
        alpha = np.ones_like(logo[:, :, 0]) * 255
    padding = int(logo_size * 0.1)
    roi = img.shape[0] - logo_size - padding, img.shape[1] - logo_size - padding
    for c in range(0, 3):
        img[roi[0] : roi[0] + logo_size, roi[1] : roi[1] + logo_size, c] = (
            alpha / 255.0
        ) * logo[:, :, c] + (1 - alpha / 255.0) * img[
            roi[0] : roi[0] + logo_size, roi[1] : roi[1] + logo_size, c
        ]
    return img


def split_list_by_lengths(data, length_list):
    split_data = []
    start_idx = 0
    for length in length_list:
        end_idx = start_idx + length
        sublist = data[start_idx:end_idx]
        split_data.append(sublist)
        start_idx = end_idx
    return split_data


def merge_img_sequence_from_ref(ref_video_path, image_sequence, output_file_name):
    video_clip = VideoFileClip(ref_video_path, fps_source="fps")
    fps = video_clip.fps
    duration = video_clip.duration
    total_frames = video_clip.reader.nframes
    audio_clip = video_clip.audio if video_clip.audio is not None else None
    edited_video_clip = ImageSequenceClip(image_sequence, fps=fps)

    if audio_clip is not None:
        edited_video_clip = edited_video_clip.set_audio(audio_clip)

    bitrate = get_bitrate_for_resolution(min(*edited_video_clip.size), "high")

    edited_video_clip.set_duration(duration).write_videofile(
        output_file_name, codec="libx264", bitrate=bitrate,
    )
    edited_video_clip.close()
    video_clip.close()


def scale_bbox_from_center(bbox, scale_width, scale_height, image_width, image_height):
    # Extract the coordinates of the bbox
    x1, y1, x2, y2 = bbox

    # Calculate the center point of the bbox
    center_x = (x1 + x2) / 2
    center_y = (y1 + y2) / 2

    # Calculate the new width and height of the bbox based on the scaling factors
    width = x2 - x1
    height = y2 - y1
    new_width = width * scale_width
    new_height = height * scale_height

    # Calculate the new coordinates of the bbox, considering the image boundaries
    new_x1 = center_x - new_width / 2
    new_y1 = center_y - new_height / 2
    new_x2 = center_x + new_width / 2
    new_y2 = center_y + new_height / 2

    # Adjust the coordinates to ensure the bbox remains within the image boundaries
    new_x1 = max(0, new_x1)
    new_y1 = max(0, new_y1)
    new_x2 = min(image_width - 1, new_x2)
    new_y2 = min(image_height - 1, new_y2)

    # Return the scaled bbox coordinates
    scaled_bbox = [new_x1, new_y1, new_x2, new_y2]
    return scaled_bbox


def laplacian_blending(A, B, m, num_levels=7):
    assert A.shape == B.shape
    assert B.shape == m.shape
    height = m.shape[0]
    width = m.shape[1]
    size_list = np.array([4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192])
    size = size_list[np.where(size_list > max(height, width))][0]
    GA = np.zeros((size, size, 3), dtype=np.float32)
    GA[:height, :width, :] = A
    GB = np.zeros((size, size, 3), dtype=np.float32)
    GB[:height, :width, :] = B
    GM = np.zeros((size, size, 3), dtype=np.float32)
    GM[:height, :width, :] = m
    gpA = [GA]
    gpB = [GB]
    gpM = [GM]
    for i in range(num_levels):
        GA = cv2.pyrDown(GA)
        GB = cv2.pyrDown(GB)
        GM = cv2.pyrDown(GM)
        gpA.append(np.float32(GA))
        gpB.append(np.float32(GB))
        gpM.append(np.float32(GM))
    lpA  = [gpA[num_levels-1]]
    lpB  = [gpB[num_levels-1]]
    gpMr = [gpM[num_levels-1]]
    for i in range(num_levels-1,0,-1):
        LA = np.subtract(gpA[i-1], cv2.pyrUp(gpA[i]))
        LB = np.subtract(gpB[i-1], cv2.pyrUp(gpB[i]))
        lpA.append(LA)
        lpB.append(LB)
        gpMr.append(gpM[i-1])
    LS = []
    for la,lb,gm in zip(lpA,lpB,gpMr):
        ls = la * gm + lb * (1.0 - gm)
        LS.append(ls)
    ls_ = LS[0]
    for i in range(1,num_levels):
        ls_ = cv2.pyrUp(ls_)
        ls_ = cv2.add(ls_, LS[i])
    ls_ = ls_[:height, :width, :]
    #ls_ = (ls_ - np.min(ls_)) * (255.0 / (np.max(ls_) - np.min(ls_)))
    return ls_.clip(0, 255)


def mask_crop(mask, crop):
    top, bottom, left, right = crop
    shape = mask.shape
    top = int(top)
    bottom = int(bottom)
    if top + bottom < shape[1]:
        if top > 0: mask[:top, :] = 0
        if bottom > 0: mask[-bottom:, :] = 0

    left = int(left)
    right = int(right)
    if left + right < shape[0]:
        if left > 0: mask[:, :left] = 0
        if right > 0: mask[:, -right:] = 0

    return mask

def create_image_grid(images, size=128):
    num_images = len(images)
    num_cols = int(np.ceil(np.sqrt(num_images)))
    num_rows = int(np.ceil(num_images / num_cols))
    grid = np.zeros((num_rows * size, num_cols * size, 3), dtype=np.uint8)

    for i, image in enumerate(images):
        row_idx = (i // num_cols) * size
        col_idx = (i % num_cols) * size
        image = cv2.resize(image.copy(), (size,size))
        if image.dtype != np.uint8:
            image = (image.astype('float32') * 255).astype('uint8')
        if image.ndim == 2:
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
        grid[row_idx:row_idx + size, col_idx:col_idx + size] = image

    return grid


#### UTILS.PY CODE END ###

#### APP.PY CODE END ###

import os
import spaces
import cv2
import glob
import time
import torch
import shutil
import argparse
import platform
import datetime
import subprocess
import insightface
import onnxruntime
import numpy as np
import gradio as gr
import threading
import queue
from tqdm import tqdm
import concurrent.futures
from moviepy.editor import VideoFileClip

from nsfw_checker import NSFWChecker
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, StreamerThread, ProcessBar, open_directory, split_list_by_lengths, merge_img_sequence_from_ref, create_image_grid

## ------------------------------ USER ARGS ------------------------------

parser = argparse.ArgumentParser(description="Swap-Mukham Face Swapper")
parser.add_argument("--out_dir", help="Default Output directory", default=os.getcwd())
parser.add_argument("--batch_size", help="Gpu batch size", default=32)
parser.add_argument("--cuda", action="store_true", help="Enable cuda", default=False)
parser.add_argument(
    "--colab", action="store_true", help="Enable colab mode", default=False
)
user_args = parser.parse_args()

## ------------------------------ DEFAULTS ------------------------------

USE_COLAB = user_args.colab
USE_CUDA = user_args.cuda
DEF_OUTPUT_PATH = user_args.out_dir
BATCH_SIZE = int(user_args.batch_size)
WORKSPACE = None
OUTPUT_FILE = None
CURRENT_FRAME = 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
NSFW_DETECTOR = None
FACE_ENHANCER_LIST = ["NONE"]
FACE_ENHANCER_LIST.extend(get_available_enhancer_names())
FACE_ENHANCER_LIST.extend(cv2_interpolations)

## ------------------------------ SET EXECUTION PROVIDER ------------------------------
# Note: Non CUDA users may change settings here

PROVIDER = ["CPUExecutionProvider"]

if USE_CUDA:
    available_providers = onnxruntime.get_available_providers()
    if "CUDAExecutionProvider" in available_providers:
        print("\n********** Running on CUDA **********\n")
        PROVIDER = ["CUDAExecutionProvider", "CPUExecutionProvider"]
    else:
        USE_CUDA = False
        print("\n********** CUDA unavailable running on CPU **********\n")
else:
    USE_CUDA = False
    print("\n********** Running on CPU **********\n")

device = "cuda" if USE_CUDA else "cpu"
EMPTY_CACHE = lambda: torch.cuda.empty_cache() if device == "cuda" else None

## ------------------------------ LOAD MODELS ------------------------------

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:
        batch = int(BATCH_SIZE) if device == "cuda" else 1
        FACE_SWAPPER = Inswapper(model_file=path, batch_size=batch, 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)

def load_nsfw_detector_model(path="./assets/pretrained_models/open-nsfw.onnx"):
    global NSFW_DETECTOR
    if NSFW_DETECTOR is None:
        NSFW_DETECTOR = NSFWChecker(model_path=path, providers=PROVIDER)


load_face_analyser_model()
load_face_swapper_model()

## ------------------------------ MAIN PROCESS ------------------------------


@spaces.GPU(duration=300, enable_queue=True)
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
    global OUTPUT_FILE
    global PREVIEW
    WORKSPACE, OUTPUT_FILE, PREVIEW = None, None, None

    ## ------------------------------ GUI UPDATE FUNC ------------------------------

    def ui_before():
        return (
            gr.update(visible=True, value=PREVIEW),
            gr.update(interactive=False),
            gr.update(interactive=False),
            gr.update(visible=False),
        )

    def ui_after():
        return (
            gr.update(visible=True, value=PREVIEW),
            gr.update(interactive=True),
            gr.update(interactive=True),
            gr.update(visible=False),
        )

    def ui_after_vid():
        return (
            gr.update(visible=False),
            gr.update(interactive=True),
            gr.update(interactive=True),
            gr.update(value=OUTPUT_FILE, visible=True),
        )

    start_time = time.time()
    total_exec_time = lambda start_time: divmod(time.time() - start_time, 60)
    get_finsh_text = lambda start_time: f"✔️ Completed in {int(total_exec_time(start_time)[0])} min {int(total_exec_time(start_time)[1])} sec."

    ## ------------------------------ PREPARE INPUTS & LOAD MODELS ------------------------------

    yield "### \n ⌛ Loading NSFW detector model...", *ui_before()
    load_nsfw_detector_model()

    yield "### \n ⌛ Loading face analyser model...", *ui_before()
    load_face_analyser_model()

    yield "### \n ⌛ Loading face swapper model...", *ui_before()
    load_face_swapper_model()

    if face_enhancer_name != "NONE":
        if face_enhancer_name not in cv2_interpolations:
            yield f"### \n ⌛ Loading {face_enhancer_name} model...", *ui_before()
        FACE_ENHANCER = load_face_enhancer_model(name=face_enhancer_name, device=device)
    else:
        FACE_ENHANCER = None

    if enable_face_parser:
        yield "### \n ⌛ Loading face parsing model...", *ui_before()
        load_face_parser_model()

    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)

    def swap_process(image_sequence):
        ## ------------------------------ CONTENT CHECK ------------------------------

        yield "### \n ⌛ Checking contents...", *ui_before()
        nsfw = NSFW_DETECTOR.is_nsfw(image_sequence)
        if nsfw:
            message = "NSFW Content detected !!!"
            yield f"### \n 🔞 {message}", *ui_before()
            assert not nsfw, message
            return False
        EMPTY_CACHE()

        ## ------------------------------ ANALYSE FACE ------------------------------

        yield "### \n ⌛ Analysing face data...", *ui_before()
        if condition != "Specific Face":
            source_data = source_path, age
        else:
            source_data = ((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
        )

        ## ------------------------------ SWAP FUNC ------------------------------

        yield "### \n ⌛ Generating faces...", *ui_before()
        preds = []
        matrs = []
        count = 0
        global PREVIEW
        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()
            count += 1

            if USE_CUDA:
                image_grid = create_image_grid(batch_pred, size=128)
                PREVIEW = image_grid[:, :, ::-1]
                yield f"### \n ⌛ Generating face Batch {count}", *ui_before()

        ## ------------------------------ FACE ENHANCEMENT ------------------------------

        generated_len = len(preds)
        if face_enhancer_name != "NONE":
            yield f"### \n ⌛ Upscaling faces with {face_enhancer_name}...", *ui_before()
            for idx, pred in tqdm(enumerate(preds), total=generated_len, desc=f"Upscaling with {face_enhancer_name}"):
                enhancer_model, enhancer_model_runner = FACE_ENHANCER
                pred = enhancer_model_runner(pred, enhancer_model)
                preds[idx] = cv2.resize(pred, (512,512))
        EMPTY_CACHE()

        ## ------------------------------ FACE PARSING ------------------------------

        if enable_face_parser:
            yield "### \n ⌛ Face-parsing mask...", *ui_before()
            masks = []
            count = 0
            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()
                count += 1

                if len(batch_mask) > 1:
                    image_grid = create_image_grid(batch_mask, size=128)
                    PREVIEW = image_grid[:, :, ::-1]
                    yield f"### \n ⌛ Face parsing Batch {count}", *ui_before()
            masks = np.concatenate(masks, axis=0) if len(masks) >= 1 else masks
        else:
            masks = [None] * generated_len

        ## ------------------------------ SPLIT LIST ------------------------------

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

        ## ------------------------------ PASTE-BACK ------------------------------

        yield "### \n ⌛ Pasting back...", *ui_before()
        def post_process(frame_idx, frame_img, split_preds, split_matrs, split_masks, enable_laplacian_blend, crop_mask, blur_amount, erode_amount):
            whole_img_path = frame_img
            whole_img = cv2.imread(whole_img_path)
            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(whole_img_path, whole_img)

        def concurrent_post_process(image_sequence, *args):
            with concurrent.futures.ThreadPoolExecutor() as executor:
                futures = []
                for idx, frame_img in enumerate(image_sequence):
                    future = executor.submit(post_process, idx, frame_img, *args)
                    futures.append(future)

                for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="Pasting back"):
                    result = future.result()

        concurrent_post_process(
            image_sequence,
            split_preds,
            split_matrs,
            split_masks,
            enable_laplacian_blend,
            crop_mask,
            blur_amount,
            erode_amount
        )


    ## ------------------------------ IMAGE ------------------------------

    if input_type == "Image":
        target = cv2.imread(image_path)
        output_file = os.path.join(output_path, output_name + ".png")
        cv2.imwrite(output_file, target)

        for info_update in swap_process([output_file]):
            yield info_update

        OUTPUT_FILE = output_file
        WORKSPACE = output_path
        PREVIEW = cv2.imread(output_file)[:, :, ::-1]

        yield get_finsh_text(start_time), *ui_after()

    ## ------------------------------ VIDEO ------------------------------

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

        yield "### \n ⌛ Extracting video frames...", *ui_before()
        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()
        cv2.destroyAllWindows()

        for info_update in swap_process(image_sequence):
            yield info_update

        yield "### \n ⌛ Merging sequence...", *ui_before()
        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:
            yield "### \n ⌛ Removing temporary files...", *ui_before()
            shutil.rmtree(temp_path)

        WORKSPACE = output_path
        OUTPUT_FILE = output_video_path

        yield get_finsh_text(start_time), *ui_after_vid()

    ## ------------------------------ DIRECTORY ------------------------------

    elif input_type == "Directory":
        extensions = ["jpg", "jpeg", "png", "bmp", "tiff", "ico", "webp"]
        temp_path = os.path.join(output_path, output_name)
        if os.path.exists(temp_path):
            shutil.rmtree(temp_path)
        os.mkdir(temp_path)

        file_paths =[]
        for file_path in glob.glob(os.path.join(directory_path, "*")):
            if any(file_path.lower().endswith(ext) for ext in extensions):
                img = cv2.imread(file_path)
                new_file_path = os.path.join(temp_path, os.path.basename(file_path))
                cv2.imwrite(new_file_path, img)
                file_paths.append(new_file_path)

        for info_update in swap_process(file_paths):
            yield info_update

        PREVIEW = cv2.imread(file_paths[-1])[:, :, ::-1]
        WORKSPACE = temp_path
        OUTPUT_FILE = file_paths[-1]

        yield get_finsh_text(start_time), *ui_after()

    ## ------------------------------ STREAM ------------------------------

    elif input_type == "Stream":
        pass


## ------------------------------ GRADIO FUNC ------------------------------


def update_radio(value):
    if value == "Image":
        return (
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(visible=False),
        )
    elif value == "Video":
        return (
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(visible=False),
        )
    elif value == "Directory":
        return (
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=True),
        )
    elif value == "Stream":
        return (
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=True),
        )


def swap_option_changed(value):
    if value.startswith("Age"):
        return (
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(visible=True),
        )
    elif value == "Specific Face":
        return (
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(visible=False),
        )
    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)


def video_changed(video_path):
    sliders_update = gr.Slider.update
    button_update = gr.Button.update
    number_update = gr.Number.update

    if video_path is None:
        return (
            sliders_update(minimum=0, maximum=0, value=0),
            sliders_update(minimum=1, maximum=1, value=1),
            number_update(value=1),
        )
    try:
        clip = VideoFileClip(video_path)
        fps = clip.fps
        total_frames = clip.reader.nframes
        clip.close()
        return (
            sliders_update(minimum=0, maximum=total_frames, value=0, interactive=True),
            sliders_update(
                minimum=0, maximum=total_frames, value=total_frames, interactive=True
            ),
            number_update(value=fps),
        )
    except:
        return (
            sliders_update(value=0),
            sliders_update(value=0),
            number_update(value=1),
        )


def analyse_settings_changed(detect_condition, detection_size, detection_threshold):
    yield "### \n ⌛ Applying new values..."
    global FACE_ANALYSER
    global DETECT_CONDITION
    DETECT_CONDITION = detect_condition
    FACE_ANALYSER = insightface.app.FaceAnalysis(name="buffalo_l", providers=PROVIDER)
    FACE_ANALYSER.prepare(
        ctx_id=0,
        det_size=(int(detection_size), int(detection_size)),
        det_thresh=float(detection_threshold),
    )
    yield f"### \n ✔️ Applied detect condition:{detect_condition}, detection size: {detection_size}, detection threshold: {detection_threshold}"


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


def slider_changed(show_frame, video_path, frame_index):
    if not show_frame:
        return None, None
    if video_path is None:
        return None, None
    clip = VideoFileClip(video_path)
    frame = clip.get_frame(frame_index / clip.fps)
    frame_array = np.array(frame)
    clip.close()
    return gr.Image.update(value=frame_array, visible=True), gr.Video.update(
        visible=False
    )


def trim_and_reload(video_path, output_path, output_name, start_frame, stop_frame):
    yield video_path, f"### \n ⌛ Trimming video frame {start_frame} to {stop_frame}..."
    try:
        output_path = os.path.join(output_path, output_name)
        trimmed_video = trim_video(video_path, output_path, start_frame, stop_frame)
        yield trimmed_video, "### \n ✔️ Video trimmed and reloaded."
    except Exception as e:
        print(e)
        yield video_path, "### \n ❌ Video trimming failed. See console for more info."


## ------------------------------ GRADIO GUI ------------------------------

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

with gr.Blocks(css=css) as interface:
    gr.Markdown("# 🗿 Swap Mukham")
    gr.Markdown("### Face swap app based on insightface inswapper.")
    with gr.Row():
        with gr.Row():
            with gr.Column(scale=0.4):
                with gr.Tab("📄 Swap Condition"):
                    swap_option = gr.Dropdown(
                        swap_options_list,
                        info="Choose which face or faces in the target image to swap.",
                        multiselect=False,
                        show_label=False,
                        value=swap_options_list[0],
                        interactive=True,
                    )
                    age = gr.Number(
                        value=25, label="Value", interactive=True, visible=False
                    )

                with gr.Tab("🎚️ Detection Settings"):
                    detect_condition_dropdown = gr.Dropdown(
                        detect_conditions,
                        label="Condition",
                        value=DETECT_CONDITION,
                        interactive=True,
                        info="This condition is only used when multiple faces are detected on source or specific image.",
                    )
                    detection_size = gr.Number(
                        label="Detection Size", value=DETECT_SIZE, interactive=True
                    )
                    detection_threshold = gr.Number(
                        label="Detection Threshold",
                        value=DETECT_THRESH,
                        interactive=True,
                    )
                    apply_detection_settings = gr.Button("Apply settings")

                with gr.Tab("📤 Output Settings"):
                    output_directory = gr.Text(
                        label="Output Directory",
                        value=DEF_OUTPUT_PATH,
                        interactive=True,
                    )
                    output_name = gr.Text(
                        label="Output Name", value="Result", interactive=True
                    )
                    keep_output_sequence = gr.Checkbox(
                        label="Keep output sequence", value=False, interactive=True
                    )

                with gr.Tab("🪄 Other Settings"):
                    face_scale = gr.Slider(
                        label="Face Scale",
                        minimum=0,
                        maximum=2,
                        value=1,
                        interactive=True,
                    )

                    face_enhancer_name = gr.Dropdown(
                        FACE_ENHANCER_LIST, label="Face Enhancer", value="NONE", multiselect=False, interactive=True
                    )

                    with gr.Accordion("Advanced Mask", open=False):
                        enable_face_parser_mask = gr.Checkbox(
                            label="Enable Face Parsing",
                            value=False,
                            interactive=True,
                        )

                        mask_include = gr.Dropdown(
                            mask_regions.keys(),
                            value=MASK_INCLUDE,
                            multiselect=True,
                            label="Include",
                            interactive=True,
                        )
                        mask_soft_kernel = gr.Number(
                            label="Soft Erode Kernel",
                            value=MASK_SOFT_KERNEL,
                            minimum=3,
                            interactive=True,
                            visible = False
                        )
                        mask_soft_iterations = gr.Number(
                            label="Soft Erode Iterations",
                            value=MASK_SOFT_ITERATIONS,
                            minimum=0,
                            interactive=True,

                        )


                    with gr.Accordion("Crop Mask", open=False):
                        crop_top = gr.Slider(label="Top", minimum=0, maximum=511, value=0, step=1, interactive=True)
                        crop_bott = gr.Slider(label="Bottom", minimum=0, maximum=511, value=511, step=1, interactive=True)
                        crop_left = gr.Slider(label="Left", minimum=0, maximum=511, value=0, step=1, interactive=True)
                        crop_right = gr.Slider(label="Right", minimum=0, maximum=511, value=511, step=1, interactive=True)


                    erode_amount = gr.Slider(
                            label="Mask Erode",
                            minimum=0,
                            maximum=1,
                            value=MASK_ERODE_AMOUNT,
                            step=0.05,
                            interactive=True,
                        )

                    blur_amount = gr.Slider(
                            label="Mask Blur",
                            minimum=0,
                            maximum=1,
                            value=MASK_BLUR_AMOUNT,
                            step=0.05,
                            interactive=True,
                        )

                    enable_laplacian_blend = gr.Checkbox(
                        label="Laplacian Blending",
                        value=True,
                        interactive=True,
                    )


                source_image_input = gr.Image(
                    label="Source face", type="filepath", interactive=True
                )

                with gr.Group(visible=False) as specific_face:
                    for i in range(NUM_OF_SRC_SPECIFIC):
                        idx = i + 1
                        code = "\n"
                        code += f"with gr.Tab(label='({idx})'):"
                        code += "\n\twith gr.Row():"
                        code += f"\n\t\tsrc{idx} = gr.Image(interactive=True, type='numpy', label='Source Face {idx}')"
                        code += f"\n\t\ttrg{idx} = gr.Image(interactive=True, type='numpy', label='Specific Face {idx}')"
                        exec(code)

                    distance_slider = gr.Slider(
                        minimum=0,
                        maximum=2,
                        value=0.6,
                        interactive=True,
                        label="Distance",
                        info="Lower distance is more similar and higher distance is less similar to the target face.",
                    )

                with gr.Group():
                    input_type = gr.Radio(
                        ["Image", "Video"],
                        label="Target Type",
                        value="Image",
                    )

                    with gr.Group(visible=True) as input_image_group:
                        image_input = gr.Image(
                            label="Target Image", interactive=True, type="filepath"
                        )

                    with gr.Group(visible=False) as input_video_group:
                        vid_widget = gr.Video if USE_COLAB else gr.Text
                        video_input = gr.Video(
                            label="Target Video", interactive=True
                        )
                        with gr.Accordion("✂️ Trim video", open=False):
                            with gr.Column():
                                with gr.Row():
                                    set_slider_range_btn = gr.Button(
                                        "Set frame range", interactive=True
                                    )
                                    show_trim_preview_btn = gr.Checkbox(
                                        label="Show frame when slider change",
                                        value=True,
                                        interactive=True,
                                    )

                                video_fps = gr.Number(
                                    value=30,
                                    interactive=False,
                                    label="Fps",
                                    visible=False,
                                )
                                start_frame = gr.Slider(
                                    minimum=0,
                                    maximum=1,
                                    value=0,
                                    step=1,
                                    interactive=True,
                                    label="Start Frame",
                                    info="",
                                )
                                end_frame = gr.Slider(
                                    minimum=0,
                                    maximum=1,
                                    value=1,
                                    step=1,
                                    interactive=True,
                                    label="End Frame",
                                    info="",
                                )
                            trim_and_reload_btn = gr.Button(
                                "Trim and Reload", interactive=True
                            )

                    with gr.Group(visible=False) as input_directory_group:
                        direc_input = gr.Text(label="Path", interactive=True)

            with gr.Column(scale=0.6):
                info = gr.Markdown(value="...")

                with gr.Row():
                    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
                )

                with gr.Row():
                    output_directory_button = gr.Button(
                        "📂", interactive=False, visible=False
                    )
                    output_video_button = gr.Button(
                        "🎬", interactive=False, visible=False
                    )

                with gr.Group():
                    with gr.Row():
                        gr.Markdown(
                            "### [🤝 Sponsor](https://github.com/sponsors/harisreedhar)"
                        )
                        gr.Markdown(
                            "### [👨‍💻 Source code](https://github.com/harisreedhar/Swap-Mukham)"
                        )
                        gr.Markdown(
                            "### [⚠️ Disclaimer](https://github.com/harisreedhar/Swap-Mukham#disclaimer)"
                        )
                        gr.Markdown(
                            "### [🌐 Run in Colab](https://colab.research.google.com/github/harisreedhar/Swap-Mukham/blob/main/swap_mukham_colab.ipynb)"
                        )
                        gr.Markdown(
                            "### [🤗 Acknowledgements](https://github.com/harisreedhar/Swap-Mukham#acknowledgements)"
                        )

    ## ------------------------------ GRADIO EVENTS ------------------------------

    set_slider_range_event = set_slider_range_btn.click(
        video_changed,
        inputs=[video_input],
        outputs=[start_frame, end_frame, video_fps],
    )

    trim_and_reload_event = trim_and_reload_btn.click(
        fn=trim_and_reload,
        inputs=[video_input, output_directory, output_name, start_frame, end_frame],
        outputs=[video_input, info],
    )

    start_frame_event = start_frame.release(
        fn=slider_changed,
        inputs=[show_trim_preview_btn, video_input, start_frame],
        outputs=[preview_image, preview_video],
        show_progress=True,
    )

    end_frame_event = end_frame.release(
        fn=slider_changed,
        inputs=[show_trim_preview_btn, video_input, end_frame],
        outputs=[preview_image, preview_video],
        show_progress=True,
    )

    input_type.change(
        update_radio,
        inputs=[input_type],
        outputs=[input_image_group, input_video_group, input_directory_group],
    )
    swap_option.change(
        swap_option_changed,
        inputs=[swap_option],
        outputs=[age, specific_face, source_image_input],
    )

    apply_detection_settings.click(
        analyse_settings_changed,
        inputs=[detect_condition_dropdown, detection_size, detection_threshold],
        outputs=[info],
    )

    src_specific_inputs = []
    gen_variable_txt = ",".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)]
    )
    exec(f"src_specific_inputs = ({gen_variable_txt})")
    swap_inputs = [
        input_type,
        image_input,
        video_input,
        direc_input,
        source_image_input,
        output_directory,
        output_name,
        keep_output_sequence,
        swap_option,
        age,
        distance_slider,
        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_outputs = [
        info,
        preview_image,
        output_directory_button,
        output_video_button,
        preview_video,
    ]

    swap_event = swap_button.click(
        fn=process, inputs=swap_inputs, outputs=swap_outputs, show_progress=True
    )

    cancel_button.click(
        fn=stop_running,
        inputs=None,
        outputs=[info],
        cancels=[
            swap_event,
            trim_and_reload_event,
            set_slider_range_event,
            start_frame_event,
            end_frame_event,
        ],
        show_progress=True,
    )
    output_directory_button.click(
        lambda: open_directory(path=WORKSPACE), inputs=None, outputs=None
    )
    output_video_button.click(
        lambda: open_directory(path=OUTPUT_FILE), inputs=None, outputs=None
    )

if __name__ == "__main__":
    if USE_COLAB:
        print("Running in colab mode")

    interface.launch()


#### APP.PY CODE END ###