Add device-agnostic inference support

.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

.gitignore

@@ -1,2 +1,7 @@
 *.pyc
-*.DS_Store
+*.DS_Store
+.venv/
+/models/
+/smoke_test_output/
+/output.wav
+/apollo_*_smoke.wav

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.

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*

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()

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