added elastic model

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import torch
-from transformers import AutoProcessor, MusicgenForConditionalGeneration
+from transformers import AutoProcessor
+from elastic_models.transformers import MusicgenForConditionalGeneration
 import scipy.io.wavfile
 import numpy as np
 import subprocess
@@ -51,7 +52,13 @@ def setup_flash_attention():
 def load_model():
     """Load the musicgen model and processor"""
     processor = AutoProcessor.from_pretrained("facebook/musicgen-large")
-    model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-large")
+    model = MusicgenForConditionalGeneration.from_pretrained(
+                "facebook/musicgen-large",
+                torch_dtype=torch.float16,
+                device=self.device,
+                mode="S",
+                __paged=True,
+            )
     return processor, model
 
 def generate_music(text_prompt, duration=10, temperature=1.0, top_k=250, top_p=0.0):
@@ -75,6 +82,7 @@ def generate_music(text_prompt, duration=10, temperature=1.0, top_k=250, top_p=0
                 temperature=temperature,
                 top_k=top_k,
                 top_p=top_p,
+                cache_implementation="paged"
             )
         
         # Convert to numpy array and prepare for output
@@ -159,6 +167,8 @@ with gr.Blocks(title="MusicGen Large - Music Generation") as demo:
             ["classical violin concerto"],
             ["reggae with steel drums and bass"],
             ["rock ballad with electric guitar solo"],
+            ["test example"],
+
         ],
         inputs=text_input,
         label="Example Prompts"