| """Converts Alibaba's MossFormer2_SS_16K (ClearerVoice-Studio) to Core ML, then checks the Core ML model against the |
| PyTorch one on the same input. |
| |
| pip install clearvoice coremltools |
| python convert.py |
| |
| Writes MossFormer2_SS_16K.mlpackage. Converted with clearvoice 0.1.2, torch 2.15 and coremltools 9.0 on macOS 26. |
| """ |
| import time |
|
|
| import numpy as np |
| import torch |
| import torch.backends.mps |
|
|
| torch.backends.mps.is_available = lambda: False |
| import coremltools as ct |
| from clearvoice import ClearVoice |
|
|
| |
| |
| WINDOW = 32776 |
|
|
| network = ClearVoice(task="speech_separation", model_names=["MossFormer2_SS_16K"]).models[0].model.eval().cpu() |
|
|
| |
| |
| _pad = torch.nn.functional.pad |
|
|
|
|
| def pad(x, amounts, mode="constant", value=None): |
| if all(a >= 0 for a in amounts): |
| return _pad(x, amounts, mode=mode, value=value) |
| keep = [slice(None)] * x.dim() |
| positive = [] |
| for i in range(0, len(amounts), 2): |
| dim, before, after = x.dim() - 1 - i // 2, amounts[i], amounts[i + 1] |
| keep[dim] = slice(-before if before < 0 else 0, x.shape[dim] + after if after < 0 else x.shape[dim]) |
| positive += [max(before, 0), max(after, 0)] |
| return _pad(x[tuple(keep)], positive, mode=mode, value=value) |
|
|
|
|
| |
| _einsum = torch.einsum |
|
|
|
|
| def einsum(equation, *operands): |
| a, b = operands[0], operands[1] if len(operands) > 1 else None |
| form = equation.replace(" ", "") |
| if form == "...d,hd->...hd": |
| return a.unsqueeze(-2) * b |
| if form == "...id,...jd->...ij": |
| return a @ b.transpose(-1, -2) |
| if form == "...ij,...jd->...id": |
| return a @ b |
| if form == "bgnd,bgne->bde": |
| batch, groups, n, d = a.shape |
| return a.reshape(batch, groups * n, d).transpose(1, 2) @ b.reshape(batch, groups * n, -1) |
| if form == "bgnd,bde->bgne": |
| return a @ b.unsqueeze(1) |
| if form in ("i,j->ij", "...,f->...f"): |
| return a.unsqueeze(-1) * b |
| raise ValueError(f"einsum form not handled: {equation}") |
|
|
|
|
| from clearvoice.models.mossformer2_ss import mossformer2_block |
| import rotary_embedding_torch.rotary_embedding_torch as rotary |
|
|
| torch.nn.functional.pad = pad |
| mossformer2_block.einsum = einsum |
| if hasattr(rotary, "einsum"): |
| rotary.einsum = einsum |
|
|
|
|
| class Separate(torch.nn.Module): |
| """One tensor in and one out: audio [1, WINDOW] -> voices [1, 2, WINDOW].""" |
|
|
| def __init__(self, network): |
| super().__init__() |
| self.network = network |
|
|
| def forward(self, audio): |
| return torch.stack(self.network(audio), dim=1) |
|
|
|
|
| model = Separate(network).eval() |
| example = torch.from_numpy(np.random.default_rng(0).standard_normal((1, WINDOW)).astype(np.float32) * 0.05) |
| with torch.no_grad(): |
| reference = model(example).numpy()[0] |
|
|
| started = time.time() |
| traced = torch.jit.trace(model, example, check_trace=False) |
| mlmodel = ct.convert( |
| traced, |
| inputs=[ct.TensorType(name="audio", shape=(1, WINDOW), dtype=np.float32)], |
| outputs=[ct.TensorType(name="voices", dtype=np.float32)], |
| convert_to="mlprogram", |
| compute_precision=ct.precision.FLOAT16, |
| minimum_deployment_target=ct.target.macOS15, |
| ) |
| mlmodel.short_description = "MossFormer2 (ClearerVoice-Studio) speech separation, 16 kHz, two voices, 2.05 s windows" |
| mlmodel.license = "Apache-2.0" |
| mlmodel.save("MossFormer2_SS_16K.mlpackage") |
| print(f"converted in {time.time() - started:.0f} s") |
|
|
|
|
| def si_sdr(estimate, target): |
| estimate, target = estimate - estimate.mean(), target - target.mean() |
| projection = np.dot(estimate, target) / np.dot(target, target) * target |
| return 10 * np.log10(np.sum(projection ** 2) / np.sum((estimate - projection) ** 2)) |
|
|
|
|
| loaded = ct.models.MLModel("MossFormer2_SS_16K.mlpackage", compute_units=ct.ComputeUnit.CPU_AND_GPU) |
| voices = loaded.predict({"audio": example.numpy()})["voices"][0] |
| for k in range(2): |
| print(f"voice {k + 1}: Core ML vs PyTorch {si_sdr(voices[k], reference[k]):.1f} dB SI-SDR") |
|
|
