new version

CHANGELOG.md

@@ -13,6 +13,8 @@ Now repeating the conditioning periodically if it is too short.
 
 More options when launching Gradio app locally (thanks @ashleykleynhans).
 
+Testing out PyTorch 2.0 memory efficient attention.
+
 ## [0.0.1] - 2023-06-09
 
 Initial release, with model evaluation only.

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:
@@ -60,73 +90,133 @@ 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.
+    args = parser.parse_args()
 
-    See [github.com/facebookresearch/audiocraft](https://github.com/facebookresearch/audiocraft)
-    for more details.
-    """)
+    ui(
+        username=args.username,
+        password=args.password,
+        inbrowser=args.inbrowser,
+        server_port=args.server_port,
+        share=args.share,
+        listen=args.listen
+    )
 
-demo.queue(max_size=15).launch()

audiocraft/modules/transformer.py

@@ -247,20 +247,20 @@ class StreamingMultiheadAttention(StreamingModule):
         # Complete the key/value pair using the streaming state.
         if self._streaming_state:
             pk = self._streaming_state['past_keys']
-            nk = torch.cat([pk, k], dim=1)
+            nk = torch.cat([pk, k], dim=2)
             if v is k:
                 nv = nk
             else:
                 pv = self._streaming_state['past_values']
-                nv = torch.cat([pv, v], dim=1)
+                nv = torch.cat([pv, v], dim=2)
         else:
             nk = k
             nv = v
 
-        assert nk.shape[1] == nv.shape[1]
+        assert nk.shape[2] == nv.shape[2]
         offset = 0
         if self.past_context is not None:
-            offset = max(0, nk.shape[1] - self.past_context)
+            offset = max(0, nk.shape[2] - self.past_context)
         if self._is_streaming:
             self._streaming_state['past_keys'] = nk[:, offset:]
             if v is not k:
@@ -271,6 +271,7 @@ class StreamingMultiheadAttention(StreamingModule):
                 self._streaming_state['offset'] = torch.tensor(0)
         return nk, nv
 
+
     def _apply_rope(self, query: torch.Tensor, key: torch.Tensor):
         # Apply rope embeddings to query and key tensors.
         assert self.rope is not None
@@ -325,7 +326,7 @@ class StreamingMultiheadAttention(StreamingModule):
                     q = self.q_layer_norm(q)
                     k = self.k_layer_norm(k)
                 # q, k, v = [rearrange(x, "b t (h d) -> (b h) t d", h=self.num_heads) for x in [q, k, v]]
-                q, k, v = [rearrange(x, "b t (h d) -> b t h d", h=self.num_heads) for x in [q, k, v]]
+                q, k, v = [rearrange(x, "b t (h d) -> b h t d", h=self.num_heads) for x in [q, k, v]]
             else:
                 if not _is_profiled():
                     # profiling breaks that propertysomehow.
@@ -333,7 +334,7 @@ class StreamingMultiheadAttention(StreamingModule):
                     assert value is key, "specialized implementation"
                 projected = nn.functional.linear(query, self.in_proj_weight, self.in_proj_bias)
                 if self.kv_repeat == 1:
-                    packed = rearrange(projected, "b t (p h d) -> b t p h d", p=3, h=self.num_heads)
+                    packed = rearrange(projected, "b t (p h d) -> b h p t d", p=3, h=self.num_heads)
                     q, k, v = ops.unbind(packed, dim=2)
                 else:
                     embed_dim = self.embed_dim
@@ -355,6 +356,7 @@ class StreamingMultiheadAttention(StreamingModule):
                     k = self.k_layer_norm(k)
                     q, k = [rearrange(x, "b t (h d) -> b t h d", h=self.num_heads) for x in [q, k]]
                 if self.rope:
+                    assert False, "Not supported for now"
                     q, k = self._apply_rope(q, k)
                 k, v = self._complete_kv(k, v)
                 if self.kv_repeat > 1:
@@ -364,7 +366,8 @@ class StreamingMultiheadAttention(StreamingModule):
                 q, k, v = [x.float() for x in [q, k, v]]
             if self.memory_efficient:
                 p = self.dropout if self.training else 0
-                x = ops.memory_efficient_attention(q, k, v, attn_mask, p=p)
+                x = torch.nn.functional.scaled_dot_product_attention(
+                    q, k, v, is_causal=attn_mask is not None, dropout_p=p)
             else:
                 # We include the dot product as float32, for consistency
                 # with the other implementations that include that step
@@ -385,7 +388,7 @@ class StreamingMultiheadAttention(StreamingModule):
                 w = F.dropout(w, self.dropout, training=self.training).to(v)
                 x = torch.einsum("bhqk,bkhc->bqhc", w, v)
             x = x.to(dtype)
-            x = rearrange(x, "b t h d -> b t (h d)", h=self.num_heads)
+            x = rearrange(x, "b h t d -> b t (h d)", h=self.num_heads)
             x = self.out_proj(x)
         else:
             key, value = self._complete_kv(key, value)