Add SFX model support with model selector

app.py

@@ -14,7 +14,6 @@ if hf_token:
 
 
 # Required for ZeroGPU Spaces - must have at least one @spaces.GPU function
-# This is a minimal startup check that uses negligible quota
 @spaces.GPU(duration=1)
 def _gpu_startup_check():
     return "GPU check passed"
@@ -24,18 +23,29 @@ def _gpu_startup_check():
 _gpu_startup_check()
 
 
-# Load model once at startup (on CPU)
-print("Loading Stable Audio 3 Small model...")
-model = StableAudioModel.from_pretrained(
-    "small-music",
-    device="cpu"
-)
-print("Model loaded successfully!")
+# Model cache
+MODEL_CACHE = {}
 
 
 @spaces.GPU(duration=1)
-def generate_audio(prompt, duration, steps, cfg_scale, seed):
-    print(f"Generating: prompt='{prompt}', duration={duration}s, steps={steps}, cfg={cfg_scale}, seed={seed}")
+def load_model(model_name):
+    """Load model on demand and cache it."""
+    if model_name not in MODEL_CACHE:
+        print(f"Loading {model_name} model...")
+        model = StableAudioModel.from_pretrained(
+            model_name,
+            device="cpu"
+        )
+        MODEL_CACHE[model_name] = model
+        print(f"{model_name} loaded successfully!")
+    return MODEL_CACHE[model_name]
+
+
+@spaces.GPU(duration=1)
+def generate_audio(prompt, duration, steps, cfg_scale, seed, model_name):
+    print(f"Generating with {model_name}: prompt='{prompt}', duration={duration}s, steps={steps}, cfg={cfg_scale}, seed={seed}")
+
+    model = load_model(model_name)
 
     audio = model.generate(
         prompt=prompt,
@@ -57,14 +67,19 @@ def generate_audio(prompt, duration, steps, cfg_scale, seed):
 
 
 with gr.Blocks(title="Stable Audio 3 Small") as demo:
-    gr.Markdown("# 🎵 Stable Audio 3 Small - Music Generation")
-    gr.Markdown("Generate music using Stability AI's Stable Audio 3 Small model. Runs on CPU.")
+    gr.Markdown("# 🎵 Stable Audio 3 Small - Music & SFX Generation")
+    gr.Markdown("Generate music and sound effects using Stability AI's Stable Audio 3 Small models. Runs on CPU.")
 
     with gr.Row():
         with gr.Column():
+            model_name = gr.Dropdown(
+                choices=["stabilityai/stable-audio-3-small-music", "stabilityai/stable-audio-3-small-sfx"],
+                value="stabilityai/stable-audio-3-small-music",
+                label="Model"
+            )
             prompt = gr.Textbox(
                 label="Prompt",
-                placeholder="Describe the music you want to generate...",
+                placeholder="Describe the music or sound effect you want to generate...",
                 lines=2
             )
             duration = gr.Slider(
@@ -92,7 +107,7 @@ with gr.Blocks(title="Stable Audio 3 Small") as demo:
 
     btn.click(
         fn=generate_audio,
-        inputs=[prompt, duration, steps, cfg_scale, seed],
+        inputs=[prompt, duration, steps, cfg_scale, seed, model_name],
         outputs=audio_output
     )