[Djshinshin]

Summary

• Add opt-in --chunk-seconds and --overlap-seconds options for bounded-memory long-audio inference.
• Add configurable --chunk-batch-size so a fixed number of chunks can share one model forward pass.
• Keep the batch-size default at 1 for minimum accelerator memory; larger values are an explicit throughput/memory tradeoff.
• Preserve full-file inference as the default and use the original path automatically when the input is no longer than one chunk.
• Keep the complete input on CPU and transfer only one fixed-size chunk batch at a time to CUDA, MPS, or CPU.
• Pad only the final short chunk, crop it back to its valid length, and combine segments with normalized linear crossfades.
• Preserve the exact input frame count and move every batch output back to CPU before overlap-add and writing.

Chunk and batch behavior

• Chunk duration must be positive and finite.
• Overlap may be zero but must be finite, non-negative, and no more than half the chunk duration; this prevents more than two chunks from overlapping at once.
• Chunk batch size must be a positive integer and is validated before checkpoint download.
• A short or structurally incompatible model output fails clearly instead of silently leaving uncovered or truncated samples.
• Chunking remains explicit because reducing model context can change the result relative to full-file inference.
• Larger batches may improve throughput on some accelerators, but use more memory and are not assumed to be faster on every backend.

Tests

• python -m unittest discover -s tests -v: 20/20 passed.
• Added CPU tests for identity-preserving batched overlap-add, exact shape and samples, [2, 2, 1] batch grouping, bounded transfer lengths, final-chunk padding/crop, short-input full-file behavior, invalid settings, early batch-size validation, and short model output rejection.
• python -m py_compile inference.py tests/test_inference.py: passed.
• python -m pip check: passed.
• git diff --check: passed.

Apple Silicon validation

MPS validation used an Apple M3, macOS 27.0, native arm64 Python 3.10, and PyTorch/TorchAudio 2.11.0. The repository's public six-second fixture was concatenated with itself locally to create a 12-second public test input; the checked-in fixture was not modified.

The earlier full-file versus sequential-chunk reference remains:

Mode	End-to-end time	Peak memory footprint
Full file	10.22 s	10.36 GiB
6 s chunks, 1 s overlap	12.59 s	5.39 GiB

A new local-checkpoint run compared fixed chunk batch sizes:

Chunk batch size	End-to-end time	Peak memory footprint
1	11.76 s	4.65 GiB
2	12.28 s	10.48 GiB

Both batched outputs were float32 PCM, stereo, 44.1 kHz, exactly 12 seconds, and finite. Batch sizes 1 and 2 had maximum absolute difference 1.04308128356934e-06, RMSE 3.84001256540259e-08, and correlation 0.999999999999966. The public input SHA-256 remained unchanged.

On this MPS machine, batch size 2 did not improve runtime and increased peak memory. The implementation therefore exposes batching without claiming a universal speedup. CUDA may behave differently, but CUDA hardware was not available for local validation.

Limitations and scope

• Crossfades reduce boundary discontinuities but do not prove perceptual equivalence or replace seam and musical-consistency listening tests.
• CUDA was not hardware-tested on the Apple Silicon machine; the existing device selection and model/device placement behavior is preserved.
• This PR does not add high-bitrate MP3 degradation, training changes, resampling, or model architecture changes.

Copyright and safety

Validation used only the repository's public fixture. No private music, checkpoint, WAV input/output, virtual environment, cache, or generated artifact is included in the commit.

Refs #37

[JusperLee]

Could the chunk-level inference process be switched to batch mode? This would significantly speed up inference.

[Djshinshin]

Thank you for the suggestion. I have implemented configurable chunk batching so --chunk-batch-size=N groups a fixed number of padded chunks into each model forward pass, then moves the batch output back to CPU for the existing normalized overlap-add. Memory remains bounded by the selected batch size, and 1 remains the default for the minimum-memory path.

The local CPU tests confirm batch grouping of [2, 2, 1], exact output shape and samples, final-chunk padding/cropping, and early validation before checkpoint download. The full suite now passes 20/20.

I also tested batch sizes 1 and 2 end to end on the public 12-second fixture using MPS. Both outputs were float32 stereo at 44.1 kHz, exactly 12 seconds and finite, with maximum absolute difference 1.04308128356934e-06, RMSE 3.84001256540259e-08, and correlation 0.999999999999966. On this Apple M3, batch size 2 was slightly slower (12.28 s versus 11.76 s) and used more peak memory (10.48 GiB versus 4.65 GiB), so I documented batching as a backend-specific throughput/memory tradeoff rather than claiming a universal speedup. CUDA may benefit differently, but I do not have CUDA hardware for a local measurement.

If a higher default is preferable for the project's CUDA-oriented baseline, I would be happy to adjust it following your guidance. Thank you again.