import gradio as gr
import spaces
import torch
import subprocess
import os
import ffmpeg

zero = torch.Tensor([0]).cuda()
print(zero.device) # <-- 'cpu' 🤔

@spaces.GPU
def greet(n):
    print(zero.device) # <-- 'cuda:0' 🤗
    return f"Hello {zero + n} Tensor"

def audio_video():
    print("started =========================")
    input_video = ffmpeg.input('results/result_voice.mp4')

    input_audio = ffmpeg.input('sample_data/sir.mp3')
    os.system(f"rm -rf results/final_output.mp4")
    ffmpeg.concat(input_video, input_audio, v=1, a=1).output('results/final_output.mp4').run()
    
    return "results/final_output.mp4"

def run_infrence(input_video,input_audio):
    audio = "sample_data/sir.mp3"
    video = "sample_data/spark_input.mp4"
    command = f'python3 inference.py --checkpoint_path checkpoints/wav2lip_gan.pth --face sample_data/spark.png --audio sample_data/sir.mp3'
    print("running ")
    # Execute the command
    process = subprocess.Popen(command, stdout=subprocess.PIPE, shell=True)

    # Get the output
    output, error = process.communicate()

    return audio_video()

def run():
  with gr.Blocks(css=".gradio-container {background-color: lightgray} #radio_div {background-color: #FFD8B4; font-size: 40px;}") as demo:
    gr.Markdown("<h1 style='text-align: center;'>"+ "One Shot Talking Face from Text" + "</h1><br/><br/>")
    with gr.Group():
      # with gr.Box():
        with gr.Row():
        # with gr.Row().style(equal_height=True):
            input_video = gr.Video(label="Input Video")
            input_audio = gr.Audio(label="Input Audio")
            video_out = gr.Video(show_label=True,label="Output")
        with gr.Row():
            btn = gr.Button("Generate")   

    btn.click(run_infrence,inputs=[input_video,input_audio], outputs=[video_out])
    # btn.click(run_infrence,inputs=[input_video,input_audio])
    demo.queue()
    demo.launch(server_name="0.0.0.0", server_port=7860)


if __name__ == "__main__":
    run()