Add padding to chunked inference

MACOS_ARM64.md

@@ -40,6 +40,17 @@ python inference.py \
   --chunk-batch-size=2
 ```
 
+Model output is wrong near the edges of whatever it is given — worst at the
+end, where the model runs out of audio. `--chunk-pad-seconds` infers each
+chunk with that much surrounding audio on each side and discards it from the
+output, keeping those wrong edges away from chunk boundaries. The padding is
+real neighboring audio, and the file's own edges take none: the first padded
+chunk begins at the file start and the last ends at the file end, extending
+further left instead, so both file edges come out as an unchunked
+pass would produce them. The default of 0 keeps the previous output
+exactly; 1.0 is past the measured wrong-edge extent (about 0.5 s) and gives
+clean seams. Each forward pass grows from chunk to chunk plus twice the pad.
+
 `--chunk-batch-size` controls how many chunks are transferred and processed in
 one model forward pass. The default of `1` minimizes accelerator memory; larger
 values keep memory bounded by a fixed batch but may trade additional memory for

README.md

@@ -102,14 +102,18 @@ python inference.py \
   --device=auto \
   --chunk-seconds=6 \
   --overlap-seconds=1 \
+  --chunk-pad-seconds=1 \
   --chunk-batch-size=2
 ```
 
 `--chunk-batch-size` groups a fixed number of chunks into each model forward
 pass. Its default is `1` for minimum accelerator memory; larger values may
 improve throughput on some devices but use more memory and should be measured
 on the target backend. Normalized linear crossfades preserve the input frame
-count, and the overlap must be no more than half the chunk duration. Chunking
+count, and the overlap must be no more than half the chunk duration.
+`--chunk-pad-seconds` infers each chunk with that much surrounding audio per
+side and discards it from the output; with its default of `0` the listener may
+hear artifacts at the chunk seams. Chunking
 can change the result relative to full-file inference because the model sees
 less context, so important material should still receive a listening check. See
 [Apple Silicon macOS notes](MACOS_ARM64.md) for a tested MPS setup and current

inference.py

@@ -83,22 +83,25 @@ def save_audio(file_path, audio, sample_rate):
     sf.write(output_path, audio, sample_rate, subtype="FLOAT")
 
 
-def resolve_chunking(chunk_seconds, overlap_seconds):
+def resolve_chunking(chunk_seconds, overlap_seconds, chunk_pad_seconds=0.0):
     """Convert optional chunk settings to samples and validate crossfades."""
     if chunk_seconds is None:
-        return None, 0
+        return None, 0, 0
     if not np.isfinite(chunk_seconds) or chunk_seconds <= 0:
         raise ValueError("Chunk duration must be a positive finite number.")
     if not np.isfinite(overlap_seconds) or overlap_seconds < 0:
         raise ValueError("Overlap duration must be a non-negative finite number.")
+    if not np.isfinite(chunk_pad_seconds) or chunk_pad_seconds < 0:
+        raise ValueError("Chunk pad duration must be a non-negative finite number.")
 
     chunk_samples = int(round(chunk_seconds * SAMPLE_RATE))
     overlap_samples = int(round(overlap_seconds * SAMPLE_RATE))
+    chunk_pad_samples = int(round(chunk_pad_seconds * SAMPLE_RATE))
     if chunk_samples < 1:
         raise ValueError("Chunk duration is shorter than one sample.")
     if overlap_samples * 2 > chunk_samples:
         raise ValueError("Overlap duration must not exceed half the chunk duration.")
-    return chunk_samples, overlap_samples
+    return chunk_samples, overlap_samples, chunk_pad_samples
 
 
 def validate_chunk_batch_size(chunk_batch_size):
@@ -141,8 +144,14 @@ def run_model(
     chunk_samples=None,
     overlap_samples=0,
     chunk_batch_size=1,
+    chunk_pad_samples=0,
 ):
-    """Run full-file or batched normalized overlap-add chunked inference."""
+    """Run full-file or batched normalized overlap-add chunked inference.
+
+    Each chunk is inferred with `chunk_pad_samples` of surrounding audio per
+    side, discarded from the output, because model output is wrong near the
+    edges of its input.
+    """
     chunk_batch_size = validate_chunk_batch_size(chunk_batch_size)
 
     total_samples = audio.shape[-1]
@@ -152,19 +161,29 @@ def run_model(
     output_sum = torch.zeros_like(audio, device="cpu")
     weight_sum = torch.zeros((1, 1, total_samples), dtype=audio.dtype)
     starts = chunk_starts(total_samples, chunk_samples, overlap_samples)
+    padded_chunk_samples = chunk_samples + 2 * chunk_pad_samples
 
     for batch_start in range(0, len(starts), chunk_batch_size):
         batch_starts = starts[batch_start : batch_start + chunk_batch_size]
         batch_chunks = []
         valid_lengths = []
+        batch_padded_starts = []
         for start in batch_starts:
             end = min(start + chunk_samples, total_samples)
             valid_samples = end - start
-            chunk = audio[..., start:end]
-            if valid_samples < chunk_samples:
-                chunk = F.pad(chunk, (0, chunk_samples - valid_samples))
+            # Real audio pads each side. The file's own edges take none:
+            # the first padded chunk begins at the file start and the last
+            # ends at the file end, extending further left instead —
+            # padding with silence measurably degrades the output near it.
+            padded_start = max(0, start - chunk_pad_samples)
+            if chunk_pad_samples and end == total_samples:
+                padded_start = max(0, total_samples - padded_chunk_samples)
+            chunk = audio[..., padded_start : padded_start + padded_chunk_samples]
+            if chunk.shape[-1] < padded_chunk_samples:
+                chunk = F.pad(chunk, (0, padded_chunk_samples - chunk.shape[-1]))
             batch_chunks.append(chunk)
             valid_lengths.append(valid_samples)
+            batch_padded_starts.append(padded_start)
 
         chunk_batch = torch.cat(batch_chunks, dim=0)
         batch_output = model(chunk_batch.to(device)).detach().to("cpu")
@@ -173,16 +192,19 @@ def run_model(
             raise RuntimeError(
                 "Chunk output must preserve the input batch and channel dimensions."
             )
-        if batch_output.shape[-1] < max(valid_lengths):
+        if batch_output.shape[-1] < padded_chunk_samples:
             raise RuntimeError(
                 "Chunk output is shorter than the corresponding input segment."
             )
 
-        for offset, (start, valid_samples) in enumerate(
-            zip(batch_starts, valid_lengths)
+        for offset, (start, valid_samples, padded_start) in enumerate(
+            zip(batch_starts, valid_lengths, batch_padded_starts)
         ):
             end = start + valid_samples
-            chunk_output = batch_output[offset : offset + 1, ..., :valid_samples]
+            valid_start = start - padded_start
+            chunk_output = batch_output[
+                offset : offset + 1, ..., valid_start : valid_start + valid_samples
+            ]
             index = batch_start + offset
             weights = crossfade_weights(
                 valid_samples,
@@ -207,6 +229,7 @@ def run_inference(
     chunk_seconds=None,
     overlap_seconds=1.0,
     chunk_batch_size=1,
+    chunk_pad_seconds=0.0,
 ):
     input_path = Path(input_wav)
     output_path = Path(output_wav)
@@ -215,8 +238,8 @@ def run_inference(
     if output_path.exists() and input_path.samefile(output_path):
         raise ValueError("Input and output paths must not reference the same file.")
 
-    chunk_samples, overlap_samples = resolve_chunking(
-        chunk_seconds, overlap_seconds
+    chunk_samples, overlap_samples, chunk_pad_samples = resolve_chunking(
+        chunk_seconds, overlap_seconds, chunk_pad_seconds
     )
     chunk_batch_size = validate_chunk_batch_size(chunk_batch_size)
     device = select_device(requested_device)
@@ -238,6 +261,7 @@ def run_inference(
             chunk_samples=chunk_samples,
             overlap_samples=overlap_samples,
             chunk_batch_size=chunk_batch_size,
+            chunk_pad_samples=chunk_pad_samples,
         )
     save_audio(output_path, output, sample_rate)
     return device
@@ -274,6 +298,12 @@ def main():
         default=1.0,
         help="Chunk crossfade duration (default: 1.0; at most half a chunk)",
     )
+    parser.add_argument(
+        "--chunk-pad-seconds",
+        type=float,
+        default=0.0,
+        help="Audio inferred beyond each chunk edge and discarded (default: 0; 1.0 gives clean seams)",
+    )
     parser.add_argument(
         "--chunk-batch-size",
         type=int,
@@ -290,6 +320,7 @@ def main():
         args.chunk_seconds,
         args.overlap_seconds,
         args.chunk_batch_size,
+        args.chunk_pad_seconds,
     )
     print(f"Inference completed on {device}: {args.out_wav}")
 

tests/test_inference.py

@@ -35,6 +35,29 @@ def forward(self, audio):
         return audio[..., :-1]
 
 
+class DegradedEdgesModel(IdentityModel):
+    """Identity, except both edges of every input come back wrong.
+
+    This mimics the real model, whose output within ~0.5 s of either edge of
+    its input differs from a well-placed prediction — mildly at the start,
+    strongly at the end. The marker is infinity so nothing wrong can hide:
+    with enough chunk padding it is sliced away before the crossfade
+    arithmetic, and with none it floods through the crossfades.
+    """
+
+    MARKER = float("inf")
+
+    def __init__(self, degraded_samples):
+        super().__init__()
+        self.degraded_samples = degraded_samples
+
+    def forward(self, audio):
+        audio = super().forward(audio).clone()
+        audio[..., : self.degraded_samples] = self.MARKER
+        audio[..., -self.degraded_samples :] = self.MARKER
+        return audio
+
+
 class DeviceSelectionTests(unittest.TestCase):
     @mock.patch.object(torch.backends.mps, "is_available", return_value=True)
     @mock.patch.object(torch.cuda, "is_available", return_value=True)
@@ -188,6 +211,135 @@ def test_chunked_cpu_inference_preserves_samples_and_limits_chunk_size(self):
                 hashlib.sha256(input_path.read_bytes()).hexdigest(), input_digest
             )
 
+    def test_chunk_padding_discards_degraded_input_edges(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            input_path = root / "input.wav"
+            output_path = root / "output.wav"
+            checkpoint_path = root / "checkpoint.bin"
+            checkpoint_path.touch()
+
+            sample_count = int(inference.SAMPLE_RATE * 1.8)
+            samples = np.linspace(-0.5, 0.5, sample_count, dtype=np.float32)
+            stereo = np.stack((samples, -samples), axis=1)
+            sf.write(input_path, stereo, inference.SAMPLE_RATE, subtype="FLOAT")
+
+            # 0.5 s chunks, 0.1 s overlap. The model degrades 0.05 s at
+            # each edge of its input, and chunk padding of the same 0.05 s
+            # discards exactly those samples.
+            degraded_seconds = 0.05
+            degraded_samples = int(degraded_seconds * inference.SAMPLE_RATE)
+            model = DegradedEdgesModel(
+                degraded_samples=degraded_samples
+            )
+            with mock.patch.object(
+                inference.look2hear.models.BaseModel,
+                "from_pretrain",
+                return_value=model,
+            ):
+                inference.run_inference(
+                    input_path,
+                    output_path,
+                    checkpoint_path,
+                    requested_device="cpu",
+                    chunk_seconds=0.5,
+                    overlap_seconds=0.1,
+                    chunk_pad_seconds=degraded_seconds,
+                )
+
+            output, _ = sf.read(output_path, dtype="float32", always_2d=True)
+            # Padding is real neighboring audio, so the first chunk starts
+            # at the file start and its degraded head is kept because there
+            # is nothing before the file to pad with. The check therefore
+            # starts after the file's first degraded_seconds. The same
+            # symmetry applies at the end of the file.
+            np.testing.assert_allclose(
+                output[degraded_samples:-degraded_samples],
+                stereo[degraded_samples:-degraded_samples],
+                rtol=0,
+                atol=1e-6,
+            )
+
+    def test_unpadded_chunks_leak_degraded_edges(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            input_path = root / "input.wav"
+            output_path = root / "output.wav"
+            checkpoint_path = root / "checkpoint.bin"
+            checkpoint_path.touch()
+
+            sample_count = int(inference.SAMPLE_RATE * 1.8)
+            samples = np.linspace(-0.5, 0.5, sample_count, dtype=np.float32)
+            stereo = np.stack((samples, -samples), axis=1)
+            sf.write(input_path, stereo, inference.SAMPLE_RATE, subtype="FLOAT")
+
+            # The same scenario with chunk padding off: the degraded edges
+            # reach the output through the overlap crossfades.
+            model = DegradedEdgesModel(
+                degraded_samples=int(0.05 * inference.SAMPLE_RATE)
+            )
+            with mock.patch.object(
+                inference.look2hear.models.BaseModel,
+                "from_pretrain",
+                return_value=model,
+            ):
+                inference.run_inference(
+                    input_path,
+                    output_path,
+                    checkpoint_path,
+                    requested_device="cpu",
+                    chunk_seconds=0.5,
+                    overlap_seconds=0.1,
+                    chunk_pad_seconds=0.0,
+                )
+
+            output, _ = sf.read(output_path, dtype="float32", always_2d=True)
+            self.assertFalse(np.isfinite(output).all())
+
+    def test_chunk_padding_wider_than_the_hop_reproduces_the_input(self):
+        """
+        Regression test
+
+        `padded_start` was `start - chunk_pad_samples` for every chunk after
+        the first, which goes negative once the padding exceeds the hop, and
+        makes slicing index from the end of the file
+
+        The test reproduces the bug, it needs padding wider than the hop,
+        so the file has to hold at least two chunks: 0.5 s chunks with
+        a 0.1 s overlap leave a 0.4 s hop, and 0.5 s of padding exceeds it.
+        The second chunk starts at 0.4 s, giving `padded_start = -0.1 s`.
+        """
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            input_path = root / "input.wav"
+            output_path = root / "output.wav"
+            checkpoint_path = root / "checkpoint.bin"
+            checkpoint_path.touch()
+
+            sample_count = int(inference.SAMPLE_RATE * 1.8)
+            samples = np.linspace(-0.5, 0.5, sample_count, dtype=np.float32)
+            stereo = np.stack((samples, -samples), axis=1)
+            sf.write(input_path, stereo, inference.SAMPLE_RATE, subtype="FLOAT")
+
+            model = IdentityModel()
+            with mock.patch.object(
+                inference.look2hear.models.BaseModel,
+                "from_pretrain",
+                return_value=model,
+            ):
+                inference.run_inference(
+                    input_path,
+                    output_path,
+                    checkpoint_path,
+                    requested_device="cpu",
+                    chunk_seconds=0.5,
+                    overlap_seconds=0.1,
+                    chunk_pad_seconds=0.5,
+                )
+
+            output, _ = sf.read(output_path, dtype="float32", always_2d=True)
+            np.testing.assert_allclose(output, stereo, rtol=0, atol=1e-6)
+
     def test_short_input_uses_original_full_file_path(self):
         audio = torch.zeros(1, 2, inference.SAMPLE_RATE // 4)
         model = IdentityModel()