import gradio as gr
import subprocess
import os

# Define the command template
COMMAND_TEMPLATE = "python scripts/demo.py --config_file ./config/LS3DCG.json --infer --audio_file {audio_path} --body_model_name s2g_LS3DCG --body_model_path experiments/2022-10-19-smplx_S2G-LS3DCG/ckpt-99.pth --id 0"

# Define the function to process the audio and generate the video
def process_audio(audio):
    # Save uploaded audio to a specific file path
    audio_path = "demo_audio/style.wav"
    with open(audio_path, "wb") as f:
        f.write(audio.read())

    # Run the script command with the specified audio file
    command = COMMAND_TEMPLATE.format(audio_path=audio_path)
    subprocess.run(command, shell=True)

    # Specify the path to the output video
    output_video_path = "./output_video.mp4"  # Update if output location is different

    # Check if the video file exists and return it for display
    if os.path.exists(output_video_path):
        return output_video_path
    else:
        return "Error: Video file not generated."

# Set up Gradio interface
audio_input = gr.Audio(label="Upload a .wav audio file", type="file")
video_output = gr.Video(label="Generated Video Output")

interface = gr.Interface(
    fn=process_audio,
    inputs=audio_input,
    outputs=video_output,
    title="Audio-to-Video Generator",
    description="Upload a .wav file to generate an animated video output."
)

# Launch the Gradio app
interface.launch()