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..356bda4 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,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}") diff --git a/tests/test_inference.py b/tests/test_inference.py index b9778cf..f6e1380 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,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()