Update app.py

app.py

@@ -6,15 +6,15 @@ import numpy as np
 import random
 import uuid
 
-# Load the model
+# Load the model - using 'melody' to ensure it follows your notes
 model = MusicGen.get_pretrained('facebook/musicgen-melody')
 
 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.",
+    "Phonk": "Aggressive high-pitched cowbell melodies, distorted 808 percussion, dark lo-fi atmosphere, 120 BPM.",
+    "Hyperpop": "High-energy metallic percussion, glitchy synth arps, futuristic digital textures, 160 BPM.",
+    "Neo-Soul": "Warm Rhodes piano, subtle jazzy guitar licks, laid-back swinging drum kit, organic groove, 90 BPM.",
+    "Dark Techno": "Heavy driving industrial kick drum, pulsing white noise, rhythmic hypnotic synth sequence, 130 BPM.",
+    "Lofi Hip-Hop": "Dusty boom-bap drum loop, vinyl crackle, soft flute melody, nostalgic vibe, 85 BPM.",
     "Custom (Type Below)": ""
 }
 
@@ -22,27 +22,41 @@ def get_random_style():
     choices = [s for s in STYLE_MAP.keys() if s != "Custom (Type Below)"]
     return random.choice(choices)
 
-def finish_my_song(audio_input, style_choice, custom_prompt, duration, temp, top_k, top_p, current_history):
+def finish_my_song(audio_input, style_choice, custom_prompt, duration, temp, top_k, top_p, bpm, 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]
+    # 1. Build the Smart Prompt
+    style_desc = STYLE_MAP[style_choice] if style_choice != "Custom (Type Below)" else custom_prompt
+    final_prompt = f"{style_desc} The tempo is strictly {bpm} BPM. High precision rhythmic alignment, studio quality."
+
+    # 2. Audio Pre-Processing (Normalization for rough recordings)
+    sr, data = audio_input
+    audio_data = data.astype(np.float32)
+    # Normalize volume
+    audio_data /= np.max(np.abs(audio_data)) + 1e-6 
     
-    # APPLY THE CREATIVITY SETTINGS
+    audio_tensor = torch.from_numpy(audio_data).t().unsqueeze(0)
+    if audio_tensor.shape[1] > 1:
+        audio_tensor = audio_tensor.mean(dim=1, keepdim=True)
+
+    # 3. Apply Precision Settings
     model.set_generation_params(
         duration=duration,
         temperature=temp,
         top_k=top_k,
-        top_p=top_p
+        top_p=top_p,
+        cfg_coeff=5.0 # High guidance to fix off-beat timing
     )
-    
-    sr, data = audio_input
-    audio_tensor = torch.from_numpy(data).float().t().unsqueeze(0)
-    if audio_tensor.shape[1] > 1:
-        audio_tensor = audio_tensor.mean(dim=1, keepdim=True)
 
-    wav = model.generate_with_chroma(descriptions=[final_prompt], melody_wavs=audio_tensor, sr=sr)
+    # 4. Generate with Chroma (Extracts the 'DNA' of your chords)
+    wav = model.generate_with_chroma(
+        descriptions=[final_prompt], 
+        melody_wavs=audio_tensor, 
+        sr=sr
+    )
     
+    # 5. Save and Track History
     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")
@@ -51,43 +65,19 @@ def finish_my_song(audio_input, style_choice, custom_prompt, duration, temp, top
     new_history = [full_file_path] + current_history
     return full_file_path, new_history
 
+# --- GRADIO UI ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     history_state = gr.State([])
 
-    gr.Markdown("# 🎹 The Song Finisher Pro + Advanced Controls")
+    gr.Markdown("# 🎹 The Song Finisher Pro")
+    gr.Markdown("Upload your chords/bass. The AI will quantize your timing to the BPM you set.")
     
     with gr.Row():
         with gr.Column(scale=2):
-            audio_in = gr.Audio(label="Upload Chords/Bass")
+            audio_in = gr.Audio(label="Step 1: Upload Your Rough Track")
+            
             with gr.Row():
-                style_dropdown = gr.Dropdown(choices=list(STYLE_MAP.keys()), value="Phonk", label="Style")
+                style_dropdown = gr.Dropdown(choices=list(STYLE_MAP.keys()), value="Phonk", label="Step 2: 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")
-            
-            # ADVANCED CONTROLS ACCORDION
-            with gr.Accordion("🛠️ Advanced Creativity Settings", open=False):
-                temp_slider = gr.Slider(minimum=0.1, maximum=2.0, value=1.0, step=0.1, label="Temperature (Randomness)")
-                top_k_slider = gr.Slider(minimum=0, maximum=500, value=250, step=1, label="Top-k (Variety)")
-                top_p_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.0, step=0.01, label="Top-p (Nucleus Sampling)")
-                gr.Markdown("*Tip: If the song sounds too messy, lower the Temperature to 0.7.*")
-
-            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")
-            history_gallery = gr.Files(label="Download Past Versions", interactive=False)
-
-    random_btn.click(fn=get_random_style, outputs=style_dropdown)
-    
-    submit_btn.click(
-        fn=finish_my_song, 
-        inputs=[audio_in, style_dropdown, custom_text, length, temp_slider, top_k_slider, top_p_slider, history_state], 
-        outputs=[latest_output, history_state]
-    ).then(
-        fn=lambda x: x, inputs=history_state, outputs=history_gallery
-    )
-
-demo.launch()
+            bpm_slider =