Initial

.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+static/

app.py

@@ -0,0 +1,111 @@
+import gradio as gr
+import os, json
+from generator import HijackedMusicGen
+from audiocraft.data.audio import audio_write
+from audio import predict
+from itertools import zip_longest
+
+def split_prompt(bigly_prompt, num_segments):
+    prompts = bigly_prompt.split(',,')
+    num_segments = int(num_segments)  # Assuming 'segment' comes as a string from Gradio slider
+    # repeat last prompt to fill in the rest
+    if len(prompts) < num_segments:
+        prompts += [prompts[-1]] * (num_segments - len(prompts))
+    elif len(prompts) > num_segments:
+        prompts = prompts[:num_segments]
+    return prompts
+
+loaded_model = None
+audio_files = []
+def model_interface(model_name, top_k, top_p, temperature, cfg_coef, segments, overlap, duration, optional_audio, prompt):
+    global loaded_model
+
+    if loaded_model is None or loaded_model.name != model_name:
+        loaded_model = HijackedMusicGen.get_pretrained(None, name=model_name)
+        
+    print(optional_audio)
+        
+    loaded_model.set_generation_params(
+        use_sampling=True,
+        duration=duration,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
+        cfg_coef=cfg_coef,
+    )
+    
+    extension_parameters = {"segments":segments, "overlap":overlap}
+    
+    prompts = split_prompt(prompt, segments)
+    first_prompt = prompts[0]
+    sample_rate, audio = predict(loaded_model, prompts, None, extension_parameters)
+    counter = 1
+    audio_path = "static/"
+    audio_name = first_prompt
+    while os.path.exists(audio_path + audio_name + ".wav"):
+        audio_name = f"{first_prompt}({counter})"
+        counter += 1
+        
+    file = audio_write(audio_path + audio_name, audio.squeeze(), sample_rate, strategy="loudness")
+    audio_files.append(file)
+
+    audio_list_html = "<br>".join([
+        f'''
+        <div style="border:1px solid #000; padding:10px; margin-bottom:10px;">
+            <div>{os.path.splitext(os.path.basename(file))[0]}</div>
+            <audio controls><source src="/file={file}" type="audio/wav"></audio>
+        </div>
+        ''' 
+        for file in reversed(audio_files)
+    ])
+
+    return audio_list_html
+
+slider_param = {
+    "top_k": {"minimum": 0, "maximum": 1000, "value": 0, "label": "Top K"},
+    "top_p": {"minimum": 0.0, "maximum": 1.0, "value": 0.0, "label": "Top P"},
+    "temperature": {"minimum": 0.1, "maximum": 10.0, "value": 1.0, "label": "Temperature"},
+    "cfg_coef": {"minimum": 0.0, "maximum": 10.0, "value": 4.0, "label": "CFG Coefficient"},
+    "segments": {"minimum": 1, "maximum": 10, "value": 1, "step": 1, "label": "Number of Segments"},
+    "overlap": {"minimum": 0.0, "maximum": 10.0, "value": 1.0, "label": "Segment Overlap"},
+    "duration": {"minimum": 1, "maximum": 300, "value": 10, "label": "Duration"},
+}
+
+slider_params = {
+    key: gr.components.Slider(**params) 
+    for key, params in slider_param.items()
+}
+
+with gr.Blocks() as interface:
+    with gr.Row():
+
+        with gr.Column():
+            with gr.Row():
+                model_dropdown = gr.components.Dropdown(choices=["small", "medium", "large", "melody"], label="Model Size", value="large")
+                optional_audio = gr.compoents.Audio(source="upload", type="filepath", label="Optional Audio", interactive=True)
+
+            slider_keys = list(slider_param.keys())
+            slider_pairs = list(zip_longest(slider_keys[::2], slider_keys[1::2]))
+
+            for key1, key2 in slider_pairs:
+                with gr.Row():
+                    with gr.Column():
+                        slider_params[key1] = gr.components.Slider(**slider_param[key1])
+                    if key2 is not None:
+                        with gr.Column():
+                            slider_params[key2] = gr.components.Slider(**slider_param[key2])
+
+            prompt_box = gr.components.Textbox(lines=5, placeholder="""Insert a double comma ,, to indicate this should prompt a new segment. For example: 
+    Rock Opera,,Dueling Banjos
+    This allows you to prompt each segment individually. If you only provide one prompt, every segment will use that one prompt. If you provide multiple prompts but less than the number of segments, then the last prompt will be used to fill in the rest.
+    """)
+            submit = gr.Button("Submit")
+
+        with gr.Column():
+            output = gr.outputs.HTML()
+
+    inputs_list = [model_dropdown] + list(slider_params.values()) + [optional_audio] + [prompt_box]
+    submit.click(model_interface, inputs=inputs_list, outputs=[output])
+    
+interface.queue()
+interface.launch()

audio.py

@@ -0,0 +1,58 @@
+import numpy as np
+import os, re, json, sys
+import torch, torchaudio, pathlib
+
+def load_and_process_audio(model, melody, sample_rate):
+    if melody is not None:
+        melody = torch.from_numpy(melody).to(model.device).float().t().unsqueeze(0)
+        if melody.dim() == 2:
+            melody = melody[None]
+        melody = melody[..., :int(sample_rate * model.lm.cfg.dataset.segment_duration)]
+        return melody
+    else:
+        return None
+
+#From https://colab.research.google.com/drive/154CqogsdP-D_TfSF9S2z8-BY98GN_na4?usp=sharing#scrollTo=exKxNU_Z4i5I
+#Thank you DragonForged for the link
+def extend_audio(model, prompt_waveform, prompts, prompt_sr, segments=5, overlap=2):
+    # Calculate the number of samples corresponding to the overlap
+    overlap_samples = int(overlap * prompt_sr)
+
+    device = model.device
+    prompt_waveform = prompt_waveform.to(device)
+
+    for i in range(1, segments):
+        # Grab the end of the waveform
+        end_waveform = prompt_waveform[...,-overlap_samples:]
+
+        # Process the trimmed waveform using the model
+        new_audio = model.generate_continuation(end_waveform, descriptions=[prompts[i]], prompt_sample_rate=prompt_sr, progress=True)
+            
+        # Cut the seed audio off the newly generated audio
+        new_audio = new_audio[...,overlap_samples:]
+
+        prompt_waveform = torch.cat([prompt_waveform, new_audio], dim=2)
+
+    return prompt_waveform
+
+def predict(model, prompts, melody_parameters, extension_parameters):
+    melody = None #load_and_process_audio(MODEL, **melody_parameters)
+
+    if melody is not None:
+        output = MODEL.generate_with_chroma(
+            descriptions=[prompt[0]],
+            melody_wavs=melody,
+            melody_sample_rate=melody_parameters['sample_rate'],
+            progress=False
+        )
+    else:
+        output = model.generate(descriptions=[prompts[0]], progress=True)
+
+    sample_rate = model.sample_rate
+    
+    if extension_parameters['segments'] > 1:
+        output_tensors = extend_audio(model, output, prompts, sample_rate, **extension_parameters).detach().cpu().float()
+    else:
+        output_tensors = output.detach().cpu().float()
+    
+    return sample_rate, output_tensors

generator.py

@@ -0,0 +1,31 @@
+import torch
+import time
+import typing as tp
+from audiocraft.models import MusicGen
+from audiocraft.modules.conditioners import ConditioningAttributes
+
+class HijackedMusicGen(MusicGen):
+    def __init__(self, socketio=None, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.socketio = socketio
+        self._progress_callback = self._timed_progress_callback if socketio is not None else None
+        self._last_update_time = time.time()
+        
+    def _timed_progress_callback(self, generated_tokens: int, tokens_to_generate: int):
+        current_time = time.time()
+        if current_time - self._last_update_time >= 0.1:  # 0.1 seconds have passed
+            self.socketio.emit('progress', {'generated_tokens': generated_tokens, 'tokens_to_generate': tokens_to_generate})
+            self._last_update_time = current_time
+        
+    @staticmethod
+    def get_pretrained(socketio, name: str = 'melody', device='cuda'):
+        music_gen = MusicGen.get_pretrained(name, device)
+        return HijackedMusicGen(socketio, music_gen.name, music_gen.compression_model, music_gen.lm)
+
+    @property
+    def progress_callback(self):
+        raise Exception("Progress callback is write-only")
+
+    @progress_callback.setter
+    def progress_callback(self, callback):
+        self._progress_callback = callback