Change Files

app.py

@@ -24,7 +24,7 @@ def _get_model(ckpt_name: str):
         raise ValueError(f"Invalid checkpoint {ckpt_name!r}, choose from {VALID_CKPTS}")
     if ckpt_name in _model_cache:
         return _model_cache[ckpt_name]
-    ckpt_path = os.path.join(CHECKPOINT_DIR, f"{ckpt_name}.pt")
+    ckpt_path = os.path.join(CHECKPOINT_DIR, f"{ckpt_name}.pth")
     model = UFormer(config).to(DEVICE).eval()
     state = torch.load(ckpt_path, map_location=DEVICE)
     model.load_state_dict(state)
@@ -43,6 +43,7 @@ def _overlap_add(model, x: np.ndarray, sr: int, chunk_s: float=5., hop_s: float=
     out   = np.zeros_like(x_pad)
     norm  = np.zeros((1, x_pad.shape[1]))
     n_chunks = 1 + (x_pad.shape[1] - chunk) // hop
+    print(f"Processing {n_chunks} chunks of size {chunk} with hop {hop}...")
 
     for i in range(n_chunks):
         s = i * hop
@@ -52,7 +53,8 @@ def _overlap_add(model, x: np.ndarray, sr: int, chunk_s: float=5., hop_s: float=
         out[:, s:s+chunk]  += y * win
         norm[:, s:s+chunk] += win
 
-    return (out / norm)[:, :T]
+    eps = 1e-8
+    return (out / (norm + eps))[:, :T]
 
 # ——————————————————————
 # 3) Restore function for Gradio
@@ -81,12 +83,12 @@ def restore_fn(audio_path, checkpoint):
 demo = gr.Interface(
     fn=restore_fn,
     inputs=[
-        gr.Audio(source="upload", type="filepath", label="Your Input"),
+        gr.Audio(sources="upload", type="filepath", label="Your Input"),
         gr.Dropdown(VALID_CKPTS, label="Checkpoint")
     ],
     outputs=gr.Audio(type="filepath", label="Restored Output"),
     title="🎵 Music Source Restoration",
-    description="Upload a WAV file and choose an instrument/group checkpoint to restore.",
+    description="Upload an (stereo) audio file and choose an instrument/group checkpoint to restore. Please note that these are baseline models for demonstration purposes only, and most of them don't perform really well...",
     allow_flagging="never"
 )