import os
import gradio as gr
import subprocess

title = "DaGAN Demo"
description = "reconstruction and animation demos for DaGAN"

def inference(mode, img, video):
    os.makedirs('tmp', exist_ok=True)

    cmd = f"ffmpeg -y -ss 00:00:00 -i {video} -to 00:00:05 -c copy cut_vid.mp4"
    subprocess.run(cmd.split())

    video = "cut_vid.mp4"
    os.system(f"python run.py --source_image {img} --driving_video {video} --result_video tmp/result.mp4 --mode {mode}")
    return "tmp/result.mp4"

gr.Interface(
    inference,
    [
        gr.Dropdown(["reconstruction", "animation"]),
        gr.inputs.Image(type="filepath", label="Source (only used if mode is reconstruction)"),
        gr.inputs.Image(type="mp4", label="Driving Video")
    ],
    outputs=gr.outputs.Video(type="mp4", label="Output Video"),
    title=title,
    description=description,
    theme="huggingface",
    allow_flagging=False).launch(debug=False,enable_queue=True)