Duplicate from aufklarer/WeSpeaker-ResNet34-LM-MLX

Co-authored-by: Ivan <aufklarer@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,91 @@
+---
+license: mit
+tags:
+  - mlx
+  - speaker-embedding
+  - speaker-verification
+  - speaker-diarization
+  - wespeaker
+  - resnet
+  - apple-silicon
+base_model: pyannote/wespeaker-voxceleb-resnet34-LM
+library_name: mlx
+pipeline_tag: audio-classification
+---
+
+# WeSpeaker ResNet34-LM — MLX
+
+MLX-compatible weights for [WeSpeaker ResNet34-LM](https://huggingface.co/pyannote/wespeaker-voxceleb-resnet34-LM), converted from the pyannote speaker embedding model with BatchNorm fused into Conv2d.
+
+## Model
+
+WeSpeaker ResNet34-LM is a speaker embedding model (~6.6M params) that produces 256-dimensional L2-normalized speaker embeddings from audio. Trained on VoxCeleb for speaker verification and diarization.
+
+**Architecture:**
+
+```
+Input: [B, T, 80, 1] log-mel spectrogram (80 fbank, 16kHz)
+  │
+  ├─ Conv2d(1→32, k=3, p=1) + ReLU
+  ├─ Layer1: 3× BasicBlock(32→32)
+  ├─ Layer2: 4× BasicBlock(32→64, stride=2)
+  ├─ Layer3: 6× BasicBlock(64→128, stride=2)
+  ├─ Layer4: 3× BasicBlock(128→256, stride=2)
+  │
+  ├─ Statistics Pooling: mean + std → [B, 5120]
+  ├─ Linear(5120→256) → L2 normalize
+  │
+  Output: [B, 256] speaker embedding
+```
+
+BatchNorm is fused into Conv2d at conversion time — no BN layers in the MLX model.
+
+## Usage (Swift / MLX)
+
+```swift
+import SpeechVAD
+
+// Speaker embedding
+let model = try await WeSpeakerModel.fromPretrained()
+let embedding = model.embed(audio: samples, sampleRate: 16000)
+// embedding: [Float] of length 256, L2-normalized
+
+// Compare speakers
+let similarity = WeSpeakerModel.cosineSimilarity(embeddingA, embeddingB)
+
+// Full speaker diarization pipeline
+let pipeline = try await DiarizationPipeline.fromPretrained()
+let result = pipeline.diarize(audio: samples, sampleRate: 16000)
+for seg in result.segments {
+    print("Speaker \(seg.speakerId): \(seg.startTime)s - \(seg.endTime)s")
+}
+```
+
+Part of [qwen3-asr-swift](https://github.com/ivan-digital/qwen3-asr-swift).
+
+## Conversion
+
+```bash
+python3 scripts/convert_wespeaker.py --upload
+```
+
+Converts the original pyannote/wespeaker-voxceleb-resnet34-LM checkpoint using a custom unpickler (no pyannote.audio dependency required). Key transformations:
+
+- **Fuse BatchNorm** into Conv2d: `w_fused = w × γ/√(σ²+ε)`, `b_fused = β − μ×γ/√(σ²+ε)`
+- **Transpose Conv2d** weights: `[O, I, H, W]` → `[O, H, W, I]` for MLX channels-last
+- **Rename**: strip `resnet.` prefix, `seg_1` → `embedding`
+- **Drop** `num_batches_tracked` keys
+
+## Weight Mapping
+
+| PyTorch Key | MLX Key | Shape |
+|-------------|---------|-------|
+| `resnet.conv1.weight` + `resnet.bn1.*` | `conv1.weight` | [32, 3, 3, 1] |
+| `resnet.layer{L}.{B}.conv{1,2}.weight` + `bn{1,2}.*` | `layer{L}.{B}.conv{1,2}.weight` | [O, 3, 3, I] |
+| `resnet.layer{L}.0.shortcut.0.weight` + `shortcut.1.*` | `layer{L}.0.shortcut.weight` | [O, 1, 1, I] |
+| `resnet.seg_1.weight` | `embedding.weight` | [256, 5120] |
+| `resnet.seg_1.bias` | `embedding.bias` | [256] |
+
+## License
+
+The original WeSpeaker model is released under the [MIT License](https://github.com/wenet-e2e/wespeaker/blob/master/LICENSE).

config.json

@@ -0,0 +1,19 @@
+{
+  "model_type": "wespeaker-resnet34-lm",
+  "sample_rate": 16000,
+  "n_mels": 80,
+  "embedding_dim": 256,
+  "layers": [
+    3,
+    4,
+    6,
+    3
+  ],
+  "channels": [
+    32,
+    64,
+    128,
+    256
+  ],
+  "pooling_output_dim": 5120
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f56204883f2de969f584af7893e5373575556b422190c14c206a5fd94f3d7fe6
+size 26526952