| |
| """ |
| Export the HTDemucs split models (4-stem and 6-stem) as fp16 CoreML mlpackages. |
| |
| This mirrors convert_htdemucs_6s_split_coreml.py but: |
| * parameterises the model name, so the same wrapper drives both htdemucs |
| (4 sources) and htdemucs_6s (6 sources) - the wrapper already derives the |
| output channel counts from len(model.sources); |
| * emits fp16 weight/compute precision instead of fp32. |
| |
| STFT/iSTFT stay outside the CoreML graph, exactly as in the fp32 exports, so |
| the two extra outputs keep the same interface the app already expects: |
| |
| inputs : spectral_magnitude [1, 4, 2048, 336] |
| audio_waveform [1, 2, 343980] |
| outputs: freq_output [1, stems*4, 2048, 336] |
| time_output [1, stems*2, 343980] |
| |
| Usage: |
| python convert_split_fp16_coreml.py --outdir /path/to/output |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import coremltools as ct |
| import torch |
| from demucs.pretrained import get_model |
| from einops import rearrange |
|
|
|
|
| SEGMENT_SAMPLES = 343_980 |
| NUM_BINS = 2_048 |
| NUM_FRAMES = 336 |
|
|
| |
| MODELS = [ |
| ( |
| "htdemucs", |
| "HTDemucs_SourceSeparation_F16", |
| "drums, bass, other, vocals", |
| ), |
| ( |
| "htdemucs_6s", |
| "HTDemucs_6S_Split_F16", |
| "drums, bass, other, vocals, guitar, piano", |
| ), |
| ] |
|
|
|
|
| class HTDemucsSplitExport(torch.nn.Module): |
| """Waveform/spectral split wrapper - identical logic to the fp32 6-stem |
| export, with the stem count taken from the model itself.""" |
|
|
| def __init__(self, model: torch.nn.Module): |
| super().__init__() |
| self.model = model |
|
|
| def forward( |
| self, |
| spectral_magnitude: torch.Tensor, |
| audio_waveform: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| model = self.model |
| x = spectral_magnitude |
| xt = audio_waveform |
|
|
| batch = x.shape[0] |
| spec_channels = x.shape[1] |
| freq_bins = x.shape[2] |
| frames = x.shape[3] |
|
|
| mean = x.mean(dim=(1, 2, 3), keepdim=True) |
| std = x.std(dim=(1, 2, 3), keepdim=True) |
| x = (x - mean) / (1e-5 + std) |
|
|
| time_mean = xt.mean(dim=(1, 2), keepdim=True) |
| time_std = xt.std(dim=(1, 2), keepdim=True) |
| xt = (xt - time_mean) / (1e-5 + time_std) |
|
|
| saved = [] |
| saved_t = [] |
| lengths = [] |
| lengths_t = [] |
|
|
| for idx, encode in enumerate(model.encoder): |
| lengths.append(x.shape[-1]) |
| inject = None |
|
|
| if idx < len(model.tencoder): |
| lengths_t.append(xt.shape[-1]) |
| tenc = model.tencoder[idx] |
| xt = tenc(xt) |
|
|
| if not tenc.empty: |
| saved_t.append(xt) |
| else: |
| inject = xt |
|
|
| x = encode(x, inject) |
|
|
| if idx == 0 and model.freq_emb is not None: |
| freqs = torch.arange(x.shape[-2], device=x.device) |
| emb = model.freq_emb(freqs).t()[None, :, :, None].expand_as(x) |
| x = x + model.freq_emb_scale * emb |
|
|
| saved.append(x) |
|
|
| if model.crosstransformer: |
| if model.bottom_channels: |
| _, _, freq, _ = x.shape |
| x = rearrange(x, "b c f t -> b c (f t)") |
| x = model.channel_upsampler(x) |
| x = rearrange(x, "b c (f t) -> b c f t", f=freq) |
| xt = model.channel_upsampler_t(xt) |
|
|
| x, xt = model.crosstransformer(x, xt) |
|
|
| if model.bottom_channels: |
| _, _, freq, _ = x.shape |
| x = rearrange(x, "b c f t -> b c (f t)") |
| x = model.channel_downsampler(x) |
| x = rearrange(x, "b c (f t) -> b c f t", f=freq) |
| xt = model.channel_downsampler_t(xt) |
|
|
| for idx, decode in enumerate(model.decoder): |
| skip = saved.pop(-1) |
| x, pre = decode(x, skip, lengths.pop(-1)) |
|
|
| offset = model.depth - len(model.tdecoder) |
| if idx >= offset: |
| tdec = model.tdecoder[idx - offset] |
| length_t = lengths_t.pop(-1) |
|
|
| if tdec.empty: |
| pre = pre[:, :, 0] |
| xt, _ = tdec(pre, None, length_t) |
| else: |
| skip = saved_t.pop(-1) |
| xt, _ = tdec(xt, skip, length_t) |
|
|
| stem_count = len(model.sources) |
|
|
| freq_out = x.view(batch, stem_count, -1, freq_bins, frames) |
| freq_out = freq_out * std[:, None] + mean[:, None] |
| freq_out = freq_out.reshape(batch, stem_count * spec_channels, freq_bins, frames) |
|
|
| time_out = xt.view(batch, stem_count, -1, audio_waveform.shape[-1]) |
| time_out = time_out * time_std[:, None] + time_mean[:, None] |
| time_out = time_out.reshape(batch, stem_count * 2, audio_waveform.shape[-1]) |
|
|
| return freq_out, time_out |
|
|
|
|
| def convert(model_name: str, out_name: str, source_order: str, outdir: Path) -> Path: |
| print(f"[convert] {model_name} -> {out_name} (fp16)") |
| bag = get_model(model_name) |
| model = bag.models[0] if hasattr(bag, "models") else bag |
| model.eval() |
|
|
| stem_count = len(model.sources) |
| print(f" sources: {model.sources}") |
| print(f" samplerate: {model.samplerate} segment: {int(model.segment * model.samplerate)}") |
|
|
| wrapper = HTDemucsSplitExport(model).eval() |
| spectral = torch.randn(1, 4, NUM_BINS, NUM_FRAMES) |
| waveform = torch.randn(1, 2, SEGMENT_SAMPLES) |
|
|
| with torch.no_grad(): |
| freq_out, time_out = wrapper(spectral, waveform) |
| print(f" freq_output: {tuple(freq_out.shape)} time_output: {tuple(time_out.shape)}") |
|
|
| traced = torch.jit.trace( |
| wrapper, |
| (spectral, waveform), |
| strict=False, |
| check_trace=False, |
| ) |
|
|
| mlmodel = ct.convert( |
| traced, |
| inputs=[ |
| ct.TensorType(name="spectral_magnitude", shape=spectral.shape), |
| ct.TensorType(name="audio_waveform", shape=waveform.shape), |
| ], |
| outputs=[ |
| ct.TensorType(name="freq_output"), |
| ct.TensorType(name="time_output"), |
| ], |
| minimum_deployment_target=ct.target.iOS16, |
| convert_to="mlprogram", |
| compute_precision=ct.precision.FLOAT16, |
| ) |
|
|
| mlmodel.author = "Meta Research (Demucs)" |
| mlmodel.license = "MIT License" |
| mlmodel.short_description = ( |
| f"HTDemucs {stem_count}-stem split CoreML export (fp16). STFT/iSTFT are " |
| "computed by the app; the model runs the HTDemucs neural network body." |
| ) |
| mlmodel.input_description["spectral_magnitude"] = ( |
| "Stereo complex STFT [1,4,2048,336]: L.real, L.imag, R.real, R.imag." |
| ) |
| mlmodel.input_description["audio_waveform"] = ( |
| "Stereo waveform [1,2,343980] at 44100 Hz." |
| ) |
| mlmodel.output_description["freq_output"] = ( |
| f"Frequency branch [1,{stem_count * 4},2048,336]. Source order: {source_order}." |
| ) |
| mlmodel.output_description["time_output"] = ( |
| f"Time branch [1,{stem_count * 2},343980]. Source order: {source_order}." |
| ) |
|
|
| out_path = outdir / f"{out_name}.mlpackage" |
| if out_path.exists(): |
| import shutil |
| shutil.rmtree(out_path) |
| mlmodel.save(str(out_path)) |
| return out_path |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument( |
| "--outdir", |
| type=Path, |
| required=True, |
| help="Directory to write the .mlpackage outputs into.", |
| ) |
| parser.add_argument( |
| "--only", |
| choices=[m[0] for m in MODELS], |
| help="Convert just one model (default: both).", |
| ) |
| args = parser.parse_args() |
| args.outdir.mkdir(parents=True, exist_ok=True) |
|
|
| for model_name, out_name, source_order in MODELS: |
| if args.only and model_name != args.only: |
| continue |
| path = convert(model_name, out_name, source_order, args.outdir) |
| size = sum(f.stat().st_size for f in path.rglob("*") if f.is_file()) |
| print(f"[done] {path} ({size / 1024 / 1024:.1f} MB)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
