Update barks.py

barks.py

@@ -1,3 +1,4 @@
+
 import os
 import torch
 import torchaudio
@@ -65,9 +66,9 @@ except Exception as e:
     sys.exit(1)
 
 try:
-    print("Loading Bark small model into system RAM with 4-bit quantization...")
+    print("Loading Bark small model into GPU VRAM with 4-bit quantization...")
     bark_processor = AutoProcessor.from_pretrained("suno/bark-small")
-    bark_model = BarkModel.from_pretrained("suno/bark-small", load_in_4bit=True, device_map="cpu")  # Quantize and offload to CPU
+    bark_model = BarkModel.from_pretrained("suno/bark-small", load_in_4bit=True, device_map="cuda")  # Load directly to GPU with quantization
 except Exception as e:
     print(f"ERROR: Failed to load Bark model: {e}")
     print("Ensure Bark model weights and bitsandbytes are installed.")
@@ -270,10 +271,7 @@ def generate_vocals(vocal_prompt: str, total_duration: int):
     
     try:
         print("Generating vocals with Bark...")
-        # Move Bark model to GPU
-        bark_model = bark_model.to("cuda")
-        
-        # Process vocal prompt
+        # Process vocal prompt, ensuring inputs are on cuda
         inputs = bark_processor(vocal_prompt, return_tensors="pt").to("cuda")
         
         # Generate vocals with mixed precision
@@ -293,8 +291,6 @@ def generate_vocals(vocal_prompt: str, total_duration: int):
         if len(vocal_segment) < total_duration * 1000:
             vocal_segment = vocal_segment + AudioSegment.silent(duration=(total_duration * 1000 - len(vocal_segment)))
         
-        # Move Bark model back to CPU
-        bark_model = bark_model.to("cpu")
         memory_cleanup()
         
         return vocal_segment, "✅ Vocals generated successfully."
@@ -407,10 +403,6 @@ def generate_music(instrumental_prompt: str, vocal_prompt: str, cfg_scale: float
     finally:
         memory_cleanup()
 
-# Function to clear inputs
-def clear_inputs():
-    return "", "", 3.0, 250, 0.9, 1.0, 30, 5, 1000, 120, "none", "none", "none", "none", "none"
-
 # 8) CUSTOM CSS
 css = """
 body { 
@@ -682,4 +674,4 @@ try:
     fastapi_app.redoc_url = None
     fastapi_app.openapi_url = None
 except Exception:
-    pass
+    pass