Update app.py

app.py

@@ -3,62 +3,83 @@ from audiocraft.models import MusicGen
 from audiocraft.data.audio import audio_write
 import torch
 import numpy as np
+import random
+import uuid # For unique filenames
 
-# Load the model (using 'melody' specifically for your chords/bass)
-# This will download the first time you run the app
+# Load the model
 model = MusicGen.get_pretrained('facebook/musicgen-melody')
 
-def finish_my_song(audio_input, text_description, duration):
-    if audio_input is None:
-        return None, "Please upload an audio file first!"
+STYLE_MAP = {
+    "Phonk": "Add aggressive high-pitched cowbell melodies, distorted 808 percussion, and a dark, lo-fi Memphis rap atmosphere. Heavy sidechain compression, 120 BPM, gritty and phonk-style.",
+    "Hyperpop": "Add high-energy metallic percussion, glitchy synth arpeggios, and bright bubblegum-pop textures. Use sharp, digital drums and a very polished, futuristic sound. 160 BPM.",
+    "Neo-Soul": "Add a warm Rhodes electric piano layer, subtle jazzy guitar licks, and a laid-back, swinging drum kit with soft rimshots. Soulful, organic, 90 BPM, deep groove.",
+    "Dark Techno": "Add a heavy, driving industrial kick drum, pulsing white noise textures, and a rhythmic hypnotic synth sequence. Minimalist, dark warehouse vibe, 130 BPM.",
+    "Lofi Hip-Hop": "Add a dusty, bit-crushed boom-bap drum loop, mellow vinyl crackle, and a soft flute melody in the background. Nostalgic, rainy day vibe, 85 BPM.",
+    "Custom (Type Below)": ""
+}
+
+def get_random_style():
+    choices = [s for s in STYLE_MAP.keys() if s != "Custom (Type Below)"]
+    return random.choice(choices)
 
-    # Set how long the AI should play for
+def finish_my_song(audio_input, style_choice, custom_prompt, duration, current_history):
+    if audio_input is None:
+        return None, current_history
+    
+    final_prompt = custom_prompt if style_choice == "Custom (Type Below)" else STYLE_MAP[style_choice]
     model.set_generation_params(duration=duration)
     
-    # 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
+    if audio_tensor.shape[1] > 1:
         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_tensor, 
-        sr=sr
-    )
+    wav = model.generate_with_chroma(descriptions=[final_prompt], melody_wavs=audio_tensor, sr=sr)
     
-    # Save to a temporary file for downloading
-    output_path = "ai_completion_idea"
+    # Generate a unique filename so history doesn't overwrite itself
+    unique_id = str(uuid.uuid4())[:8]
+    output_path = f"finished_beat_{unique_id}"
     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"
+    full_file_path = f"{output_path}.wav"
+    
+    # Add new file to the top of history
+    new_history = [full_file_path] + current_history
+    return full_file_path, new_history
 
-# 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.")
+    # State to hold the list of file paths
+    history_state = gr.State([])
+
+    gr.Markdown("# 🎹 The Song Finisher Pro + History")
     
     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")
+        with gr.Column(scale=2):
+            audio_in = gr.Audio(label="Upload Chords/Bass")
+            with gr.Row():
+                style_dropdown = gr.Dropdown(choices=list(STYLE_MAP.keys()), value="Phonk", label="Style")
+                random_btn = gr.Button("🎲 Random")
+            
+            custom_text = gr.Textbox(label="Custom Prompt", placeholder="Type here...")
+            length = gr.Slider(minimum=5, maximum=30, value=15, step=5, label="Seconds")
+            submit_btn = gr.Button("🔥 Generate Arrangement", variant="primary")
+            
+            latest_output = gr.Audio(label="Latest Generation", type="filepath")
 
+        with gr.Column(scale=1):
+            gr.Markdown("### 📜 Session History")
+            # This gallery shows all previous versions
+            history_gallery = gr.Files(label="Download Past Versions", interactive=False)
+
+    # Logic
+    random_btn.click(fn=get_random_style, outputs=style_dropdown)
+    
     submit_btn.click(
         fn=finish_my_song, 
-        inputs=[audio_in, prompt, length], 
-        outputs=[audio_out]
+        inputs=[audio_in, style_dropdown, custom_text, length, history_state], 
+        outputs=[latest_output, history_state]
+    ).then(
+        fn=lambda x: x, inputs=history_state, outputs=history_gallery
     )
 
 demo.launch()