Model loading once

app.py

@@ -3,42 +3,44 @@ from audiocraft.models import AudioGen
 from audiocraft.data.audio import audio_write
 import os
 import gradio as gr
+import spaces
 
 model = AudioGen.get_pretrained('facebook/audiogen-medium')
 model.set_generation_params(duration=5)  # generate 5 seconds.
 
-def generate_audio(descriptions):
-    if not os.path.exists('audio_files'):
-        os.makedirs('audio_files')
-    
-    wav = model.generate([descriptions])  # generates 3 samples.
-    results = []
-
-    for idx, one_wav in enumerate(wav):
-        filename = f'{descriptions}.wav'
-        file_path = os.path.join('audio_files', filename)
-        audio_write(file_path, one_wav.cpu(), model.sample_rate, strategy="loudness", loudness_compressor=True, add_suffix=False)
-        print(f"Generated audio for '{descriptions}'")
-        results.append(file_path)
-    
-    return results[0]
+OUTPUT_DIR = "audio_files"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+@spaces.GPU
+def generate_audio(descriptions: str):
+    safe_name = "_".join(descriptions.split())
+    output_path = os.path.join(OUTPUT_DIR, safe_name)
+
+    wav = model.generate([descriptions]) 
+    audio_write(
+        output_path,
+        wav[0].cpu(),
+        model.sample_rate,
+        strategy="loudness",
+        loudness_compressor=True,
+        add_suffix=False,
+    )
+
+    final_path = f"{output_path}.wav"
+    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 descriptions of the audio to generate")
+    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():    
+        with gr.Row():
             output = gr.Audio(label="Generated Audio")
 
-        generate_button.click(fn=generate_audio, inputs=descriptions, outputs=[output])
-        interface.queue().launch()
-    
-ui_full()
+        generate_button.click(fn=generate_audio, inputs=descriptions, outputs=output)
+    return interface
+
+demo = ui_full()