Running htdemucs on a phone: memory, GPUs and a 9.5x faster iSTFT

Music source separation takes a finished track and pulls it back apart: vocals on one stem, bass on another, drums on a third, everything else on a fourth. Producers use it to remix. Musicians use it to practise against a song with their own instrument removed. Until fairly recently you did it on a workstation, or you uploaded your audio to somebody else's server and hoped.

We wanted it on the phone, with the audio never leaving the device. Demixr is the result, and it runs htdemucs — an off-the-shelf, well-understood model. Getting that model to produce correct output was a weekend. Getting it to run inside a phone's memory budget, and then fast enough that anyone would wait for it, took the rest of the work.

This is what was actually hard.

The model, and where it comes from

We trained nothing. htdemucs is Hybrid Transformer Demucs, from Meta's Demucs project — a waveform-and-spectrogram hybrid with a transformer in the middle of the U-Net, described in Hybrid Transformers for Music Source Separation (Rouard, Massa and Défossez, 2022). The reference implementation is github.com/adefossez/demucs, MIT licensed, weights included.

There are three variants worth knowing about, and the differences matter more for a phone than they do for a workstation:

Variant	Stems	What it is	Shipped?
`htdemucs`	4 — vocals, bass, drums, other	The default. One pass, one model.	Yes
`htdemucs_ft`	4	Fine-tuned per source. Better separation, four times the work — it runs a dedicated model per stem.	No
`htdemucs_6s`	6 — adds guitar and piano	Experimental upstream, and upstream says so.	No — see below

htdemucs_ft is the obvious upgrade and the one we cannot afford: four sequential passes over a song, on a device already fighting for memory, turns a two-minute wait into an eight-minute one. That trade is fine on a laptop with a fan. It is not a mobile app.

If you want to try this yourself, the exports other people have already published are the fastest way in — Hugging Face has around seventy repositories matching htdemucs, including ONNX, Core ML and MLX conversions of the same weights. We exported our own, for reasons that become clear in the section on where the STFT lives, but starting from somebody else's export is a much shorter route to hearing four stems come out of a phone.

How the model is actually used

A separation is not one call into a model. The model handles a few seconds of audio at a time, and everything around it — getting to the right sample rate, cutting the song into chunks, turning waveforms into spectrograms and back, stitching the pieces together without a seam — is host code. That host code is where all of this article's problems live.

flowchart TD
  A["Song on disk<br/>(any format, any rate)"] --> B["Decode + resample<br/>to 44.1 kHz stereo"]
  B --> C["Cut into overlapping chunks"]
  C --> D["STFT<br/>waveform to spectrogram"]
  D --> E["htdemucs<br/>4 masks out"]
  E --> F["Apply masks<br/>per stem"]
  F --> G["Inverse STFT<br/>spectrogram to waveform"]
  G --> H["Overlap-add<br/>stitch chunks"]
  H --> I["Add the time branch"]
  I --> J["4 stems written to disk"]

One pass. The model is step five of nine; the other eight are ours.

Worth saying plainly, because it is the thing that surprises people: the neural network was never the slow part. Every performance win below came from somewhere in that diagram other than the box marked htdemucs.

The first version was Android-only, and we deleted it

The original engine was PyTorch Lite behind a native MethodChannel, with an Oboe resampler in C++ and a CMake build. It worked. It also meant an Android-shaped hole in the codebase: macOS and iOS got nothing, and every change had to be made twice.

We replaced it with a single ONNX engine shared across all three platforms, and then deleted the native path entirely — the plugin, the WAV helpers, the C++ directory, the externalNativeBuild block, the PyTorch dependency. MainActivity went back to being a plain FlutterActivity.

	Before	After
Runtime	PyTorch Lite	ONNX Runtime
Platforms	Android only	Android, iOS, macOS
Bridge	Custom `MethodChannel` plugin	`flutter_onnxruntime`, vendored
Resampling	Oboe, C++	Dart
Native build	CMake + `externalNativeBuild`	None

That is a satisfying diff to write and a nervous one to ship, so it was gated on evidence: 4-stem and 6-stem output had to match the old path to within one least-significant bit on macOS before the old path could go. It did.

Then Android ran out of memory at five gigabytes

The cross-platform engine was correct everywhere and viable nowhere on mobile. On device, separation climbed to roughly 5GB of resident memory and the process was killed.

Two things were responsible, and neither is obvious from the ONNX Runtime documentation.

The first was graph optimisation. ORT rewrites the model graph on load — fusing operators, folding constants — and to do that it holds intermediate representations in memory. For a model this size, on a device with this little headroom, the optimisation pass cost more than it saved. Turning it off required vendoring flutter_onnxruntime so we could reach the session options at all.

The second was the memory arena. ORT's default allocator keeps freed blocks around to reuse, which is exactly right on a server and exactly wrong on a phone: across a long run of chunks, peak usage ratchets upward and never comes back down. Disabling the arena keeps the peak flat at about 2.2GB no matter how many chunks the song takes.

xychart-beta
  title "Peak resident memory during one separation"
  x-axis ["ORT defaults", "Arena + graph opt off"]
  y-axis "Gigabytes" 0 --> 6
  bar [5.0, 2.2]

Both numbers are configuration, not code. Neither is signposted.

The two settings, in the terms the runtime uses for them:

Session option	Value	Why
Graph optimisation level	`ORT_DISABLE_ALL`	The fusion pass needs more memory than the fused graph saves, on a model this size.
CPU memory arena	disabled	Freed blocks are returned rather than pooled. Peak stops ratcheting across chunks.

Flat is the word that matters. A peak that grows with song length is a bug that only your users with long songs will find.

iOS needed permission to use the memory it had

The same build on iPhone was killed by jetsam — iOS's memory watchdog — while sitting well under the device's physical RAM. iOS caps what a normal app may hold, and a 2.2GB working set is over that line.

The fix is a single entitlement, com.apple.developer.kernel.increased-memory-limit, which raises the ceiling for apps that genuinely need it. Worth knowing it exists before you spend a day assuming your allocations are wrong.

We rejected the GPU, with data, and then accepted a different one

The obvious next move was to get off the CPU. We measured before committing, and the first answer was no: a PyTorch-MPS path showed a headline 12.6× in isolation that did not survive contact with the rest of the pipeline, and the integration cost was large. That investigation was written up and rejected on the evidence rather than quietly abandoned, which matters later when somebody asks why the obvious thing wasn't done.

The second answer was yes. A spike with ExecuTorch and Core ML ran the htdemucs core roughly 4× faster in a single partition — no graph splitting, no fallback to CPU for unsupported operators, which is usually where mobile GPU acceleration falls over.

Exporting the model the right way mattered as much as the backend. Putting the STFT inside the graph — so the export takes audio in and gives stems out, rather than expecting the host to prepare spectrograms — produced a 4.8× speedup at a correlation of 1.0 against the reference implementation. Same numbers out, four times sooner.

Fixing the bottleneck moved the bottleneck

With the core on the GPU, profiling pointed somewhere unglamorous: the post-processing stage. Masking and the inverse STFT had been exported as a second .pte, and inside it the inverse STFT was implemented as a dense-DFT ConvTranspose1d. That is an O(N²) matrix multiply standing in for an operation with a well-known O(N log N) algorithm, and it cost 1.75 seconds per chunk — ten times the core it was feeding.

pie showData title Time per chunk, before the rewrite
  "Inverse STFT (dense DFT)" : 1750
  "Everything else, including the model" : 175

The neural network is the small slice. This is what profiling after a win is for.

So we took it out of the model and wrote it in Dart: a radix-2 FFT inverse STFT reproducing htdemucs' own masking and spectral inversion step for step — build the conjugate-symmetric spectrum per frame, inverse FFT, window and overlap-add at the −1536 centre offset, divide by the overlap-add envelope, add the time branch.

xychart-beta
  title "Inverse STFT, one run"
  x-axis ["ConvTranspose1d in-graph", "Radix-2 FFT in Dart"]
  y-axis "Milliseconds" 0 --> 6000
  bar [5252, 550]

The results were better than the speedup suggested:

Measure	Before	After	Change
Inverse STFT	5252 ms	550 ms	9.5× faster
End to end, 12-second clip	16 s	11 s	31% faster
Post-processing model in the download	129 MB	0 MB	Removed
Correlation vs NumPy mirror	—	1.0	Relative error 4.7e-7

The download got smaller because the code got faster. Replacing a model artefact with fifty lines of arithmetic you understand is usually a good trade, and it is the kind of trade that only shows up if you profile the slow thing instead of assuming the neural network is the slow thing.

The first run felt broken, and it wasn't

One problem survived all of this. The first separation after install took about ten seconds longer than every subsequent one, and the progress bar sat still for all of it.

That was Core ML compiling the model for the specific device — roughly 10 seconds on a Mac, 20 on an iPhone. We assumed it was cached on disk after the first time. It isn't: two sequential cold loads both took about 9.6 seconds.

sequenceDiagram
  participant U as User
  participant A as App
  participant C as Core ML
  U->>A: Install, model downloads
  A->>C: Compile for this device
  Note over A,C: 9631 ms, behind "Optimising the model for your device"
  C-->>A: Compiled, kept resident
  U->>A: Separate a song
  A->>C: Run
  Note over A,C: 0 ms of compile
  C-->>A: Stems

The compile cannot be avoided. It can be moved to a moment the user is already waiting.

Since the compile can't be avoided, it can at least be moved somewhere the user is already waiting. The model is now warmed up immediately after it downloads, behind the honest label "Optimising the model for your device", and kept resident so later runs reuse it. On an already-downloaded model, the app warms it in the background at launch.

Measured: first warm-up 9631ms, reuse 0ms.

What shipped, and what didn't

The app ships two engines. GPU via ExecuTorch is the default — a Core ML build on Apple platforms, Vulkan on Android, the same weights either way. CPU via ONNX stays available as a smaller download and a fallback.

	CPU engine	GPU engine
Runtime	ONNX Runtime	ExecuTorch
Backend	CPU	Core ML (Apple), Vulkan (Android)
Weights	The same htdemucs weights either way
Speed, macOS	baseline	8.4×
Speed, iPhone	baseline	2.5×
Download size	Smaller	Larger
Role	Fallback	Default

One thing did not ship. There is a six-stem htdemucs variant that separates guitar and piano as well, and we cut it: the guitar and piano quality was poor enough in testing that shipping it would have been a worse product with a longer feature list. The enum values stay in the code for compatibility with libraries saved by earlier builds — unreachable, harmless, and cheaper than a migration.

What we'd tell anyone attempting this

Measure the peak, not the average. A memory profile that grows slowly across chunks looks fine on a 12-second test clip and fails on a real song.

Read the allocator's defaults. Both memory fixes were configuration, not code. Neither is signposted, and both were worth more than any optimisation we wrote by hand.

Profile after every win. Moving the core to the GPU didn't make the app fast; it made the inverse STFT the problem. We'd have never looked there otherwise.

Decide where the STFT lives, early. In the graph it is one export and one call, and it is the difference between a single GPU partition and a graph that falls back to CPU halfway through. Out of the graph it is yours to optimise. Both are defensible; drifting between them is not.

Write down the things you reject. The GPU investigation that ended in "no" was as useful as the one that ended in "yes", because it stopped the question being reopened from scratch three weeks later.

We put a music-separation model on a phone. The model was the easy part.