revert

app.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 """DialogueSidon — two-speaker dialogue separation demo.
-
 Loads exported torch.export components from sarulab-speech/DialogueSidon on
 Hugging Face Hub and runs diffusion-based speaker separation.
 Inputs up to 120 s are processed in one shot; longer inputs use chunked
@@ -12,7 +11,6 @@ import json
 import os
 import subprocess
 import tempfile
-import time
 
 try:
     import spaces
@@ -96,20 +94,20 @@ def extract_fbank_features(
         mean = feat.mean(0, keepdim=True)
         var = feat.var(0, keepdim=True)
         feat = (feat - mean) / torch.sqrt(var + 1e-5)
-        features.append(feat.to(device=device, dtype=torch.float32))
+        features.append(feat.to(device))
 
     input_features, attention_mask = _pad_batch(features)
     b, t, c = input_features.shape
     t = (t // stride) * stride
     input_features = input_features[:, :t, :]
     attention_mask = attention_mask[:, :t]
-    input_features = input_features.reshape(b, t // stride, c * stride).contiguous()
-    attention_mask = attention_mask[:, 1::stride].contiguous()
+    input_features = input_features.reshape(b, t // stride, c * stride)
+    attention_mask = attention_mask[:, 1::stride]
     return {"input_features": input_features, "attention_mask": attention_mask}
 
 
 # ---------------------------------------------------------------------------
-# Model loading (cached per device)
+# Model loading (cached)
 # ---------------------------------------------------------------------------
 
 _cache: dict = {}
@@ -126,7 +124,6 @@ def load_models(device: torch.device) -> dict:
     with open(paths["metadata.json"]) as fp:
         meta = json.load(fp)
 
-    print(f"Loading models onto {device} ...")
     ssl_encoder = torch.export.load(paths["ssl_encoder.pt2"]).module().to(device)
     diffusion_head = torch.export.load(paths["diffusion_head.pt2"]).module().to(device)
     vae_decoder = torch.export.load(paths["vae_decoder.pt2"]).module().to(device)
@@ -186,10 +183,9 @@ def _separate_chunk(
     latent_dim = models["latent_dim"]
 
     noisy_ssl = extract_fbank_features([wav.view(-1)], device)
-    input_features = noisy_ssl["input_features"].to(device=device, dtype=torch.float32).contiguous()
-    attention_mask = noisy_ssl["attention_mask"].to(device=device, dtype=torch.int64).contiguous()
-
-    features, pred0, pred1 = models["ssl_encoder"](input_features)
+    features, pred0, pred1 = models["ssl_encoder"](
+        noisy_ssl["input_features"], noisy_ssl["attention_mask"]
+    )
 
     predicted_latents = torch.cat([pred0, pred1], dim=-1)
     conditioning = torch.cat([_normalize(predicted_latents, models), features], dim=-1)
@@ -207,7 +203,7 @@ def _separate_chunk(
         ).prev_sample
 
     latents = _denormalize(latents, models)
-    spk1 = models["vae_decoder"](latents[:, :, :latent_dim].transpose(1, 2)).squeeze(0)
+    spk1 = models["vae_decoder"](latents[:, :, :latent_dim].transpose(1, 2)).squeeze(0)  # [1, T]
     spk2 = models["vae_decoder"](latents[:, :, latent_dim:].transpose(1, 2)).squeeze(0)
     return torch.cat([spk1, spk2], dim=0)  # [2, T]
 
@@ -241,6 +237,7 @@ def separate(
     models = load_models(device)
     out_sr = models["sample_rate"]
 
+    # resample to 16 kHz
     if sample_rate != SAMPLE_RATE_IN:
         wav_16k = torchaudio.functional.resample(wav, sample_rate, SAMPLE_RATE_IN)
     else:
@@ -251,11 +248,13 @@ def separate(
     total_samples = wav_16k.shape[-1]
 
     if total_samples <= chunk_samples:
+        # single-shot inference
         max_val = wav_16k.abs().max().clamp_min(1e-6)
         wav_norm = torch.nn.functional.pad(0.9 * wav_16k / max_val, (160, 160))
         separated = _separate_chunk(wav_norm, num_steps, models, device)
         return separated, out_sr
 
+    # chunked streaming inference
     overlap_samples_in = int(OVERLAP_SECONDS * SAMPLE_RATE_IN)
     hop_samples = chunk_samples - overlap_samples_in
     starts = list(range(0, total_samples, hop_samples))
@@ -263,14 +262,15 @@ def separate(
     stitched: torch.Tensor | None = None
     prev_end_in = 0
 
-    for start in starts:
+    for idx, start in enumerate(starts):
         end = min(start + chunk_samples, total_samples)
         chunk = wav_16k[:, start:end]
         max_val = chunk.abs().max().clamp_min(1e-6)
         chunk_norm = torch.nn.functional.pad(0.9 * chunk / max_val, (160, 160))
 
-        pred = _separate_chunk(chunk_norm, num_steps, models, device)
+        pred = _separate_chunk(chunk_norm, num_steps, models, device)  # [2, T_out]
 
+        # match output length to input length (resampling ratio)
         target_out = max(1, round((end - start) * out_sr / SAMPLE_RATE_IN))
         if pred.shape[-1] > target_out:
             pred = pred[:, :target_out]
@@ -284,14 +284,11 @@ def separate(
             continue
 
         overlap_in = max(0, prev_end_in - start)
-        overlap_out = max(
-            0,
-            min(
-                round(overlap_in * out_sr / SAMPLE_RATE_IN),
-                stitched.shape[-1],
-                pred.shape[-1],
-            ),
-        )
+        overlap_out = max(0, min(
+            round(overlap_in * out_sr / SAMPLE_RATE_IN),
+            stitched.shape[-1],
+            pred.shape[-1],
+        ))
 
         if overlap_out > 0:
             pred, _ = _maybe_swap(stitched[:, -overlap_out:], pred, overlap_out)
@@ -315,7 +312,8 @@ def extract_audio_from_video(video_path: str) -> tuple[torch.Tensor, int]:
     with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
         tmp_path = tmp.name
     subprocess.run(
-        ["ffmpeg", "-y", "-i", video_path, "-ac", "1", "-ar", str(SAMPLE_RATE_IN), "-vn", tmp_path],
+        ["ffmpeg", "-y", "-i", video_path, "-ac", "1", "-ar", str(SAMPLE_RATE_IN),
+         "-vn", tmp_path],
         check=True,
         stdout=subprocess.DEVNULL,
         stderr=subprocess.DEVNULL,
@@ -331,9 +329,13 @@ def create_stereo_video(
     spk2: np.ndarray,
     out_sr: int,
 ) -> str:
-    """Mux separated speakers (L=spk1, R=spk2) back into a video file."""
-    stereo = np.stack([spk1, spk2], axis=0)
+    """Mux separated speakers (L=spk1, R=spk2) back into a video file.
+    Returns path to the output video (caller is responsible for cleanup).
+    """
+    # Write stereo audio to a temp wav
+    stereo = np.stack([spk1, spk2], axis=0)  # [2, T]
     stereo_tensor = torch.from_numpy(stereo).float()
+    # Normalise to avoid clipping
     peak = stereo_tensor.abs().max().clamp_min(1e-6)
     stereo_tensor = stereo_tensor / peak * 0.9
 
@@ -347,13 +349,13 @@ def create_stereo_video(
     subprocess.run(
         [
             "ffmpeg", "-y",
-            "-i", video_path,
-            "-i", audio_path,
-            "-c:v", "copy",
-            "-c:a", "aac",
+            "-i", video_path,       # original video (video stream)
+            "-i", audio_path,       # new stereo audio
+            "-c:v", "copy",         # copy video stream unchanged
+            "-c:a", "aac",          # encode audio as AAC
             "-b:a", "192k",
-            "-map", "0:v:0",
-            "-map", "1:a:0",
+            "-map", "0:v:0",        # take video from first input
+            "-map", "1:a:0",        # take audio from second input
             "-shortest",
             out_path,
         ],
@@ -365,45 +367,6 @@ def create_stereo_video(
     return out_path
 
 
-# ---------------------------------------------------------------------------
-# RTF report & mel helpers
-# ---------------------------------------------------------------------------
-
-def _extract_mel_db(wav_np: np.ndarray, sr: int, n_mels: int = 80) -> np.ndarray:
-    wav_t = torch.from_numpy(wav_np.copy()).float()
-    if wav_t.ndim != 1:
-        wav_t = wav_t.reshape(-1)
-    peak = wav_t.abs().max().clamp_min(1e-6)
-    wav_t = (wav_t / peak).unsqueeze(0)
-    mel_tf = torchaudio.transforms.MelSpectrogram(
-        sample_rate=sr, n_fft=1024, hop_length=256, n_mels=n_mels
-    )
-    mel = mel_tf(wav_t)
-    mel_db = torchaudio.transforms.AmplitudeToDB(stype="power", top_db=80)(mel)
-    return mel_db[0].numpy()
-
-
-def _rtf_report(duration_s: float, proc_time_s: float, mel_stats: dict, device: str = "unknown") -> str:
-    rtf = proc_time_s / max(duration_s, 1e-6)
-    lines = [
-        "=== DialogueSidon Separation Report ===",
-        "",
-        f"Device          : {device}",
-        f"Input duration  : {duration_s:.2f} s",
-        f"Processing time : {proc_time_s:.2f} s",
-        f"RTF             : {rtf:.3f}x  ({'faster' if rtf < 1.0 else 'slower'} than real-time)",
-        "",
-        "--- Mel Spectrogram Statistics (80 bins) ---",
-        f"{'Source':<12} {'Mean (dB)':>12} {'Peak (dB)':>12} {'Std (dB)':>12}",
-        "-" * 52,
-    ]
-    for name, stats in mel_stats.items():
-        lines.append(
-            f"{name:<12} {stats['mean']:>12.1f} {stats['peak']:>12.1f} {stats['std']:>12.1f}"
-        )
-    return "\n".join(lines)
-
-
 # ---------------------------------------------------------------------------
 # Gradio interface
 # ---------------------------------------------------------------------------
@@ -413,7 +376,8 @@ def get_device() -> torch.device:
 
 
 def _wav_to_numpy_output(wav_tensor: torch.Tensor, sr: int) -> tuple[int, np.ndarray]:
-    arr = wav_tensor.detach().cpu().numpy()
+    arr = wav_tensor.cpu().numpy()
+    # Convert to int16 for Gradio audio output
     arr = np.clip(arr / max(np.abs(arr).max(), 1e-6) * 0.9, -1.0, 1.0)
     return sr, (arr * 32767).astype(np.int16)
 
@@ -421,13 +385,11 @@ def _wav_to_numpy_output(wav_tensor: torch.Tensor, sr: int) -> tuple[int, np.nda
 def run_separation_audio(
     input_audio: tuple[int, np.ndarray] | None,
     num_steps: int,
-) -> tuple[tuple[int, np.ndarray], tuple[int, np.ndarray], str]:
+) -> tuple[tuple[int, np.ndarray], tuple[int, np.ndarray]]:
     if input_audio is None:
         raise gr.Error("Please upload an audio file.")
 
     sr, audio_np = input_audio
-    duration_s = audio_np.shape[0] / sr
-
     wav = torch.from_numpy(audio_np.copy()).float()
 
     if wav.ndim == 1:
@@ -437,68 +399,39 @@ def run_separation_audio(
             wav = wav.T
         wav = wav.mean(dim=0, keepdim=True)
 
-    scale = float(np.iinfo(audio_np.dtype).max) if audio_np.dtype in (np.int16, np.int32) else 1.0
-    if scale != 1.0:
-        wav = wav / scale
+    if audio_np.dtype in (np.int16, np.int32):
+        wav = wav / float(np.iinfo(audio_np.dtype).max)
 
     device = get_device()
+    separated, out_sr = separate(wav, sr, num_steps, device)
 
-    t0 = time.perf_counter()
-    separated, out_sr = separate(wav, sr, int(num_steps), device)
-    proc_time = time.perf_counter() - t0
-
-    spk1_out = _wav_to_numpy_output(separated[0], out_sr)
-    spk2_out = _wav_to_numpy_output(separated[1], out_sr)
-
-    audio_mono = audio_np.mean(axis=-1) if audio_np.ndim == 2 else audio_np
-    mel_in = _extract_mel_db((audio_mono / scale).astype(np.float32), sr)
-    mel_s1 = _extract_mel_db(spk1_out[1].astype(np.float32) / 32767.0, out_sr)
-    mel_s2 = _extract_mel_db(spk2_out[1].astype(np.float32) / 32767.0, out_sr)
-
-    mel_stats = {
-        "Input": {"mean": float(mel_in.mean()), "peak": float(mel_in.max()), "std": float(mel_in.std())},
-        "Speaker 1": {"mean": float(mel_s1.mean()), "peak": float(mel_s1.max()), "std": float(mel_s1.std())},
-        "Speaker 2": {"mean": float(mel_s2.mean()), "peak": float(mel_s2.max()), "std": float(mel_s2.std())},
-    }
-    report = _rtf_report(duration_s, proc_time, mel_stats, str(device))
-
-    return spk1_out, spk2_out, report
+    return (
+        _wav_to_numpy_output(separated[0], out_sr),
+        _wav_to_numpy_output(separated[1], out_sr),
+    )
 
 
 def run_separation_video(
     video_path: str | None,
     num_steps: int,
-) -> tuple[str | None, tuple[int, np.ndarray], tuple[int, np.ndarray], str]:
+) -> tuple[str | None, tuple[int, np.ndarray], tuple[int, np.ndarray]]:
     if video_path is None:
         raise gr.Error("Please upload a video file.")
 
     wav, sr = extract_audio_from_video(video_path)
-    duration_s = wav.shape[-1] / sr
     device = get_device()
+    separated, out_sr = separate(wav, sr, num_steps, device)
 
-    t0 = time.perf_counter()
-    separated, out_sr = separate(wav, sr, int(num_steps), device)
-    proc_time = time.perf_counter() - t0
+    spk1_np = separated[0].cpu().numpy()
+    spk2_np = separated[1].cpu().numpy()
 
-    spk1_np = separated[0].detach().cpu().numpy()
-    spk2_np = separated[1].detach().cpu().numpy()
     out_video = create_stereo_video(video_path, spk1_np, spk2_np, out_sr)
 
-    spk1_out = _wav_to_numpy_output(separated[0], out_sr)
-    spk2_out = _wav_to_numpy_output(separated[1], out_sr)
-
-    mel_in = _extract_mel_db(wav[0].detach().cpu().numpy(), sr)
-    mel_s1 = _extract_mel_db(spk1_out[1].astype(np.float32) / 32767.0, out_sr)
-    mel_s2 = _extract_mel_db(spk2_out[1].astype(np.float32) / 32767.0, out_sr)
-
-    mel_stats = {
-        "Input": {"mean": float(mel_in.mean()), "peak": float(mel_in.max()), "std": float(mel_in.std())},
-        "Speaker 1": {"mean": float(mel_s1.mean()), "peak": float(mel_s1.max()), "std": float(mel_s1.std())},
-        "Speaker 2": {"mean": float(mel_s2.mean()), "peak": float(mel_s2.max()), "std": float(mel_s2.std())},
-    }
-    report = _rtf_report(duration_s, proc_time, mel_stats, str(device))
-
-    return out_video, spk1_out, spk2_out, report
+    return (
+        out_video,
+        _wav_to_numpy_output(separated[0], out_sr),
+        _wav_to_numpy_output(separated[1], out_sr),
+    )
 
 
 if HAS_SPACES:
@@ -513,10 +446,8 @@ with gr.Blocks(title="DialogueSidon — Dialogue Separation") as demo:
         Upload a degraded or noisy audio/video recording of a two-speaker conversation.
         DialogueSidon jointly **separates** the two speakers and **restores** clean, high-quality speech
         from the mixture — handling background noise, reverberation, and channel degradation in one pass.
-
         Inputs up to 120 s are processed in one shot. Longer inputs are processed in 120 s chunks
         with 10 s overlap crossfade and automatic speaker re-alignment across chunks.
-
         **Model**: [sarulab-speech/DialogueSidon](https://huggingface.co/sarulab-speech/DialogueSidon)
         """
     )
@@ -527,6 +458,9 @@ with gr.Blocks(title="DialogueSidon — Dialogue Separation") as demo:
     )
 
     with gr.Tabs():
+        # ------------------------------------------------------------------
+        # Audio tab
+        # ------------------------------------------------------------------
         with gr.Tab("Audio"):
             with gr.Row():
                 with gr.Column():
@@ -540,20 +474,20 @@ with gr.Blocks(title="DialogueSidon — Dialogue Separation") as demo:
                     audio_spk1 = gr.Audio(label="Speaker 1", type="numpy")
                     audio_spk2 = gr.Audio(label="Speaker 2", type="numpy")
 
-            audio_report = gr.Textbox(
-                label="RTF Report", lines=10, max_lines=15, interactive=False
-            )
             audio_btn.click(
                 fn=run_separation_audio,
                 inputs=[audio_input, num_steps],
-                outputs=[audio_spk1, audio_spk2, audio_report],
+                outputs=[audio_spk1, audio_spk2],
             )
             gr.Examples(
-                examples=[["LDC2026S02_mono.wav", 100]],
-                inputs=[audio_input, num_steps],
+                examples=[["LDC2026S02_mono.wav",100]],
+                inputs=[audio_input,num_steps],
                 label="Ex1"
             )
 
+        # ------------------------------------------------------------------
+        # Video tab
+        # ------------------------------------------------------------------
         with gr.Tab("Video"):
             with gr.Row():
                 with gr.Column():
@@ -569,13 +503,10 @@ with gr.Blocks(title="DialogueSidon — Dialogue Separation") as demo:
                     video_spk1 = gr.Audio(label="Speaker 1 (audio only)", type="numpy")
                     video_spk2 = gr.Audio(label="Speaker 2 (audio only)", type="numpy")
 
-            video_report = gr.Textbox(
-                label="RTF Report", lines=10, max_lines=15, interactive=False
-            )
             video_btn.click(
                 fn=run_separation_video,
                 inputs=[video_input, num_steps],
-                outputs=[video_output, video_spk1, video_spk2, video_report],
+                outputs=[video_output, video_spk1, video_spk2],
             )
 
     gr.Markdown("---\n**License**: CC-BY-NC 4.0 — non-commercial use only.")