Update app.py

app.py

@@ -1,137 +1,211 @@
-import gradio as gr
-import note_seq
 import numpy as np
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-tokenizer = AutoTokenizer.from_pretrained("TristanBehrens/js-fakes-4bars")
-model = AutoModelForCausalLM.from_pretrained("TristanBehrens/js-fakes-4bars")
-
-NOTE_LENGTH_16TH_120BPM = 0.25 * 60 / 120
-BAR_LENGTH_120BPM = 4.0 * 60 / 120
-SAMPLE_RATE=44100
-
-def token_sequence_to_note_sequence(token_sequence, use_program=True, use_drums=True, instrument_mapper=None, only_piano=False):
-    if isinstance(token_sequence, str):
-        token_sequence = token_sequence.split()
-    note_sequence = empty_note_sequence()
-
-    # Render all notes.
-    current_program = 1
-    current_is_drum = False
-    current_instrument = 0
-    track_count = 0
-    for token_index, token in enumerate(token_sequence):
-
-        if token == "PIECE_START":
-            pass
-        elif token == "PIECE_END":
-            print("The end.")
-            break
-        elif token == "TRACK_START":
-            current_bar_index = 0
-            track_count += 1
-            pass
-        elif token == "TRACK_END":
-            pass
-        elif token == "KEYS_START":
-            pass
-        elif token == "KEYS_END":
-            pass
-        elif token.startswith("KEY="):
-            pass
-        elif token.startswith("INST"):
-            instrument = token.split("=")[-1]
-            if instrument != "DRUMS" and use_program:
-                if instrument_mapper is not None:
-                    if instrument in instrument_mapper:
-                        instrument = instrument_mapper[instrument]
-                current_program = int(instrument)
-                current_instrument = track_count
-                current_is_drum = False
-            if instrument == "DRUMS" and use_drums:
-                current_instrument = 0
-                current_program = 0
-                current_is_drum = True
-        elif token == "BAR_START":
-            current_time = current_bar_index * BAR_LENGTH_120BPM
-            current_notes = {}
-        elif token == "BAR_END":
-            current_bar_index += 1
-            pass
-        elif token.startswith("NOTE_ON"):
-            pitch = int(token.split("=")[-1])
-            note = note_sequence.notes.add()
-            note.start_time = current_time
-            note.end_time = current_time + 4 * NOTE_LENGTH_16TH_120BPM
-            note.pitch = pitch
-            note.instrument = current_instrument
-            note.program = current_program
-            note.velocity = 80
-            note.is_drum = current_is_drum
-            current_notes[pitch] = note
-        elif token.startswith("NOTE_OFF"):
-            pitch = int(token.split("=")[-1])
-            if pitch in current_notes:
-                note = current_notes[pitch]
-                note.end_time = current_time
-        elif token.startswith("TIME_DELTA"):
-            delta = float(token.split("=")[-1]) * NOTE_LENGTH_16TH_120BPM
-            current_time += delta
-        elif token.startswith("DENSITY="):
-            pass
-        elif token == "[PAD]":
-            pass
+import torch
+import gradio as gr
+import spaces
+from queue import Queue
+from threading import Thread
+from typing import Optional
+from transformers import MusicgenForConditionalGeneration, MusicgenProcessor, set_seed
+from transformers.generation.streamers import BaseStreamer
+
+model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+processor = MusicgenProcessor.from_pretrained("facebook/musicgen-small")
+
+title = "9🌍MusicHub - Text to Music Stream Generator"
+description = """  Facebook MusicGen-Small Model - Generate and stream music with model https://huggingface.co/facebook/musicgen-small  """
+article = """
+## How It Works:
+MusicGen is an auto-regressive transformer-based model, meaning generates audio codes (tokens) in a causal fashion.
+At each decoding step, the model generates a new set of audio codes, conditional on the text input and all previous audio codes. From the 
+frame rate of the [EnCodec model](https://huggingface.co/facebook/encodec_32khz) used to decode the generated codes to audio waveform.
+"""
+
+
+class MusicgenStreamer(BaseStreamer):
+    def __init__(
+        self,
+        model: MusicgenForConditionalGeneration,
+        device: Optional[str] = None,
+        play_steps: Optional[int] = 10,
+        stride: Optional[int] = None,
+        timeout: Optional[float] = None,
+    ):
+        """
+        Streamer that stores playback-ready audio in a queue, to be used by a downstream application as an iterator. This is
+        useful for applications that benefit from acessing the generated audio in a non-blocking way (e.g. in an interactive
+        Gradio demo).
+        Parameters:
+            model (`MusicgenForConditionalGeneration`):
+                The MusicGen model used to generate the audio waveform.
+            device (`str`, *optional*):
+                The torch device on which to run the computation. If `None`, will default to the device of the model.
+            play_steps (`int`, *optional*, defaults to 10):
+                The number of generation steps with which to return the generated audio array. Using fewer steps will 
+                mean the first chunk is ready faster, but will require more codec decoding steps overall. This value 
+                should be tuned to your device and latency requirements.
+            stride (`int`, *optional*):
+                The window (stride) between adjacent audio samples. Using a stride between adjacent audio samples reduces
+                the hard boundary between them, giving smoother playback. If `None`, will default to a value equivalent to 
+                play_steps // 6 in the audio space.
+            timeout (`int`, *optional*):
+                The timeout for the audio queue. If `None`, the queue will block indefinitely. Useful to handle exceptions
+                in `.generate()`, when it is called in a separate thread.
+        """
+        self.decoder = model.decoder
+        self.audio_encoder = model.audio_encoder
+        self.generation_config = model.generation_config
+        self.device = device if device is not None else model.device
+
+        # variables used in the streaming process
+        self.play_steps = play_steps
+        if stride is not None:
+            self.stride = stride
+        else:
+            hop_length = np.prod(self.audio_encoder.config.upsampling_ratios)
+            self.stride = hop_length * (play_steps - self.decoder.num_codebooks) // 6
+        self.token_cache = None
+        self.to_yield = 0
+
+        # varibles used in the thread process
+        self.audio_queue = Queue()
+        self.stop_signal = None
+        self.timeout = timeout
+
+    def apply_delay_pattern_mask(self, input_ids):
+        # build the delay pattern mask for offsetting each codebook prediction by 1 (this behaviour is specific to MusicGen)
+        _, decoder_delay_pattern_mask = self.decoder.build_delay_pattern_mask(
+            input_ids[:, :1],
+            pad_token_id=self.generation_config.decoder_start_token_id,
+            max_length=input_ids.shape[-1],
+        )
+        # apply the pattern mask to the input ids
+        input_ids = self.decoder.apply_delay_pattern_mask(input_ids, decoder_delay_pattern_mask)
+
+        # revert the pattern delay mask by filtering the pad token id
+        input_ids = input_ids[input_ids != self.generation_config.pad_token_id].reshape(
+            1, self.decoder.num_codebooks, -1
+        )
+
+        # append the frame dimension back to the audio codes
+        input_ids = input_ids[None, ...]
+
+        # send the input_ids to the correct device
+        input_ids = input_ids.to(self.audio_encoder.device)
+
+        output_values = self.audio_encoder.decode(
+            input_ids,
+            audio_scales=[None],
+        )
+        audio_values = output_values.audio_values[0, 0]
+        return audio_values.cpu().float().numpy()
+
+    def put(self, value):
+        batch_size = value.shape[0] // self.decoder.num_codebooks
+        if batch_size > 1:
+            raise ValueError("MusicgenStreamer only supports batch size 1")
+
+        if self.token_cache is None:
+            self.token_cache = value
+        else:
+            self.token_cache = torch.concatenate([self.token_cache, value[:, None]], dim=-1)
+
+        if self.token_cache.shape[-1] % self.play_steps == 0:
+            audio_values = self.apply_delay_pattern_mask(self.token_cache)
+            self.on_finalized_audio(audio_values[self.to_yield : -self.stride])
+            self.to_yield += len(audio_values) - self.to_yield - self.stride
+
+    def end(self):
+        """Flushes any remaining cache and appends the stop symbol."""
+        if self.token_cache is not None:
+            audio_values = self.apply_delay_pattern_mask(self.token_cache)
         else:
-            #print(f"Ignored token {token}.")
-            pass
-
-    # Make the instruments right.
-    instruments_drums = []
-    for note in note_sequence.notes:
-        pair = [note.program, note.is_drum]
-        if pair not in instruments_drums:
-            instruments_drums += [pair]
-        note.instrument = instruments_drums.index(pair)
-
-    if only_piano:
-        for note in note_sequence.notes:
-            if not note.is_drum:
-                note.instrument = 0
-                note.program = 0
-
-    return note_sequence
-
-def empty_note_sequence(qpm=120.0, total_time=0.0):
-    note_sequence = note_seq.protobuf.music_pb2.NoteSequence()
-    note_sequence.tempos.add().qpm = qpm
-    note_sequence.ticks_per_quarter = note_seq.constants.STANDARD_PPQ
-    note_sequence.total_time = total_time
-    return note_sequence
-
-def process(text):
-  input_ids = tokenizer.encode(text, return_tensors="pt")
-  generated_ids = model.generate(input_ids, max_length=500)
-  generated_sequence = tokenizer.decode(generated_ids[0])
-
-  # Convert text of notes to audio
-  note_sequence = token_sequence_to_note_sequence(generated_sequence)
-  synth = note_seq.midi_synth.synthesize
-  array_of_floats = synth(note_sequence, sample_rate=SAMPLE_RATE)
-  note_plot = note_seq.plot_sequence(note_sequence, False)
-  array_of_floats /=1.414
-  array_of_floats *= 32767
-  int16_data = array_of_floats.astype(np.int16)
-  return SAMPLE_RATE, int16_data
-  
-title = "Music generation with GPT-2"
-
-iface = gr.Interface(
-  fn=process, 
-  inputs=[gr.inputs.Textbox(default="PIECE_START")],
-  outputs=['audio'],
-  title=title,
-  examples=[["PIECE_START"], ["PIECE_START STYLE=JSFAKES GENRE=JSFAKES TRACK_START INST=48 BAR_START NOTE_ON=61"]],
-  article="This demo is inspired in the notebook from https://huggingface.co/TristanBehrens/js-fakes-4bars"
+            audio_values = np.zeros(self.to_yield)
+
+        self.on_finalized_audio(audio_values[self.to_yield :], stream_end=True)
+
+    def on_finalized_audio(self, audio: np.ndarray, stream_end: bool = False):
+        """Put the new audio in the queue. If the stream is ending, also put a stop signal in the queue."""
+        self.audio_queue.put(audio, timeout=self.timeout)
+        if stream_end:
+            self.audio_queue.put(self.stop_signal, timeout=self.timeout)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        value = self.audio_queue.get(timeout=self.timeout)
+        if not isinstance(value, np.ndarray) and value == self.stop_signal:
+            raise StopIteration()
+        else:
+            return value
+
+
+sampling_rate = model.audio_encoder.config.sampling_rate
+frame_rate = model.audio_encoder.config.frame_rate
+
+target_dtype = np.int16
+max_range = np.iinfo(target_dtype).max
+
+
+@spaces.GPU
+def generate_audio(text_prompt, audio_length_in_s=10.0, play_steps_in_s=2.0, seed=0):
+    max_new_tokens = int(frame_rate * audio_length_in_s)
+    play_steps = int(frame_rate * play_steps_in_s)
+
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    if device != model.device:
+        model.to(device)
+        if device == "cuda:0":
+            model.half()
+
+    inputs = processor(
+        text=text_prompt,
+        padding=True,
+        return_tensors="pt",
+    )
+
+    streamer = MusicgenStreamer(model, device=device, play_steps=play_steps)
+
+    generation_kwargs = dict(
+        **inputs.to(device),
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+    )
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    set_seed(seed)
+    for new_audio in streamer:
+        print(f"Sample of length: {round(new_audio.shape[0] / sampling_rate, 2)} seconds")
+        new_audio = (new_audio * max_range).astype(np.int16)
+        yield (sampling_rate, new_audio)
+
+
+demo = gr.Interface(
+    fn=generate_audio,
+    inputs=[
+        gr.Text(label="Prompt", value="80s pop track with synth and instrumentals"),
+        gr.Slider(10, 30, value=15, step=5, label="Audio length in seconds"),
+        gr.Slider(0.5, 2.5, value=0.5, step=0.5, label="Streaming interval in seconds", info="Lower = shorter chunks, lower latency, more codec steps"),
+        gr.Slider(0, 10, value=5, step=1, label="Seed for random generations"),
+    ],
+    outputs=[
+        gr.Audio(label="Generated Music", streaming=True, autoplay=True)
+    ],
+    examples = [
+        ["Country acoustic guitar fast line dance singer like  Kenny Chesney and Garth brooks and Luke Combs and Chris Stapleton. bpm: 100", 30, 0.5, 5],
+        ["Electronic Dance track with pulsating bass and high energy synths. bpm: 126", 30, 0.5, 5],
+        ["Rap Beats with deep bass and snappy snares. bpm: 80", 30, 0.5, 5],
+        ["Lo-Fi track with smooth beats and chill vibes. bpm: 100", 30, 0.5, 5],
+        ["Global Groove track with international instruments and dance rhythms. bpm: 128", 30, 0.5, 5],
+        ["Relaxing Meditation music with ambient pads and soothing melodies. bpm: 80", 30, 0.5, 5],
+        ["Rave Dance track with hard-hitting beats and euphoric synths. bpm: 128", 30, 0.5, 5]
+    ],
+
+    title=title,
+    description=description,
+    article=article,
+    cache_examples=False,
 )
 
-iface.launch(debug=True)
+demo.queue().launch()