| |
| """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() |
|
|
