Update figs

app.py

@@ -43,12 +43,21 @@ with open(cfg_f, 'r') as f:
 # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 device = "cuda"
 model  = SEMamba(cfg).to(device)
-sdict  = torch.load(ckpt, map_location=device)
-model.load_state_dict(sdict["generator"])
-model.eval()
+#sdict  = torch.load(ckpt, map_location=device)
+#model.load_state_dict(sdict["generator"])
+#model.eval()
 
 @spaces.GPU
 def enhance(filepath):
+    # Load model based on selection
+    ckpt_path = {
+        "VCTK-Demand": "ckpts/SEMamba_advanced.pth",
+        "VCTK+DNS": "ckpts/vd.pth"
+    }[model_name]
+
+    print("Loading:", ckpt_path)
+    model.load_state_dict(torch.load(ckpt_path, map_location=device)["generator"])
+    model.eval()
     with torch.no_grad():
         # load & resample
         wav, orig_sr = librosa.load(filepath, sr=None)
@@ -107,13 +116,33 @@ def enhance(filepath):
 
     return "enhanced.wav", fig
 
+#with gr.Blocks() as demo:
+#    gr.Markdown(ABOUT)
+#    input_audio = gr.Audio(label="Input Audio", type="filepath", interactive=True)
+#    enhance_btn = gr.Button("Enhance")
+#    output_audio = gr.Audio(label="Enhanced Audio", type="filepath")
+#    plot_output = gr.Plot(label="Spectrograms")
+#
+#    enhance_btn.click(fn=enhance, inputs=input_audio, outputs=[output_audio, plot_output])
+#
+#demo.queue().launch()
+
 with gr.Blocks() as demo:
     gr.Markdown(ABOUT)
     input_audio = gr.Audio(label="Input Audio", type="filepath", interactive=True)
+    model_choice = gr.Radio(
+        label="Choose Model",
+        choices=["VCTK-Demand", "VCTK+DNS"],
+        value="VCTK-Demand"
+    )
     enhance_btn = gr.Button("Enhance")
     output_audio = gr.Audio(label="Enhanced Audio", type="filepath")
     plot_output = gr.Plot(label="Spectrograms")
 
-    enhance_btn.click(fn=enhance, inputs=input_audio, outputs=[output_audio, plot_output])
+    enhance_btn.click(
+        fn=enhance,
+        inputs=[input_audio, model_choice],
+        outputs=[output_audio, plot_output]
+    )
 
 demo.queue().launch()