Update app.py

app.py

@@ -9,17 +9,12 @@ from transformers import MusicgenForConditionalGeneration, MusicgenProcessor, se
 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")
 
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
-if device != model.device:
-    model.to(device)
-    if device == "cuda:0":
-        model.half()
-
 title = "MusicGen Streaming"
 
 description = """
@@ -49,7 +44,7 @@ particularly when the chunk size is chosen to be small. In practice, the chunk s
 smaller chunk size will mean that the first chunk is ready faster, but should not be chosen so small that the model generates slower 
 than the time it takes to play the audio.
 
-For details on how the streaming class works, check out the source code for the [MusicgenStreamer](https://huggingface.co/spaces/sanchit-gandhi/musicgen-streaming/blob/main/app.py#L56).
+For details on how the streaming class works, check out the source code for the [MusicgenStreamer](https://huggingface.co/spaces/sanchit-gandhi/musicgen-streaming/blob/main/app.py#L50).
 """
 
 
@@ -180,10 +175,17 @@ 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,