Proposal: Apple Silicon/MPS support for lossy EDM and pop-dance music restoration

[Djshinshin]

Hello Apollo authors,

Thank you for releasing Apollo and its pretrained checkpoints. I am interested in contributing practical testing and implementation work specifically for lossy electronic dance music, pop music, and harmony-rich pop-dance music restoration.

My background combines:

• DJ and dance music production/mastering workflows, especially EDM, club music, pop music, K-pop remixes, harmony-rich pop-dance arrangements, kick/sub-bass and vocal clarity evaluation
• professional voice acting, with listening experience focused on consonant articulation, sibilance, breath, resonance, intelligibility, and vocal timbre

I would like to explore the following contributions:

1. Apple Silicon / MPS support

   • replace hard-coded .cuda() calls with a device abstraction
   • add MPS, CUDA and CPU fallback
   • test complex STFT/iSTFT compatibility and document unsupported operations
   • add chunked inference for long stereo music on limited unified memory

2. High-bitrate lossy EDM, pop and pop-dance evaluation and training

   • extend simulated degradation beyond the published 24-128 kbps range to 192/256/320 kbps
   • evaluate whether the model remains identity-preserving when codec artifacts are already subtle
   • include transient-heavy electronic dance music, mainstream pop, and harmony-rich pop-dance music
   • evaluate dense vocals, layered synthesizers, chord clarity, harmonic overtones, kick/sub-bass interaction and mono compatibility

3. Continual fine-tuning with a safe promotion gate

   • keep a stable inference checkpoint while a challenger checkpoint trains separately
   • train only from legitimate clean WAV/FLAC targets paired with locally generated MP3 degradations
   • never train MP3-only material against itself as ground truth
   • promote a checkpoint only after objective metrics and level-matched human listening tests improve without added artifacts

4. Perceptual validation

   • codec chirp/warbling and pre-echo
   • transient smearing and low-frequency tearing
   • vocal consonants, sibilance, breath and resonance
   • chord definition, harmonic balance and separation between layered synths and vocals
   • loudness-matched A/B, true peak and mono checks

Would the maintainers welcome a PR that first adds device-agnostic inference and an Apple Silicon test path? If so, is there a preferred PyTorch version or branch to target?

Also, would extending the official data preprocessing configuration to include 192/256/320 kbps MP3 and dedicated lossy EDM, pop and harmony-rich pop-dance evaluation profiles be aligned with the project direction?

Any test data used for this work will remain local or come from properly licensed/open datasets; copyrighted music will not be uploaded to the repository.

Thank you.

[JusperLee]

Thank you for the detailed proposal and for your willingness to contribute implementation and practical testing.

Apple Silicon/MPS support would be valuable. We would welcome a focused PR that:

• replaces hard-coded .cuda() calls with device-agnostic handling;
• supports CUDA, MPS, and CPU fallback while preserving the existing CUDA behavior;
• checks MPS compatibility for the main inference pipeline, especially complex-valued operations and STFT/iSTFT;
• documents any unsupported operations or known limitations.

Please target the current main branch and the PyTorch version currently used by the project. If proper MPS support requires raising the minimum PyTorch version, please explain the compatibility implications in the PR.

Chunked inference for long stereo audio would also be useful, although it may be better handled in a separate PR because overlap, boundary artifacts, and output consistency need to be considered carefully.

Regarding 192/256/320 kbps MP3, you are welcome to experiment with fine-tuning or training for these bitrate settings. Extending the degradation pipeline and training configuration to cover high-bitrate MP3 is aligned with the project direction. It would be helpful to retain clean WAV/FLAC audio as the training target and generate the lossy inputs locally, as you proposed.

For high-bitrate inputs, please also pay attention to identity preservation, since codec artifacts may already be subtle and the model should avoid unnecessarily modifying clean or near-clean content. Evaluation on transient-heavy EDM, pop, and harmony-rich arrangements, together with objective metrics and loudness-matched listening tests, would be particularly useful.

Thank you again for the thoughtful proposal.

[Djshinshin]

Thank you again for the clear guidance. I have now opened a focused draft PR for the first stage:

PR #38 — Add device-agnostic inference support
#38

The draft currently:

• removes the hard-coded .cuda() calls and CUDA_VISIBLE_DEVICES setting from inference.py;
• adds --device auto|cuda|mps|cpu, with auto preserving the CUDA → MPS → CPU priority;
• keeps the model and input tensor on the same device and moves the output back to CPU before writing;
• validates the audio sample rate, non-empty shape, and finite values on CPU before checkpoint download or accelerator transfer;
• limits remote checkpoint download to the official JusperLee/Apollo repository, while custom checkpoints must be explicitly provided as existing trusted local files;
• prevents output from overwriting the input through either the same path or a hard-link alias;
• documents the tested Apple Silicon setup, PyTorch version difference, and known limitations.

Apple Silicon validation was performed on:

• Apple M3 with 24 GB unified memory
• macOS 27.0
• native arm64 Python 3.10
• PyTorch/TorchAudio 2.11.0

The upstream environment records PyTorch 2.0.0+cu118 and TorchAudio 2.0.1+cu118 as the original Linux/CUDA baseline. The PR documents this difference explicitly. CUDA remains the first auto choice, but it was not hardware-tested on the Apple Silicon machine.

Using the repository's public six-second stereo fixture, the updated entry point completed end-to-end inference on both MPS and CPU:

• MPS: 5.78 seconds
• CPU: 316.46 seconds
• both outputs: float32 PCM, stereo, 44.1 kHz, exactly six seconds, with no NaN or infinite values
• maximum CPU/MPS absolute difference: 0.000249214470387
• RMSE: 2.02574076096e-05
• correlation: 0.999999991468
• the input SHA-256 remained unchanged

The current local validation also includes 14/14 passing unit tests, py_compile, pip check, and git diff --check. The two Copilot review threads have been addressed and resolved.

As suggested, this PR does not include long-audio chunking, overlap/crossfade, high-bitrate MP3 degradation, or training changes. Those remain separate follow-up work.

I also added a minimal workflow named CPU unit tests. It only runs Python compilation and mock-based CPU unit tests; it does not download the checkpoint or claim MPS/CUDA coverage. GitHub currently marks the first workflow run from the fork as action_required, so it appears to need maintainer approval before it can execute. If this small workflow fits the project's CI approach, would you please approve the run? If you would prefer not to add a workflow in this PR, I can remove it and retain the test code and documented local results.

The PR remains in draft while awaiting that CI decision and any maintainer feedback. Validation used only the repository's public fixture; no private music, checkpoint, WAV output, virtual environment, or generated artifact was committed or uploaded.

Thank you.

[JusperLee]

Thank you for the detailed update and for keeping the first PR focused.

The scope of PR #38 aligns well with the proposed first stage. Keeping long-audio chunking, high-bitrate degradation, and training changes in separate follow-up PRs is appropriate.

The lightweight CPU unit-test workflow fits the project's CI approach, and I have approved the workflow run. Please keep the workflow in the PR.

The documented Apple Silicon results and the distinction between local MPS validation and the existing CUDA baseline are also helpful. Once the workflow has completed successfully and the implementation is ready, please mark the PR as ready for review. We can continue the code-specific discussion in PR #38.

Thank you again for the careful implementation and validation.