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=25,  # Default chunk duration
+        duration=20,  # Default chunk duration
         two_step_cfg=False  # Disable two-step CFG for stability
     )
 except Exception as e:
@@ -58,6 +58,15 @@ def print_resource_usage(stage: str):
     print(f"CPU Memory Used: {psutil.virtual_memory().percent}%")
     print("---------------")
 
+# Check available GPU memory
+def check_vram_availability(required_gb=6.0):
+    total_vram = torch.cuda.get_device_properties(0).total_memory / (1024**3)
+    allocated_vram = torch.cuda.memory_allocated() / (1024**3)
+    available_vram = total_vram - allocated_vram
+    if available_vram < required_gb:
+        print(f"WARNING: Low VRAM available ({available_vram:.2f} GB). Consider reducing total_duration or num_variations.")
+    return available_vram >= required_gb
+
 # 4) GENRE PROMPT FUNCTIONS
 def set_red_hot_chili_peppers_prompt(bpm):
     rhythm = "strong rhythmic steps" if bpm > 120 else "groovy rhythmic flow"
@@ -158,6 +167,10 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
         output_files = []
         variation_segments = []
 
+        # Check VRAM availability
+        if not check_vram_availability(required_gb=6.0):
+            return None, "⚠️ Insufficient VRAM for generation. Reduce total_duration or num_variations."
+
         # 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)
@@ -170,8 +183,8 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
 
             if use_chunks:
                 # Chunked generation
-                num_chunks = max(1, min(chunk_count, total_duration // 25))  # Cap at 25 seconds per chunk
-                chunk_duration = total_duration / num_chunks
+                num_chunks = max(1, min(chunk_count, total_duration // 20))  # Cap at 20 seconds per chunk
+                chunk_duration = min(total_duration / num_chunks, 20)  # Max 20 seconds
                 overlap_duration = min(1.0, crossfade_duration / 1000.0)
                 generation_duration = chunk_duration + overlap_duration
                 audio_chunks = []
@@ -218,6 +231,7 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
 
                     torch.cuda.empty_cache()
                     gc.collect()
+                    torch.cuda.synchronize()
                     time.sleep(0.5)
                     print_resource_usage(f"After Chunk {i+1} Generation (Variation {var+1})")
 
@@ -273,6 +287,7 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
 
                 torch.cuda.empty_cache()
                 gc.collect()
+                torch.cuda.synchronize()
                 time.sleep(0.5)
                 print_resource_usage(f"After Full Track Generation (Variation {var+1})")
 
@@ -312,6 +327,7 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
     finally:
         torch.cuda.empty_cache()
         gc.collect()
+        torch.cuda.synchronize()
 
 # Function to toggle interactivity of chunk-related sliders
 def toggle_chunk_interactivity(use_chunks):
@@ -322,7 +338,7 @@ def toggle_chunk_interactivity(use_chunks):
     )
 
 def clear_inputs():
-    return "", 3.0, 250, 0.9, 1.0, 30, 1000, 1000, 1, 2, True, 120
+    return "", 3.0, 250, 0.9, 1.0, 20, 1000, 1000, 1, 2, True, 120
 
 # 7) CUSTOM CSS
 css = """
@@ -461,7 +477,7 @@ with gr.Blocks(css=css) as demo:
             label="Total Duration (seconds)",
             minimum=10,
             maximum=90,
-            value=30,
+            value=20,
             step=1,
             info="Total duration of the track (10 to 90 seconds)."
         )
@@ -479,7 +495,7 @@ with gr.Blocks(css=css) as demo:
             maximum=2000,
             value=1000,
             step=100,
-            info="Crossfade duration between variations (only used if chunking is enabled)."
+            info="Crossfade duration between variations."
         )
         num_variations = gr.Slider(
             label="Number of Variations",
@@ -495,7 +511,7 @@ with gr.Blocks(css=css) as demo:
             maximum=8,
             value=2,
             step=1,
-            info="Number of chunks to split the track into (only used if chunking is enabled)."
+            info="Number of chunks to split the track into (only used if chunking is enabled, max 20 seconds per chunk)."
         )
         use_chunks = gr.Checkbox(
             label="Generate in Chunks",