#!/usr/bin/env python3 """Convert PalabraAI ReDimNet2-B6 speaker embedding model to Core ML. The generated model accepts one mono 16 kHz waveform with a fixed sample count and returns an L2-normalized speaker embedding. The Swift pipeline pads shorter chunks and center-crops longer chunks to this length before inference. """ from __future__ import annotations import argparse from pathlib import Path import coremltools as ct import numpy as np import torch import torch.nn.functional as F class NormalizedEmbedding(torch.nn.Module): def __init__(self, model: torch.nn.Module) -> None: super().__init__() self.model = model def forward(self, audio: torch.Tensor) -> torch.Tensor: embedding = self.model(audio) return torch.nn.functional.normalize(embedding, p=2, dim=1) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Convert ReDimNet2-B6 to Core ML") parser.add_argument("--output", default="Models/speaker/ReDimNet2-B6.mlpackage") parser.add_argument("--model-name", default="b6", choices=["b0", "b1", "b2", "b3", "b4", "b5", "b6"]) parser.add_argument("--train-type", default="lm", choices=["ft_lm", "ptn", "lm"]) parser.add_argument("--dataset", default="vb2+vox2_v0") parser.add_argument("--seconds", type=float) parser.add_argument("--sample-count", type=int, default=160320, help="Default is 10.02s at 16 kHz, aligned to ReDimNet2 frame stride") parser.add_argument("--sample-rate", type=int, default=16000) parser.add_argument("--minimum-deployment-target", default="macOS14") return parser.parse_args() def deployment_target(name: str) -> ct.target: try: return getattr(ct.target, name) except AttributeError as error: available = ", ".join(sorted(k for k in dir(ct.target) if k.startswith("macOS"))) raise SystemExit(f"Unknown Core ML deployment target {name!r}. Available macOS targets: {available}") from error def main() -> None: args = parse_args() output = Path(args.output) output.parent.mkdir(parents=True, exist_ok=True) sample_count = int(round(args.seconds * args.sample_rate)) if args.seconds is not None else args.sample_count if sample_count <= 0: raise SystemExit("--seconds must produce a positive sample count") print( "Loading ReDimNet2 from PalabraAI/redimnet2 " f"model={args.model_name} train_type={args.train_type} dataset={args.dataset}" ) torch.set_grad_enabled(False) model = torch.hub.load( "PalabraAI/redimnet2", "redimnet2", model_name=args.model_name, train_type=args.train_type, dataset=args.dataset, pretrained=True, trust_repo=True, ).eval() install_coreml_trace_patches() wrapped = NormalizedEmbedding(model).eval() example = torch.zeros(1, sample_count, dtype=torch.float32) set_fixed_aligned_frames(model, example) print(f"Tracing fixed waveform input: [1, {sample_count}]") traced = torch.jit.trace(wrapped, example, strict=False) traced = torch.jit.freeze(traced.eval()) traced = torch.jit.optimize_for_inference(traced) print("Converting to Core ML") mlmodel = ct.convert( traced, convert_to="mlprogram", inputs=[ ct.TensorType( name="audio", shape=example.shape, dtype=np.float32, ) ], outputs=[ct.TensorType(name="embedding")], minimum_deployment_target=deployment_target(args.minimum_deployment_target), compute_precision=ct.precision.FLOAT16, ) mlmodel.short_description = "ReDimNet2-B6 speaker embedding model converted from PalabraAI/redimnet2." mlmodel.input_description["audio"] = ( f"Mono 16 kHz waveform padded/cropped to {sample_count} samples." ) mlmodel.output_description["embedding"] = "L2-normalized speaker embedding." mlmodel.user_defined_metadata["source"] = "https://github.com/PalabraAI/redimnet2" mlmodel.user_defined_metadata["model_name"] = args.model_name mlmodel.user_defined_metadata["train_type"] = args.train_type mlmodel.user_defined_metadata["dataset"] = args.dataset mlmodel.user_defined_metadata["sample_rate"] = str(args.sample_rate) mlmodel.user_defined_metadata["sample_count"] = str(sample_count) mlmodel.save(output) print(f"Wrote {output}") def install_coreml_trace_patches() -> None: from redimnet2 import redimnet2 as redimnet2_module from redimnet2.layers import attention as attention_module from redimnet2.layers import redim_structural def to1d_forward(self, x: torch.Tensor) -> torch.Tensor: return torch.flatten(x.permute(0, 2, 1, 3), start_dim=1, end_dim=2) def redimnet2_forward(self, inp): if not self.is_subnet: aligned_frames = getattr(self, "_coreml_aligned_frames", None) if aligned_frames is None: aligned_frames = (inp.shape[-1] // self.time_stride) * self.time_stride inp = inp[:, :, :, :aligned_frames] x = self.stem(inp) if self.agg_gnorm: x = self.stem_gnorm(x) outputs_1d = [x] else: outputs_1d = list(inp) x = self.stem(inp) if self.agg_gnorm: x = self.stem_gnorm(x) outputs_1d.append(x) for stage_ind in range(self.num_stages): outputs_1d.extend(self.run_stage(outputs_1d, stage_ind)) x = self.fin_wght1d(outputs_1d) outputs_1d.append(x) x = self.fin_to2d(x) x = self.head(x) if self.return_all_outputs: return x, outputs_1d return x def wrap_forward(self, x: torch.Tensor) -> torch.Tensor: if self.pad_right_samples is not None: x = torch.nn.functional.pad(x, (0, self.pad_right_samples), mode="constant", value=None) x = self.spec(x) if x.ndim == 3: x = x.unsqueeze(1) if self.return_all_outputs: out, all_outs_1d = self.backbone(x) else: out = self.backbone(x) if out.ndim == 4: out = torch.flatten(out, start_dim=1, end_dim=2) if self.before_pool_offset is not None: out = out[:, :, self.before_pool_offset:] out = self.bn(self.pool(out)) out = self.linear(out) if self.bn2 is not None: out = self.bn2(out) if self.return_all_outputs: return out, all_outs_1d return out def attention_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: bsz, tgt_len, _ = hidden_states.size() def shape(tensor: torch.Tensor, seq_len: int) -> torch.Tensor: return tensor.reshape(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() query_states = self.q_proj(hidden_states) * self.scaling key_states = shape(self.k_proj(hidden_states), -1) value_states = shape(self.v_proj(hidden_states), -1) query_states = shape(query_states, tgt_len) query_states = torch.flatten(query_states, start_dim=0, end_dim=1) key_states = torch.flatten(key_states, start_dim=0, end_dim=1) value_states = torch.flatten(value_states, start_dim=0, end_dim=1) attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) attn_weights = F.softmax(attn_weights, dim=-1) attn_probs = F.dropout(attn_weights, p=self.dropout, training=self.training) attn_output = torch.bmm(attn_probs, value_states) attn_output = attn_output.reshape(bsz, self.num_heads, tgt_len, self.head_dim) attn_output = attn_output.transpose(1, 2) attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim) return self.out_proj(attn_output) redim_structural.to1d.forward = to1d_forward redimnet2_module.ReDimNet2.forward = redimnet2_forward redimnet2_module.ReDimNet2Wrap.forward = wrap_forward attention_module.MultiHeadAttention.forward = attention_forward def set_fixed_aligned_frames(model: torch.nn.Module, example: torch.Tensor) -> None: with torch.no_grad(): spec = model.spec(example) frames = int(spec.shape[-1]) model.backbone._coreml_aligned_frames = (frames // model.backbone.time_stride) * model.backbone.time_stride if __name__ == "__main__": main()