Skip to content

Add device-agnostic inference support - #38

Merged
JusperLee merged 2 commits into
JusperLee:mainfrom
Djshinshin:codex/mps-device-support
Aug 14, 2026
Merged

Add device-agnostic inference support#38
JusperLee merged 2 commits into
JusperLee:mainfrom
Djshinshin:codex/mps-device-support

Conversation

@Djshinshin

@Djshinshin Djshinshin commented Aug 13, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Remove the hard-coded .cuda() calls and CUDA_VISIBLE_DEVICES setting from inference.py.
  • Add --device auto|cuda|mps|cpu; auto preserves the existing CUDA-first behavior, then selects MPS, then CPU.
  • Keep the model and input tensor on the same device and move output safely back to CPU before writing.
  • Download only the official JusperLee/Apollo remote checkpoint; custom checkpoints must be existing local files from trusted sources.
  • Validate sample rate, non-empty shape, and finite samples on CPU before checkpoint download or accelerator transfer.
  • Prevent output from overwriting the input through either the same path or a hard-link alias.

Compatibility

The upstream environment records PyTorch 2.0.0+cu118 and TorchAudio 2.0.1+cu118 as the original Linux/CUDA baseline. Apple Silicon validation used an Apple M3 with macOS 27.0, native arm64 Python 3.10, and PyTorch/TorchAudio 2.11.0.

CUDA remains the first auto choice and the CUDA inference path is preserved, but CUDA was not hardware-tested on the Apple Silicon validation machine. MPS support can also vary with PyTorch and macOS versions.

Tests

  • python -m unittest discover -s tests -v: 14/14 passed, including trusted checkpoint resolution, rejection of arbitrary/dotted remote IDs, validation ordering, and hard-link protection.
  • python -m py_compile inference.py tests/test_inference.py: passed.
  • python -m pip check: passed.
  • git diff --check: passed.
  • Added a minimal GitHub Actions workflow named CPU unit tests; it runs compile and mock-based unit tests only and does not claim MPS or CUDA coverage.
  • Ran the official public asserts/input_wav.wav fixture end to end on MPS and CPU after the review fixes.
  • Both outputs decoded as float32 PCM, stereo, 44.1 kHz, exactly 6 seconds, with no NaN or infinite values.
  • The input SHA-256 remained 3c9a053913c3016b493ddc0b92e21a4e682e8fb5f872dafc945154155bddf771.
  • CPU/MPS comparison: maximum absolute difference 0.000249214470387, RMSE 2.02574076096e-05, correlation 0.999999991468.
  • End-to-end reference time: MPS 5.78 s; CPU 316.46 s on the tested Apple M3.

Security and trust boundary

The upstream loader uses torch.load(..., weights_only=False). To avoid downloading arbitrary serialized objects, the CLI accepts the exact official JusperLee/Apollo remote ID or an explicitly existing local file only. Missing local files, directories, and all other remote IDs fail clearly. The tested official checkpoint SHA-256 is documented as a reproducibility record, not enforced as a permanent hash gate, so future legitimate official updates are not blocked.

Limitations

  • test.py and the training configuration retain their existing CUDA assumptions; this PR is limited to the public inference entry point.
  • Long-audio chunking and crossfade are outside this PR.
  • High-bitrate MP3 degradation or evaluation work is outside this PR.
  • 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.

Copyright and safety

Validation used only the repository's official public fixture. No private music was read or uploaded, and no checkpoint, WAV output, virtual environment, cache, or other generated artifact is included in the commits.

Refs #37

@JusperLee
JusperLee marked this pull request as ready for review August 13, 2026 16:51
@JusperLee
JusperLee requested a lite review from Copilot August 13, 2026 16:51

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR updates the public inference.py entry point to support device-agnostic inference (CUDA/MPS/CPU), removes hard-coded CUDA assumptions, and documents a tested Apple Silicon workflow while adding unit coverage for the new device/IO behavior.

Changes:

  • Add --device auto|cuda|mps|cpu and device selection logic; keep model/input on the same device and write output back on CPU.
  • Resolve checkpoints from either a local path or by downloading pytorch_model.bin from a Hugging Face repo ID.
  • Add unit tests for device selection and CPU inference, plus Apple Silicon documentation and a macOS arm64 requirements file.

Reviewed changes

Copilot reviewed 5 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file
File Description
inference.py Device selection, HF/local checkpoint resolution, SoundFile-based IO, and safer CPU writeback.
tests/test_inference.py Unit tests covering select_device() and CPU inference behavior.
README.md Updated inference command and documentation for --device / --checkpoint.
MACOS_ARM64.md Apple Silicon/MPS setup notes and validation details.
requirements-macos-arm64.txt Pinned macOS arm64 inference/smoke-test dependencies.
.gitignore Ignore common local inference outputs and venv/models dirs.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread inference.py
Comment thread inference.py Outdated
@Djshinshin
Djshinshin marked this pull request as draft August 13, 2026 17:23
@JusperLee
JusperLee marked this pull request as ready for review August 14, 2026 04:01
@JusperLee
JusperLee merged commit d8b3d20 into JusperLee:main Aug 14, 2026
1 check failed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants