diff --git a/.github/workflows/cpu-unit-tests.yml b/.github/workflows/cpu-unit-tests.yml new file mode 100644 index 0000000..4248fbb --- /dev/null +++ b/.github/workflows/cpu-unit-tests.yml @@ -0,0 +1,25 @@ +name: CPU unit tests + +on: + pull_request: + push: + branches: [main] + +jobs: + unit-tests: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.10" + cache: pip + - name: Install CPU test dependencies + run: | + python -m pip install --upgrade pip + python -m pip install torch==2.11.0 --index-url https://download.pytorch.org/whl/cpu + python -m pip install numpy==1.26.4 huggingface-hub==1.27.0 soundfile==0.14.0 + - name: Compile and run unit tests + run: | + python -m py_compile inference.py tests/test_inference.py + python -m unittest discover -s tests -v diff --git a/.gitignore b/.gitignore index e3b2e50..b828eb1 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,7 @@ *.pyc -*.DS_Store \ No newline at end of file +*.DS_Store +.venv/ +/models/ +/smoke_test_output/ +/output.wav +/apollo_*_smoke.wav diff --git a/MACOS_ARM64.md b/MACOS_ARM64.md new file mode 100644 index 0000000..6802116 --- /dev/null +++ b/MACOS_ARM64.md @@ -0,0 +1,86 @@ +# Apple Silicon macOS inference + +The upstream `look2hear.yml` records a Linux CUDA baseline with +`torch==2.0.0+cu118` and `torchaudio==2.0.1+cu118`; it cannot be installed +unchanged on macOS arm64. The MPS smoke test for this change used native arm64 +Python 3.10, PyTorch 2.11.0, and TorchAudio 2.11.0. The inference entry point +uses SoundFile for WAV I/O so that it works with both the older baseline and +newer TorchAudio releases whose I/O depends on TorchCodec. + +Create and activate an isolated environment, then install the macOS runtime +requirements: + +```bash +python3.10 -m venv .venv +source .venv/bin/activate +python -m pip install -r requirements-macos-arm64.txt +``` + +Run the repository's public six-second test WAV on MPS: + +```bash +python inference.py \ + --device=mps \ + --checkpoint=models/pytorch_model.bin \ + --in_wav=asserts/input_wav.wav \ + --out_wav=apollo_mps_smoke.wav +``` + +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 +because the upstream loader uses PyTorch deserialization. Only use local +checkpoints from sources you trust. The official checkpoint used for local +validation had this SHA-256 digest (recorded for reproducibility, not enforced +as a gate so that future official updates remain possible): + +```text +99d9af7f1ff20e63c393035513a655392818d66b4d7fc23d658175c1f15e8d76 +``` + +## Local validation reference + +The public six-second WAV was tested on an Apple M3 with 24 GB unified memory +and macOS 27.0. These are reference measurements, not performance guarantees: + +| Device | End-to-end time | Maximum RSS | Peak memory footprint | +| --- | ---: | ---: | ---: | +| MPS | 5.78 s | 412.0 MiB | 4.83 GiB | +| CPU | 316.46 s | 3.76 GiB | 3.27 GiB | + +Both outputs were stereo, 44.1 kHz, six seconds long, and contained no NaN or +infinite values. Comparing the float32 CPU and MPS WAV outputs gave a maximum +absolute difference of `0.000249214470387`, RMSE of `2.02574076096e-05`, and +correlation of `0.999999991468`. The input SHA-256 remained unchanged: + +```text +3c9a053913c3016b493ddc0b92e21a4e682e8fb5f872dafc945154155bddf771 +``` + +## Device behavior + +- `auto` preserves CUDA as the first choice, then selects MPS, then CPU. +- `cuda`, `mps`, and `cpu` explicitly select one backend and fail clearly if + the requested accelerator is unavailable. +- The model and input are moved to the same device. Output is detached and + moved to CPU before it is written. +- Apollo validates sample rate, basic shape, and finite samples on CPU before + downloading the checkpoint or moving the input to an accelerator. +- Apollo expects 44.1 kHz audio. Input and output paths must not identify the + same file, including through a hard link. + +## Known limitations + +- MPS validation covered the public stereo, 44.1 kHz, six-second test WAV on + PyTorch 2.11.0. MPS operator support can vary with PyTorch and macOS versions. +- PyTorch 2.11.0 emitted non-fatal STFT/iSTFT output-resize deprecation warnings + during the MPS run. +- CUDA remains the first `auto` choice and follows the original inference + behavior, but CUDA was not hardware-tested on the Apple Silicon validation + 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. +- 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 371b9f4..f0a2181 100644 --- a/README.md +++ b/README.md @@ -76,9 +76,20 @@ python train.py --conf_dir=configs/apollo.yml To evaluate the Apollo model, run the following command: ```bash -python inference.py --in_wav=assets/input.wav --out_wav=assets/output.wav +python inference.py \ + --in_wav=asserts/input_wav.wav \ + --out_wav=output.wav \ + --device=auto ``` +`--device=auto` prefers CUDA, then Apple Silicon MPS, and finally CPU. The +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 +[Apple Silicon macOS notes](MACOS_ARM64.md) for a tested MPS setup and current +limitations. + ## 📊 Results *Here, you can include a brief overview of the performance metrics or results that Apollo achieves using different bitrates* diff --git a/inference.py b/inference.py index 872c77c..720d214 100644 --- a/inference.py +++ b/inference.py @@ -1,30 +1,139 @@ -import os -import torch -import torchaudio import argparse +from pathlib import Path + +from huggingface_hub import hf_hub_download +import numpy as np +import soundfile as sf +import torch + import look2hear.models + +SAMPLE_RATE = 44_100 +DEVICE_CHOICES = ("auto", "cuda", "mps", "cpu") +OFFICIAL_CHECKPOINT = "JusperLee/Apollo" +LOCAL_CHECKPOINT_SUFFIXES = {".bin", ".ckpt", ".pt", ".pth", ".safetensors"} + + +def select_device(requested): + """Resolve an inference device, preferring CUDA, then MPS, then CPU.""" + if requested == "auto": + if torch.cuda.is_available(): + return torch.device("cuda") + if torch.backends.mps.is_available(): + return torch.device("mps") + return torch.device("cpu") + + if requested == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is not available.") + if requested == "mps" and not torch.backends.mps.is_available(): + raise RuntimeError("MPS was requested but is not available.") + return torch.device(requested) + + +def resolve_checkpoint(reference): + """Resolve the official remote checkpoint or an existing local file.""" + reference = str(reference) + if reference == OFFICIAL_CHECKPOINT: + return Path( + hf_hub_download( + repo_id=OFFICIAL_CHECKPOINT, filename="pytorch_model.bin" + ) + ) + + path = Path(reference).expanduser() + if path.is_file(): + return path + if path.exists(): + raise IsADirectoryError(f"Checkpoint must be a file, not a directory: {path}") + if ( + path.suffix.lower() in LOCAL_CHECKPOINT_SUFFIXES + or path.is_absolute() + or reference.startswith(("./", "../", "~/")) + ): + raise FileNotFoundError(f"Local checkpoint not found: {path}") + raise ValueError( + "Unsupported checkpoint source. Use JusperLee/Apollo or an existing " + "local checkpoint file." + ) + + def load_audio(file_path): - audio, samplerate = torchaudio.load(file_path) - return audio.unsqueeze(0).cuda() # [1, 1, samples] + audio, sample_rate = sf.read( + file_path, dtype="float32", always_2d=True + ) + if sample_rate != SAMPLE_RATE: + raise ValueError( + f"Apollo expects {SAMPLE_RATE} Hz audio, got {sample_rate} Hz." + ) + if audio.shape[0] == 0 or audio.shape[1] == 0: + raise ValueError("Input audio must contain at least one sample and channel.") + if not np.isfinite(audio).all(): + raise ValueError("Input audio contains NaN or infinite values.") + # SoundFile uses [samples, channels]; Apollo uses [batch, channels, samples]. + audio = torch.from_numpy(np.ascontiguousarray(audio.T)).unsqueeze(0) + return audio, sample_rate -def save_audio(file_path, audio, samplerate=44100): - audio = audio.squeeze(0).cpu() - torchaudio.save(file_path, audio, samplerate) -def main(input_wav, output_wav): - os.environ['CUDA_VISIBLE_DEVICES'] = "0" +def save_audio(file_path, audio, sample_rate): + output_path = Path(file_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + audio = audio.detach().squeeze(0).to("cpu").numpy().T + sf.write(output_path, audio, sample_rate, subtype="FLOAT") - model = look2hear.models.BaseModel.from_pretrain("JusperLee/Apollo", sr=44100, win=20, feature_dim=256, layer=6).cuda() - test_data = load_audio(input_wav) - with torch.no_grad(): - out = model(test_data) - save_audio(output_wav, out) -if __name__ == "__main__": +def run_inference(input_wav, output_wav, checkpoint, requested_device="auto"): + 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.") + + 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), + sr=SAMPLE_RATE, + win=20, + feature_dim=256, + layer=6, + ).to(device).eval() + with torch.inference_mode(): + output = model(test_data) + save_audio(output_path, output, sample_rate) + return device + + +def main(): parser = argparse.ArgumentParser(description="Audio Inference Script") - parser.add_argument("--in_wav", type=str, required=True, help="Path to input wav file") - parser.add_argument("--out_wav", type=str, required=True, help="Path to output wav file") + parser.add_argument( + "--in_wav", type=Path, required=True, help="Path to input WAV file" + ) + parser.add_argument( + "--out_wav", type=Path, required=True, help="Path to output WAV file" + ) + parser.add_argument( + "--checkpoint", + default=OFFICIAL_CHECKPOINT, + help="JusperLee/Apollo or an existing local checkpoint file", + ) + parser.add_argument( + "--device", + choices=DEVICE_CHOICES, + default="auto", + help="Inference device (auto prefers CUDA, then MPS, then CPU)", + ) args = parser.parse_args() - main(args.in_wav, args.out_wav) + device = run_inference( + args.in_wav, args.out_wav, args.checkpoint, args.device + ) + print(f"Inference completed on {device}: {args.out_wav}") + + +if __name__ == "__main__": + main() diff --git a/requirements-macos-arm64.txt b/requirements-macos-arm64.txt new file mode 100644 index 0000000..1883e34 --- /dev/null +++ b/requirements-macos-arm64.txt @@ -0,0 +1,8 @@ +# Native Apple Silicon inference/smoke-test environment. +# The upstream look2hear.yml is a linux-64 CUDA environment. +torch==2.11.0 +torchaudio==2.11.0 +numpy==1.26.4 +huggingface-hub==1.27.0 +omegaconf==2.3.0 +soundfile==0.14.0 diff --git a/tests/test_inference.py b/tests/test_inference.py new file mode 100644 index 0000000..7e0e64f --- /dev/null +++ b/tests/test_inference.py @@ -0,0 +1,178 @@ +import hashlib +import os +from pathlib import Path +import tempfile +import unittest +from unittest import mock + +import numpy as np +import soundfile as sf +import torch + +import inference + + +class IdentityModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.anchor = torch.nn.Parameter(torch.zeros(())) + self.input_device = None + + def forward(self, audio): + self.input_device = audio.device + 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) + def test_auto_prefers_cuda_over_mps(self, cuda_available, mps_available): + self.assertEqual(inference.select_device("auto"), torch.device("cuda")) + + @mock.patch.object(torch.backends.mps, "is_available", return_value=True) + @mock.patch.object(torch.cuda, "is_available", return_value=False) + def test_auto_uses_mps_when_cuda_is_unavailable( + self, cuda_available, mps_available + ): + self.assertEqual(inference.select_device("auto"), torch.device("mps")) + + @mock.patch.object(torch.backends.mps, "is_available", return_value=False) + @mock.patch.object(torch.cuda, "is_available", return_value=False) + def test_auto_falls_back_to_cpu(self, cuda_available, mps_available): + self.assertEqual(inference.select_device("auto"), torch.device("cpu")) + + @mock.patch.object(torch.cuda, "is_available", return_value=False) + def test_explicit_unavailable_cuda_fails(self, cuda_available): + with self.assertRaisesRegex(RuntimeError, "CUDA was requested"): + inference.select_device("cuda") + + @mock.patch.object(torch.backends.mps, "is_available", return_value=False) + def test_explicit_unavailable_mps_fails(self, mps_available): + with self.assertRaisesRegex(RuntimeError, "MPS was requested"): + inference.select_device("mps") + + +class InferenceTests(unittest.TestCase): + def test_official_checkpoint_downloads_from_trusted_repository(self): + with mock.patch.object( + inference, + "hf_hub_download", + return_value="/cache/pytorch_model.bin", + ) as download: + path = inference.resolve_checkpoint(inference.OFFICIAL_CHECKPOINT) + + self.assertEqual(path, Path("/cache/pytorch_model.bin")) + download.assert_called_once_with( + repo_id=inference.OFFICIAL_CHECKPOINT, + filename="pytorch_model.bin", + ) + + def test_existing_local_checkpoint_is_accepted(self): + with tempfile.TemporaryDirectory() as directory: + checkpoint = Path(directory) / "model.pth" + checkpoint.touch() + self.assertEqual(inference.resolve_checkpoint(checkpoint), checkpoint) + + def test_arbitrary_hugging_face_repository_is_rejected(self): + for reference in ("someone/model", "someone/model.v2"): + with self.subTest(reference=reference), mock.patch.object( + inference, "hf_hub_download" + ) as download: + with self.assertRaisesRegex(ValueError, "Unsupported checkpoint"): + inference.resolve_checkpoint(reference) + download.assert_not_called() + + def test_missing_local_checkpoint_has_clear_error(self): + with self.assertRaisesRegex(FileNotFoundError, "Local checkpoint not found"): + inference.resolve_checkpoint("models/missing.pth") + + def test_checkpoint_directory_is_rejected(self): + with tempfile.TemporaryDirectory() as directory: + with self.assertRaisesRegex(IsADirectoryError, "must be a file"): + inference.resolve_checkpoint(directory) + + def test_cpu_inference_preserves_input_and_audio_shape(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() + + samples = np.linspace(-0.5, 0.5, inference.SAMPLE_RATE) + stereo = np.stack((samples, -samples), axis=1).astype(np.float32) + 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", + ) + + output, sample_rate = sf.read( + output_path, dtype="float32", always_2d=True + ) + self.assertEqual(device, torch.device("cpu")) + self.assertEqual(model.input_device, torch.device("cpu")) + self.assertEqual(sample_rate, inference.SAMPLE_RATE) + self.assertEqual(output.shape, stereo.shape) + self.assertTrue(np.isfinite(output).all()) + self.assertEqual( + hashlib.sha256(input_path.read_bytes()).hexdigest(), input_digest + ) + + def test_refuses_to_overwrite_input(self): + path = Path("same.wav") + with self.assertRaisesRegex(ValueError, "must be different"): + inference.run_inference(path, path, "checkpoint.bin", "cpu") + + def test_wrong_sample_rate_fails_before_download_or_device_transfer(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), 48_000) + + with mock.patch.object( + inference, "hf_hub_download" + ) as download, mock.patch.object(torch.Tensor, "to") as transfer: + with self.assertRaisesRegex(ValueError, "expects 44100 Hz"): + inference.run_inference( + input_path, + root / "output.wav", + inference.OFFICIAL_CHECKPOINT, + "cpu", + ) + download.assert_not_called() + transfer.assert_not_called() + + def test_refuses_hard_link_output_alias(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + input_path = root / "input.wav" + output_path = root / "output.wav" + sf.write( + input_path, + np.zeros((100, 2), dtype=np.float32), + inference.SAMPLE_RATE, + ) + try: + os.link(input_path, output_path) + except OSError as error: + self.skipTest(f"hard links are unavailable: {error}") + + with self.assertRaisesRegex(ValueError, "same file"): + inference.run_inference( + input_path, output_path, "checkpoint.bin", "cpu" + ) + + +if __name__ == "__main__": + unittest.main()