Update app.py

app.py

@@ -108,49 +108,78 @@ def demo_fn(speech_upl: str, noise_type: str, snr: int, mic_input: Optional[str]
     snr = int(snr)
     noise_fn = NOISES[noise_type]
     meta = AudioMetaData(-1, -1, -1, -1, "")
-    max_s = 10  # limit to 10 seconds
+
+    # --- Load audio (full file, no hard 10s crop) ---
     if speech_upl is not None:
         sample, meta = load_audio(speech_upl, sr)
-        max_len = max_s * sr
-        if sample.shape[-1] > max_len:
-            start = torch.randint(0, sample.shape[-1] - max_len, ()).item()
-            sample = sample[..., start : start + max_len]
     else:
         sample, meta = load_audio("samples/p232_013_clean.wav", sr)
-        sample = sample[..., : max_s * sr]
+
+    # Mix to mono if multi-channel
     if sample.dim() > 1 and sample.shape[0] > 1:
         assert (
             sample.shape[1] > sample.shape[0]
         ), f"Expecting channels first, but got {sample.shape}"
         sample = sample.mean(dim=0, keepdim=True)
+
     logger.info(f"Loaded sample with shape {sample.shape}")
+
+    # Add noise if requested
     if noise_fn is not None:
         noise, _ = load_audio(noise_fn, sr)  # type: ignore
         logger.info(f"Loaded noise with shape {noise.shape}")
         _, _, sample = mix_at_snr(sample, noise, snr)
+
     logger.info("Start denoising audio")
-    enhanced = enhance(model, df, sample)
+
+    # --- Process in chunks instead of single forward pass ---
+    chunk_size = sr * 10   # 10 seconds per chunk
+    enhanced_chunks = []
+
+    for i in range(0, sample.shape[-1], chunk_size):
+        chunk = sample[..., i:i + chunk_size]
+        if chunk.shape[-1] == 0:
+            continue
+        logger.info(f"Enhancing chunk {i//chunk_size + 1}")
+        enhanced_chunk = enhance(model, df, chunk)
+
+        # Apply short fade-in to smooth first chunk
+        if i == 0:
+            lim = torch.linspace(0.0, 1.0, int(sr * 0.15)).unsqueeze(0)
+            lim = torch.cat((lim, torch.ones(1, enhanced_chunk.shape[1] - lim.shape[1])), dim=1)
+            enhanced_chunk = enhanced_chunk * lim
+
+        enhanced_chunks.append(enhanced_chunk)
+
+    # Concatenate all enhanced chunks into one
+    enhanced = torch.cat(enhanced_chunks, dim=-1)
     logger.info("Denoising finished")
-    lim = torch.linspace(0.0, 1.0, int(sr * 0.15)).unsqueeze(0)
-    lim = torch.cat((lim, torch.ones(1, enhanced.shape[1] - lim.shape[1])), dim=1)
-    enhanced = enhanced * lim
+
+    # Resample back if needed
     if meta.sample_rate != sr:
         enhanced = resample(enhanced, sr, meta.sample_rate)
         sample = resample(sample, sr, meta.sample_rate)
         sr = meta.sample_rate
+
+    # Save noisy & enhanced wavs
     noisy_wav = tempfile.NamedTemporaryFile(suffix="noisy.wav", delete=False).name
     save_audio(noisy_wav, sample, sr)
     enhanced_wav = tempfile.NamedTemporaryFile(suffix="enhanced.wav", delete=False).name
     save_audio(enhanced_wav, enhanced, sr)
-    logger.info(f"saved audios: {noisy_wav}, {enhanced_wav}")
+    logger.info(f"Saved audios: {noisy_wav}, {enhanced_wav}")
+
+    # Plot spectrograms
     ax_noisy.clear()
     ax_enh.clear()
     noisy_im = spec_im(sample, sr=sr, figure=fig_noisy, ax=ax_noisy)
     enh_im = spec_im(enhanced, sr=sr, figure=fig_enh, ax=ax_enh)
+
+    # Cleanup temp files
     filter = [speech_upl, noisy_wav, enhanced_wav]
     if mic_input is not None and mic_input != "":
         filter.append(mic_input)
     cleanup_tmp(filter)
+
     return noisy_wav, noisy_im, enhanced_wav, enh_im