import time
import requests
import numpy as np
import gradio as gr
from loguru import logger

musicGen_url = 'http://13.125.95.210:8080/'

def musicGen(query):
    tmp_audio_name = str(round(time.time() * 1000)) + '0'
    tmp_video_name = str(round(time.time() * 1000)) + '1'
    res = requests.get(url=musicGen_url, params={'query': query}).json()
    audio_url = res['audio_url']
    video_url = res['video_url']
    logger.info(f'request query = {query} \naudio url = {audio_url} \nvideo url = {video_url}')
    audio_content = requests.get(audio_url).content
    video_content = requests.get(video_url).content
    with open(f'/tmp/{tmp_audio_name}', 'wb') as a:
        a.write(audio_content)
    with open(f'/tmp/{tmp_video_name}', 'wb') as v:
        v.write(video_content)

    return f'/tmp/{tmp_video_name}', f'/tmp/{tmp_audio_name}'

demo = gr.Interface(title="The inference server has been shutdown", fn=musicGen, inputs='text', outputs=['video', 'audio'])

if __name__ == "__main__":
    demo.launch()