Model for GPU adjusting ui

app.py

@@ -1,19 +1,22 @@
 import spaces
+import os
 import torchaudio
 from audiocraft.models import AudioGen
 from audiocraft.data.audio import audio_write
-import os
 import gradio as gr
 
-
-model = AudioGen.get_pretrained('facebook/audiogen-medium')
-model.set_generation_params(duration=5)  # generate 5 seconds.
-
 OUTPUT_DIR = "audio_files"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
+model = None  
+
 @spaces.GPU
 def generate_audio(descriptions: str):
+    global model
+    if model is None:
+        model = AudioGen.get_pretrained('facebook/audiogen-medium')
+        model.set_generation_params(duration=5)  # 5 seconds
+
     safe_name = "_".join(descriptions.split())
     output_path = os.path.join(OUTPUT_DIR, safe_name)
 
@@ -31,17 +34,13 @@ def generate_audio(descriptions: str):
     print(f"Generated audio for '{descriptions}' -> {final_path}")
     return final_path
 
-def ui_full():
-    with gr.Blocks() as interface:
-        gr.Markdown("# AudioGen Demo")
-        with gr.Row():
-            descriptions = gr.Textbox(lines=2, label="Enter a description of the audio")
-        with gr.Row():
-            generate_button = gr.Button("Generate Audio")
-        with gr.Row():
-            output = gr.Audio(label="Generated Audio")
-
-        generate_button.click(fn=generate_audio, inputs=descriptions, outputs=output)
-    return interface
+with gr.Blocks() as demo:
+    gr.Markdown("# AudioGen Demo")
+    with gr.Row():
+        descriptions = gr.Textbox(lines=2, label="Enter a description of the audio")
+    with gr.Row():
+        generate_button = gr.Button("Generate Audio")
+    with gr.Row():
+        output = gr.Audio(label="Generated Audio")
 
-demo = ui_full()
+    generate_button.click(fn=generate_audio, inputs=descriptions, outputs=output)