import gradio as gr
import os 
import subprocess
import tempfile
from glob import glob

def run_inference(image_input, audio_input, progress=gr.Progress(track_tqdm=True)):

    # Create a temporary folder for downloaded and processed images
    temp_dir = tempfile.mkdtemp()

    try:
        # Start the subprocess with Popen to capture logs
        process = subprocess.Popen(
            [
                "python", "inference.py",
                "--config", "configs/inference.yaml",
                "--input_image", image_input,
                "--input_audio", audio_input,
                "--output_dir", temp_dir,
            ],
            stdout=subprocess.PIPE,  # Capture standard output
            stderr=subprocess.PIPE,  # Capture standard error
            text=True,  # Decode output to text (instead of bytes)
        )

        # Stream logs from the subprocess in real-time
        for line in process.stdout:
            print(line, end="")  # Print logs to the console (or handle them as needed)

        # Wait for the subprocess to finish and check for errors
        process.wait()
        if process.returncode != 0:
            error_message = process.stderr.read()
            raise gr.Error(f"Inference failed with error: {error_message}")

        # Collect the output video
        output_video = glob(os.path.join(temp_dir, "*.mp4"))
        return output_video[0] if output_video else "No video generated."

    except Exception as e:
        raise gr.Error(f"Error during inference: {str(e)}")


with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("# MEMO")
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(label="Image Input", type="filepath")
                audio_input = gr.Audio(label="Audio Input", type="filepath")
                submit_btn = gr.Button("Submit")
            with gr.Column():
                output_result = gr.Video(label="Result")

    submit_btn.click(
        fn =run_inference,
        inputs =[image_input, audio_input],
        outputs = [output_result]
    )

demo.queue().launch()