[zbitouzakaria]

What

A new --chunk-pad-seconds option (default 0) for chunked inference: each chunk's model input is extended with that much surrounding audio on each side, and the model's output for the extension is discarded. Chunk layout, overlap, crossfade, and the number of forward passes are unchanged; each pass grows from chunk to chunk + 2 × pad samples. --chunk-pad-seconds 0 reproduces the current output bit-exactly.

Why

The model's output is wrong near the edges of its input. Restoring the same audio twice with different chunk positions and subtracting the two outputs: they match bit for bit except within ~0.5 s of an input edge, where the difference reaches the level of the music itself at the trailing edge (about −25 dB at the leading edge). The extent is ~0.5 s whether the chunk is 4, 6, or 10 s long.

Chunked inference crossfades exactly those regions into every seam. On the repository's 6 s fixture tiled to 20 s, with 6 s chunks and a 1 s overlap, SNR against a single unchunked pass collapses at every chunk boundary:

The remaining dips common to both traces are numerical differences between the CUDA reference and these MPS renders; padding 1.0 and 2.0 agree with each other past 120 dB at those spots. These still need to be investigated.

How

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 one ends at the file end, extending further left instead. Both file edges therefore come out exactly as an unchunked pass produces them.

1s is the recommended value: it clears the measured ~0.5 s edge extent with margin, and costs one third more samples per forward pass at 6 s chunks.

Tests

• 2 tests: test chunk padding effectively discard the model's degraded input edges (an identity model returning inf over the outer 0.05 s of every input, run through run_inference with --chunk-pad-seconds 0.05), and that when --chunk-pad-seconds 0 with the same model inf values make it to the output.
• Local check: --chunk-pad-seconds 0 compared sample-for-sample against main's output on the released checkpoint (20 s render): identical.
• python -m unittest discover -s tests -v: 22/22 passed.
• python -m py_compile inference.py tests/test_inference.py: passed.
• python -m pip check: passed.
• git diff --check: passed.

[JusperLee]

[P1] Handle padding that reaches before the file start

--chunk-pad-seconds is accepted without an upper bound, so chunk_pad_samples can exceed the hop size. For a non-first chunk where start < chunk_pad_samples, padded_start becomes negative. Tensor slicing then interprets that index relative to the end of the audio, and the remainder is zero-padded, while the crop below still assumes a full chunk_pad_samples of left context.

For example, with total=300, chunk=100, overlap=20, and pad=90, the second chunk starts at 80 and gets padded_start=-10. An identity model no longer reproduces the source around that chunk, even though all CLI arguments are valid.

Please derive the crop offset from the actual padded start after applying the edge handling, e.g. valid_start = start - padded_start, and add a regression test where pad > chunk - overlap.

@zbitouzakaria

[zbitouzakaria]

Thanks for the review @JusperLee ! I added the regression test (that does fail at the moment).

The fix I have in mind:

Before inference

Replace
chunk = audio[..., padded_start : padded_start + padded_chunk_samples]
with
chunk = audio[..., max(0, padded_start) : padded_start + padded_chunk_samples]

And save the padded_start values in an array.
The 0 right padding always makes sure the padded chunks have the same length

After inference

Derive the crop offset using valid_start = start - max(0, padded_start)

Working on the fix, let me know if you have any issues with this.

[zbitouzakaria]

The fix is up.

python -m unittest discover -s tests -v: 23/23 passed.
python -m py_compile inference.py tests/test_inference.py: passed.
python -m pip check: passed.
git diff --check: passed.