Update barks.py

barks.py

@@ -1,4 +1,4 @@
-```python
+
 import os
 import torch
 import torchaudio
@@ -36,8 +36,8 @@ if device != "cuda":
     sys.exit(1)
 print(f"CUDA is available. Using GPU: {torch.cuda.get_device_name(0)}")
 
-# Initialize accelerator with enhanced CPU offloading
-accelerator = Accelerator(mixed_precision="fp16", cpu_offload=True)
+# Initialize accelerator without cpu_offload
+accelerator = Accelerator(mixed_precision="fp16")
 
 # Pre-run memory cleanup
 def aggressive_memory_cleanup():
@@ -51,7 +51,7 @@ aggressive_memory_cleanup()
 
 # 2) LOAD MODELS
 try:
-    print("Loading MusicGen medium model into VRAM...")
+    print("Loading MusicGen medium model into system RAM...")
     local_model_path = "./models/musicgen-medium"
     if not os.path.exists(local_model_path):
         print(f"ERROR: Local model path {local_model_path} does not exist.")
@@ -92,7 +92,7 @@ def check_vram_availability(required_gb=3.0):  # Lowered threshold
     available_vram = total_vram - allocated_vram
     if available_vram < required_gb:
         print(f"WARNING: Low VRAM available ({available_vram:.2f} GB < {required_gb:.2f} GB required).")
-        print("Reduce total_duration, chunk_duration, or enable more CPU offloading.")
+        print("Reduce total_duration, chunk_duration, or skip vocals.")
     print(f"Total VRAM: {total_vram:.2f} GB, Available: {available_vram:.2f} GB")
     return available_vram >= required_gb
 
@@ -274,10 +274,10 @@ def generate_vocals(vocal_prompt: str, total_duration: int):
     try:
         print("Generating vocals with Bark...")
         # Move Bark model to GPU
-        bark_model = accelerator.prepare(bark_model)
+        bark_model = bark_model.to("cuda")
         
         # Process vocal prompt
-        inputs = bark_processor(vocal_prompt, return_tensors="pt").to(accelerator.device)
+        inputs = bark_processor(vocal_prompt, return_tensors="pt").to("cuda")
         
         # Generate vocals with mixed precision
         with torch.no_grad(), autocast():
@@ -330,7 +330,7 @@ def generate_music(instrumental_prompt: str, vocal_prompt: str, cfg_scale: float
         np.random.seed(seed)
 
         # Move MusicGen to GPU
-        musicgen_model = accelerator.prepare(musicgen_model)
+        musicgen_model = musicgen_model.to("cuda")
 
         for i in range(num_chunks):
             chunk_prompt = instrumental_prompt
@@ -567,7 +567,7 @@ with gr.Blocks(css=css) as demo:
                 maximum=1.0,
                 value=0.9,
                 step=0.05,
-                pair_with="Keeps tokens with cumulative probability above p."
+                info="Keeps tokens with cumulative probability above p."
             )
             temperature = gr.Slider(
                 label="Temperature 🔥",
@@ -693,4 +693,3 @@ try:
     fastapi_app.openapi_url = None
 except Exception:
     pass
-```