Update demos/musicgen_colab.py

demos/musicgen_colab.py

@@ -200,11 +200,11 @@ class Predictor:
                         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()
-                    return output, outputs_diffusion #Return the task id.  <-- Corrected return
+                    return task_id, (output, outputs_diffusion) #Return the task id.
                 else:
-                  return output, None  # <-- Corrected return
+                  return task_id, (output, None)
             except Exception as e:
-                return task_id, e #Should never be used, kept for consistency.
+                return task_id, e
         else:
             # Use the multiprocessing queue (multi-process mode)
             self.current_task_id += 1
@@ -238,9 +238,7 @@ _default_model_name = "facebook/musicgen-melody"
 def predict_full(model, model_path, depth, use_mbd, text, melody, duration, topk, topp, temperature, cfg_coef):
     # Initialize Predictor *INSIDE* the function
     predictor = Predictor(model, depth)
-
-    # Call predict() - this will return either (wav, diffusion_wav) or a task_id
-    prediction_result = predictor.predict(
+    task_id, (wav, diffusion_wav) = predictor.predict( # Unpack directly!
         text=text,
         melody=melody,
         duration=duration,
@@ -250,16 +248,7 @@ def predict_full(model, model_path, depth, use_mbd, text, melody, duration, topk
         temperature=temperature,
         cfg_coef=cfg_coef,
     )
-
-    # Handle daemon and non-daemon cases
-    if predictor.is_daemon:
-        wav, diffusion_wav = prediction_result  # Direct unpacking (daemon mode)
-    else:
-        # Get the result using the task_id (multi-process mode)
-        task_id = prediction_result
-        wav, diffusion_wav = predictor.get_result(task_id)
-
-    # Save and return audio files (rest of the function remains the same)
+    # Save and return audio files
     wav_paths = []
     video_paths = []
     # Save standard output
@@ -274,7 +263,6 @@ def predict_full(model, model_path, depth, use_mbd, text, melody, duration, topk
         video_paths.append(video_path)
         file_cleaner.add(file.name)
         file_cleaner.add(video_path)
-
     # Save MBD output if used
     if diffusion_wav is not None:
         with NamedTemporaryFile("wb", suffix=".wav", delete=False) as file:
@@ -288,15 +276,13 @@ def predict_full(model, model_path, depth, use_mbd, text, melody, duration, topk
             video_paths.append(video_path)
             file_cleaner.add(file.name)
             file_cleaner.add(video_path)
-
     # Shutdown predictor to prevent hanging processes!
-    if not predictor.is_daemon:  # Important!
+    if not predictor.is_daemon: # Important!
         predictor.shutdown()
-
     if use_mbd:
-        return video_paths[0], wav_paths[0], video_paths[1], wav_paths[1]
+         return video_paths[0], wav_paths[0], video_paths[1], wav_paths[1]
     return video_paths[0], wav_paths[0], None, None
-    
+
 def toggle_audio_src(choice):
     if choice == "mic":
         return gr.update(sources="microphone", value=None, label="Microphone")