import gradio as gr
import os
from run import run_faceswap

def process_face_swap(source_image, target_image, target_video):
    source_path = "face.png"
    target_path = "target"

    # Save source image
    source_image.save(source_path)

    # Use whichever is provided (image or video)
    if target_image is not None:
        ext = ".jpg"
        target_image.save(target_path + ext)
        used_path = target_path + ext
    elif target_video is not None:
        ext = ".mp4"
        target_video.save(target_path + ext)
        used_path = target_path + ext
    else:
        raise ValueError("Please upload either an image or a video.")

    output_path = run_faceswap(source_path, used_path)
    return output_path

with gr.Blocks() as demo:
    gr.Markdown("## Lightweight Roop Face Swap")
    gr.Markdown("Upload face.png and target image or video. Output will be result.jpg or result.mp4")

    with gr.Row():
        source_image = gr.Image(label="Source Face (face.png)", type="pil")

    with gr.Row():
        with gr.Tab("Target Image"):
            target_image = gr.Image(label="Upload Target Image (.jpg)", type="pil")
        with gr.Tab("Target Video"):
            target_video = gr.Video(label="Upload Target Video (.mp4)")

    output_file = gr.File(label="Swapped Output")

    submit_btn = gr.Button("Submit")

    submit_btn.click(
        fn=process_face_swap,
        inputs=[source_image, target_image, target_video],
        outputs=output_file
    )

demo.launch()