Update app.py

app.py

@@ -2,26 +2,24 @@
 
 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 = "AI Radio"
 
-class MusicgenStreamer(BaseStreamer):
+class MusicgenStreamer:
     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.device = device or 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
+        self.stride = stride or 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.timeout = None, 0, Queue(), timeout
+        self.stop_signal = object()
 
     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])
@@ -32,15 +30,15 @@ class MusicgenStreamer(BaseStreamer):
     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
+        self.token_cache = torch.cat([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.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)
+        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)
@@ -52,29 +50,26 @@ class MusicgenStreamer(BaseStreamer):
 
     def __next__(self):
         value = self.audio_queue.get(timeout=self.timeout)
-        if not isinstance(value, np.ndarray) and value == self.stop_signal:
+        if value is 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()
+    max_new_tokens = int(model.audio_encoder.config.frame_rate * audio_length_in_s)
+    play_steps = int(model.audio_encoder.config.frame_rate * play_steps_in_s)
     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()
+    Thread(target=model.generate, kwargs=dict(**inputs.to(device), streamer=streamer, max_new_tokens=max_new_tokens)).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
+        print(f"Sample of length: {round(new_audio.shape[0] / model.audio_encoder.config.sampling_rate, 2)} seconds")
+        yield model.audio_encoder.config.sampling_rate, new_audio
 
 demo = gr.Interface(
     fn=generate_audio,