Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import librosa
 import numpy as np
 import torch
+import random
 
 from diffusers import SpectrogramDiffusionPipeline, MidiProcessor
 
@@ -12,6 +13,11 @@ pipe.enable_xformers_memory_efficient_attention()
 
 processor = MidiProcessor()
 
+COLORS = [
+    ["#ff0000", "#00ff00"],
+    ["#00ff00", "#0000ff"],
+    ["#0000ff", "#ff0000"],
+]
 
 def predict(audio_file_pth):
 
@@ -19,7 +25,7 @@ def predict(audio_file_pth):
         output = pipe(processor(audio_file_pth.name)[:2])
         audio = output.audios[0]
 
-    return (16000, audio.ravel())
+    return gr.make_waveform((16000, audio.ravel()), bars_color=random.choice(COLORS), bar_count=75)
 
 
 title = "Music Spectrogram Diffusion: Multi-instrument Music Synthesis with Spectrogram Diffusion"
@@ -39,7 +45,7 @@ gr.Interface(
         gr.File(label="Upload MIDI", file_count="single", file_types=[".mid"]),
     ],
     outputs=[
-        gr.Audio(label="Synthesised Music", type="numpy"),
+        gr.Video(label="Synthesised Audio"),
     ],
     title=title,
     description=description,