Update demos/musicgen_app.py

demos/musicgen_app.py

@@ -7,7 +7,7 @@ import subprocess as sp
 import sys
 import time
 import typing as tp
-from tempfile import NamedTemporaryFile
+from tempfile import NamedTemporaryFile, gettempdir
 
 from einops import rearrange
 import torch
@@ -18,11 +18,11 @@ from audiocraft.data.audio import audio_write
 from audiocraft.models.encodec import InterleaveStereoCompressionModel
 from audiocraft.models import MusicGen, MultiBandDiffusion
 import multiprocessing as mp
+import warnings
 
 # --- Utility Functions and Classes ---
-import warnings # <--- Import the warnings module!
 
-class FileCleaner:  # Unchanged
+class FileCleaner:
     def __init__(self, file_lifetime: float = 3600):
         self.file_lifetime = file_lifetime
         self.files = []
@@ -42,15 +42,38 @@ class FileCleaner:  # Unchanged
                 break
 file_cleaner = FileCleaner()
 
-def make_waveform(*args, **kwargs):  # Unchanged
-    be = time.time()
-    with warnings.catch_warnings():
-        warnings.simplefilter('ignore')
-        out = gr.make_waveform(*args, **kwargs)
-        print("Make a video took", time.time() - be)
-        return out
 
-# --- Worker Process --- (Modified for conditional use)
+def convert_wav_to_mp4(wav_path, output_path=None):
+    """Converts a WAV file to a waveform MP4 video using ffmpeg."""
+    if output_path is None:
+        # Create output path in the same directory as the input
+        output_path = Path(wav_path).with_suffix(".mp4")
+    try:
+        command = [
+            "ffmpeg",
+            "-y",  # Overwrite output file if it exists
+            "-i", str(wav_path),
+            "-filter_complex",
+            "[0:a]showwaves=s=1280x202:mode=line,format=yuv420p[v]",  # Waveform filter
+            "-map", "[v]",
+            "-map", "0:a",
+            "-c:v", "libx264",  # Video codec
+            "-c:a", "aac",       # Audio codec
+            "-preset", "fast", # Important, don't do veryslow.
+            str(output_path),
+        ]
+
+        process = sp.run(command, capture_output=True, text=True, check=True)
+        return str(output_path)
+
+    except sp.CalledProcessError as e:
+        print(f"Error in ffmpeg conversion: {e}")
+        print(f"ffmpeg stdout: {e.stdout}")
+        print(f"ffmpeg stderr: {e.stderr}")
+        raise  # Re-raise the exception to be caught by Gradio
+
+
+# --- Worker Process ---
 
 def model_worker(model_name: str, task_queue: mp.Queue, result_queue: mp.Queue):
     """
@@ -257,8 +280,12 @@ def predict_full(model, model_path, use_mbd, text, melody, duration, topk, topp,
             loudness_headroom_db=16, loudness_compressor=True, add_suffix=False
         )
         wav_paths.append(file.name)
-        video_paths.append(make_waveform(file.name)) # Make and clean up video
+        # Make and clean up video:
+        video_path = convert_wav_to_mp4(file.name)
+        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:
@@ -268,8 +295,12 @@ def predict_full(model, model_path, use_mbd, text, melody, duration, topk, topp,
                 loudness_headroom_db=16, loudness_compressor=True, add_suffix=False
             )
             wav_paths.append(file.name)
-            video_paths.append(make_waveform(file.name)) # Make and clean up video
+            # Make and clean up video:
+            video_path = convert_wav_to_mp4(file.name)
+            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!