Update app.py

app.py

@@ -4,123 +4,176 @@ from audiocraft.data.audio import audio_write
 import torch
 import numpy as np
 import uuid
+import torchaudio
 
 # Load Model
 model = MusicGen.get_pretrained('facebook/musicgen-melody')
 
-# --- STYLE ENGINE (Upgraded for Bass House/Jayms Style) ---
-STYLE_MAP = {
-    "Bass House (Jayms Style)": "Aggressive Bass House, metallic FM wobble bass, shuffling tech-house drums, growl fills, 126 BPM, heavy sidechain, high-energy club mix.",
-    "Future House (Mainstage)": "Future House, bouncy metallic donk bass, bright lead synths, 4x4 shuffling beat, 126 BPM, Oliver Heldens style, clean production.",
-    "Tech House (Deep)": "Minimal Tech House, rolling sub-bass, dry percussion, crisp hi-hats, 125 BPM, hypnotizing groove, dark atmosphere.",
-    "Dubstep (Heavy)": "Heavy Dubstep drop, robotic growl bass, half-time riddim drums, sub-bass impact, 140 BPM, distorted textures.",
-    "Hip Hop (90s Boom Bap)": "Golden era Hip Hop, dusty vinyl drums, jazz piano chop, upright bass, 90 BPM, swing groove, lo-fi.",
-    "Trap (Dark)": "Hard dark trap, rapid hi-hat triplets, distorted 808 glides, ominous bell melody, 140 BPM, aggressive.",
-    "RnB (Silk)": "Smooth contemporary RnB, velvet synth pads, deep sub-bass, snapped snare, 95 BPM, soulful atmosphere.",
-    "Custom (Write Your Own)": ""
+# --- STRUCTURE DEFINITIONS ---
+# These act as "presets" for the song's roadmap
+STRUCTURE_PRESETS = {
+    "Radio Hit (Intro-Verse-Chorus-Drop)": ["Intro", "Verse", "Build-Up", "Drop/Chorus", "Outro"],
+    "Club Extended (Intro-Build-Drop-Break-Drop)": ["Intro", "Build-Up", "Drop", "Breakdown", "Drop 2", "Outro"],
+    "Hip Hop Classic (Intro-Verse-Hook-Verse-Hook)": ["Intro", "Verse 1", "Hook", "Verse 2", "Hook", "Outro"],
+    "Custom (Single Prompt)": ["Full Track"]
 }
 
-def generate_song(audio_input, style_choice, custom_text, total_duration, bpm, temp, cfg_scale, top_k, current_history):
-    if audio_input is None:
-        return None, current_history
-    
-    # 1. Prompt Engineering
-    base_prompt = STYLE_MAP[style_choice] if style_choice != "Custom (Write Your Own)" else custom_text
-    # We force "High Fidelity" and "Quantized" to ensure professional quality
-    final_prompt = f"{base_prompt} strictly {bpm} BPM, professional mastering, high fidelity, quantized, clean mix."
+# --- STYLE PROMPTS FOR SECTIONS ---
+# How each section should sound based on the genre
+SECTION_PROMPTS = {
+    "Bass House (Jayms Style)": {
+        "Intro": "Atmospheric intro, filtered bass, simple hi-hats, tension building.",
+        "Verse": "Minimal groove, shuffling hi-hats, deep sub-bass, metallic plucks, clean vocal chops.",
+        "Build-Up": "Rising snare roll, pitch riser, white noise sweep, accelerating energy, hype vocals.",
+        "Drop": "Explosive Bass House drop, metallic FM wobble bass, heavy kick, sidechain, festival energy.",
+        "Breakdown": "Atmospheric pads, filtered chords, no drums, emotional melody.",
+        "Outro": "Drum loop fading out, DJ friendly intro/outro, simple bass."
+    },
+    "Future House (Mainstage)": {
+        "Intro": "Bright piano chords, filtered kick, uplifting atmosphere.",
+        "Verse": "Plucky melody, snapping fingers, deep house bass, clean production.",
+        "Build-Up": "uplifting risers, snare build, hands in the air energy.",
+        "Drop": "Bouncy Future House drop, metallic donk bass, shuffling beat, catchy lead.",
+        "Outro": "Stripped back drums, fading melody."
+    },
+    "Dubstep (Heavy)": {
+        "Intro": "Dark cinematic drone, ominous atmosphere, slow percussion.",
+        "Build-Up": "Intense drum roll, siren rising, pre-drop vocal scream.",
+        "Drop": "Heavy dubstep drop, mechanical growl bass, half-time drums, sub-bass impact.",
+        "Breakdown": "Orchestral strings, emotional piano, dark ambience.",
+        "Outro": "Dark fading textures, slow drum hit."
+    }
+}
+
+def generate_structured_song(audio_input, genre_style, structure_preset, total_duration, bpm, temp, cfg_scale, top_k, history):
+    if audio_input is None: return None, history
 
-    # 2. Audio Processing (Normalization)
+    # 1. SETUP AUDIO INPUT (The DNA)
     sr, data = audio_input
     audio_data = data.astype(np.float32)
-    max_val = np.max(np.abs(audio_data))
-    if max_val > 0:
-        audio_data = audio_data / max_val
+    if np.max(np.abs(audio_data)) > 0: audio_data /= np.max(np.abs(audio_data))
+    
+    # Convert to Tensor
+    audio_tensor = torch.from_numpy(audio_data).t()
+    if audio_tensor.shape[0] > 1: audio_tensor = audio_tensor.mean(dim=0, keepdim=True)
+    else: audio_tensor = audio_tensor.unsqueeze(0)
+    
+    if sr != 32000:
+        resampler = torchaudio.transforms.Resample(sr, 32000)
+        audio_tensor = resampler(audio_tensor)
+
+    # 2. DETERMINE SECTIONS
+    sections = STRUCTURE_PRESETS[structure_preset]
+    if structure_preset == "Custom (Single Prompt)":
+        # Fallback to simple generation
+        pass 
+
+    # Calculate time per section
+    # e.g., 60s total / 4 sections = 15s per section
+    sec_duration = total_duration / len(sections)
+    
+    full_song = audio_tensor # Start with the user's input as context (optional, or start fresh)
+    # Actually, let's treat the user input as the "Intro" DNA and generate from there.
+    
+    current_context = audio_tensor.unsqueeze(0) # [1, 1, T]
     
-    # Convert to Mono for the model
-    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. SETTINGS TUNING
-    # For Bass House, we need high CFG (strictness) to get that specific sound
-    model.set_generation_params(
-        duration=min(total_duration, 30), 
-        temperature=temp, 
-        cfg_coef=cfg_scale, 
-        top_k=int(top_k)
-    )
-
-    # 4. GENERATION LOOP
-    # Phase 1: The Foundation
-    combined_wav = model.generate_with_chroma(
-        descriptions=[final_prompt], 
-        melody_wavs=audio_tensor, 
-        melody_sample_rate=sr
-    )
+    print(f"Generating {len(sections)} sections...")
     
-    # Phase 2: The Extension (Recursive Stitching)
-    overlap = 5 
-    current_len = combined_wav.shape[-1] / model.sample_rate
-
-    while current_len < total_duration:
-        prompt_segment = combined_wav[..., -int(overlap * model.sample_rate):]
-        remaining = total_duration - current_len
-        chunk_len = min(30, remaining + overlap)
+    generated_parts = []
+
+    # 3. GENERATION LOOP (The Architect)
+    for i, section_name in enumerate(sections):
+        # Get the specific prompt for this section (e.g., "Verse" description)
+        if genre_style in SECTION_PROMPTS and section_name in SECTION_PROMPTS[genre_style]:
+            specific_desc = SECTION_PROMPTS[genre_style][section_name]
+        else:
+            # Fallback if specific section text missing
+            specific_desc = f"{genre_style} style, {section_name} section, dynamic change."
         
-        # Ensure parameters stick during the loop
-        model.set_generation_params(duration=chunk_len, temperature=temp, cfg_coef=cfg_scale, top_k=int(top_k))
+        final_prompt = f"{specific_desc} strictly {bpm} BPM, high fidelity."
+        print(f"Rendering {section_name}: {final_prompt}")
+
+        # Set params for this chunk
+        # Note: We generate a bit more than needed to allow for crossfading overlap
+        model.set_generation_params(
+            duration=sec_duration + 2, # +2s for overlap/continuity
+            temperature=temp,
+            cfg_coef=cfg_scale,
+            top_k=int(top_k)
+        )
+
+        # Generate Continuation from previous part
+        # We use the LAST 10 seconds of the previous audio as the "seed" for the new part
+        # to ensure flow, but we change the text prompt to change the energy.
         
+        # Grab last 10s of context
+        context_len = current_context.shape[-1]
+        trim_len = int(10 * 32000)
+        if context_len > trim_len:
+            input_seed = current_context[..., -trim_len:]
+        else:
+            input_seed = current_context
+
         next_chunk = model.generate_continuation(
-            prompt_segment, 
-            model.sample_rate, 
-            descriptions=[final_prompt]
+            prompt=input_seed,
+            prompt_sample_rate=32000,
+            descriptions=[final_prompt],
+            progress=True
         )
         
-        combined_wav = torch.cat([combined_wav, next_chunk[..., int(overlap * model.sample_rate):]], dim=-1)
-        current_len = combined_wav.shape[-1] / model.sample_rate
+        # Add to list
+        # Remove the input_seed part from the output to avoid duplication?
+        # generate_continuation returns the FULL audio (seed + new). 
+        # We only want the NEW part.
+        new_audio = next_chunk[..., input_seed.shape[-1]:] 
+        
+        generated_parts.append(new_audio)
+        
+        # Update context for next loop
+        current_context = next_chunk
+
+    # 4. STITCHING
+    # Concatenate all parts
+    final_mix = torch.cat(generated_parts, dim=-1)
 
     # 5. EXPORT
     uid = str(uuid.uuid4())[:8]
-    filename = f"track_{style_choice.split(' ')[0]}_{uid}"
-    audio_write(filename, combined_wav[0].cpu(), model.sample_rate, strategy="loudness")
+    filename = f"structured_song_{uid}"
+    audio_write(filename, final_mix[0].cpu(), model.sample_rate, strategy="loudness")
     
-    path = f"{filename}.wav"
-    return path, [path] + current_history
+    return f"{filename}.wav", [f"{filename}.wav"] + history
+
+# --- UI ---
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="zinc")) as demo:
+    history = gr.State([])
+    gr.Markdown("# 🏗️ Infinite Song Architect: Structure Mode")
+    gr.Markdown("Define a Verse-Chorus structure, and the AI will build each section dynamically.")
 
-# --- UI DESIGN ---
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="cyan", secondary_hue="slate")) as demo:
-    history_state = gr.State([])
-    
-    gr.Markdown("# 🎛️ Infinite Song Architect (Bass House Edition)")
-    gr.Markdown("Upload a loop to generate a track in the style of Jayms, Oliver Heldens, or Tchami.")
-    
     with gr.Row():
-        with gr.Column(scale=2):
-            audio_in = gr.Audio(label="🎧 Input Audio (Rough Idea / Melody)", type="numpy")
+        with gr.Column():
+            audio_in = gr.Audio(label="🎧 Input DNA (Chords/Bass)", type="numpy")
             
             with gr.Row():
-                style_drop = gr.Dropdown(choices=list(STYLE_MAP.keys()), value="Bass House (Jayms Style)", label="🎵 Target Style")
-                custom_text = gr.Textbox(label="✍️ Custom Prompt", placeholder="e.g. 'Aggressive wobble bass...'", visible=True)
+                genre_drop = gr.Dropdown(list(SECTION_PROMPTS.keys()), value="Bass House (Jayms Style)", label="Sound Palette")
+                struct_drop = gr.Dropdown(list(STRUCTURE_PRESETS.keys()), value="Radio Hit (Intro-Verse-Chorus-Drop)", label="Song Structure")
             
-            with gr.Accordion("🎚️ Mixing Console (Tweak the AI)", open=True):
-                with gr.Row():
-                    bpm_slider = gr.Slider(120, 140, 126, step=1, label="BPM (126 is Standard)")
-                    len_slider = gr.Slider(15, 120, 30, step=15, label="Duration (Seconds)")
-                with gr.Row():
-                    cfg_slider = gr.Slider(1.0, 15.0, 8.0, step=0.5, label="Style Strictness (Higher = More Metallic Bass)")
-                    temp_slider = gr.Slider(0.1, 1.5, 0.65, step=0.05, label="Groove/Swing (Lower = More Consistent)")
-                    top_k_slider = gr.Slider(10, 250, 250, step=10, label="Audio Stability")
-
-            btn = gr.Button("🎹 Render Track", variant="primary", size="lg")
-            out_audio = gr.Audio(label="💿 Final Master")
+            with gr.Row():
+                bpm_slider = gr.Slider(60, 180, 126, label="BPM")
+                len_slider = gr.Slider(30, 120, 60, step=10, label="Total Length")
             
-        with gr.Column(scale=1):
-            history_list = gr.Files(label="📂 Session Library")
-
-    btn.click(
-        fn=generate_song,
-        inputs=[audio_in, style_drop, custom_text, len_slider, bpm_slider, temp_slider, cfg_slider, top_k_slider, history_state],
-        outputs=[out_audio, history_list]
-    )
+            with gr.Accordion("Fine Tune", open=False):
+                cfg_slider = gr.Slider(1, 15, 8, label="Strictness")
+                temp_slider = gr.Slider(0.1, 1.5, 0.6, label="Creativity")
+                top_k_slider = gr.Slider(10, 250, 50, label="Stability")
+
+            btn = gr.Button("🏗️ BUILD FULL SONG", variant="primary")
+            out_audio = gr.Audio(label="Full Track Output")
+
+        with gr.Column():
+            history_list = gr.Files(label="History")
+
+    btn.click(generate_structured_song, 
+              [audio_in, genre_drop, struct_drop, len_slider, bpm_slider, temp_slider, cfg_slider, top_k_slider, history], 
+              [out_audio, history_list])
 
 demo.launch()