Update app.py

app.py

@@ -2,34 +2,63 @@ import gradio as gr
 from audiocraft.models import MusicGen
 from audiocraft.data.audio import audio_write
 import torch
+import numpy as np
 
-# Load the model - 'melody' is key for your use case
+# Load the model (using 'melody' specifically for your chords/bass)
+# This will download the first time you run the app
 model = MusicGen.get_pretrained('facebook/musicgen-melody')
 
 def finish_my_song(audio_input, text_description, duration):
-    # Set the length of the generated idea
+    if audio_input is None:
+        return None, "Please upload an audio file first!"
+
+    # Set how long the AI should play for
     model.set_generation_params(duration=duration)
     
-    # Process your chords/bass + your text prompt
-    # melody_wavs expects the audio you want the AI to "follow"
+    # Get the sampling rate and the audio data from your upload
+    sr, data = audio_input
+    
+    # Convert to the format the AI needs
+    audio_tensor = torch.from_numpy(data).float().t().unsqueeze(0)
+    if audio_tensor.shape[1] > 1: # Convert stereo to mono if needed
+        audio_tensor = audio_tensor.mean(dim=1, keepdim=True)
+
+    # Generate the song based on your description + your audio
     wav = model.generate_with_chroma(
         descriptions=[text_description],
-        melody_wavs=audio_input[1], 
-        sr=audio_input[0]
+        melody_wavs=audio_tensor, 
+        sr=sr
     )
     
-    # Return the generated audio
-    return wav[0].cpu().numpy()
+    # Save to a temporary file for downloading
+    output_path = "ai_completion_idea"
+    audio_write(output_path, wav[0].cpu(), model.sample_rate, strategy="loudness")
+    
+    # Return the file path so Gradio shows a player and a download button
+    return f"{output_path}.wav"
 
-interface = gr.Interface(
-    fn=finish_my_song,
-    inputs=[
-        gr.Audio(label="Upload your Chords & Bass"),
-        gr.Textbox(label="Describe what to add (e.g., drums, lead synth, mood)"),
-        gr.Slider(minimum=10, maximum=30, value=15, label="Seconds to generate")
-    ],
-    outputs=gr.Audio(label="Full Song Idea"),
-    title="The Song Finisher"
-)
+# Create a sleek Dark Mode interface
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🎵 The Song Finisher AI")
+    gr.Markdown("Upload your **Chords and Bass**, describe the vibe, and let AI build the rest.")
+    
+    with gr.Row():
+        with gr.Column():
+            audio_in = gr.Audio(label="Step 1: Upload Chords/Bass (WAV or MP3)")
+            prompt = gr.Textbox(
+                label="Step 2: Describe the Vibe", 
+                placeholder="e.g., Add heavy trap drums, a wide synth lead, and a dark atmosphere..."
+            )
+            length = gr.Slider(minimum=5, maximum=30, value=15, step=5, label="Seconds to Generate")
+            submit_btn = gr.Button("Finish My Song", variant="primary")
+        
+        with gr.Column():
+            audio_out = gr.Audio(label="Step 3: Listen & Download", type="filepath")
+
+    submit_btn.click(
+        fn=finish_my_song, 
+        inputs=[audio_in, prompt, length], 
+        outputs=[audio_out]
+    )
 
-interface.launch()
+demo.launch()