Update app.py

app.py

@@ -1,24 +1,37 @@
 import gradio as gr
-from transformers import pipeline
-import soundfile as sf
-import io
-import numpy as np
+import torch
+from transformers import AutoProcessor, MusicgenForConditionalGeneration
 
-model = pipeline("text-to-music", model="facebook/musicgen-small")
+# Load Meta's AudioGen model
+model_id = "facebook/musicgen-small"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def generate_music(text_input):
-    music_output = model(text_input)
+processor = AutoProcessor.from_pretrained(model_id)
+model = MusicgenForConditionalGeneration.from_pretrained(model_id).to(device)
 
-    audio_buffer = io.BytesIO() 
-    sf.write(audio_buffer, music_output, 22050, format='WAV')  
-    audio_buffer.seek(0)
-    
-    return audio_buffer 
+def generate_audio(text_prompt, duration=5):
+    """Generate music based on a text prompt."""
+    inputs = processor(
+        text=[text_prompt], 
+        return_tensors="pt"
+    ).to(device)
 
+    with torch.no_grad():
+        audio_waveform = model.generate(**inputs, max_new_tokens=duration * 50)
+
+    return audio_waveform[0].cpu().numpy(), 32000  # Sample rate
+
+# Gradio UI
 iface = gr.Interface(
-    fn=generate_music,
-    inputs=gr.Textbox(label="Enter Text for Music"),
-    outputs=gr.Audio(label="Generated Music")
+    fn=generate_audio,
+    inputs=[
+        gr.Textbox(label="Music Prompt", placeholder="e.g., A calm piano melody"),
+        gr.Slider(1, 30, value=5, step=1, label="Duration (seconds)")
+    ],
+    outputs=gr.Audio(label="Generated Music"),
+    title="Text-to-Music Generator",
+    description="Enter a text prompt to generate music using Meta's MusicGen model.",
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()