Update barks.py

barks.py

@@ -1,4 +1,3 @@
-
 import os
 import torch
 import torchaudio
@@ -36,18 +35,16 @@ if device != "cuda":
     sys.exit(1)
 print(f"CUDA is available. Using GPU: {torch.cuda.get_device_name(0)}")
 
-# Initialize accelerator without cpu_offload
+# Initialize accelerator
 accelerator = Accelerator(mixed_precision="fp16")
 
-# Pre-run memory cleanup
-def aggressive_memory_cleanup():
+# Simplified memory cleanup to avoid freezes
+def memory_cleanup():
     torch.cuda.empty_cache()
     gc.collect()
-    torch.cuda.ipc_collect()
-    torch.cuda.synchronize()
-    print("Performed aggressive memory cleanup.")
+    print("Performed memory cleanup.")
 
-aggressive_memory_cleanup()
+memory_cleanup()
 
 # 2) LOAD MODELS
 try:
@@ -86,7 +83,7 @@ def print_resource_usage(stage: str):
     print("---------------")
 
 # Check available GPU memory
-def check_vram_availability(required_gb=3.0):  # Lowered threshold
+def check_vram_availability(required_gb=3.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
@@ -298,7 +295,7 @@ def generate_vocals(vocal_prompt: str, total_duration: int):
         
         # Move Bark model back to CPU
         bark_model = bark_model.to("cpu")
-        aggressive_memory_cleanup()
+        memory_cleanup()
         
         return vocal_segment, "✅ Vocals generated successfully."
     except Exception as e:
@@ -370,12 +367,12 @@ def generate_music(instrumental_prompt: str, vocal_prompt: str, cfg_scale: float
             os.remove(temp_wav_path)
             audio_segments.append(segment)
 
-            aggressive_memory_cleanup()
+            memory_cleanup()
             print_resource_usage(f"After Chunk {i+1} Generation")
 
         # Move MusicGen back to CPU
         musicgen_model = musicgen_model.to("cpu")
-        aggressive_memory_cleanup()
+        memory_cleanup()
 
         print("Combining instrumental chunks...")
         final_segment = audio_segments[0]
@@ -415,7 +412,7 @@ def generate_music(instrumental_prompt: str, vocal_prompt: str, cfg_scale: float
     except Exception as e:
         return None, f"❌ Generation failed: {e}"
     finally:
-        aggressive_memory_cleanup()
+        memory_cleanup()
 
 # Function to clear inputs
 def clear_inputs():
@@ -692,4 +689,4 @@ try:
     fastapi_app.redoc_url = None
     fastapi_app.openapi_url = None
 except Exception:
-    pass
+    pass