Update app.py

roll back i messed up ><

app.py

@@ -47,7 +47,7 @@ def print_resource_usage(stage: str):
     print(f"GPU Memory Reserved: {torch.cuda.memory_reserved() / (1024**3):.2f} GB")
     print("---------------")
 
-# 4) GENRE PROMPT FUNCTIONS (Unchanged)
+# 4) GENRE PROMPT FUNCTIONS (Updated for consistency, more instruments, and popular styles)
 def set_rock_prompt():
     return "Hard rock with dynamic electric guitars, heavy steady drums, deep groovy bass, subtle organ layers, and a hint of Red Hot Chili Peppers' funky rock energy, maintaining a cohesive structure throughout"
 
@@ -77,34 +77,23 @@ def set_deep_house_prompt():
 
 # 5) AUDIO PROCESSING FUNCTIONS
 def apply_chorus(segment):
-    delayed = segment - 6
+    # Enhanced chorus effect for richer sound
+    delayed = segment - 6  # Reduced gain to -6 dB for a subtler effect
     delayed = delayed.set_frame_rate(segment.frame_rate)
     return segment.overlay(delayed, position=20)
 
-def apply_reverb(segment):
-    # Simulate reverb by overlaying multiple delayed copies with decreasing amplitude
-    reverb_segment = segment
-    for delay_ms, gain_db in [(50, -10), (100, -15), (150, -20)]:
-        delayed = segment - gain_db
-        delayed = delayed.set_frame_rate(segment.frame_rate)
-        reverb_segment = reverb_segment.overlay(delayed, position=delay_ms)
-    return reverb_segment
-
 def apply_eq(segment):
-    # Use only low-pass and high-pass filters, as equalizer is not available in pydub
-    segment = segment.low_pass_filter(8000)
-    segment = segment.high_pass_filter(80)
+    # Adjusted EQ for a more balanced sound
+    segment = segment.low_pass_filter(8000)  # Raised cutoff to 8kHz for brighter highs
+    segment = segment.high_pass_filter(80)  # Lowered cutoff to 80Hz for deeper bass
     return segment
 
 def apply_limiter(segment, max_db=-3.0):
+    # Apply limiter with a higher threshold to preserve dynamics
     if segment.dBFS > max_db:
         segment = segment - (segment.dBFS - max_db)
     return segment
 
-def apply_final_gain(segment, target_db=-12.0):
-    gain_adjustment = target_db - segment.dBFS
-    return segment + gain_adjustment
-
 # 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):
     global musicgen_model
@@ -113,22 +102,30 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
     try:
         start_time = time.time()
 
+        # Ensure total duration is within reasonable bounds (up to 90 seconds for longer tracks)
         total_duration = min(max(total_duration, 10), 90)
-        chunk_duration = total_duration  # Single chunk to minimize overhead
-        num_chunks = 1  # Single chunk generation
+        chunk_duration = 15
+        num_chunks = max(2, (total_duration + chunk_duration - 1) // chunk_duration)
+        chunk_duration = total_duration / num_chunks
+
+        # Generate slightly longer chunks for overlap
+        overlap_duration = min(1.0, crossfade_duration / 1000.0)
+        generation_duration = chunk_duration + overlap_duration
 
         audio_chunks = []
         sample_rate = musicgen_model.sample_rate
 
+        # Generate audio in chunks with a consistent prompt
         for i in range(num_chunks):
-            chunk_prompt = instrumental_prompt
+            chunk_prompt = instrumental_prompt  # Use the same prompt for all chunks
+
             print(f"Generating chunk {i+1}/{num_chunks} on GPU (prompt: {chunk_prompt})...")
             musicgen_model.set_generation_params(
-                duration=chunk_duration,
+                duration=generation_duration,
                 use_sampling=True,
-                top_k=250,
-                top_p=0.9,
-                temperature=0.9,
+                top_k=top_k,
+                top_p=top_p,
+                temperature=temperature,
                 cfg_coef=cfg_scale
             )
 
@@ -152,38 +149,36 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
             if audio_chunk.shape[0] != 2:
                 raise ValueError(f"Expected stereo audio with shape (2, samples), got shape {audio_chunk.shape}")
 
-            # Process in memory using pydub without intermediate file I/O
-            audio_array = audio_chunk.numpy()
-            audio_array = (audio_array * 32767).astype(np.int16)  # Convert to 16-bit PCM
-            segment = AudioSegment(
-                audio_array.tobytes(),
-                frame_rate=sample_rate,
-                sample_width=2,  # 16-bit
-                channels=2
-            )
-            audio_chunks.append(segment)
+            temp_wav_path = f"temp_chunk_{i}.wav"
+            chunk_path = f"chunk_{i}.mp3"
+            torchaudio.save(temp_wav_path, audio_chunk, sample_rate, bits_per_sample=24)
+            segment = AudioSegment.from_wav(temp_wav_path)
+            segment.export(chunk_path, format="mp3", bitrate="320k")
+            os.remove(temp_wav_path)
+            audio_chunks.append(chunk_path)
 
             torch.cuda.empty_cache()
             gc.collect()
             time.sleep(0.5)
             print_resource_usage(f"After Chunk {i+1} Generation")
 
+        # Combine chunks with crossfade
         print("Combining audio chunks...")
-        final_segment = audio_chunks[0]
+        final_segment = AudioSegment.from_mp3(audio_chunks[0])
         for i in range(1, len(audio_chunks)):
-            next_segment = audio_chunks[i]
-            next_segment = next_segment + 1
+            next_segment = AudioSegment.from_mp3(audio_chunks[i])
+            next_segment = next_segment + 1  # Reduced gain boost to +1 dB
             final_segment = final_segment.append(next_segment, crossfade=crossfade_duration)
 
+        # Trim to exact total duration
         final_segment = final_segment[:total_duration * 1000]
 
+        # Post-process with improved dynamics
         print("Post-processing final track...")
         final_segment = apply_eq(final_segment)
         final_segment = apply_chorus(final_segment)
-        final_segment = apply_reverb(final_segment)
-        final_segment = apply_limiter(final_segment, max_db=-3.0)
-        final_segment = final_segment.normalize(headroom=-6.0)
-        final_segment = apply_final_gain(final_segment, target_db=-12.0)
+        final_segment = apply_limiter(final_segment, max_db=-3.0)  # Apply limiter only once with higher threshold
+        final_segment = final_segment.normalize(headroom=-6.0)  # Increased headroom to -6 dB
 
         mp3_path = "output_cleaned.mp3"
         final_segment.export(
@@ -194,6 +189,9 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
         )
         print(f"Saved final audio to {mp3_path}")
 
+        for chunk_path in audio_chunks:
+            os.remove(chunk_path)
+
         print_resource_usage("After Final Generation")
         print(f"Total Generation Time: {time.time() - start_time:.2f} seconds")
 
@@ -205,7 +203,7 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
         gc.collect()
 
 def clear_inputs():
-    return "", 3.0, 250, 0.9, 0.9, 45, 750
+    return "", 3.0, 300, 0.95, 1.0, 30, 500
 
 # 7) CUSTOM CSS (Unchanged)
 css = """
@@ -376,7 +374,7 @@ with gr.Blocks(css=css) as demo:
             label="Top-K Sampling",
             minimum=10,
             maximum=500,
-            value=250,
+            value=300,
             step=10,
             info="Limits sampling to the top k most likely tokens. Higher values increase diversity."
         )
@@ -384,7 +382,7 @@ with gr.Blocks(css=css) as demo:
             label="Top-P Sampling (Nucleus Sampling)",
             minimum=0.0,
             maximum=1.0,
-            value=0.9,
+            value=0.95,
             step=0.1,
             info="Keeps tokens with cumulative probability above p. Higher values increase diversity."
         )
@@ -392,7 +390,7 @@ with gr.Blocks(css=css) as demo:
             label="Temperature",
             minimum=0.1,
             maximum=2.0,
-            value=0.9,
+            value=1.0,
             step=0.1,
             info="Controls randomness. Higher values make output more diverse but less predictable."
         )
@@ -400,7 +398,7 @@ with gr.Blocks(css=css) as demo:
             label="Total Duration (seconds)",
             minimum=10,
             maximum=90,
-            value=45,
+            value=30,
             step=1,
             info="Total duration of the track (10 to 90 seconds)."
         )
@@ -408,7 +406,7 @@ with gr.Blocks(css=css) as demo:
             label="Crossfade Duration (ms)",
             minimum=100,
             maximum=2000,
-            value=750,
+            value=500,
             step=100,
             info="Crossfade duration between chunks for smoother transitions."
         )