import gradio as gr
import subprocess
import os

def run_talkshow_model(audio_file):
    # Path to the TalkSHOW demo script
    demo_script = 'scripts/demo.py'
    
    # Configuration and model parameters
    config_file = './config/LS3DCG.json'
    body_model_name = 's2g_LS3DCG'
    body_model_path = 'experiments/2022-10-19-smplx_S2G-LS3DCG/ckpt-99.pth'
    
    # Path of the uploaded audio file
    audio_file_path = audio_file
    
    # Path where the output .mp4 video will be saved
    output_video_path = './output_video/result.mp4'
    
    # Run the demo.py script with the necessary arguments
    command = [
        'python', demo_script,
        '--config_file', config_file,
        '--infer',
        '--audio_file', audio_file_path,
        '--body_model_name', body_model_name,
        '--body_model_path', body_model_path,
        '--id', '0',
        '--output', output_video_path  # Assuming demo.py has an argument to specify output
    ]
    
    try:
        # Run the subprocess and capture output
        subprocess.run(command, check=True, capture_output=True, text=True)
        
        # Check if the .mp4 file is generated
        if os.path.exists(output_video_path):
            return output_video_path  # Return the path of the generated video
        else:
            return "Error: Output video not generated."
    
    except subprocess.CalledProcessError as e:
        return f"Error running the model: {e.stderr}"  # Return the error message

# Set up the Gradio interface
interface = gr.Interface(
    fn=run_talkshow_model, 
    inputs=gr.Audio(source="upload", type="filepath"), 
    outputs=gr.Video(),  # Use gr.Video to output the generated .mp4 video
    title="TalkSHOW: Audio to Mesh"
)

# Launch the interface
if __name__ == "__main__":
    interface.launch()