Update app.py

app.py

@@ -1,7 +1,10 @@
-import gradio as gr
 import numpy as np
-import torch
 import random
+import subprocess
+import tempfile
+import torch
+
+import gradio as gr
 
 from transformers import AutoProcessor, MusicgenForConditionalGeneration
 
@@ -11,7 +14,7 @@ COLORS = [
     ["#0000ff", "#ff0000"],
 ]
 
-path = "facebook/musicgen-large"
+path = "facebook/musicgen-small"
 processor = AutoProcessor.from_pretrained(path)
 model = MusicgenForConditionalGeneration.from_pretrained(path, torch_dtype=torch.float16).to("cuda")
 
@@ -25,7 +28,12 @@ def predict(text):
     with torch.autocast("cuda"):
         outputs = model.generate(**inputs, do_sample=True, guidance_scale=3, max_new_tokens=512)
 
-    return (32000, outputs[0][0].cpu().numpy().astype(np.float16)), gr.make_waveform((32000, outputs[0].cpu().numpy().astype(np.float16).ravel()), bars_color=random.choice(COLORS), bar_count=75)
+    with tempfile.NamedTemporaryFile() as tmp:
+        video_path = gr.make_waveform((32000, outputs[0].cpu().numpy().astype(np.float16).ravel()), bars_color=random.choice(COLORS), bar_count=75)
+        command = f'ffmpeg -y -i {video_path} -vf "scale=250:150" {tmp.name}'
+        subprocess.run(command.split()]) 
+
+        return (32000, outputs[0][0].cpu().numpy().astype(np.float16)), tmp.name 
 
 
 title = "MusicGen"