Update demos/musicgen_app.py

demos/musicgen_app.py

@@ -1,5 +1,4 @@
 import spaces  # <--- IMPORTANT: Add this import
-
 import argparse
 import logging
 import os
@@ -20,7 +19,6 @@ from audiocraft.models.encodec import InterleaveStereoCompressionModel
 from audiocraft.models import MusicGen, MultiBandDiffusion
 import multiprocessing as mp
 
-
 # --- Utility Functions and Classes ---
 
 class FileCleaner:  # Unchanged
@@ -51,20 +49,20 @@ def make_waveform(*args, **kwargs):  # Unchanged
         print("Make a video took", time.time() - be)
         return out
 
-# --- Worker Process ---
-#This stays the same, since the worker is designed for this purpose
+# --- Worker Process --- (Modified for conditional use)
+
 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.
+    Persistent worker process (used when NOT running as a daemon).
     """
     try:
         device = 'cuda' if torch.cuda.is_available() else 'cpu'
         model = MusicGen.get_pretrained(model_name, device=device)
-        mbd = MultiBandDiffusion.get_mbd_musicgen(device=device)  # Load MBD here too
+        mbd = MultiBandDiffusion.get_mbd_musicgen(device=device)
 
         while True:
             task = task_queue.get()
-            if task is None:  # Sentinel value to exit
+            if task is None:
                 break
 
             task_id, text, melody, duration, use_diffusion, gen_params = task
@@ -104,75 +102,134 @@ def model_worker(model_name: str, task_queue: mp.Queue, result_queue: mp.Queue):
                         assert outputs_diffusion.shape[1] == 1  # output is mono
                         outputs_diffusion = rearrange(outputs_diffusion, '(s b) c t -> b (s c) t', s=2)
                     outputs_diffusion = outputs_diffusion.detach().cpu()
-                    result_queue.put((task_id, (output, outputs_diffusion))) # Send BOTH results.
+                    result_queue.put((task_id, (output, outputs_diffusion)))
                 else:
-                    result_queue.put((task_id, (output, None)))  # Send back the result
+                    result_queue.put((task_id, (output, None)))
 
             except Exception as e:
-                result_queue.put((task_id, e))  # Send back the exception
+                result_queue.put((task_id, e))
 
     except Exception as e:
-        result_queue.put((-1,e)) #Fatal error on loading.
+        result_queue.put((-1, e))
 
-# --- Gradio Interface Functions ---
+
+# --- Predictor Class (Modified for conditional process creation) ---
 
 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()
-        self.process = mp.Process(target=model_worker, args=(model_name, self.task_queue, self.result_queue))
-        self.process.start()
-        self.current_task_id = 0
-        self._check_initialization()
-
+        self.model_name = model_name
+        self.is_daemon = mp.current_process().daemon
+        if self.is_daemon:
+            # Running in a daemonic process (e.g., on Spaces)
+            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+            self.model = MusicGen.get_pretrained(self.model_name, device=self.device)
+            self.mbd = MultiBandDiffusion.get_mbd_musicgen(device=self.device)  # Load MBD here too
+            self.current_task_id = 0 # Initialize task ID
+        else:
+            # Running in a non-daemonic process (e.g., locally)
+            self.task_queue = mp.Queue()
+            self.result_queue = mp.Queue()
+            self.process = mp.Process(
+                target=model_worker, args=(self.model_name, self.task_queue, self.result_queue)
+            )
+            self.process.start()
+            self.current_task_id = 0
+            self._check_initialization()
 
     def _check_initialization(self):
-        """Check if the worker process initialized successfully."""
-        # Give it some time to either load or report failure.
-        time.sleep(2)
-        try:
-            task_id, result = self.result_queue.get(timeout=3) # Get result from model_worker
-
-            if isinstance(result, Exception):
-                if task_id == -1:
-                    raise RuntimeError("Model loading failed in worker process.") from result
-        except:
-            pass # Expected if model loads fast enough
+        """Check if the worker process initialized successfully (only in non-daemon mode)."""
+        if not self.is_daemon:
+            time.sleep(2)
+            try:
+                task_id, result = self.result_queue.get(timeout=3)
+                if isinstance(result, Exception):
+                    if task_id == -1:
+                        raise RuntimeError("Model loading failed in worker process.") from result
+            except:
+                pass
 
     def predict(self, text, melody, duration, use_diffusion, **gen_params):
-        """
-        Submits a prediction task to the worker process.
-        """
-        self.current_task_id += 1
-        task = (self.current_task_id, text, melody, duration, use_diffusion, gen_params)
-        self.task_queue.put(task)
-        return self.current_task_id
+        """Submits a prediction task."""
+        if self.is_daemon:
+             # Directly perform the prediction (single-process mode)
+            self.current_task_id +=1
+            task_id = self.current_task_id
+            try:
+                self.model.set_generation_params(duration=duration, **gen_params)
+                target_sr = self.model.sample_rate
+                target_ac = 1
+                processed_melody = None
+                if melody:
+                    sr, melody_data = melody
+                    melody_tensor = torch.from_numpy(melody_data).to(self.device).float().t()
+                    if melody_tensor.ndim == 1:
+                        melody_tensor = melody_tensor.unsqueeze(0)
+                    melody_tensor = melody_tensor[..., :int(sr * duration)]
+                    processed_melody = convert_audio(melody_tensor, sr, target_sr, target_ac)
+
+                if processed_melody is not None:
+                    output, tokens = self.model.generate_with_chroma(
+                        descriptions=[text],
+                        melody_wavs=[processed_melody],
+                        melody_sample_rate=target_sr,
+                        progress=True,
+                        return_tokens=True
+                    )
+                else:
+                    output, tokens = self.model.generate([text], progress=True, return_tokens=True)
+
+                output = output.detach().cpu()
+
+                if use_diffusion:
+                    if isinstance(self.model.compression_model, InterleaveStereoCompressionModel):
+                        left, right = self.model.compression_model.get_left_right_codes(tokens)
+                        tokens = torch.cat([left, right])
+                    outputs_diffusion = self.mbd.tokens_to_wav(tokens)
+                    if isinstance(self.model.compression_model, InterleaveStereoCompressionModel):
+                        assert outputs_diffusion.shape[1] == 1
+                        outputs_diffusion = rearrange(outputs_diffusion, '(s b) c t -> b (s c) t', s=2)
+                    outputs_diffusion = outputs_diffusion.detach().cpu()
+                    return task_id, (output, outputs_diffusion) #Return the task id.
+                else:
+                  return task_id, (output, None)
+
+
+            except Exception as e:
+                return task_id, e
+
+        else:
+            # Use the multiprocessing queue (multi-process mode)
+            self.current_task_id += 1
+            task = (self.current_task_id, text, melody, duration, use_diffusion, gen_params)
+            self.task_queue.put(task)
+            return self.current_task_id
 
     def get_result(self, task_id):
-        """
-        Retrieves the result of a prediction task.  Blocks until the result is available.
-        """
-        while True: # Loop to get the correct task
-            result_task_id, result = self.result_queue.get()
-            if result_task_id == task_id:
-                if isinstance(result, Exception):
-                    raise result  # Re-raise the exception in the main process
-                return result # (wav, diffusion_wav) or (wav, None)
+        """Retrieves the result of a prediction task."""
+        if self.is_daemon:
+            # Results are returned directly by 'predict' in daemon mode
+            result_id, result = task_id, task_id #predictor return (task_id, results)
+        else:
+            # Get result from the queue (multi-process mode)
+            while True:
+                result_task_id, result = self.result_queue.get()
+                if result_task_id == task_id:
+                    break  # Found the correct result
+
+        if isinstance(result, Exception):
+            raise result
+        return result
 
     def shutdown(self):
-        """
-        Shuts down the worker process.
-        """
-        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
+        """Shuts down the worker process (if running)."""
+        if not self.is_daemon and self.process.is_alive():
+            self.task_queue.put(None)
+            self.process.join()
 
-# NO GLOBAL PREDICTOR ANYMORE
 
-_default_model_name = 'facebook/musicgen-melody'
+_default_model_name = "facebook/musicgen-melody"
 
-@spaces.GPU(duration=60)  # <--- IMPORTANT: Add this decorator
+@spaces.GPU(duration=60)  # Use the decorator for Spaces
 def predict_full(model, model_path, use_mbd, text, melody, duration, topk, topp, temperature, cfg_coef):
     # Initialize Predictor *INSIDE* the function
     predictor = Predictor(model)