Update app.py

app.py

@@ -13,6 +13,16 @@ from scipy.io import wavfile
 import torch
 import torchaudio
 
+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 get_spectro(prompt):
+    image = pipe(prompt).images[0] 
+    return image
+
 def wav_bytes_from_spectrogram_image(image):
     """
     Reconstruct a WAV audio clip from a spectrogram image. Also returns the duration in seconds.
@@ -133,4 +143,4 @@ def waveform_from_spectrogram(
     return waveform
 
 
-gr.Interface(fn=wav_bytes_from_spectrogram_image, inputs=[gr.Image(source="upload", type="pil")], outputs=[gr.Audio()]).launch()
+gr.Interface(fn=get_spectro, inputs=[gr.Textbox()], outputs=[gr.Image()]).launch()