Update demos/musicgen_app.py

demos/musicgen_app.py

@@ -186,7 +186,7 @@ class Predictor:
                         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
+                        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 task_id, (output, outputs_diffusion) #Return the task id.
@@ -234,7 +234,7 @@ def predict_full(model, model_path, use_mbd, text, melody, duration, topk, topp,
     # Initialize Predictor *INSIDE* the function
     predictor = Predictor(model)
 
-    task_id = predictor.predict(
+    task_id, (wav, diffusion_wav) = predictor.predict( # Unpack directly!
         text=text,
         melody=melody,
         duration=duration,
@@ -245,8 +245,6 @@ def predict_full(model, model_path, use_mbd, text, melody, duration, topk, topp,
         cfg_coef=cfg_coef,
     )
 
-    wav, diffusion_wav = predictor.get_result(task_id)
-
     # Save and return audio files
     wav_paths = []
     video_paths = []
@@ -272,7 +270,9 @@ def predict_full(model, model_path, use_mbd, text, melody, duration, topk, topp,
             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 not predictor.is_daemon: # Important!
+        predictor.shutdown()
 
     if use_mbd:
          return video_paths[0], wav_paths[0], video_paths[1], wav_paths[1]
@@ -293,6 +293,7 @@ def toggle_diffusion(choice):
         return [gr.update(visible=False)] * 2
 # --- Gradio UI ---
 
+
 def ui_full(launch_kwargs):
     with gr.Blocks() as interface:
         gr.Markdown(