From 99b31871fba0e86bad72137ac7601f5085d8e4a7 Mon Sep 17 00:00:00 2001 From: Zakaria Zbitou Date: Wed, 19 Aug 2026 02:21:00 +0200 Subject: [PATCH 1/3] Pad chunk inference and discard the model's degraded edges --- MACOS_ARM64.md | 11 ++++ README.md | 6 ++- inference.py | 53 ++++++++++++++++---- tests/test_inference.py | 108 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 166 insertions(+), 12 deletions(-) diff --git a/MACOS_ARM64.md b/MACOS_ARM64.md index dc33865..849adb7 100644 --- a/MACOS_ARM64.md +++ b/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 diff --git a/README.md b/README.md index 829dc5d..2cc6ece 100644 --- a/README.md +++ b/README.md @@ -102,6 +102,7 @@ python inference.py \ --device=auto \ --chunk-seconds=6 \ --overlap-seconds=1 \ + --chunk-pad-seconds=1 \ --chunk-batch-size=2 ``` @@ -109,7 +110,10 @@ python inference.py \ 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 diff --git a/inference.py b/inference.py index ea5f5c0..e50d476 100644 --- a/inference.py +++ b/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,6 +161,7 @@ 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] @@ -160,9 +170,16 @@ def run_model( 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 = start - chunk_pad_samples if start else 0 + 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) @@ -173,7 +190,7 @@ 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." ) @@ -182,7 +199,12 @@ def run_model( zip(batch_starts, valid_lengths) ): end = start + valid_samples - chunk_output = batch_output[offset : offset + 1, ..., :valid_samples] + valid_start = chunk_pad_samples if start else 0 + if chunk_pad_samples and end == total_samples: + valid_start = start - max(0, total_samples - padded_chunk_samples) + 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}") diff --git a/tests/test_inference.py b/tests/test_inference.py index b9778cf..147be58 100644 --- a/tests/test_inference.py +++ b/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,91 @@ 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_short_input_uses_original_full_file_path(self): audio = torch.zeros(1, 2, inference.SAMPLE_RATE // 4) model = IdentityModel() From ab97f7cf5d16bc2f58c069ee48621124bc024644 Mon Sep 17 00:00:00 2001 From: Zakaria Zbitou Date: Sun, 23 Aug 2026 18:56:21 +0200 Subject: [PATCH 2/3] Add a failing test for chunk padding wider than the hop This commit is expected to fail, so the bug can be confirmed before any fix lands. Check out this commit and run: python -m unittest discover -s tests The suite reports one failure: with 0.5 s chunks, a 0.1 s overlap and 0.5 s of padding, an identity model differs from its own input on 27.8% of samples. The fix follows in the next commit. --- tests/test_inference.py | 45 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/tests/test_inference.py b/tests/test_inference.py index 147be58..c6360b5 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -296,6 +296,51 @@ def test_unpadded_chunks_leak_degraded_edges(self): 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: chunk padding wider than the hop silenced whole chunks. + + `padded_start` was `start - chunk_pad_samples` for every chunk after + the first, which goes negative once the padding exceeds the hop. A + negative slice start is read from the end of the audio, so the slice + came back empty and `F.pad` filled the whole chunk with zeros, which + the crossfades then blended into the output. + + Reproducing it needs padding wider than the hop and a chunk after the + first, 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() From 10b4427640f3e6c1f4a1778ef0526bbc0691f586 Mon Sep 17 00:00:00 2001 From: Zakaria Zbitou Date: Sun, 23 Aug 2026 19:57:11 +0200 Subject: [PATCH 3/3] Clamp the padded chunk start and derive the crop offset from it --- inference.py | 12 ++++++------ tests/test_inference.py | 13 ++++++------- 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/inference.py b/inference.py index e50d476..356bda4 100644 --- a/inference.py +++ b/inference.py @@ -167,6 +167,7 @@ def run_model( 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 @@ -174,7 +175,7 @@ def run_model( # 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 = start - chunk_pad_samples if start else 0 + 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] @@ -182,6 +183,7 @@ def run_model( 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") @@ -195,13 +197,11 @@ def run_model( "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 - valid_start = chunk_pad_samples if start else 0 - if chunk_pad_samples and end == total_samples: - valid_start = start - max(0, total_samples - padded_chunk_samples) + valid_start = start - padded_start chunk_output = batch_output[ offset : offset + 1, ..., valid_start : valid_start + valid_samples ] diff --git a/tests/test_inference.py b/tests/test_inference.py index c6360b5..f6e1380 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -297,16 +297,15 @@ def test_unpadded_chunks_leak_degraded_edges(self): self.assertFalse(np.isfinite(output).all()) def test_chunk_padding_wider_than_the_hop_reproduces_the_input(self): - """Regression: chunk padding wider than the hop silenced whole chunks. + """ + Regression test `padded_start` was `start - chunk_pad_samples` for every chunk after - the first, which goes negative once the padding exceeds the hop. A - negative slice start is read from the end of the audio, so the slice - came back empty and `F.pad` filled the whole chunk with zeros, which - the crossfades then blended into the output. + the first, which goes negative once the padding exceeds the hop, and + makes slicing index from the end of the file - Reproducing it needs padding wider than the hop and a chunk after the - first, so the file has to hold at least two chunks: 0.5 s chunks with + 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`. """