Update app.py

quality of life updates

app.py

@@ -11,6 +11,7 @@ from pydub import AudioSegment
 from audiocraft.models import MusicGen
 from torch.cuda.amp import autocast
 import warnings
+import random
 
 # Suppress warnings for cleaner output
 warnings.filterwarnings("ignore")
@@ -71,7 +72,13 @@ 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_foo_fighters_prompt():
-    return "Alternative rock with punchy guitar riffs, tight drums, melodic hooks, and a Foo Fighters-inspired anthemic energy with gritty verses."
+    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."
 
 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."
@@ -126,34 +133,89 @@ 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 = 1):
+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):
     global musicgen_model
     if not instrumental_prompt.strip():
         return None, "⚠️ Please enter a valid instrumental prompt!"
     try:
         start_time = time.time()
         total_duration = min(max(total_duration, 10), 90)
-        chunk_duration = 15
-        num_chunks = max(1, total_duration // chunk_duration)
-        chunk_duration = total_duration / num_chunks
-        overlap_duration = min(1.0, crossfade_duration / 1000.0)
-        generation_duration = chunk_duration + overlap_duration
-
-        output_files = []
         sample_rate = musicgen_model.sample_rate
+        output_files = []
 
         for var in range(num_variations):
             print(f"Generating variation {var+1}/{num_variations}...")
-            audio_chunks = []
             seed = 42 + var  # Use different seeds for variations
             torch.manual_seed(seed)
             np.random.seed(seed)
 
-            for i in range(num_chunks):
-                chunk_prompt = instrumental_prompt
-                print(f"Generating chunk {i+1}/{num_chunks} for variation {var+1} on GPU (prompt: {chunk_prompt})...")
+            if use_chunks:
+                # Chunked generation
+                chunk_duration = 15
+                num_chunks = max(1, total_duration // chunk_duration)
+                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
+                    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,
+                        cfg_coef=cfg_scale
+                    )
+
+                    print_resource_usage(f"Before Chunk {i+1} Generation (Variation {var+1})")
+
+                    with torch.no_grad():
+                        with autocast():
+                            audio_chunk = musicgen_model.generate([chunk_prompt], progress=True)[0]
+
+                    audio_chunk = audio_chunk.cpu().to(dtype=torch.float32)
+                    if audio_chunk.dim() == 1:
+                        audio_chunk = torch.stack([audio_chunk, audio_chunk], dim=0)
+                    elif audio_chunk.dim() == 2 and audio_chunk.shape[0] == 1:
+                        audio_chunk = torch.cat([audio_chunk, audio_chunk], dim=0)
+                    elif audio_chunk.dim() == 2 and audio_chunk.shape[0] != 2:
+                        audio_chunk = audio_chunk[:1, :]
+                        audio_chunk = torch.cat([audio_chunk, audio_chunk], dim=0)
+                    elif audio_chunk.dim() > 2:
+                        audio_chunk = audio_chunk.view(2, -1)
+
+                    if audio_chunk.shape[0] != 2:
+                        raise ValueError(f"Expected stereo audio with shape (2, samples), got shape {audio_chunk.shape}")
+
+                    temp_wav_path = f"temp_chunk_{var}_{i}.wav"
+                    chunk_path = f"chunk_{var}_{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 (Variation {var+1})")
+
+                print(f"Combining audio chunks for variation {var+1}...")
+                final_segment = AudioSegment.from_mp3(audio_chunks[0])
+                for i in range(1, len(audio_chunks)):
+                    next_segment = AudioSegment.from_mp3(audio_chunks[i])
+                    next_segment = next_segment + 1
+                    final_segment = final_segment.append(next_segment, crossfade=crossfade_duration)
+
+                final_segment = final_segment[:total_duration * 1000]
+            else:
+                # Single-shot generation
+                print(f"Generating full track for variation {var+1} on GPU (prompt: {instrumental_prompt})...")
                 musicgen_model.set_generation_params(
-                    duration=generation_duration,
+                    duration=total_duration,
                     use_sampling=True,
                     top_k=top_k,
                     top_p=top_p,
@@ -161,11 +223,11 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
                     cfg_coef=cfg_scale
                 )
 
-                print_resource_usage(f"Before Chunk {i+1} Generation (Variation {var+1})")
+                print_resource_usage(f"Before Full Track Generation (Variation {var+1})")
 
                 with torch.no_grad():
                     with autocast():
-                        audio_chunk = musicgen_model.generate([chunk_prompt], progress=True)[0]
+                        audio_chunk = musicgen_model.generate([instrumental_prompt], progress=True)[0]
 
                 audio_chunk = audio_chunk.cpu().to(dtype=torch.float32)
                 if audio_chunk.dim() == 1:
@@ -181,27 +243,15 @@ 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}")
 
-                temp_wav_path = f"temp_chunk_{var}_{i}.wav"
-                chunk_path = f"chunk_{var}_{i}.mp3"
+                temp_wav_path = f"temp_full_{var}.wav"
                 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")
+                final_segment = AudioSegment.from_wav(temp_wav_path)
                 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 (Variation {var+1})")
-
-            print(f"Combining audio chunks for variation {var+1}...")
-            final_segment = AudioSegment.from_mp3(audio_chunks[0])
-            for i in range(1, len(audio_chunks)):
-                next_segment = AudioSegment.from_mp3(audio_chunks[i])
-                next_segment = next_segment + 1
-                final_segment = final_segment.append(next_segment, crossfade=crossfade_duration)
-
-            final_segment = final_segment[:total_duration * 1000]
+                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)
@@ -220,8 +270,9 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
             print(f"Saved final audio to {mp3_path}")
             output_files.append(mp3_path)
 
-            for chunk_path in audio_chunks:
-                os.remove(chunk_path)
+            if use_chunks:
+                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")
@@ -234,8 +285,12 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
         torch.cuda.empty_cache()
         gc.collect()
 
+# Function to toggle crossfade_duration interactivity
+def toggle_crossfade_interactivity(use_chunks):
+    return gr.update(interactive=use_chunks)
+
 def clear_inputs():
-    return "", 3.0, 250, 0.9, 1.0, 30, 500, 1
+    return "", 3.0, 250, 0.9, 1.0, 30, 500, 1, True
 
 # 7) CUSTOM CSS
 css = """
@@ -384,7 +439,7 @@ with gr.Blocks(css=css) as demo:
             maximum=2000,
             value=500,
             step=100,
-            info="Crossfade duration between chunks."
+            info="Crossfade duration between chunks (only used if chunking is enabled)."
         )
         num_variations = gr.Slider(
             label="Number of Variations",
@@ -394,6 +449,11 @@ with gr.Blocks(css=css) as demo:
             step=1,
             info="Number of different versions to generate with varying random seeds."
         )
+        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)."
+        )
         with gr.Row(elem_classes="action-buttons"):
             gen_btn = gr.Button("Generate Music")
             clr_btn = gr.Button("Clear Inputs")
@@ -402,6 +462,9 @@ 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])
@@ -418,13 +481,13 @@ with gr.Blocks(css=css) as demo:
     deep_house_btn.click(set_deep_house_prompt, inputs=None, outputs=[instrumental_prompt])
     gen_btn.click(
         generate_music,
-        inputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, crossfade_duration, num_variations],
+        inputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, crossfade_duration, num_variations, use_chunks],
         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]
+        outputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, crossfade_duration, num_variations, use_chunks]
     )
 
 # 9) TURN OFF OPENAPI/DOCS