Upload musicgen_app.py

demos/musicgen_app.py

@@ -8,30 +8,31 @@
 # also released under the MIT license.
 
 import argparse
-from concurrent.futures import ProcessPoolExecutor
 import logging
 import os
-from pathlib import Path
 import subprocess as sp
 import sys
-from tempfile import NamedTemporaryFile
 import time
 import typing as tp
 import warnings
+from concurrent.futures import ProcessPoolExecutor
+from pathlib import Path
+from tempfile import NamedTemporaryFile
 
-from einops import rearrange
-import torch
 import gradio as gr
+import torch
+from einops import rearrange
 
-from audiocraft.data.audio_utils import convert_audio
 from audiocraft.data.audio import audio_write
+from audiocraft.data.audio_utils import convert_audio
+from audiocraft.models import MultiBandDiffusion, MusicGen
 from audiocraft.models.encodec import InterleaveStereoCompressionModel
-from audiocraft.models import MusicGen, MultiBandDiffusion
-
 
 MODEL = None  # Last used model
-SPACE_ID = os.environ.get('SPACE_ID', '')
-IS_BATCHED = "facebook/MusicGen" in SPACE_ID or 'musicgen-internal/musicgen_dev' in SPACE_ID
+SPACE_ID = os.environ.get("SPACE_ID", "")
+IS_BATCHED = (
+    "facebook/MusicGen" in SPACE_ID or "musicgen-internal/musicgen_dev" in SPACE_ID
+)
 print(IS_BATCHED)
 MAX_BATCH_SIZE = 12
 BATCHED_DURATION = 15
@@ -43,8 +44,8 @@ _old_call = sp.call
 
 def _call_nostderr(*args, **kwargs):
     # Avoid ffmpeg vomiting on the logs.
-    kwargs['stderr'] = sp.DEVNULL
-    kwargs['stdout'] = sp.DEVNULL
+    kwargs["stderr"] = sp.DEVNULL
+    kwargs["stdout"] = sp.DEVNULL
     _old_call(*args, **kwargs)
 
 
@@ -86,17 +87,19 @@ def make_waveform(*args, **kwargs):
     # Further remove some warnings.
     be = time.time()
     with warnings.catch_warnings():
-        warnings.simplefilter('ignore')
+        warnings.simplefilter("ignore")
         out = gr.make_waveform(*args, **kwargs)
         print("Make a video took", time.time() - be)
         return out
 
 
-def load_model(version='facebook/musicgen-melody'):
+def load_model(version="facebook/musicgen-melody"):
     global MODEL
     print("Loading model", version)
     if MODEL is None or MODEL.name != version:
+        # Clear PyTorch CUDA cache and delete model
         del MODEL
+        torch.cuda.empty_cache()
         MODEL = None  # in case loading would crash
         MODEL = MusicGen.get_pretrained(version)
 
@@ -108,9 +111,16 @@ def load_diffusion():
         MBD = MultiBandDiffusion.get_mbd_musicgen()
 
 
-def _do_predictions(texts, melodies, duration, progress=False, gradio_progress=None, **gen_kwargs):
+def _do_predictions(
+    texts, melodies, duration, progress=False, gradio_progress=None, **gen_kwargs
+):
     MODEL.set_generation_params(duration=duration, **gen_kwargs)
-    print("new batch", len(texts), texts, [None if m is None else (m[0], m[1].shape) for m in melodies])
+    print(
+        "new batch",
+        len(texts),
+        texts,
+        [None if m is None else (m[0], m[1].shape) for m in melodies],
+    )
     be = time.time()
     processed_melodies = []
     target_sr = 32000
@@ -119,10 +129,13 @@ def _do_predictions(texts, melodies, duration, progress=False, gradio_progress=N
         if melody is None:
             processed_melodies.append(None)
         else:
-            sr, melody = melody[0], torch.from_numpy(melody[1]).to(MODEL.device).float().t()
+            sr, melody = (
+                melody[0],
+                torch.from_numpy(melody[1]).to(MODEL.device).float().t(),
+            )
             if melody.dim() == 1:
                 melody = melody[None]
-            melody = melody[..., :int(sr * duration)]
+            melody = melody[..., : int(sr * duration)]
             melody = convert_audio(melody, sr, target_sr, target_ac)
             processed_melodies.append(melody)
 
@@ -133,15 +146,17 @@ def _do_predictions(texts, melodies, duration, progress=False, gradio_progress=N
                 melody_wavs=processed_melodies,
                 melody_sample_rate=target_sr,
                 progress=progress,
-                return_tokens=USE_DIFFUSION
+                return_tokens=USE_DIFFUSION,
             )
         else:
-            outputs = MODEL.generate(texts, progress=progress, return_tokens=USE_DIFFUSION)
+            outputs = MODEL.generate(
+                texts, progress=progress, return_tokens=USE_DIFFUSION
+            )
     except RuntimeError as e:
         raise gr.Error("Error while generating " + e.args[0])
     if USE_DIFFUSION:
         if gradio_progress is not None:
-            gradio_progress(1, desc='Running MultiBandDiffusion...')
+            gradio_progress(1, desc="Running MultiBandDiffusion...")
         tokens = outputs[1]
         if isinstance(MODEL.compression_model, InterleaveStereoCompressionModel):
             left, right = MODEL.compression_model.get_left_right_codes(tokens)
@@ -149,7 +164,9 @@ def _do_predictions(texts, melodies, duration, progress=False, gradio_progress=N
         outputs_diffusion = MBD.tokens_to_wav(tokens)
         if isinstance(MODEL.compression_model, InterleaveStereoCompressionModel):
             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 = rearrange(
+                outputs_diffusion, "(s b) c t -> b (s c) t", s=2
+            )
         outputs = torch.cat([outputs[0], outputs_diffusion], dim=0)
     outputs = outputs.detach().cpu().float()
     pending_videos = []
@@ -157,8 +174,14 @@ def _do_predictions(texts, melodies, duration, progress=False, gradio_progress=N
     for output in outputs:
         with NamedTemporaryFile("wb", suffix=".wav", delete=False) as file:
             audio_write(
-                file.name, output, MODEL.sample_rate, strategy="loudness",
-                loudness_headroom_db=16, loudness_compressor=True, add_suffix=False)
+                file.name,
+                output,
+                MODEL.sample_rate,
+                strategy="loudness",
+                loudness_headroom_db=16,
+                loudness_compressor=True,
+                add_suffix=False,
+            )
             pending_videos.append(pool.submit(make_waveform, file.name))
             out_wavs.append(file.name)
             file_cleaner.add(file.name)
@@ -173,12 +196,24 @@ def _do_predictions(texts, melodies, duration, progress=False, gradio_progress=N
 def predict_batched(texts, melodies):
     max_text_length = 512
     texts = [text[:max_text_length] for text in texts]
-    load_model('facebook/musicgen-stereo-melody')
+    load_model("facebook/musicgen-stereo-melody")
     res = _do_predictions(texts, melodies, BATCHED_DURATION)
     return res
 
 
-def predict_full(model, model_path, decoder, text, melody, duration, topk, topp, temperature, cfg_coef, progress=gr.Progress()):
+def predict_full(
+    model,
+    model_path,
+    decoder,
+    text,
+    melody,
+    duration,
+    topk,
+    topp,
+    temperature,
+    cfg_coef,
+    progress=gr.Progress(),
+):
     global INTERRUPTING
     global USE_DIFFUSION
     INTERRUPTING = False
@@ -188,8 +223,10 @@ def predict_full(model, model_path, decoder, text, melody, duration, topk, topp,
         if not Path(model_path).exists():
             raise gr.Error(f"Model path {model_path} doesn't exist.")
         if not Path(model_path).is_dir():
-            raise gr.Error(f"Model path {model_path} must be a folder containing "
-                           "state_dict.bin and compression_state_dict_.bin.")
+            raise gr.Error(
+                f"Model path {model_path} must be a folder containing "
+                "state_dict.bin and compression_state_dict_.bin."
+            )
         model = model_path
     if temperature < 0:
         raise gr.Error("Temperature must be >= 0.")
@@ -215,12 +252,20 @@ def predict_full(model, model_path, decoder, text, melody, duration, topk, topp,
         progress((min(max_generated, to_generate), to_generate))
         if INTERRUPTING:
             raise gr.Error("Interrupted.")
+
     MODEL.set_custom_progress_callback(_progress)
 
     videos, wavs = _do_predictions(
-        [text], [melody], duration, progress=True,
-        top_k=topk, top_p=topp, temperature=temperature, cfg_coef=cfg_coef,
-        gradio_progress=progress)
+        [text],
+        [melody],
+        duration,
+        progress=True,
+        top_k=topk,
+        top_p=topp,
+        temperature=temperature,
+        cfg_coef=cfg_coef,
+        gradio_progress=progress,
+    )
     if USE_DIFFUSION:
         return videos[0], wavs[0], videos[1], wavs[1]
     return videos[0], wavs[0], None, None
@@ -255,42 +300,86 @@ def ui_full(launch_kwargs):
                 with gr.Row():
                     text = gr.Text(label="Input Text", interactive=True)
                     with gr.Column():
-                        radio = gr.Radio(["file", "mic"], value="file",
-                                         label="Condition on a melody (optional) File or Mic")
-                        melody = gr.Audio(source="upload", type="numpy", label="File",
-                                          interactive=True, elem_id="melody-input")
+                        radio = gr.Radio(
+                            ["file", "mic"],
+                            value="file",
+                            label="Condition on a melody (optional) File or Mic",
+                        )
+                        melody = gr.Audio(
+                            sources=["upload"],
+                            type="numpy",
+                            label="File",
+                            interactive=True,
+                            elem_id="melody-input",
+                        )
                 with gr.Row():
                     submit = gr.Button("Submit")
                     # Adapted from https://github.com/rkfg/audiocraft/blob/long/app.py, MIT license.
                     _ = gr.Button("Interrupt").click(fn=interrupt, queue=False)
                 with gr.Row():
-                    model = gr.Radio(["facebook/musicgen-melody", "facebook/musicgen-medium", "facebook/musicgen-small",
-                                      "facebook/musicgen-large", "facebook/musicgen-melody-large",
-                                      "facebook/musicgen-stereo-small", "facebook/musicgen-stereo-medium",
-                                      "facebook/musicgen-stereo-melody", "facebook/musicgen-stereo-large",
-                                      "facebook/musicgen-stereo-melody-large"],
-                                     label="Model", value="facebook/musicgen-stereo-melody", interactive=True)
+                    # model = gr.Radio(["facebook/musicgen-melody", "facebook/musicgen-medium", "facebook/musicgen-small",
+                    #                   "facebook/musicgen-large", "facebook/musicgen-melody-large",
+                    #                   "facebook/musicgen-stereo-small", "facebook/musicgen-stereo-medium",
+                    #                   "facebook/musicgen-stereo-melody", "facebook/musicgen-stereo-large",
+                    #                   "facebook/musicgen-stereo-melody-large"],
+                    #                  label="Model", value="facebook/musicgen-stereo-melody", interactive=True)
                     model_path = gr.Text(label="Model Path (custom models)")
                 with gr.Row():
-                    decoder = gr.Radio(["Default", "MultiBand_Diffusion"],
-                                       label="Decoder", value="Default", interactive=True)
+                    decoder = gr.Radio(
+                        ["Default", "MultiBand_Diffusion"],
+                        label="Decoder",
+                        value="Default",
+                        interactive=True,
+                    )
                 with gr.Row():
-                    duration = gr.Slider(minimum=1, maximum=120, value=10, label="Duration", interactive=True)
+                    duration = gr.Slider(
+                        minimum=1,
+                        maximum=60,
+                        value=10,
+                        label="Duration",
+                        interactive=True,
+                    )
                 with gr.Row():
                     topk = gr.Number(label="Top-k", value=250, interactive=True)
                     topp = gr.Number(label="Top-p", value=0, interactive=True)
-                    temperature = gr.Number(label="Temperature", value=1.0, interactive=True)
-                    cfg_coef = gr.Number(label="Classifier Free Guidance", value=3.0, interactive=True)
+                    temperature = gr.Number(
+                        label="Temperature", value=1.0, interactive=True
+                    )
+                    cfg_coef = gr.Number(
+                        label="Classifier Free Guidance", value=3.0, interactive=True
+                    )
             with gr.Column():
                 output = gr.Video(label="Generated Music")
-                audio_output = gr.Audio(label="Generated Music (wav)", type='filepath')
+                audio_output = gr.Audio(label="Generated Music (wav)", type="filepath")
                 diffusion_output = gr.Video(label="MultiBand Diffusion Decoder")
-                audio_diffusion = gr.Audio(label="MultiBand Diffusion Decoder (wav)", type='filepath')
-        submit.click(toggle_diffusion, decoder, [diffusion_output, audio_diffusion], queue=False,
-                     show_progress=False).then(predict_full, inputs=[model, model_path, decoder, text, melody, duration, topk, topp,
-                                                                     temperature, cfg_coef],
-                                               outputs=[output, audio_output, diffusion_output, audio_diffusion])
-        radio.change(toggle_audio_src, radio, [melody], queue=False, show_progress=False)
+                audio_diffusion = gr.Audio(
+                    label="MultiBand Diffusion Decoder (wav)", type="filepath"
+                )
+        submit.click(
+            toggle_diffusion,
+            decoder,
+            [diffusion_output, audio_diffusion],
+            queue=False,
+            show_progress=False,
+        ).then(
+            predict_full,
+            inputs=[
+                model,
+                model_path,
+                decoder,
+                text,
+                melody,
+                duration,
+                topk,
+                topp,
+                temperature,
+                cfg_coef,
+            ],
+            outputs=[output, audio_output, diffusion_output, audio_diffusion],
+        )
+        radio.change(
+            toggle_audio_src, radio, [melody], queue=False, show_progress=False
+        )
 
         gr.Examples(
             fn=predict_full,
@@ -299,41 +388,41 @@ def ui_full(launch_kwargs):
                     "An 80s driving pop song with heavy drums and synth pads in the background",
                     "./assets/bach.mp3",
                     "facebook/musicgen-stereo-melody",
-                    "Default"
+                    "Default",
                 ],
                 [
                     "A cheerful country song with acoustic guitars",
                     "./assets/bolero_ravel.mp3",
                     "facebook/musicgen-stereo-melody",
-                    "Default"
+                    "Default",
                 ],
                 [
                     "90s rock song with electric guitar and heavy drums",
                     None,
                     "facebook/musicgen-stereo-medium",
-                    "Default"
+                    "Default",
                 ],
                 [
                     "a light and cheerly EDM track, with syncopated drums, aery pads, and strong emotions",
                     "./assets/bach.mp3",
                     "facebook/musicgen-stereo-melody",
-                    "Default"
+                    "Default",
                 ],
                 [
                     "lofi slow bpm electro chill with organic samples",
                     None,
                     "facebook/musicgen-stereo-medium",
-                    "Default"
+                    "Default",
                 ],
                 [
                     "Punk rock with loud drum and power guitar",
                     None,
                     "facebook/musicgen-stereo-medium",
-                    "MultiBand_Diffusion"
+                    "MultiBand_Diffusion",
                 ],
             ],
             inputs=[text, melody, model, decoder],
-            outputs=[output]
+            outputs=[output],
         )
         gr.Markdown(
             """
@@ -403,20 +492,37 @@ def ui_batched(launch_kwargs):
         with gr.Row():
             with gr.Column():
                 with gr.Row():
-                    text = gr.Text(label="Describe your music", lines=2, interactive=True)
+                    text = gr.Text(
+                        label="Describe your music", lines=2, interactive=True
+                    )
                     with gr.Column():
-                        radio = gr.Radio(["file", "mic"], value="file",
-                                         label="Condition on a melody (optional) File or Mic")
-                        melody = gr.Audio(source="upload", type="numpy", label="File",
-                                          interactive=True, elem_id="melody-input")
+                        radio = gr.Radio(
+                            ["file", "mic"],
+                            value="file",
+                            label="Condition on a melody (optional) File or Mic",
+                        )
+                        melody = gr.Audio(
+                            source="upload",
+                            type="numpy",
+                            label="File",
+                            interactive=True,
+                            elem_id="melody-input",
+                        )
                 with gr.Row():
                     submit = gr.Button("Generate")
             with gr.Column():
                 output = gr.Video(label="Generated Music")
-                audio_output = gr.Audio(label="Generated Music (wav)", type='filepath')
-        submit.click(predict_batched, inputs=[text, melody],
-                     outputs=[output, audio_output], batch=True, max_batch_size=MAX_BATCH_SIZE)
-        radio.change(toggle_audio_src, radio, [melody], queue=False, show_progress=False)
+                audio_output = gr.Audio(label="Generated Music (wav)", type="filepath")
+        submit.click(
+            predict_batched,
+            inputs=[text, melody],
+            outputs=[output, audio_output],
+            batch=True,
+            max_batch_size=MAX_BATCH_SIZE,
+        )
+        radio.change(
+            toggle_audio_src, radio, [melody], queue=False, show_progress=False
+        )
         gr.Examples(
             fn=predict_batched,
             examples=[
@@ -442,7 +548,7 @@ def ui_batched(launch_kwargs):
                 ],
             ],
             inputs=[text, melody],
-            outputs=[output]
+            outputs=[output],
         )
         gr.Markdown("""
         ### More details
@@ -476,50 +582,46 @@ def ui_batched(launch_kwargs):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        '--listen',
+        "--listen",
         type=str,
-        default='0.0.0.0' if 'SPACE_ID' in os.environ else '127.0.0.1',
-        help='IP to listen on for connections to Gradio',
+        default="0.0.0.0" if "SPACE_ID" in os.environ else "127.0.0.1",
+        help="IP to listen on for connections to Gradio",
     )
     parser.add_argument(
-        '--username', type=str, default='', help='Username for authentication'
+        "--username", type=str, default="", help="Username for authentication"
     )
     parser.add_argument(
-        '--password', type=str, default='', help='Password for authentication'
+        "--password", type=str, default="", help="Password for authentication"
     )
     parser.add_argument(
-        '--server_port',
+        "--server_port",
         type=int,
         default=0,
-        help='Port to run the server listener on',
-    )
-    parser.add_argument(
-        '--inbrowser', action='store_true', help='Open in browser'
-    )
-    parser.add_argument(
-        '--share', action='store_true', help='Share the gradio UI'
+        help="Port to run the server listener on",
     )
+    parser.add_argument("--inbrowser", action="store_true", help="Open in browser")
+    parser.add_argument("--share", action="store_true", help="Share the gradio UI")
 
     args = parser.parse_args()
 
     launch_kwargs = {}
-    launch_kwargs['server_name'] = args.listen
+    launch_kwargs["server_name"] = args.listen
 
     if args.username and args.password:
-        launch_kwargs['auth'] = (args.username, args.password)
+        launch_kwargs["auth"] = (args.username, args.password)
     if args.server_port:
-        launch_kwargs['server_port'] = args.server_port
+        launch_kwargs["server_port"] = args.server_port
     if args.inbrowser:
-        launch_kwargs['inbrowser'] = args.inbrowser
+        launch_kwargs["inbrowser"] = args.inbrowser
     if args.share:
-        launch_kwargs['share'] = args.share
+        launch_kwargs["share"] = args.share
 
     logging.basicConfig(level=logging.INFO, stream=sys.stderr)
 
     # Show the interface
-    if IS_BATCHED:
-        global USE_DIFFUSION
-        USE_DIFFUSION = False
-        ui_batched(launch_kwargs)
-    else:
-        ui_full(launch_kwargs)
+    # if IS_BATCHED:
+    #     global USE_DIFFUSION
+    #     USE_DIFFUSION = False
+    #     ui_batched(launch_kwargs)
+    # else:
+    ui_full(launch_kwargs)