Add batched chunk inference for long audio

MACOS_ARM64.md

@@ -26,6 +26,28 @@ python inference.py \
   --out_wav=apollo_mps_smoke.wav
 ```
 
+Full-file inference is the default. To bound accelerator memory for longer
+audio, opt in to chunking and an overlap crossfade:
+
+```bash
+python inference.py \
+  --device=mps \
+  --checkpoint=models/pytorch_model.bin \
+  --in_wav=long_input.wav \
+  --out_wav=long_output.wav \
+  --chunk-seconds=6 \
+  --overlap-seconds=1 \
+  --chunk-batch-size=2
+```
+
+`--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
+throughput. The last chunk is padded for inference and cropped back to its
+original size; normalized linear crossfades preserve the exact input frame
+count. The overlap may be zero but must not exceed half the chunk duration.
+Inputs no longer than the requested chunk use the original full-file path.
+
 If `--checkpoint` is omitted, inference downloads the official
 `JusperLee/Apollo` `pytorch_model.bin` from Hugging Face. Custom checkpoints
 must be existing local files; arbitrary remote repositories are rejected
@@ -57,6 +79,43 @@ correlation of `0.999999991468`. The input SHA-256 remained unchanged:
 3c9a053913c3016b493ddc0b92e21a4e682e8fb5f872dafc945154155bddf771
 ```
 
+## Chunked validation reference
+
+A separate MPS check concatenated the same public fixture twice to create a
+12-second input, then compared full-file inference with 6-second chunks and a
+1-second overlap. The derived test input SHA-256 was
+`3e4a91f95975054bc092637176bb60bcb4bbd1469b8f8e92c14eb6e4ad0431c2`;
+the repository fixture itself was not modified.
+
+| 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 |
+
+The chunked run reduced the measured peak memory footprint by 48.0%. Both
+outputs were float32 PCM, stereo, 44.1 kHz, exactly 12 seconds long, and finite.
+Compared with full-file output, the chunked output had maximum absolute
+difference `0.7703997492790222`, RMSE `0.027596083317049432`, and correlation
+`0.9835697382563144`. At the two tested chunk starts, output sample differences
+from the full-file result were at most `3.3527612686157227e-08`. These objective
+checks do not establish that the two versions are perceptually equivalent.
+
+The same public 12-second input was also used to compare chunk batch sizes on
+MPS. In this Apple M3 run, batching did not improve MPS performance and should
+therefore remain an explicit memory/throughput choice rather than a universal
+speed claim:
+
+| 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 outputs were float32 PCM, stereo, 44.1 kHz, exactly 12 seconds long, and
+finite. Batch sizes 1 and 2 differed by at most `1.04308128356934e-06`, with
+RMSE `3.84001256540259e-08` and correlation `0.999999999999966`. CUDA may have
+different throughput characteristics, but it was not available for local
+hardware validation.
+
 ## Device behavior
 
 - `auto` preserves CUDA as the first choice, then selects MPS, then CPU.
@@ -80,7 +139,12 @@ correlation of `0.999999991468`. The input SHA-256 remained unchanged:
   machine.
 - Training and dataset evaluation still contain CUDA-specific configuration;
   this change only makes the public inference entry point device-agnostic.
-- There is no long-audio chunking or crossfade in this change. Long inputs may
-  require substantially more memory.
+- Chunking reduces accelerator memory demand but also limits model context and
+  can produce a different result from full-file inference. Crossfades reduce
+  boundary discontinuities but do not replace listening tests for seam or
+  musical-consistency artifacts.
+- Larger chunk batches can reduce the number of model forward passes but use
+  more accelerator memory. They are not guaranteed to improve every backend;
+  the tested MPS configuration was slower with batch size 2.
 - A successful technical smoke test is not a listening-quality review and does
   not prove that information missing from a lossy source was truly recovered.

README.md

@@ -90,7 +90,28 @@ python inference.py \
 default downloads the official `JusperLee/Apollo` checkpoint from Hugging Face.
 For a custom checkpoint, `--checkpoint` accepts only an existing local file;
 arbitrary remote repositories are rejected because the upstream checkpoint
-loader deserializes the file with PyTorch. See
+loader deserializes the file with PyTorch.
+
+Full-file inference remains the default. For long audio, enable bounded-memory
+inference explicitly:
+
+```bash
+python inference.py \
+  --in_wav=long_input.wav \
+  --out_wav=long_output.wav \
+  --device=auto \
+  --chunk-seconds=6 \
+  --overlap-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
+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
 limitations.
 

inference.py

@@ -5,6 +5,7 @@
 import numpy as np
 import soundfile as sf
 import torch
+import torch.nn.functional as F
 
 import look2hear.models
 
@@ -82,18 +83,145 @@ def save_audio(file_path, audio, sample_rate):
     sf.write(output_path, audio, sample_rate, subtype="FLOAT")
 
 
-def run_inference(input_wav, output_wav, checkpoint, requested_device="auto"):
+def resolve_chunking(chunk_seconds, overlap_seconds):
+    """Convert optional chunk settings to samples and validate crossfades."""
+    if chunk_seconds is None:
+        return None, 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.")
+
+    chunk_samples = int(round(chunk_seconds * SAMPLE_RATE))
+    overlap_samples = int(round(overlap_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
+
+
+def validate_chunk_batch_size(chunk_batch_size):
+    """Require a fixed positive number of chunks per model forward pass."""
+    if (
+        isinstance(chunk_batch_size, bool)
+        or not isinstance(chunk_batch_size, int)
+        or chunk_batch_size < 1
+    ):
+        raise ValueError("Chunk batch size must be a positive integer.")
+    return chunk_batch_size
+
+
+def chunk_starts(total_samples, chunk_samples, overlap_samples):
+    """Return starts that cover the signal without a redundant final chunk."""
+    hop_samples = chunk_samples - overlap_samples
+    starts = [0]
+    while starts[-1] + chunk_samples < total_samples:
+        starts.append(starts[-1] + hop_samples)
+    return starts
+
+
+def crossfade_weights(length, overlap_samples, fade_in, fade_out, dtype):
+    """Create linear edge weights for normalized overlap-add."""
+    weights = torch.ones(length, dtype=dtype)
+    fade_samples = min(overlap_samples, length)
+    if fade_samples:
+        ramp = torch.linspace(0.0, 1.0, fade_samples, dtype=dtype)
+        if fade_in:
+            weights[:fade_samples] = ramp
+        if fade_out:
+            weights[-fade_samples:] = torch.flip(ramp, dims=(0,))
+    return weights.view(1, 1, -1)
+
+
+def run_model(
+    model,
+    audio,
+    device,
+    chunk_samples=None,
+    overlap_samples=0,
+    chunk_batch_size=1,
+):
+    """Run full-file or batched normalized overlap-add chunked inference."""
+    chunk_batch_size = validate_chunk_batch_size(chunk_batch_size)
+
+    total_samples = audio.shape[-1]
+    if chunk_samples is None or total_samples <= chunk_samples:
+        return model(audio.to(device)).detach().to("cpu")
+
+    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)
+
+    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 = []
+        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))
+            batch_chunks.append(chunk)
+            valid_lengths.append(valid_samples)
+
+        chunk_batch = torch.cat(batch_chunks, dim=0)
+        batch_output = model(chunk_batch.to(device)).detach().to("cpu")
+        expected_shape = (len(batch_starts), audio.shape[1])
+        if batch_output.ndim != 3 or batch_output.shape[:2] != expected_shape:
+            raise RuntimeError(
+                "Chunk output must preserve the input batch and channel dimensions."
+            )
+        if batch_output.shape[-1] < max(valid_lengths):
+            raise RuntimeError(
+                "Chunk output is shorter than the corresponding input segment."
+            )
+
+        for offset, (start, valid_samples) in enumerate(
+            zip(batch_starts, valid_lengths)
+        ):
+            end = start + valid_samples
+            chunk_output = batch_output[offset : offset + 1, ..., :valid_samples]
+            index = batch_start + offset
+            weights = crossfade_weights(
+                valid_samples,
+                overlap_samples,
+                fade_in=index > 0,
+                fade_out=index < len(starts) - 1,
+                dtype=audio.dtype,
+            )
+            output_sum[..., start:end] += chunk_output * weights
+            weight_sum[..., start:end] += weights
+
+    if torch.any(weight_sum <= 0):
+        raise RuntimeError("Chunk overlap-add left uncovered output samples.")
+    return output_sum / weight_sum
+
+
+def run_inference(
+    input_wav,
+    output_wav,
+    checkpoint,
+    requested_device="auto",
+    chunk_seconds=None,
+    overlap_seconds=1.0,
+    chunk_batch_size=1,
+):
     input_path = Path(input_wav)
     output_path = Path(output_wav)
     if input_path.resolve() == output_path.resolve():
         raise ValueError("Input and output paths must be different.")
     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_batch_size = validate_chunk_batch_size(chunk_batch_size)
     device = select_device(requested_device)
     test_data, sample_rate = load_audio(input_path)
     checkpoint_path = resolve_checkpoint(checkpoint)
-    test_data = test_data.to(device)
 
     model = look2hear.models.BaseModel.from_pretrain(
         str(checkpoint_path),
@@ -103,7 +231,14 @@ def run_inference(input_wav, output_wav, checkpoint, requested_device="auto"):
         layer=6,
     ).to(device).eval()
     with torch.inference_mode():
-        output = model(test_data)
+        output = run_model(
+            model,
+            test_data,
+            device,
+            chunk_samples=chunk_samples,
+            overlap_samples=overlap_samples,
+            chunk_batch_size=chunk_batch_size,
+        )
     save_audio(output_path, output, sample_rate)
     return device
 
@@ -127,10 +262,34 @@ def main():
         default="auto",
         help="Inference device (auto prefers CUDA, then MPS, then CPU)",
     )
+    parser.add_argument(
+        "--chunk-seconds",
+        type=float,
+        default=None,
+        help="Optional chunk duration for long audio; disabled by default",
+    )
+    parser.add_argument(
+        "--overlap-seconds",
+        type=float,
+        default=1.0,
+        help="Chunk crossfade duration (default: 1.0; at most half a chunk)",
+    )
+    parser.add_argument(
+        "--chunk-batch-size",
+        type=int,
+        default=1,
+        help="Chunks per model forward pass (default: 1; higher uses more memory)",
+    )
     args = parser.parse_args()
 
     device = run_inference(
-        args.in_wav, args.out_wav, args.checkpoint, args.device
+        args.in_wav,
+        args.out_wav,
+        args.checkpoint,
+        args.device,
+        args.chunk_seconds,
+        args.overlap_seconds,
+        args.chunk_batch_size,
     )
     print(f"Inference completed on {device}: {args.out_wav}")