import gradio as gr
import torch
from transformers import pipeline

from transformers import AutoProcessor, MusicgenForConditionalGeneration
import scipy
import wave
import io

def generate_music(text):
    processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
    model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
    
    inputs = processor(
        text=text,
        padding=True,
        return_tensors="pt",
    )
    audio_values = model.generate(**inputs, do_sample=True
, guidance_scale=3
, max_new_tokens=1350)
    sampling_rate = model.config.audio_encoder.sampling_rate
    scipy.io.wavfile.write("musicgen_out.wav", rate=sampling_rate, data=audio_values[0, 0].numpy())
    file_path= "musicgen_out.wav"
    
    try:
        Audio = gr.Audio(file_path)
        return Audio
    except Exception as e:
        return str(e)


# theme='Taithrah/Minimal'
# theme = 'freddyaboulton/test-blue'
theme = gr.themes.Monochrome(
    primary_hue="indigo",
    secondary_hue="violet",
    neutral_hue="purple",
).set(
    body_background_fill='*background_fill_secondary',
    link_text_color='*secondary_700',
    link_text_color_dark='*secondary_600'
)





# 创建 Gradio Blocks 界面
iface = gr.Blocks(theme = theme)

with iface as demo:
    # 创建一个列容器
    with gr.Row():
        with gr.Column(scale=1):
            gr.Image("logo.jpg", height = 138, width = 108, show_download_button = False, show_label = False, show_share_button = False, min_width = 108)
        with gr.Column(scale=5):
            company_intro = gr.Markdown("## Byte beats\n\n We simplify soundtrack creation, empowering filmmakers globally, thereby reducing production expenses and shortening the process.\n\n\n")
            
    with gr.Row():
        # 在左侧列中添加文本输入框
        with gr.Column(scale=1):
            input_text = gr.Textbox(lines=6, placeholder="Please enter your script or upload video below...")
            video = gr.UploadButton("upload your video")
            example_demo = gr.Markdown("## Example Demo\n\n")

            
        # 在右侧列中添加音频输出组件
        with gr.Column(scale=3):
            output_audio = gr.Audio(label="Generated Music")
            generate_button = gr.Button("Generate Music")
            generate_button.click(fn=generate_music, inputs=input_text, outputs=output_audio)
    
        
# 启动界面
demo.queue(max_size=4).launch(share = True)


# demo = gr.Interface(
#   fn=generate_music,
#   inputs='text',
#   outputs='audio',
# )

# demo.queue(max_size=4).launch(share = True)