Update for app_batched

app.py

@@ -13,7 +13,7 @@ import gradio as gr
 import os
 from audiocraft.models import MusicGen
 from audiocraft.data.audio import audio_write
-from audiocraft.utils.extend import generate_music_segments, add_settings_to_image, sanitize_file_name
+from audiocraft.utils.extend import generate_music_segments, add_settings_to_image
 import numpy as np
 import random
 

app_batched.py

@@ -6,7 +6,12 @@ This source code is licensed under the license found in the
 LICENSE file in the root directory of this source tree.
 """
 
+import argparse
+from concurrent.futures import ProcessPoolExecutor
+import subprocess as sp
 from tempfile import NamedTemporaryFile
+import time
+import warnings
 import torch
 import gradio as gr
 from audiocraft.data.audio_utils import convert_audio
@@ -16,6 +21,29 @@ from audiocraft.models import MusicGen
 
 MODEL = None
 
+_old_call = sp.call
+
+
+def _call_nostderr(*args, **kwargs):
+    # Avoid ffmpeg vomitting on the logs.
+    kwargs['stderr'] = sp.DEVNULL
+    kwargs['stdout'] = sp.DEVNULL
+    _old_call(*args, **kwargs)
+
+
+sp.call = _call_nostderr
+pool = ProcessPoolExecutor(3)
+pool.__enter__()
+
+
+def make_waveform(*args, **kwargs):
+    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
+
 
 def load_model():
     print("Loading model")
@@ -28,11 +56,13 @@ def predict(texts, melodies):
         MODEL = load_model()
 
     duration = 12
+    max_text_length = 512
+    texts = [text[:max_text_length] for text in texts]
     MODEL.set_generation_params(duration=duration)
 
-    print(texts, 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
     target_ac = 1
     for melody in melodies:
@@ -40,8 +70,6 @@ def predict(texts, melodies):
             processed_melodies.append(None)
         else:
             sr, melody = melody[0], torch.from_numpy(melody[1]).to(MODEL.device).float().t()
-            duration = min(duration, melody.shape[-1] / sr)
-            MODEL.set_generation_params(duration=duration)
             if melody.dim() == 1:
                 melody = melody[None]
             melody = melody[..., :int(sr * duration)]
@@ -52,7 +80,7 @@ def predict(texts, melodies):
         descriptions=texts,
         melody_wavs=processed_melodies,
         melody_sample_rate=target_sr,
-        progress=True
+        progress=False
     )
 
     outputs = outputs.detach().cpu().float()
@@ -62,73 +90,132 @@ def predict(texts, melodies):
             audio_write(
                 file.name, output, MODEL.sample_rate, strategy="loudness",
                 loudness_headroom_db=16, loudness_compressor=True, add_suffix=False)
-            waveform_video = gr.make_waveform(file.name)
-            out_files.append(waveform_video)
-    return [out_files]
-
-
-with gr.Blocks() as demo:
-    gr.Markdown(
-        """
-        # MusicGen
-
-        This is the demo for [MusicGen](https://github.com/facebookresearch/audiocraft), a simple and controllable model for music generation
-        presented at: ["Simple and Controllable Music Generation"](https://huggingface.co/papers/2306.05284).
-        <br/>
-        <a href="https://huggingface.co/spaces/musicgen/MusicGen?duplicate=true" style="display: inline-block;margin-top: .5em;margin-right: .25em;" target="_blank">
-        <img style="margin-bottom: 0em;display: inline;margin-top: -.25em;" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
-        for longer sequences, more control and no queue.</p>
-        """
-    )
-    with gr.Row():
-        with gr.Column():
-            with gr.Row():
-                text = gr.Text(label="Describe your music", lines=2, interactive=True)
-                melody = gr.Audio(source="upload", type="numpy", label="Condition on a melody (optional)", interactive=True)
-            with gr.Row():
-                submit = gr.Button("Generate")
-        with gr.Column():
-            output = gr.Video(label="Generated Music")
-    submit.click(predict, inputs=[text, melody], outputs=[output], batch=True, max_batch_size=12)
-    gr.Examples(
-        fn=predict,
-        examples=[
-            [
-                "An 80s driving pop song with heavy drums and synth pads in the background",
-                "./assets/bach.mp3",
-            ],
-            [
-                "A cheerful country song with acoustic guitars",
-                "./assets/bolero_ravel.mp3",
+            out_files.append(pool.submit(make_waveform, file.name))
+    res = [[out_file.result() for out_file in out_files]]
+    print("batch finished", len(texts), time.time() - be)
+    return res
+
+
+def ui(**kwargs):
+    with gr.Blocks() as demo:
+        gr.Markdown(
+            """
+            # MusicGen
+
+            This is the demo for [MusicGen](https://github.com/facebookresearch/audiocraft), a simple and controllable model for music generation
+            presented at: ["Simple and Controllable Music Generation"](https://huggingface.co/papers/2306.05284).
+            <br/>
+            <a href="https://huggingface.co/spaces/musicgen/MusicGen?duplicate=true" style="display: inline-block;margin-top: .5em;margin-right: .25em;" target="_blank">
+            <img style="margin-bottom: 0em;display: inline;margin-top: -.25em;" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
+            for longer sequences, more control and no queue.</p>
+            """
+        )
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    text = gr.Text(label="Describe your music", lines=2, interactive=True)
+                    melody = gr.Audio(source="upload", type="numpy", label="Condition on a melody (optional)", interactive=True)
+                with gr.Row():
+                    submit = gr.Button("Generate")
+            with gr.Column():
+                output = gr.Video(label="Generated Music")
+        submit.click(predict, inputs=[text, melody], outputs=[output], batch=True, max_batch_size=8)
+        gr.Examples(
+            fn=predict,
+            examples=[
+                [
+                    "An 80s driving pop song with heavy drums and synth pads in the background",
+                    "./assets/bach.mp3",
+                ],
+                [
+                    "A cheerful country song with acoustic guitars",
+                    "./assets/bolero_ravel.mp3",
+                ],
+                [
+                    "90s rock song with electric guitar and heavy drums",
+                    None,
+                ],
+                [
+                    "a light and cheerly EDM track, with syncopated drums, aery pads, and strong emotions bpm: 130",
+                    "./assets/bach.mp3",
+                ],
+                [
+                    "lofi slow bpm electro chill with organic samples",
+                    None,
+                ],
             ],
-            [
-                "90s rock song with electric guitar and heavy drums",
-                None,
-            ],
-            [
-                "a light and cheerly EDM track, with syncopated drums, aery pads, and strong emotions bpm: 130",
-                "./assets/bach.mp3",
-            ],
-            [
-                "lofi slow bpm electro chill with organic samples",
-                None,
-            ],
-        ],
-        inputs=[text, melody],
-        outputs=[output]
+            inputs=[text, melody],
+            outputs=[output]
+        )
+        gr.Markdown("""
+        ### More details
+
+        The model will generate 12 seconds of audio based on the description you provided.
+        You can optionaly provide a reference audio from which a broad melody will be extracted.
+        The model will then try to follow both the description and melody provided.
+        All samples are generated with the `melody` model.
+
+        You can also use your own GPU or a Google Colab by following the instructions on our repo.
+
+        See [github.com/facebookresearch/audiocraft](https://github.com/facebookresearch/audiocraft)
+        for more details.
+        """)
+
+        # Show the interface
+        launch_kwargs = {}
+        username = kwargs.get('username')
+        password = kwargs.get('password')
+        server_port = kwargs.get('server_port', 0)
+        inbrowser = kwargs.get('inbrowser', False)
+        share = kwargs.get('share', False)
+        server_name = kwargs.get('listen')
+
+        launch_kwargs['server_name'] = server_name
+
+        if username and password:
+            launch_kwargs['auth'] = (username, password)
+        if server_port > 0:
+            launch_kwargs['server_port'] = server_port
+        if inbrowser:
+            launch_kwargs['inbrowser'] = inbrowser
+        if share:
+            launch_kwargs['share'] = share
+        demo.queue(max_size=60).launch(**launch_kwargs)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--listen',
+        type=str,
+        default='127.0.0.1',
+        help='IP to listen on for connections to Gradio',
+    )
+    parser.add_argument(
+        '--username', type=str, default='', help='Username for authentication'
+    )
+    parser.add_argument(
+        '--password', type=str, default='', help='Password for authentication'
+    )
+    parser.add_argument(
+        '--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'
     )
-    gr.Markdown("""
-    ### More details
-
-    The model will generate 12 seconds of audio based on the description you provided.
-    You can optionaly provide a reference audio from which a broad melody will be extracted.
-    The model will then try to follow both the description and melody provided.
-    All samples are generated with the `melody` model.
-  
-    You can also use your own GPU or a Google Colab by following the instructions on our repo.
 
-    See [github.com/facebookresearch/audiocraft](https://github.com/facebookresearch/audiocraft)
-    for more details.
-    """)
+    args = parser.parse_args()
 
-demo.queue(max_size=15).launch()
+    ui(
+        username=args.username,
+        password=args.password,
+        inbrowser=args.inbrowser,
+        server_port=args.server_port,
+        share=args.share,
+        listen=args.listen
+    )