kind of working

app.py

@@ -59,6 +59,9 @@ def load_model(version='melody'):
 
 
 def _do_predictions(texts, melodies, duration, **gen_kwargs):
+    if duration > MODEL.lm.cfg.dataset.segment_duration:
+        raise gr.Error("MusicGen currently supports durations of up to 30 seconds!")
+
     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])
     be = time.time()
@@ -76,7 +79,7 @@ def _do_predictions(texts, melodies, duration, **gen_kwargs):
             melody = convert_audio(melody, sr, target_sr, target_ac)
             processed_melodies.append(melody)
 
-    if processed_melodies.any():
+    if any(m is not None for m in processed_melodies):
         outputs = MODEL.generate_with_chroma(
             descriptions=texts,
             melody_wavs=processed_melodies,
@@ -110,12 +113,10 @@ def predict_batched(texts, melodies):
 def predict_full(model, text, melody, duration, topk, topp, temperature, cfg_coef):
     topk = int(topk)
     load_model(model)
-    if duration > MODEL.lm.cfg.dataset.segment_duration:
-        raise gr.Error("MusicGen currently supports durations of up to 30 seconds!")
 
     outs = _do_predictions(
         [text], [melody], duration,
-        topk=topk, top_p=topp, temperature=temperature, cfg_coef=cfg_coef)
+        top_k=topk, top_p=topp, temperature=temperature, cfg_coef=cfg_coef)
     return outs[0]
 
 
@@ -138,7 +139,7 @@ def ui_full(launch_kwargs):
                 with gr.Row():
                     model = gr.Radio(["melody", "medium", "small", "large"], label="Model", value="melody", interactive=True)
                 with gr.Row():
-                    duration = gr.Slider(minimum=1, maximum=30, value=10, label="Duration", interactive=True)
+                    duration = gr.Slider(minimum=1, maximum=120, 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)
@@ -184,7 +185,12 @@ def ui_full(launch_kwargs):
             ### More details
 
             The model will generate a short music extract based on the description you provided.
-            You can generate up to 30 seconds of audio.
+            The model can generate up to 30 seconds of audio in one pass. It is now possible
+            to extend the generation by feeding back the end of the previous chunk of audio.
+            This can take a long time, and the model might lose consistency. The model might also
+            decide at arbitrary positions that the song ends.
+
+            **WARNING:** Choosing long durations will take a long time to generate (2min might take ~10min).
 
             We present 4 model variations:
             1. Melody -- a music generation model capable of generating music condition on text and melody inputs. **Note**, you can also use text only.

audiocraft/models/musicgen.py

@@ -45,6 +45,7 @@ class MusicGen:
         self.device = next(iter(lm.parameters())).device
         self.generation_params: dict = {}
         self.set_generation_params(duration=15)  # 15 seconds by default
+        self._progress_callback: tp.Optional[tp.Callable[[int, int], None]] = None
         if self.device.type == 'cpu':
             self.autocast = TorchAutocast(enabled=False)
         else:
@@ -127,6 +128,9 @@ class MusicGen:
             'two_step_cfg': two_step_cfg,
         }
 
+    def set_custom_progress_callback(self, progress_callback: tp.Optional[tp.Callable[[int, int], None]] = None):
+        self._progress_callback = progress_callback
+
     def generate_unconditional(self, num_samples: int, progress: bool = False) -> torch.Tensor:
         """Generate samples in an unconditional manner.
 
@@ -274,6 +278,10 @@ class MusicGen:
         current_gen_offset: int = 0
 
         def _progress_callback(generated_tokens: int, tokens_to_generate: int):
+            generated_tokens += current_gen_offset
+            if self._progress_callback is not None:
+                self._progress_callback(generated_tokens, total_gen_len)
+            else:
             print(f'{current_gen_offset + generated_tokens: 6d} / {total_gen_len: 6d}', end='\r')
 
         if prompt_tokens is not None:
@@ -296,11 +304,17 @@ class MusicGen:
             # melody conditioning etc.
             ref_wavs = [attr.wav['self_wav'] for attr in attributes]
             all_tokens = []
-            if prompt_tokens is not None:
+            if prompt_tokens is None:
+                prompt_length = 0
+            else:
                 all_tokens.append(prompt_tokens)
+                prompt_length = prompt_tokens.shape[-1]
+
+
+            stride_tokens = int(self.frame_rate * self.extend_stride)
 
-            time_offset = 0.
-            while time_offset < self.duration:
+            while current_gen_offset + prompt_length < total_gen_len:
+                time_offset = current_gen_offset / self.frame_rate
                 chunk_duration = min(self.duration - time_offset, self.max_duration)
                 max_gen_len = int(chunk_duration * self.frame_rate)
                 for attr, ref_wav in zip(attributes, ref_wavs):
@@ -321,14 +335,13 @@ class MusicGen:
                     gen_tokens = self.lm.generate(
                         prompt_tokens, attributes,
                         callback=callback, max_gen_len=max_gen_len, **self.generation_params)
-                stride_tokens = int(self.frame_rate * self.extend_stride)
                 if prompt_tokens is None:
                     all_tokens.append(gen_tokens)
                 else:
                     all_tokens.append(gen_tokens[:, :, prompt_tokens.shape[-1]:])
-                prompt_tokens = gen_tokens[:, :, stride_tokens]
+                prompt_tokens = gen_tokens[:, :, stride_tokens:]
+                prompt_length = prompt_tokens.shape[-1]
                 current_gen_offset += stride_tokens
-                time_offset += self.extend_stride
 
             gen_tokens = torch.cat(all_tokens, dim=-1)