Update app.py

app.py

@@ -42,7 +42,7 @@ try:
         sys.exit(1)
     musicgen_model = MusicGen.get_pretrained(local_model_path, device=device)
     musicgen_model.set_generation_params(
-        duration=15,  # Default chunk duration
+        duration=25,  # Default chunk duration
         two_step_cfg=False  # Disable two-step CFG for stability
     )
 except Exception as e:
@@ -59,53 +59,67 @@ def print_resource_usage(stage: str):
     print("---------------")
 
 # 4) GENRE PROMPT FUNCTIONS
-def set_red_hot_chili_peppers_prompt():
-    return "Funk rock with groovy basslines, syncopated guitar riffs, energetic drums, and a Red Hot Chili Peppers-inspired vibe with dynamic vocal energy and funky breakdowns."
+def set_red_hot_chili_peppers_prompt(bpm):
+    rhythm = "strong rhythmic steps" if bpm > 120 else "groovy rhythmic flow"
+    return f"Funk rock with groovy basslines, syncopated guitar riffs, energetic drums, and a Red Hot Chili Peppers-inspired vibe with dynamic vocal energy and funky breakdowns, {rhythm} at {bpm} BPM."
 
-def set_nirvana_grunge_prompt():
-    return "Grunge with raw distorted guitar riffs, heavy drums, melodic basslines, and a Nirvana-inspired angst-filled sound with quiet-loud dynamics."
+def set_nirvana_grunge_prompt(bpm):
+    rhythm = "intense rhythmic steps" if bpm > 120 else "grungy rhythmic pulse"
+    return f"Grunge with raw distorted guitar riffs, heavy drums, melodic basslines, and a Nirvana-inspired angst-filled sound with quiet-loud dynamics, {rhythm} at {bpm} BPM."
 
-def set_pearl_jam_grunge_prompt():
-    return "Grunge with soulful guitar leads, driving rhythms, deep bass, and a Pearl Jam-inspired emotional intensity with soaring choruses."
+def set_pearl_jam_grunge_prompt(bpm):
+    rhythm = "soulful rhythmic steps" if bpm > 120 else "driving rhythmic flow"
+    return f"Grunge with soulful guitar leads, driving rhythms, deep bass, and a Pearl Jam-inspired emotional intensity with soaring choruses, {rhythm} at {bpm} BPM."
 
-def set_soundgarden_grunge_prompt():
-    return "Grunge with heavy, sludgy guitar riffs, complex drum patterns, and a Soundgarden-inspired dark, psychedelic edge with powerful vocals."
+def set_soundgarden_grunge_prompt(bpm):
+    rhythm = "heavy rhythmic steps" if bpm > 120 else "sludgy rhythmic groove"
+    return f"Grunge with heavy, sludgy guitar riffs, complex drum patterns, and a Soundgarden-inspired dark, psychedelic edge with powerful vocals, {rhythm} at {bpm} BPM."
 
-def set_foo_fighters_prompt():
+def set_foo_fighters_prompt(bpm):
     styles = ["anthemic", "gritty", "melodic", "fast-paced", "driving"]
     tempos = ["upbeat", "mid-tempo", "high-energy"]
     moods = ["energetic", "introspective", "rebellious", "uplifting"]
     style = random.choice(styles)
     tempo = random.choice(tempos)
     mood = random.choice(moods)
-    return f"Alternative rock with {style} guitar riffs, {tempo} drums, melodic hooks, and a Foo Fighters-inspired {mood} vibe with powerful choruses."
+    rhythm = "powerful rhythmic steps" if bpm > 120 else "catchy rhythmic groove"
+    return f"Alternative rock with {style} guitar riffs, {tempo} drums, melodic hooks, and a Foo Fighters-inspired {mood} vibe with powerful choruses, {rhythm} at {bpm} BPM."
 
-def set_smashing_pumpkins_prompt():
-    return "Alternative rock with dreamy guitar textures, heavy distortion, dynamic drums, and a Smashing Pumpkins-inspired blend of melancholy and aggression."
+def set_smashing_pumpkins_prompt(bpm):
+    rhythm = "dynamic rhythmic steps" if bpm > 120 else "dreamy rhythmic flow"
+    return f"Alternative rock with dreamy guitar textures, heavy distortion, dynamic drums, and a Smashing Pumpkins-inspired blend of melancholy and aggression, {rhythm} at {bpm} BPM."
 
-def set_radiohead_prompt():
-    return "Experimental rock with atmospheric synths, intricate guitar layers, complex rhythms, and a Radiohead-inspired blend of introspective and innovative soundscapes."
+def set_radiohead_prompt(bpm):
+    rhythm = "complex rhythmic steps" if bpm > 120 else "intricate rhythmic pulse"
+    return f"Experimental rock with atmospheric synths, intricate guitar layers, complex rhythms, and a Radiohead-inspired blend of introspective and innovative soundscapes, {rhythm} at {bpm} BPM."
 
-def set_classic_rock_prompt():
-    return "Classic rock with bluesy electric guitars, steady drums, groovy bass, and a Led Zeppelin-inspired raw energy with dynamic solos."
+def set_classic_rock_prompt(bpm):
+    rhythm = "bluesy rhythmic steps" if bpm > 120 else "steady rhythmic groove"
+    return f"Classic rock with bluesy electric guitars, steady drums, groovy bass, and a Led Zeppelin-inspired raw energy with dynamic solos, {rhythm} at {bpm} BPM."
 
-def set_alternative_rock_prompt():
-    return "Alternative rock with distorted guitar riffs, punchy drums, melodic basslines, and a Pixies-inspired quirky, energetic vibe."
+def set_alternative_rock_prompt(bpm):
+    rhythm = "quirky rhythmic steps" if bpm > 120 else "energetic rhythmic flow"
+    return f"Alternative rock with distorted guitar riffs, punchy drums, melodic basslines, and a Pixies-inspired quirky, energetic vibe, {rhythm} at {bpm} BPM."
 
-def set_post_punk_prompt():
-    return "Post-punk with jangly guitars, driving basslines, sharp drums, and a Joy Division-inspired moody, atmospheric sound."
+def set_post_punk_prompt(bpm):
+    rhythm = "sharp rhythmic steps" if bpm > 120 else "moody rhythmic pulse"
+    return f"Post-punk with jangly guitars, driving basslines, sharp drums, and a Joy Division-inspired moody, atmospheric sound, {rhythm} at {bpm} BPM."
 
-def set_indie_rock_prompt():
-    return "Indie rock with jangly guitars, heartfelt vocals, steady drums, and an Arctic Monkeys-inspired blend of witty lyrics and catchy riffs."
+def set_indie_rock_prompt(bpm):
+    rhythm = "catchy rhythmic steps" if bpm > 120 else "jangly rhythmic flow"
+    return f"Indie rock with jangly guitars, heartfelt vocals, steady drums, and an Arctic Monkeys-inspired blend of witty lyrics and catchy riffs, {rhythm} at {bpm} BPM."
 
-def set_funk_rock_prompt():
-    return "Funk rock with slap bass, funky guitar chords, upbeat drums, and a Rage Against the Machine-inspired mix of groove and aggression."
+def set_funk_rock_prompt(bpm):
+    rhythm = "aggressive rhythmic steps" if bpm > 120 else "funky rhythmic groove"
+    return f"Funk rock with slap bass, funky guitar chords, upbeat drums, and a Rage Against the Machine-inspired mix of groove and aggression, {rhythm} at {bpm} BPM."
 
-def set_detroit_techno_prompt():
-    return "Detroit techno with deep pulsing synths, driving basslines, crisp hi-hats, and a Juan Atkins-inspired rhythmic groove."
+def set_detroit_techno_prompt(bpm):
+    rhythm = "pulsing rhythmic steps" if bpm > 120 else "deep rhythmic groove"
+    return f"Detroit techno with deep pulsing synths, driving basslines, crisp hi-hats, and a Juan Atkins-inspired rhythmic groove, {rhythm} at {bpm} BPM."
 
-def set_deep_house_prompt():
-    return "Deep house with warm analog synth chords, soulful vocal chops, deep basslines, and a Larry Heard-inspired laid-back groove."
+def set_deep_house_prompt(bpm):
+    rhythm = "soulful rhythmic steps" if bpm > 120 else "laid-back rhythmic flow"
+    return f"Deep house with warm analog synth chords, soulful vocal chops, deep basslines, and a Larry Heard-inspired laid-back groove, {rhythm} at {bpm} BPM."
 
 # 5) AUDIO PROCESSING FUNCTIONS
 def apply_chorus(segment):
@@ -133,7 +147,7 @@ def apply_fade(segment, fade_in_duration=2000, fade_out_duration=2000):
     return segment
 
 # 6) GENERATION & I/O FUNCTIONS
-def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p: float, temperature: float, total_duration: int, crossfade_duration: int, num_variations: int, use_chunks: bool):
+def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p: float, temperature: float, total_duration: int, crossfade_duration: int, variation_crossfade_duration: int, num_variations: int, chunk_count: int, use_chunks: bool, bpm: int):
     global musicgen_model
     if not instrumental_prompt.strip():
         return None, "⚠️ Please enter a valid instrumental prompt!"
@@ -142,6 +156,11 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
         total_duration = min(max(total_duration, 10), 90)
         sample_rate = musicgen_model.sample_rate
         output_files = []
+        variation_segments = []
+
+        # Adjust temperature based on BPM for energy
+        adjusted_temperature = temperature + (bpm - 120) / 600.0  # Scale temperature slightly
+        adjusted_temperature = min(max(adjusted_temperature, 0.1), 2.0)
 
         for var in range(num_variations):
             print(f"Generating variation {var+1}/{num_variations}...")
@@ -151,22 +170,21 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
 
             if use_chunks:
                 # Chunked generation
-                chunk_duration = 15
-                num_chunks = max(1, total_duration // chunk_duration)
+                num_chunks = max(1, min(chunk_count, total_duration // 25))  # Cap at 25 seconds per chunk
                 chunk_duration = total_duration / num_chunks
                 overlap_duration = min(1.0, crossfade_duration / 1000.0)
                 generation_duration = chunk_duration + overlap_duration
                 audio_chunks = []
 
                 for i in range(num_chunks):
-                    chunk_prompt = instrumental_prompt
+                    chunk_prompt = f"{instrumental_prompt}, at {bpm} BPM"
                     print(f"Generating chunk {i+1}/{num_chunks} for variation {var+1} on GPU (prompt: {chunk_prompt})...")
                     musicgen_model.set_generation_params(
                         duration=generation_duration,
                         use_sampling=True,
                         top_k=top_k,
                         top_p=top_p,
-                        temperature=temperature,
+                        temperature=adjusted_temperature,
                         cfg_coef=cfg_scale
                     )
 
@@ -211,6 +229,10 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
                     final_segment = final_segment.append(next_segment, crossfade=crossfade_duration)
 
                 final_segment = final_segment[:total_duration * 1000]
+                variation_segments.append(final_segment)
+
+                for chunk_path in audio_chunks:
+                    os.remove(chunk_path)
             else:
                 # Single-shot generation
                 print(f"Generating full track for variation {var+1} on GPU (prompt: {instrumental_prompt})...")
@@ -219,7 +241,7 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
                     use_sampling=True,
                     top_k=top_k,
                     top_p=top_p,
-                    temperature=temperature,
+                    temperature=adjusted_temperature,
                     cfg_coef=cfg_scale
                 )
 
@@ -227,7 +249,7 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
 
                 with torch.no_grad():
                     with autocast():
-                        audio_chunk = musicgen_model.generate([instrumental_prompt], progress=True)[0]
+                        audio_chunk = musicgen_model.generate([f"{instrumental_prompt}, at {bpm} BPM"], progress=True)[0]
 
                 audio_chunk = audio_chunk.cpu().to(dtype=torch.float32)
                 if audio_chunk.dim() == 1:
@@ -247,50 +269,60 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
                 torchaudio.save(temp_wav_path, audio_chunk, sample_rate, bits_per_sample=24)
                 final_segment = AudioSegment.from_wav(temp_wav_path)
                 os.remove(temp_wav_path)
+                variation_segments.append(final_segment)
 
                 torch.cuda.empty_cache()
                 gc.collect()
                 time.sleep(0.5)
                 print_resource_usage(f"After Full Track Generation (Variation {var+1})")
 
-            print(f"Post-processing final track for variation {var+1}...")
-            final_segment = apply_eq(final_segment)
-            final_segment = apply_chorus(final_segment)
-            final_segment = apply_limiter(final_segment, max_db=-6.0)
-            final_segment = final_segment.normalize(headroom=-9.0)
-            final_segment = apply_final_gain(final_segment, target_db=-18.0)
-
-            mp3_path = f"output_cleaned_variation_{var+1}.mp3"
-            final_segment.export(
-                mp3_path,
-                format="mp3",
-                bitrate="320k",
-                tags={"title": f"GhostAI Instrumental Variation {var+1}", "artist": "GhostAI"}
-            )
-            print(f"Saved final audio to {mp3_path}")
-            output_files.append(mp3_path)
-
-            if use_chunks:
-                for chunk_path in audio_chunks:
-                    os.remove(chunk_path)
+        # Combine variations with crossfade
+        print("Combining variations with crossfade...")
+        combined_segment = variation_segments[0]
+        for i in range(1, len(variation_segments)):
+            next_segment = variation_segments[i]
+            next_segment = next_segment + 1
+            combined_segment = combined_segment.append(next_segment, crossfade=variation_crossfade_duration)
+
+        # Post-process combined track
+        print("Post-processing final track...")
+        combined_segment = apply_eq(combined_segment)
+        combined_segment = apply_chorus(combined_segment)
+        combined_segment = apply_limiter(combined_segment, max_db=-6.0)
+        combined_segment = combined_segment.normalize(headroom=-9.0)
+        combined_segment = apply_final_gain(combined_segment, target_db=-18.0)
+
+        mp3_path = "output_cleaned.mp3"
+        combined_segment.export(
+            mp3_path,
+            format="mp3",
+            bitrate="320k",
+            tags={"title": "GhostAI Instrumental", "artist": "GhostAI"}
+        )
+        print(f"Saved final audio to {mp3_path}")
+        output_files.append(mp3_path)
 
         print_resource_usage("After Final Generation")
         print(f"Total Generation Time: {time.time() - start_time:.2f} seconds")
 
-        # Return the first variation for Gradio display; others are saved to disk
-        return output_files[0], f"✅ Done! Generated {num_variations} variations."
+        # Return the combined track for Gradio display
+        return mp3_path, f"✅ Done! Generated {num_variations} variations."
     except Exception as e:
         return None, f"❌ Generation failed: {e}"
     finally:
         torch.cuda.empty_cache()
         gc.collect()
 
-# Function to toggle crossfade_duration interactivity
-def toggle_crossfade_interactivity(use_chunks):
-    return gr.update(interactive=use_chunks)
+# Function to toggle interactivity of chunk-related sliders
+def toggle_chunk_interactivity(use_chunks):
+    return (
+        gr.update(interactive=use_chunks),  # crossfade_duration
+        gr.update(interactive=use_chunks),  # variation_crossfade_duration
+        gr.update(interactive=use_chunks)   # chunk_count
+    )
 
 def clear_inputs():
-    return "", 3.0, 250, 0.9, 1.0, 30, 500, 1, True
+    return "", 3.0, 250, 0.9, 1.0, 30, 1000, 1000, 1, 2, True, 120
 
 # 7) CUSTOM CSS
 css = """
@@ -434,13 +466,21 @@ with gr.Blocks(css=css) as demo:
             info="Total duration of the track (10 to 90 seconds)."
         )
         crossfade_duration = gr.Slider(
-            label="Crossfade Duration (ms)",
+            label="Chunk Crossfade Duration (ms)",
             minimum=100,
             maximum=2000,
-            value=500,
+            value=1000,
             step=100,
             info="Crossfade duration between chunks (only used if chunking is enabled)."
         )
+        variation_crossfade_duration = gr.Slider(
+            label="Variation Crossfade Duration (ms)",
+            minimum=100,
+            maximum=2000,
+            value=1000,
+            step=100,
+            info="Crossfade duration between variations (only used if chunking is enabled)."
+        )
         num_variations = gr.Slider(
             label="Number of Variations",
             minimum=1,
@@ -449,10 +489,26 @@ with gr.Blocks(css=css) as demo:
             step=1,
             info="Number of different versions to generate with varying random seeds."
         )
+        chunk_count = gr.Slider(
+            label="Chunk Count",
+            minimum=1,
+            maximum=8,
+            value=2,
+            step=1,
+            info="Number of chunks to split the track into (only used if chunking is enabled)."
+        )
         use_chunks = gr.Checkbox(
             label="Generate in Chunks",
             value=True,
-            info="Enable to generate in 15-second chunks (safer for GPU memory). Disable for single-shot generation (higher VRAM usage)."
+            info="Enable to generate in chunks (safer for GPU memory). Disable for single-shot generation (higher VRAM usage)."
+        )
+        bpm = gr.Slider(
+            label="Tempo (BPM)",
+            minimum=60,
+            maximum=180,
+            value=120,
+            step=1,
+            info="Beats per minute to influence the track's tempo."
         )
         with gr.Row(elem_classes="action-buttons"):
             gen_btn = gr.Button("Generate Music")
@@ -462,32 +518,37 @@ with gr.Blocks(css=css) as demo:
         out_audio = gr.Audio(label="Generated Stereo Instrumental Track", type="filepath")
         status = gr.Textbox(label="Status", interactive=False)
     
-    # Toggle crossfade_duration interactivity
-    use_chunks.change(fn=toggle_crossfade_interactivity, inputs=use_chunks, outputs=crossfade_duration)
-
-    rhcp_btn.click(set_red_hot_chili_peppers_prompt, inputs=None, outputs=[instrumental_prompt])
-    nirvana_btn.click(set_nirvana_grunge_prompt, inputs=None, outputs=[instrumental_prompt])
-    pearl_jam_btn.click(set_pearl_jam_grunge_prompt, inputs=None, outputs=[instrumental_prompt])
-    soundgarden_btn.click(set_soundgarden_grunge_prompt, inputs=None, outputs=[instrumental_prompt])
-    foo_fighters_btn.click(set_foo_fighters_prompt, inputs=None, outputs=[instrumental_prompt])
-    smashing_pumpkins_btn.click(set_smashing_pumpkins_prompt, inputs=None, outputs=[instrumental_prompt])
-    radiohead_btn.click(set_radiohead_prompt, inputs=None, outputs=[instrumental_prompt])
-    classic_rock_btn.click(set_classic_rock_prompt, inputs=None, outputs=[instrumental_prompt])
-    alternative_rock_btn.click(set_alternative_rock_prompt, inputs=None, outputs=[instrumental_prompt])
-    post_punk_btn.click(set_post_punk_prompt, inputs=None, outputs=[instrumental_prompt])
-    indie_rock_btn.click(set_indie_rock_prompt, inputs=None, outputs=[instrumental_prompt])
-    funk_rock_btn.click(set_funk_rock_prompt, inputs=None, outputs=[instrumental_prompt])
-    detroit_techno_btn.click(set_detroit_techno_prompt, inputs=None, outputs=[instrumental_prompt])
-    deep_house_btn.click(set_deep_house_prompt, inputs=None, outputs=[instrumental_prompt])
+    # Toggle chunk-related sliders' interactivity
+    use_chunks.change(
+        fn=toggle_chunk_interactivity,
+        inputs=use_chunks,
+        outputs=[crossfade_duration, variation_crossfade_duration, chunk_count]
+    )
+
+    # Bind genre buttons with BPM
+    rhcp_btn.click(set_red_hot_chili_peppers_prompt, inputs=bpm, outputs=instrumental_prompt)
+    nirvana_btn.click(set_nirvana_grunge_prompt, inputs=bpm, outputs=instrumental_prompt)
+    pearl_jam_btn.click(set_pearl_jam_grunge_prompt, inputs=bpm, outputs=instrumental_prompt)
+    soundgarden_btn.click(set_soundgarden_grunge_prompt, inputs=bpm, outputs=instrumental_prompt)
+    foo_fighters_btn.click(set_foo_fighters_prompt, inputs=bpm, outputs=instrumental_prompt)
+    smashing_pumpkins_btn.click(set_smashing_pumpkins_prompt, inputs=bpm, outputs=instrumental_prompt)
+    radiohead_btn.click(set_radiohead_prompt, inputs=bpm, outputs=instrumental_prompt)
+    classic_rock_btn.click(set_classic_rock_prompt, inputs=bpm, outputs=instrumental_prompt)
+    alternative_rock_btn.click(set_alternative_rock_prompt, inputs=bpm, outputs=instrumental_prompt)
+    post_punk_btn.click(set_post_punk_prompt, inputs=bpm, outputs=instrumental_prompt)
+    indie_rock_btn.click(set_indie_rock_prompt, inputs=bpm, outputs=instrumental_prompt)
+    funk_rock_btn.click(set_funk_rock_prompt, inputs=bpm, outputs=instrumental_prompt)
+    detroit_techno_btn.click(set_detroit_techno_prompt, inputs=bpm, outputs=instrumental_prompt)
+    deep_house_btn.click(set_deep_house_prompt, inputs=bpm, outputs=instrumental_prompt)
     gen_btn.click(
         generate_music,
-        inputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, crossfade_duration, num_variations, use_chunks],
+        inputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, crossfade_duration, variation_crossfade_duration, num_variations, chunk_count, use_chunks, bpm],
         outputs=[out_audio, status]
     )
     clr_btn.click(
         clear_inputs,
         inputs=None,
-        outputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, crossfade_duration, num_variations, use_chunks]
+        outputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, crossfade_duration, variation_crossfade_duration, num_variations, chunk_count, use_chunks, bpm]
     )
 
 # 9) TURN OFF OPENAPI/DOCS