import gradio as gr
import os

def generate_talking_head(image, audio):
    output_dir = "./results"
    os.makedirs(output_dir, exist_ok=True)

    # SadTalker ke inference command
    os.system(
        f"python inference.py "
        f"--driven_audio {audio} "
        f"--source_image {image} "
        f"--result_dir {output_dir}"
    )

    output_path = os.path.join(output_dir, "output.mp4")
    if os.path.exists(output_path):
        return output_path
    else:
        return None

iface = gr.Interface(
    fn=generate_talking_head,
    inputs=[gr.Image(type="filepath", label="Upload Image"),
            gr.Audio(type="filepath", label="Upload Audio")],
    outputs=gr.Video(label="Generated Video")
)

iface.launch()