add: option for model and duration separately

app.py

@@ -5,35 +5,42 @@ import numpy as np
 import gradio as gr
 from sonics import HFAudioClassifier
 
-# Model configurations
-MODEL_IDS = {
-    "SpecTTTra-α (5s)": "awsaf49/sonics-spectttra-alpha-5s",
-    "SpecTTTra-β (5s)": "awsaf49/sonics-spectttra-beta-5s",
-    "SpecTTTra-γ (5s)": "awsaf49/sonics-spectttra-gamma-5s",
-    "SpecTTTra-α (120s)": "awsaf49/sonics-spectttra-alpha-120s",
-    "SpecTTTra-β (120s)": "awsaf49/sonics-spectttra-beta-120s",
-    "SpecTTTra-γ (120s)": "awsaf49/sonics-spectttra-gamma-120s",
-}
+# Restructured model configurations for separate selectors
+MODEL_TYPES = ["SpecTTTra-α", "SpecTTTra-β", "SpecTTTra-γ"]
+DURATIONS = ["5s", "120s"]
+
+# Mapping for model IDs
+def get_model_id(model_type, duration):
+    model_map = {
+        "SpecTTTra-α-5s": "awsaf49/sonics-spectttra-alpha-5s",
+        "SpecTTTra-β-5s": "awsaf49/sonics-spectttra-beta-5s",
+        "SpecTTTra-γ-5s": "awsaf49/sonics-spectttra-gamma-5s",
+        "SpecTTTra-α-120s": "awsaf49/sonics-spectttra-alpha-120s",
+        "SpecTTTra-β-120s": "awsaf49/sonics-spectttra-beta-120s",
+        "SpecTTTra-γ-120s": "awsaf49/sonics-spectttra-gamma-120s",
+    }
+    key = f"{model_type}-{duration}"
+    return model_map[key]
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model_cache = {}
 
-
-def load_model(model_name):
+def load_model(model_type, duration):
     """Load model if not already cached"""
-    if model_name not in model_cache:
-        model_id = MODEL_IDS[model_name]
+    model_key = f"{model_type}-{duration}"
+    if model_key not in model_cache:
+        model_id = get_model_id(model_type, duration)
         model = HFAudioClassifier.from_pretrained(model_id)
         model = model.to(device)
         model.eval()
-        model_cache[model_name] = model
-    return model_cache[model_name]
+        model_cache[model_key] = model
+    return model_cache[model_key]
 
 
-def process_audio(audio_path, model_name):
+def process_audio(audio_path, model_type, duration):
     """Process audio file and return prediction"""
     try:
-        model = load_model(model_name)
+        model = load_model(model_type, duration)
         max_time = model.config.audio.max_time
 
         # Load and process audio
@@ -69,11 +76,11 @@ def process_audio(audio_path, model_name):
         return {"Error": str(e)}
 
 
-def predict(audio_file, model_name):
+def predict(audio_file, model_type, duration):
     """Gradio interface function"""
     if audio_file is None:
         return {"Message": "Please upload an audio file"}
-    return process_audio(audio_file, model_name)
+    return process_audio(audio_file, model_type, duration)
 
 
 # Updated CSS with better color scheme for resource links
@@ -146,6 +153,15 @@ css = """
     margin-top: 30px;
     padding: 15px;
 }
+
+/* Selectors wrapper for side-by-side appearance */
+.selectors-wrapper {
+    display: flex;
+    gap: 10px;
+}
+.selectors-wrapper > div {
+    flex: 1;
+}
 """
 
 # Create Gradio interface
@@ -199,12 +215,21 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean()) as demo:
                 elem_id="audio_input"
             )
             
-            model_dropdown = gr.Dropdown(
-                choices=list(MODEL_IDS.keys()),
-                value="SpecTTTra-γ (5s)",
-                label="Select Model",
-                elem_id="model_dropdown"
-            )
+            # Add CSS class to create a wrapper for side-by-side dropdowns
+            with gr.Row(elem_classes="selectors-wrapper"):
+                model_dropdown = gr.Dropdown(
+                    choices=MODEL_TYPES,
+                    value="SpecTTTra-γ",
+                    label="Select Model",
+                    elem_id="model_dropdown"
+                )
+                
+                duration_dropdown = gr.Dropdown(
+                    choices=DURATIONS,
+                    value="5s",
+                    label="Select Duration",
+                    elem_id="duration_dropdown"
+                )
             
             submit_btn = gr.Button(
                 "✨ Analyze Audio",
@@ -240,10 +265,10 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean()) as demo:
     with gr.Accordion("Example Audio Files", open=True):
         gr.Examples(
             examples=[
-                ["example/real_song.mp3", "SpecTTTra-γ (5s)"],
-                ["example/fake_song.mp3", "SpecTTTra-γ (5s)"],
+                ["example/real_song.mp3", "SpecTTTra-γ", "5s"],
+                ["example/fake_song.mp3", "SpecTTTra-γ", "5s"],
             ],
-            inputs=[audio_input, model_dropdown],
+            inputs=[audio_input, model_dropdown, duration_dropdown],
             outputs=[output],
             fn=predict,
             cache_examples=True,
@@ -260,7 +285,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean()) as demo:
     )
 
     # Prediction handling
-    submit_btn.click(fn=predict, inputs=[audio_input, model_dropdown], outputs=[output])
+    submit_btn.click(fn=predict, inputs=[audio_input, model_dropdown, duration_dropdown], outputs=[output])
 
 if __name__ == "__main__":
     demo.launch()