import gradio as gr
from diffusers import DiffusionPipeline
import scipy.io.wavfile


def load_model(model_id):
    pipeline = DiffusionPipeline.from_pretrained(model_id)
    pipeline = pipeline.to("cuda")
    return pipeline

def denoise(length_sec,model):
    pipeline = load_model(model)
    audios = pipeline(audio_length_in_s=length_sec).audios
    for audio in audios:
        scipy.io.wavfile.write("test.wav", pipeline.unet.sample_rate, audio.transpose())
    return "test.wav"




block = gr.Blocks()

with block:
    gr.HTML(
        """
            <div style="text-align: center; max-width: 700px; margin: 0 auto;">
              <div
                style="
                  display: inline-flex;
                  align-items: center;
                  gap: 0.8rem;
                  font-size: 1.75rem;
                "
              >
                <h1 style="font-weight: 900; margin-bottom: 7px;">
                  Dance Diffusion
                </h1>
              </div>
              <p style="margin-bottom: 10px; font-size: 94%">
                Dance Diffusion is the first in a suite of generative audio tools for producers and musicians to be released by Harmonai
              </p>
            </div>
        """
    )
    with gr.Group():
        with gr.Box():
            length = gr.Slider(1.0, 6.0, value=3.0, step=0.5, label="Audio length in seconds")
            model = gr.Dropdown(choices=["harmonai/maestro-150k", "harmonai/jmann-small-190k", "harmonai/honk-140k", "harmonai/unlocked-250k","harmonai/jmann-large-580k","harmonai/glitch-440k"], value="harmonai/maestro-150k",type="value", label="Model")
            out = gr.Audio(label="Output", type="filepath")
            btn = gr.Button("Submit").style(full_width=True)
        
        btn.click(denoise, inputs=[length,model], outputs=out) 
        gr.HTML('''
        <div class="footer" style="text-align: center; max-width: 700px; margin: 0 auto;">
                    <p>Model by <a href="https://huggingface.co/harmonai" style="text-decoration: underline;" target="_blank">Harmonai</a>
                    </p>
        </div>
        ''')

block.launch(debug=True)