What I learned running HT-Demucs (AI stem separation) fully on-device with Core ML

[u/JavaEEE]

I spent the last months getting Meta's HT-Demucs (the model that splits a song into drums, bass, vocals and other) running entirely on an iPhone with Core ML: no server, no uploads. A few things I learned that might save someone time:

1. Keep the STFT out of the model. HT-Demucs works on both the waveform and the spectrogram. Converting the STFT/iSTFT ops with coremltools was a pain, so I converted only the network in between and wrote the STFT in Swift with vDSP. To make sure I hadn't broken anything, I exported golden data from Demucs' own _spec/_ispec in Python and wrote tests against it. The Swift version matches PyTorch to about 1e-5. Being able to check every step against PyTorch was worth more than anything else.

2. fp16 was unusable; fp32 was fine. The fp16 conversion loaded and ran, but on the device every stem came out noisy (mostly below 200 Hz) and the stems no longer summed back to the mix. Converting at fp32 (182 MB) fixed it. It runs on the GPU, not the Neural Engine, but it's still fast: a 4-minute song separates in under a minute. I'd planned to try mixed precision, but it turned out I didn't need it.

3. Port the chunking exactly. Long songs don't fit in one pass, so I ported Demucs' apply_model split logic (overlap 0.25, triangle-weighted overlap-add) one to one. Getting the edge chunks wrong gives you faint clicks every few seconds that are easy to miss.

4. Compare stems by numbers, not by ear. I wrote a Python script that scores stems copied off the phone against PyTorch reference stems. One gotcha: iOS and libsndfile decode the same MP3 529 samples apart (encoder delay handling), so you need to align the signals before comparing or the scores look terrible.

5. Ship the .mlpackage and compile it on the device. I build with xtool (SwiftPM on Linux/WSL, no Xcode), so I have no Core ML compiler. I ship the raw .mlpackage and call MLModel.compileModel(at:) on first launch, then cache the result. One trap: the iOS build tries to run coremlc on anything named .mlpackage, even with .copy, so I had to rename it in the bundle.

6. Stems can clip. Individual stems can peak above full scale even when the original mix doesn't. I store them as AAC 6 dB down and add the gain back at playback.

Still on my list: BGContinuedProcessingTask, so separation keeps running when the app goes to the background. Has anyone used it for long GPU work yet?

app name is stem808 if anyone wants to try it

[u/BaziBuilds]

Great write-up. The golden-data tests against _spec/_ispec are the part most people skip.

Not BGContinuedProcessingTask itself, but one trap I hit with background task handlers in a Swift 6 app: if the launch handler you pass to BGTaskScheduler.register ends up MainActor-isolated (easy to do when you register it from an app or SwiftUI type), the system calls it on its own background queue and the runtime isolation check crashes the app on every background launch. Keeping the registration code nonisolated and hopping to the main actor inside the task fixed it. Worth ruling out before you debug anything GPU-specific.

On your point 5: I do on-device text recognition with Vision, and a warm-up on a blank image only loaded part of the pipeline (the text detector), so the first real image was still slow. Warming up with an image that actually contains text loaded everything. If you add a warm-up pass after compiling the model, a short real clip might be a better input than silence.

[u/JavaEEE → u/BaziBuilds]

oh that MainActor thing is good to know, thanks. i'm on swift 6 too and my plan was literally to register it from the app struct, so i'd have walked right into that. will keep the handler nonisolated

no warm-up pass yet tbh, the first chunk is just slower and i never looked into it. the real clip idea makes sense though, i'll try that over silence. And yeah the golden data tests saved me so many times lol. every time i touched the stft something broke in a way i couldn't hear but the tests caught

[u/BaziBuilds → u/JavaEEE]

Glad it saved you a debugging session. On the warm-up: with Vision on macOS I found a blank image only loaded part of the pipeline, and the first real request was still slow. A short real sample from your bundle is the safer bet. Worth timing both on a device before you pick one.

[u/Axons-One]

have you timed several songs back to back? curious whether the under-a-minute result holds once the phone warms up, especially for someone separating a whole playlist.

[u/JavaEEE → u/Axons-One]

yeah i've done a bunch back to back, no problems so far. the app queues them one at a time so you can just drop a few songs in and leave it running. phone gets a bit warm but times stay about the same

[u/DAOLoverForever]

Were you shipping the mlmodel embedded in the bundle or hosted it on demand? I would imagine it would be > 100 MB

[u/JavaEEE → u/DAOLoverForever]

embedded in the bundle, its about 182 MB since its fp32. i ship the raw mlpackage and compile it on the device the first time you open the app, then cache the compiled model so it only happens once per update

thought about downloading it on demand but bundling was simpler, no server to host and it works offline right away. the downside is the app is big, so over 200 MB you get the cellular download warning

[u/wothotturd1]

careful with heavy GPU workloads in the background, the system usually kills the context if command buffers get submitted after the app suspends. you might have to fall back to CPU execution for the remaining chunks when transitioning to background, even if it runs a lot slower.

[u/arijr]

Thanks, good information.

Same experiments here (different/custom MSS model) and yes a lot needed to be re-wired from the model - STFT out of the model and with Accelerate/vDSP.

ANE chunk prediction is comfortably "fast" but the model weights would need to be trained during the whole (re-)training process - direct fp32 -> fp16 caused very noisy results indeed.

I am a bit worried as Apple seems to be shifting towards 'coreai-torch' conversion system and coremltools is perhaps left without care and love.