import gradio as gr

import torch

from spectro import wav_bytes_from_spectrogram_image
from diffusers import StableDiffusionPipeline

model_id = "riffusion/riffusion-model-v1"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

def predict(prompt):
    spec = pipe(prompt).images[0]
    print(spec)
    wav = wav_bytes_from_spectrogram_image(spec)
    with open("output.wav", "wb") as f:
        f.write(wav[0].getbuffer())
    return 'output.wav'

gr.Interface(
    predict,
    inputs="text",
    outputs=[gr.Audio(type='filepath')],
    title="Riffusion Text-to-Music",
    description="Describe a musical prompt, generate music by getting a Riffusion spectrogram and its corresponding sound"
).queue(max_size=32, concurrency_count=20).launch(debug=True)