#######################################################################################
#
# MIT License
#
# Copyright (c) [2025] [leonelhs@gmail.com]
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
#######################################################################################
#
# Source code is based on or inspired by several projects.
# For more details and proper attribution, please refer to the following resources:
#
# - [hyxue] - [https://huggingface.co/spaces/hyxue/HiFiFace-inference-demo]
# - [maum-ai] [https://github.com/maum-ai/hififace]
#

import gradio as gr
import torch
from huggingface_hub import hf_hub_download

from benchmark.app_image import ImageSwap
from models.model import HifiFaceST, HifiFaceWGM

REPO_ID = "leonelhs/HiFiFace"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

gen_st_path = hf_hub_download(repo_id=REPO_ID,
                              filename="hififace_pretrained/standard_model/generator_320000.pth")

gen_wgm_path = hf_hub_download(repo_id=REPO_ID,
                               filename="hififace_pretrained/with_gaze_and_mouth/generator_190000.pth")

fade_detector_path = hf_hub_download(repo_id=REPO_ID,
                                     filename="face_detector/face_detector_scrfd_10g_bnkps.onnx")

identity_extractor_config = {
    "f_3d_checkpoint_path": hf_hub_download(repo_id=REPO_ID, filename="Deep3DFaceRecon/epoch_20.pth"),
    "f_id_checkpoint_path": hf_hub_download(repo_id=REPO_ID, filename="arcface/ms1mv3_arcface_r100_fp16_backbone.pth")
}

class ConfigPath:
    face_detector_weights = fade_detector_path
    model_path = ""
    model_idx = 80000
    ffmpeg_device = device
    device = device

cfg = ConfigPath()

model_standard = HifiFaceST(identity_extractor_config, device=device, generator_path=gen_st_path)

model_wgm = HifiFaceWGM(identity_extractor_config, device=device, generator_path=gen_wgm_path)

image_infer_standard = ImageSwap(cfg, model_standard)
image_infer_wgm = ImageSwap(cfg, model_wgm)

MODELS = {
    "Standard model": "standard",
    "Eye and mouth hm loss": "eyeandmouth",
}

def inference_image(source_face, target_face, method="standard", shape_rate=1.0, id_rate=1.0, iterations=1):
    if method == "standard":
        return target_face, image_infer_standard.inference(source_face, target_face, shape_rate, id_rate, int(iterations))
    return target_face, image_infer_wgm.inference(source_face, target_face, shape_rate, id_rate, int(iterations))


with gr.Blocks(title="FaceSwap") as app:
    gr.Markdown("## HiFiFace image swap")
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Row():
                source_image = gr.Image(type="numpy", label="Face image")
                target_image = gr.Image(type="numpy", label="Body image")
            mod = gr.Dropdown(choices=list(MODELS.items()), label="Model generator", value="standard")
            image_btn = gr.Button("Swap image")
            with gr.Accordion("Fine tunes", open=False):
                structure_sim = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, step=0.1, label="3d similarity")
                id_sim = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, step=0.1, label="id similarity")
                iters = gr.Slider(minimum=1, maximum=10, value=1, step=1, label="iters")
        with gr.Column(scale=1):
            with gr.Row():
                output_image = gr.ImageSlider(label="Swapped image", type="pil")

    image_btn.click(
        fn=inference_image,
        inputs=[source_image, target_image, mod, structure_sim, id_sim, iters],
        outputs=output_image,
    )

app.launch(share=False, debug=True, show_error=True, mcp_server=True, pwa=True)
app.queue()