Update app.py

app.py

@@ -1,33 +1,39 @@
 import gradio as gr
+import torch
 import torchaudio
-from transformers import MusicgenForConditionalGeneration, AutoProcessor
+from transformers import MusicgenForConditionalGeneration, MusicgenProcessor
 
-model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-melody")
-processor = AutoProcessor.from_pretrained("facebook/musicgen-melody")
+# Load melody-capable model
+model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-melody", torch_dtype=torch.float32)
+processor = MusicgenProcessor.from_pretrained("facebook/musicgen-melody")
 
 def generate_music(prompt, melody):
     if melody is None:
         return None
 
-    # Load audio and ensure correct sample rate
-    waveform, sr = torchaudio.load(melody)
-    if sr != 16000:
-        resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=16000)
-        waveform = resampler(waveform)
+    # Load melody audio file
+    melody_waveform, melody_sr = torchaudio.load(melody)
+    if melody_sr != 16000:
+        resampler = torchaudio.transforms.Resample(orig_freq=melody_sr, new_freq=16000)
+        melody_waveform = resampler(melody_waveform)
 
-    inputs = processor(audio=waveform, sampling_rate=16000, text=[prompt], return_tensors="pt")
-    audio_values = model.generate(**inputs, max_new_tokens=256)
-    return (audio_values[0].numpy(), model.config.audio_encoder.sampling_rate)
+    # Trim or pad to 30 seconds
+    melody_waveform = melody_waveform[:, :16000 * 30]
 
-interface = gr.Interface(
+    inputs = processor(audio=melody_waveform, sampling_rate=16000, text=[prompt], return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=1024)
+    audio_array = outputs[0].cpu().numpy()
+    return (audio_array, model.config.audio_encoder.sampling_rate)
+
+demo = gr.Interface(
     fn=generate_music,
     inputs=[
-        gr.Textbox(label="Prompt", placeholder="e.g., ambient techno with atmospheric pads"),
-        gr.Audio(source="upload", type="filepath", label="Melody Input (.wav or .mp3)")
+        gr.Textbox(label="Prompt", placeholder="e.g., Funky jazz with synths"),
+        gr.Audio(source="upload", type="filepath", label="Melody Input (WAV or MP3)")
     ],
-    outputs=gr.Audio(label="Generated Music"),
+    outputs=gr.Audio(label="Generated Track"),
     title="🎵 MusicGen-Melody AI Generator",
-    description="Upload a melody and describe the vibe to generate full music using Meta’s MusicGen-Melody model."
+    description="Upload a melody and describe the vibe. Generates music using Meta’s MusicGen-Melody model."
 )
 
-interface.launch()
+demo.launch()