face-parsing-bisenet-bf16

mlx-community/face-parsing-bisenet-bf16

MLX port of the BiSeNet face-parsing model (ResNet18 backbone, 19-class CelebAMask-HQ segmentation — 79999_iter.pth, from zllrunning/face-parsing.PyTorch, as redistributed by TMElyralab/MuseTalk in face-parse-bisent/). Used by MuseTalk to build the blend mask when compositing a regenerated mouth back into a frame.

Converted from PyTorch to MLX-native bf16 safetensors. Weights only — the model definition lives in the Swift consumer (xocialize/musetalk-mlx-swift, MuseTalk.BiSeNet).

Parity

Versus the PyTorch reference (CPU): argmax (parse map) agreement 100.0000% on a fixed input; feat_out logit max_abs 8.7e-2 at bf16 (8.3e-6 at fp32) — the logit noise is washed out by the argmax, so the segmentation is identical. bf16 chosen because the output is a class argmax (robust) and it halves the footprint (~27 MB).

Layout

model.safetensors   160 tensors, bf16, MLX-native keys
config.json         architecture / n_classes (19) / input 512 / key convention

Keys are MLX-native: the diffusers/torch downsample.0/.1 Sequential is renamed to downsample.conv/.bn, num_batches_tracked dropped, conv weights transposed (O,I,H,W)→(O,H,W,I). The 19 classes follow the CelebAMask-HQ label order (0 background; 1 skin; … 11–13 lips/mouth; …).

Use with MLX (Swift)

import MuseTalk   // xocialize/musetalk-mlx-swift
let net = BiSeNet()
net.train(false)  // BatchNorm: running stats
try MuseTalkWeights.load(net, from: weightsURL)   // model.safetensors
let logits = net(faceNCHW)   // (B, 19, H, W); argmax over axis 1 -> parse map

License

MIT (mirrors upstream face-parsing.PyTorch / MuseTalk).