Update barks.py

barks.py

@@ -48,13 +48,13 @@ memory_cleanup()
 
 # 2) LOAD MODELS
 try:
-    print("Loading MusicGen medium model into system RAM...")
+    print("Loading MusicGen medium model into GPU VRAM...")
     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.")
         print("Please download the MusicGen medium model weights and place them in the correct directory.")
         sys.exit(1)
-    musicgen_model = MusicGen.get_pretrained(local_model_path, device="cpu")  # Load to CPU initially
+    musicgen_model = MusicGen.get_pretrained(local_model_path, device="cuda")  # Load directly to GPU
     musicgen_model.set_generation_params(
         duration=5,  # Lower default chunk duration
         two_step_cfg=False  # Disable two-step CFG for stability
@@ -67,7 +67,7 @@ except Exception as e:
 try:
     print("Loading Bark small model into system RAM 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
+    bark_model = BarkModel.from_pretrained("suno/bark-small", load_in_4bit=True, device_map="cpu")  # Quantize and offload to CPU
 except Exception as e:
     print(f"ERROR: Failed to load Bark model: {e}")
     print("Ensure Bark model weights and bitsandbytes are installed.")
@@ -326,9 +326,6 @@ def generate_music(instrumental_prompt: str, vocal_prompt: str, cfg_scale: float
         torch.manual_seed(seed)
         np.random.seed(seed)
 
-        # Move MusicGen to GPU
-        musicgen_model = musicgen_model.to("cuda")
-
         for i in range(num_chunks):
             chunk_prompt = instrumental_prompt
             print(f"Generating chunk {i+1}/{num_chunks} on GPU (prompt: {chunk_prompt})...")
@@ -370,10 +367,6 @@ def generate_music(instrumental_prompt: str, vocal_prompt: str, cfg_scale: float
             memory_cleanup()
             print_resource_usage(f"After Chunk {i+1} Generation")
 
-        # Move MusicGen back to CPU
-        musicgen_model = musicgen_model.to("cpu")
-        memory_cleanup()
-
         print("Combining instrumental chunks...")
         final_segment = audio_segments[0]
         for i in range(1, len(audio_segments)):