Update demos/musicgen_app.py

demos/musicgen_app.py

@@ -1,4 +1,5 @@
-import spaces
+import spaces  # <--- IMPORTANT: Add this import
+
 import argparse
 import logging
 import os
@@ -19,6 +20,7 @@ from audiocraft.models.encodec import InterleaveStereoCompressionModel
 from audiocraft.models import MusicGen, MultiBandDiffusion
 import multiprocessing as mp
 
+
 # --- Utility Functions and Classes ---
 
 class FileCleaner:  # Unchanged
@@ -50,7 +52,7 @@ def make_waveform(*args, **kwargs):  # Unchanged
         return out
 
 # --- Worker Process ---
-@spaces.GPU(required=True)
+#This stays the same, since the worker is designed for this purpose
 def model_worker(model_name: str, task_queue: mp.Queue, result_queue: mp.Queue):
     """
     Persistent worker process that loads the model and handles prediction tasks.
@@ -115,6 +117,7 @@ def model_worker(model_name: str, task_queue: mp.Queue, result_queue: mp.Queue):
 # --- Gradio Interface Functions ---
 
 class Predictor:
+    #This stays the same, this is the intended design
     def __init__(self, model_name: str):
         self.task_queue = mp.Queue()
         self.result_queue = mp.Queue()
@@ -161,22 +164,19 @@ class Predictor:
         """
         Shuts down the worker process.
         """
-        self.task_queue.put(None)  # Send sentinel value to stop the worker
-        self.process.join()  # Wait for the process to terminate
+        if self.process.is_alive():
+            self.task_queue.put(None)  # Send sentinel value to stop the worker
+            self.process.join()  # Wait for the process to terminate
 
+# NO GLOBAL PREDICTOR ANYMORE
 
-# Global predictor instance
-_predictor = None
-
-def get_predictor(model_name:str = 'facebook/musicgen-melody'):
-    global _predictor
-    if _predictor is None:
-        _predictor = Predictor(model_name)
-    return _predictor
+_default_model_name = 'facebook/musicgen-melody'
 
+@spaces.GPU(duration=60)  # <--- IMPORTANT: Add this decorator
 def predict_full(model, model_path, use_mbd, text, melody, duration, topk, topp, temperature, cfg_coef):
+    # Initialize Predictor *INSIDE* the function
+    predictor = Predictor(model)
 
-    predictor = get_predictor(model)
     task_id = predictor.predict(
         text=text,
         melody=melody,
@@ -214,6 +214,8 @@ def predict_full(model, model_path, use_mbd, text, melody, duration, topk, topp,
             wav_paths.append(file.name)
             video_paths.append(make_waveform(file.name)) # Make and clean up video
             file_cleaner.add(file.name)
+    # Shutdown predictor to prevent hanging processes!
+    predictor.shutdown()
 
     if use_mbd:
          return video_paths[0], wav_paths[0], video_paths[1], wav_paths[1]
@@ -316,7 +318,7 @@ def ui_full(launch_kwargs):
                     "facebook/musicgen-melody",
                     "Default"
                 ],
-                [
+                 [
                     "lofi slow bpm electro chill with organic samples",
                     None,
                     "facebook/musicgen-medium",