Upload

app.py

@@ -56,45 +56,48 @@ model.eval()
 
    
 @spaces.GPU
-def enhance_and_plot(audio):
+# --- Inference ---
+def enhance(audio):
     if audio is None: return None, None
     orig_sr, wav_np = audio
     if orig_sr != sr:
         wav_np = librosa.resample(wav_np, orig_sr, sr)
     wav = torch.from_numpy(wav_np).float().to(device)
-    norm = torch.sqrt(len(wav)/torch.sum(wav**2))
+    norm = torch.sqrt(len(wav) / torch.sum(wav ** 2))
     wav = (wav * norm).unsqueeze(0)
 
     amp, pha, _ = mag_phase_stft(wav, **stft_cfg, compress_factor=model_cfg["compress_factor"])
     amp_g, pha_g = model(amp, pha)
     out = mag_phase_istft(amp_g, pha_g, **stft_cfg, compress_factor=model_cfg["compress_factor"])
     out = (out / norm).squeeze().cpu().numpy()
-
     if orig_sr != sr:
         out = librosa.resample(out, sr, orig_sr)
 
-    D = librosa.stft(out, n_fft=512, hop_length=256)
+    # draw spectrum
+    D = librosa.stft(out, n_fft=1024, hop_length=512)
     S_db = librosa.amplitude_to_db(np.abs(D), ref=np.max)
     fig, ax = plt.subplots()
-    img = librosa.display.specshow(S_db, sr=orig_sr, hop_length=256, x_axis='time', y_axis='hz', ax=ax)
-    plt.colorbar(img, ax=ax, format="%+2.0f dB")
-    ax.set_title("Enhanced Output Spectrum")
-
+    librosa.display.specshow(S_db, sr=orig_sr, hop_length=512, x_axis='time', y_axis='hz', ax=ax)
+    ax.set_title("Enhanced Spectrogram")
+    plt.colorbar(format="%+2.0f dB")
     return (orig_sr, out), fig
 
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🎧 SEMamba Speech Enhancement")
-    gr.Markdown("Upload or record a noisy audio sample to enhance it and view the spectrogram.")
-
-    with gr.Row():
-        with gr.Column():
-            input_audio = gr.Audio(sources=["upload", "microphone"], type="numpy", label="Noisy Input")
-            btn = gr.Button("Enhance")
-
-        with gr.Column():
-            output_audio = gr.Audio(label="Enhanced Output", type="numpy")
-            spectrum     = gr.Plot(label="Spectrogram")
-
-    btn.click(enhance_and_plot, inputs=input_audio, outputs=[output_audio, spectrum])
-
-demo.launch()
+# --- Interface ---
+se_demo = gr.Interface(
+    fn=enhance,
+    inputs=gr.Audio(sources=["upload", "microphone"], type="numpy", label="Input Audio"),
+    outputs=[
+        gr.Audio(label="Enhanced Audio", type="numpy"),
+        gr.Plot(label="Spectrogram")
+    ],
+    title="<a href='https://github.com/RoyChao19477/SEMamba' target='_blank'>SEMamba</a>: Speech Enhancement",
+    description="SEMamba is a state-space model for real-world noisy speech enhancement. Upload or record a noisy sample to hear the result and view the spectrogram.",
+    article="<p style='text-align: center'><a href='https://arxiv.org/abs/2405.15144' target='_blank'>SEMamba: Mamba for Long-Context Speech Enhancement (SLT 2024)</a></p>",
+    examples=[
+        ["examples/noisy_sample_16k.wav"]
+    ],
+    cache_examples=True
+)
+
+# --- Launch ---
+se_demo.launch()