From 0a91c5fe084bb4643410c52f72a12001871f8ec6 Mon Sep 17 00:00:00 2001 From: Djshinshin <129614060+Djshinshin@users.noreply.github.com> Date: Fri, 14 Aug 2026 14:46:38 +0800 Subject: [PATCH 1/2] Add chunked inference for long audio --- MACOS_ARM64.md | 46 ++++++++++++++- README.md | 20 ++++++- inference.py | 124 ++++++++++++++++++++++++++++++++++++++-- tests/test_inference.py | 95 ++++++++++++++++++++++++++++++ 4 files changed, 278 insertions(+), 7 deletions(-) diff --git a/MACOS_ARM64.md b/MACOS_ARM64.md index 6802116..bc6c417 100644 --- a/MACOS_ARM64.md +++ b/MACOS_ARM64.md @@ -26,6 +26,25 @@ 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 +``` + +Only one chunk is transferred to the selected accelerator at a time. 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 +76,27 @@ 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. + ## Device behavior - `auto` preserves CUDA as the first choice, then selects MPS, then CPU. @@ -80,7 +120,9 @@ 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. - 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. diff --git a/README.md b/README.md index f0a2181..c1ab71d 100644 --- a/README.md +++ b/README.md @@ -86,7 +86,25 @@ 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 +``` + +Chunked inference processes one segment on the selected device at a time and +uses normalized linear crossfades. The overlap must be no more than half the +chunk duration, and the output keeps the input frame count. 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. diff --git a/inference.py b/inference.py index 720d214..c8b270c 100644 --- a/inference.py +++ b/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,7 +83,97 @@ 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 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): + """Run full-file or normalized overlap-add chunked inference.""" + 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 index, start in enumerate(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)) + + chunk_output = model(chunk.to(device)).detach().to("cpu") + if chunk_output.ndim != 3 or chunk_output.shape[:2] != audio.shape[:2]: + raise RuntimeError( + "Chunk output must preserve the input batch and channel dimensions." + ) + if chunk_output.shape[-1] < valid_samples: + raise RuntimeError( + "Chunk output is shorter than the corresponding input segment." + ) + chunk_output = chunk_output[..., :valid_samples] + + 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, +): input_path = Path(input_wav) output_path = Path(output_wav) if input_path.resolve() == output_path.resolve(): @@ -90,10 +181,12 @@ def run_inference(input_wav, output_wav, checkpoint, requested_device="auto"): 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 + ) 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 +196,13 @@ 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, + ) save_audio(output_path, output, sample_rate) return device @@ -127,10 +226,27 @@ 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)", + ) 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, ) print(f"Inference completed on {device}: {args.out_wav}") diff --git a/tests/test_inference.py b/tests/test_inference.py index 7e0e64f..7857820 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -17,12 +17,22 @@ def __init__(self): super().__init__() self.anchor = torch.nn.Parameter(torch.zeros(())) self.input_device = None + self.input_devices = [] + self.input_lengths = [] def forward(self, audio): self.input_device = audio.device + self.input_devices.append(audio.device) + self.input_lengths.append(audio.shape[-1]) return audio +class ShortOutputModel(IdentityModel): + def forward(self, audio): + super().forward(audio) + return audio[..., :-1] + + 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) @@ -129,6 +139,91 @@ def test_cpu_inference_preserves_input_and_audio_shape(self): hashlib.sha256(input_path.read_bytes()).hexdigest(), input_digest ) + def test_chunked_cpu_inference_preserves_samples_and_limits_chunk_size(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") + input_digest = hashlib.sha256(input_path.read_bytes()).hexdigest() + + model = IdentityModel() + with mock.patch.object( + inference.look2hear.models.BaseModel, + "from_pretrain", + return_value=model, + ): + device = inference.run_inference( + input_path, + output_path, + checkpoint_path, + requested_device="cpu", + chunk_seconds=0.5, + overlap_seconds=0.1, + ) + + output, sample_rate = sf.read( + output_path, dtype="float32", always_2d=True + ) + self.assertEqual(device, torch.device("cpu")) + self.assertEqual(sample_rate, inference.SAMPLE_RATE) + self.assertEqual(output.shape, stereo.shape) + np.testing.assert_allclose(output, stereo, rtol=0, atol=1e-6) + self.assertGreater(len(model.input_lengths), 1) + self.assertEqual( + set(model.input_lengths), {int(inference.SAMPLE_RATE * 0.5)} + ) + self.assertEqual(set(model.input_devices), {torch.device("cpu")}) + self.assertEqual( + hashlib.sha256(input_path.read_bytes()).hexdigest(), input_digest + ) + + def test_short_input_uses_original_full_file_path(self): + audio = torch.zeros(1, 2, inference.SAMPLE_RATE // 4) + model = IdentityModel() + + output = inference.run_model( + model, + audio, + torch.device("cpu"), + chunk_samples=inference.SAMPLE_RATE, + overlap_samples=inference.SAMPLE_RATE // 4, + ) + + self.assertEqual(model.input_lengths, [audio.shape[-1]]) + self.assertTrue(torch.equal(output, audio)) + + def test_chunk_settings_reject_invalid_durations(self): + invalid_settings = ( + (0, 0, "positive finite"), + (float("inf"), 0, "positive finite"), + (1, -0.1, "non-negative finite"), + (1, float("nan"), "non-negative finite"), + (1, 0.6, "half the chunk"), + ) + for chunk_seconds, overlap_seconds, message in invalid_settings: + with self.subTest( + chunk_seconds=chunk_seconds, overlap_seconds=overlap_seconds + ), self.assertRaisesRegex(ValueError, message): + inference.resolve_chunking(chunk_seconds, overlap_seconds) + + def test_chunking_rejects_short_model_output(self): + audio = torch.zeros(1, 2, 100) + with self.assertRaisesRegex(RuntimeError, "shorter"): + inference.run_model( + ShortOutputModel(), + audio, + torch.device("cpu"), + chunk_samples=60, + overlap_samples=10, + ) + def test_refuses_to_overwrite_input(self): path = Path("same.wav") with self.assertRaisesRegex(ValueError, "must be different"): From 32f75791a82dfc5a2296773e55b083fd92806954 Mon Sep 17 00:00:00 2001 From: Djshinshin <129614060+Djshinshin@users.noreply.github.com> Date: Fri, 14 Aug 2026 15:08:29 +0800 Subject: [PATCH 2/2] Add batched chunk inference --- MACOS_ARM64.md | 34 +++++++++++++--- README.md | 15 ++++--- inference.py | 89 ++++++++++++++++++++++++++++++----------- tests/test_inference.py | 42 +++++++++++++++++++ 4 files changed, 145 insertions(+), 35 deletions(-) diff --git a/MACOS_ARM64.md b/MACOS_ARM64.md index bc6c417..dc33865 100644 --- a/MACOS_ARM64.md +++ b/MACOS_ARM64.md @@ -36,14 +36,17 @@ python inference.py \ --in_wav=long_input.wav \ --out_wav=long_output.wav \ --chunk-seconds=6 \ - --overlap-seconds=1 + --overlap-seconds=1 \ + --chunk-batch-size=2 ``` -Only one chunk is transferred to the selected accelerator at a time. 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. +`--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 @@ -97,6 +100,22 @@ difference `0.7703997492790222`, RMSE `0.027596083317049432`, and correlation 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. @@ -124,5 +143,8 @@ checks do not establish that the two versions are perceptually equivalent. 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. diff --git a/README.md b/README.md index c1ab71d..8995fbf 100644 --- a/README.md +++ b/README.md @@ -97,14 +97,17 @@ python inference.py \ --out_wav=long_output.wav \ --device=auto \ --chunk-seconds=6 \ - --overlap-seconds=1 + --overlap-seconds=1 \ + --chunk-batch-size=2 ``` -Chunked inference processes one segment on the selected device at a time and -uses normalized linear crossfades. The overlap must be no more than half the -chunk duration, and the output keeps the input frame count. 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 +`--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. diff --git a/inference.py b/inference.py index c8b270c..ea5f5c0 100644 --- a/inference.py +++ b/inference.py @@ -101,6 +101,17 @@ def resolve_chunking(chunk_seconds, overlap_seconds): 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 @@ -123,8 +134,17 @@ def crossfade_weights(length, overlap_samples, fade_in, fade_out, dtype): return weights.view(1, 1, -1) -def run_model(model, audio, device, chunk_samples=None, overlap_samples=0): - """Run full-file or normalized overlap-add chunked inference.""" +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") @@ -133,33 +153,46 @@ def run_model(model, audio, device, chunk_samples=None, overlap_samples=0): weight_sum = torch.zeros((1, 1, total_samples), dtype=audio.dtype) starts = chunk_starts(total_samples, chunk_samples, overlap_samples) - for index, start in enumerate(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)) - - chunk_output = model(chunk.to(device)).detach().to("cpu") - if chunk_output.ndim != 3 or chunk_output.shape[:2] != audio.shape[:2]: + 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 chunk_output.shape[-1] < valid_samples: + if batch_output.shape[-1] < max(valid_lengths): raise RuntimeError( "Chunk output is shorter than the corresponding input segment." ) - chunk_output = chunk_output[..., :valid_samples] - - 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 + + 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.") @@ -173,6 +206,7 @@ def run_inference( requested_device="auto", chunk_seconds=None, overlap_seconds=1.0, + chunk_batch_size=1, ): input_path = Path(input_wav) output_path = Path(output_wav) @@ -184,6 +218,7 @@ def run_inference( 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) @@ -202,6 +237,7 @@ def run_inference( device, chunk_samples=chunk_samples, overlap_samples=overlap_samples, + chunk_batch_size=chunk_batch_size, ) save_audio(output_path, output, sample_rate) return device @@ -238,6 +274,12 @@ def main(): 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( @@ -247,6 +289,7 @@ def main(): args.device, args.chunk_seconds, args.overlap_seconds, + args.chunk_batch_size, ) print(f"Inference completed on {device}: {args.out_wav}") diff --git a/tests/test_inference.py b/tests/test_inference.py index 7857820..b9778cf 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -19,11 +19,13 @@ def __init__(self): self.input_device = None self.input_devices = [] self.input_lengths = [] + self.input_batch_sizes = [] def forward(self, audio): self.input_device = audio.device self.input_devices.append(audio.device) self.input_lengths.append(audio.shape[-1]) + self.input_batch_sizes.append(audio.shape[0]) return audio @@ -166,6 +168,7 @@ def test_chunked_cpu_inference_preserves_samples_and_limits_chunk_size(self): requested_device="cpu", chunk_seconds=0.5, overlap_seconds=0.1, + chunk_batch_size=2, ) output, sample_rate = sf.read( @@ -179,6 +182,7 @@ def test_chunked_cpu_inference_preserves_samples_and_limits_chunk_size(self): self.assertEqual( set(model.input_lengths), {int(inference.SAMPLE_RATE * 0.5)} ) + self.assertEqual(model.input_batch_sizes, [2, 2, 1]) self.assertEqual(set(model.input_devices), {torch.device("cpu")}) self.assertEqual( hashlib.sha256(input_path.read_bytes()).hexdigest(), input_digest @@ -224,6 +228,44 @@ def test_chunking_rejects_short_model_output(self): overlap_samples=10, ) + def test_chunking_rejects_invalid_batch_size(self): + audio = torch.zeros(1, 2, 100) + for batch_size in (0, -1, 1.5, True): + with self.subTest(batch_size=batch_size), self.assertRaisesRegex( + ValueError, "positive integer" + ): + inference.run_model( + IdentityModel(), + audio, + torch.device("cpu"), + chunk_samples=60, + overlap_samples=10, + chunk_batch_size=batch_size, + ) + + def test_invalid_batch_size_fails_before_checkpoint_download(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + input_path = root / "input.wav" + sf.write( + input_path, + np.zeros((100, 2), dtype=np.float32), + inference.SAMPLE_RATE, + ) + + with mock.patch.object(inference, "hf_hub_download") as download: + with self.assertRaisesRegex(ValueError, "positive integer"): + inference.run_inference( + input_path, + root / "output.wav", + inference.OFFICIAL_CHECKPOINT, + "cpu", + chunk_seconds=1, + overlap_seconds=0, + chunk_batch_size=0, + ) + download.assert_not_called() + def test_refuses_to_overwrite_input(self): path = Path("same.wav") with self.assertRaisesRegex(ValueError, "must be different"):