app.py

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 spec, 'output.wav'

gr.Interface(
    predict,
    inputs="text",
    outputs=[gr.Image(), gr.Audio(type='filepath')],
    title="Riffusion Text-to-Music",
    description="""<p style="text-align: center;">Describe a musical prompt, generate music by getting a Riffusion spectrogram and its corresponding sound.
              <br />if you want to skip the queue, or get faster inference, you can duplicate this space: <a href="https://huggingface.co/spaces/fffiloni/spectrogram-to-music?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
              </p>
"""
).queue().launch(debug=True)