Update app.py

app.py

@@ -12,13 +12,9 @@ import gradio as gr
 import spaces
 
 
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
-
-model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small", low_cpu_mem_usage=True)
 processor = MusicgenProcessor.from_pretrained("facebook/musicgen-small")
 
-if device == "cuda:0":
-    model.to(device).half();
 
 class MusicgenStreamer(BaseStreamer):
     def __init__(
@@ -143,25 +139,26 @@ class MusicgenStreamer(BaseStreamer):
 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):
-    inputs = processor(
-        text=text_prompt,
-        padding=True,
-        return_tensors="pt",
-    )
-
     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();
 
-    if device == "cuda:0":
-        model.to(device).half();
+    inputs = processor(
+        text=text_prompt,
+        padding=True,
+        return_tensors="pt",
+    )
 
     streamer = MusicgenStreamer(model, device=device, play_steps=play_steps)