Update app.py

app.py

@@ -5,16 +5,17 @@ import torch
 
 from diffusers import SpectrogramDiffusionPipeline, MidiProcessor
 
-pipe = SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion")
-pipe = pipe.to("cuda")
+pipe = SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion", torch_dtype=torch.float16).to("cuda")
+pipe.enable_xformers_memory_efficient_attention()
+
 processor = MidiProcessor()
 
 
 def predict(audio_file_pth):
-    # audio = tuple (sample_rate, frames) or (sample_rate, (frames, channels))
 
-    output = pipe(processor(audio_file_pth.name)[:2])
-    audio = output.audios[0]
+    with torch.inference_mode():
+        output = pipe(processor(audio_file_pth.name)[:2])
+        audio = output.audios[0]
 
     return (16000, audio.ravel())
 
@@ -33,7 +34,7 @@ examples = []
 gr.Interface(
     fn=predict,
     inputs=[
-        gr.File(file_count="single", file_types=[".mid"]),
+        gr.File(label="Upload MIDI", file_count="single", file_types=[".mid"]),
     ],
     outputs=[
         gr.Audio(label="Synthesised Music", type="numpy"),