Create app.py

app.py

@@ -0,0 +1,98 @@
+from queue import Queue
+from threading import Thread
+from typing import Optional
+import numpy as np
+import torch
+from transformers import MusicgenForConditionalGeneration, MusicgenProcessor, set_seed
+from transformers.generation.streamers import BaseStreamer
+import gradio as gr
+import spaces
+
+model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+processor = MusicgenProcessor.from_pretrained("facebook/musicgen-small")
+
+title, description, article = "MusicGen Streaming (by Sanchit Gandhi)"
+
+class MusicgenStreamer(BaseStreamer):
+    def __init__(self, model, device=None, play_steps=10, stride=None, timeout=None):
+        self.decoder, self.audio_encoder, self.generation_config = model.decoder, model.audio_encoder, model.generation_config
+        self.device = device if device else model.device
+        self.play_steps = play_steps
+        self.stride = stride if stride else np.prod(self.audio_encoder.config.upsampling_ratios) * (play_steps - self.decoder.num_codebooks) // 6
+        self.token_cache, self.to_yield, self.audio_queue, self.stop_signal, self.timeout = None, 0, Queue(), None, timeout
+
+    def apply_delay_pattern_mask(self, input_ids):
+        _, 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])
+        input_ids = self.decoder.apply_delay_pattern_mask(input_ids, mask)
+        input_ids = input_ids[input_ids != self.generation_config.pad_token_id].reshape(1, self.decoder.num_codebooks, -1)[None, ...]
+        return self.audio_encoder.decode(input_ids.to(self.audio_encoder.device), audio_scales=[None]).audio_values[0, 0].cpu().float().numpy()
+
+    def put(self, value):
+        if value.shape[0] // self.decoder.num_codebooks > 1:
+            raise ValueError("MusicgenStreamer only supports batch size 1")
+        self.token_cache = torch.concatenate([self.token_cache, value[:, None]], dim=-1) if self.token_cache else value
+        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):
+        audio_values = self.apply_delay_pattern_mask(self.token_cache) if self.token_cache else np.zeros(self.to_yield)
+        self.on_finalized_audio(audio_values[self.to_yield :], stream_end=True)
+
+    def on_finalized_audio(self, audio, stream_end=False):
+        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()
+        return value
+
+sampling_rate, frame_rate = model.audio_encoder.config.sampling_rate, model.audio_encoder.config.frame_rate
+
+@spaces.GPU()
+def generate_audio(text_prompt, audio_length_in_s=10.0, play_steps_in_s=2.0, seed=0):
+    max_new_tokens, play_steps = int(frame_rate * audio_length_in_s), 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")
+        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=1.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=[
+        ["An 80s driving pop song with heavy drums and synth pads in the background", 30, 1.5, 5],
+        ["A cheerful country song with acoustic guitars", 30, 1.5, 5],
+        ["90s rock song with electric guitar and heavy drums", 30, 1.5, 5],
+        ["a light and cheerly EDM track, with syncopated drums, aery pads, and strong emotions bpm: 130", 30, 1.5, 5],
+        ["lofi slow bpm electro chill with organic samples", 30, 1.5, 5],
+    ],
+    title=title,
+    cache_examples=False,
+)
+
+demo.queue().launch()
+This compressed version maintains the core functionality while removing redundant comments and