Rohanraje10/unet-mobile-bald

UNet-Mobile: Face → Bald-Face Translation

MobileNetV3-Small encoder + Depthwise-Separable UNet decoder
Designed for real-time on-device inference on Apple Neural Engine (A17+).

Model Architecture

Input: [B, 4, 256, 256] — RGB face crop + binary hair mask
                                    │
                    ┌───────────────┤ MobileNetV3-Small Encoder
                    │               │ (pretrained ImageNet)
    s0 [16,128,128] ├── stage0 ────┤ stride-2 stem conv
    s1 [16, 64, 64] ├── stage1 ────┤ InvertedResidual
    s2 [24, 32, 32] ├── stage2 ────┤ InvertedResidual ×2
    s3 [40, 16, 16] ├── stage3 ────┤ InvertedResidual ×3
    s4 [96,  8,  8] └── stage4 ────┘ InvertedResidual ×5 (bottleneck)
                    │
                    │    DWS Decoder (bilinear↑ + DW conv + PW conv)
    d3 [64, 16, 16] ├── dec3(s4+s3)
    d2 [48, 32, 32] ├── dec2(d3+s2)
    d1 [32, 64, 64] ├── dec1(d2+s1)
    d0 [24,128,128] ├── dec0(d1+s0)
                    │
Output: [B, 3, 256, 256] — RGB bald face (via bilinear↑ + 1×1 conv + Tanh)

Specs

Metric	Value
Parameters	891K
fp16 size	1.7 MB
CoreML package	~3-4 MB (6-bit palettized)
Target latency	<40 ms on A17 ANE
Input	4-ch (RGB + hair mask), 256×256
Output	3-ch RGB, 256×256

Training

Loss Functions

• L1 reconstruction (λ=10.0) — pixel-level fidelity
• LPIPS perceptual (λ=5.0, AlexNet) — perceptual quality
• PatchGAN adversarial (λ=1.0, hinge loss) — sharpness & realism
• ArcFace identity (λ=10.0, cosine similarity) — preserve face identity
• R1 gradient penalty (γ=10.0, lazy every 16 steps) — discriminator regularization

Optimizer & Schedule

• AdamW β₁=0.5, β₂=0.999 (GAN-standard)
• TTUR: Generator lr=2e-4, Discriminator lr=4e-4
• Cosine annealing with 5-epoch linear warmup → η_min = lr/100
• 200 epochs, 8×A100-80GB, fp16 mixed precision

Dataset

~80K pairs from FFHQ + CelebA-HQ:

• Input: Face crop (256×256 RGB) + hair segmentation mask (256×256 binary)
• Target: LaMa-inpainted bald face (256×256 RGB)

Generated via gen_bald_pairs.py pipeline:

1. Face parsing → hair region mask (BiSeNet)
2. LaMa inpainting → bald face generation
3. Quality filtering → final pairs

Usage

Training

# Install dependencies
pip install torch torchvision accelerate lpips onnxruntime-gpu trackio huggingface_hub

# Single GPU (debug)
python train.py --data_dir /path/to/data --batch_size 16 --num_epochs 2

# 8×A100 distributed
bash launch_train.sh

Inference (PyTorch)

import torch
from model import UNetMobile

model = UNetMobile(in_channels=4, out_channels=3, pretrained=False)
state_dict = torch.load('generator_latest.pt', map_location='cpu')
model.load_state_dict(state_dict)
model.eval()

# Input: [1, 4, 256, 256] — RGB in [-1,1] + mask in [0,1]
input_4ch = torch.cat([face_rgb_normalized, hair_mask], dim=1)
with torch.no_grad():
    bald_face = model(input_4ch)  # [-1, 1]
    bald_face = (bald_face + 1) / 2  # → [0, 1]

CoreML Export

python export_coreml.py \
    --checkpoint generator_latest.pt \
    --output UNetMobileBald.mlpackage \
    --compress 6bit

iOS Integration

let config = MLModelConfiguration()
config.computeUnits = .all  // CPU + GPU + ANE
let model = try UNetMobileBald(configuration: config)

// Input: MLMultiArray [1, 4, 256, 256], pixel values 0-255
let prediction = try model.prediction(input_image: inputArray)
let output = prediction.output_image  // [1, 3, 256, 256], values 0-255

ANE Compatibility Notes

• All activations: Hardswish (ANE-native, hardware-accelerated)
• Upsampling: bilinear (no transposed convolutions)
• Normalization: BatchNorm (folds into conv at export)
• No GELU, SiLU, LayerNorm, or dynamic shapes
• Fixed spatial dims: multiples of 16 for optimal ANE tiling

Files

File	Description
model.py	UNet-Mobile architecture
losses.py	L1 + LPIPS + PatchGAN + ArcFace losses
dataset.py	Dataset loader for bald-pair data
train.py	Full training script (Accelerate DDP)
export_coreml.py	PyTorch → CoreML mlprogram conversion
launch_train.sh	8×A100 launch script
accelerate_config.yaml	Accelerate distributed config

Citation

Based on:

• MobileNetV3 — Howard et al., ICCV 2019 (arXiv:1905.02244)
• SimSwap loss recipe — Chen et al., ACM MM 2020 (arXiv:2106.06340)
• FaceDancer loss weights — Rosberg et al., WACV 2023 (arXiv:2210.10473)
• LPIPS perceptual loss — Zhang et al., CVPR 2018 (arXiv:1801.03924)
• ArcFace identity — Deng et al., CVPR 2019 (arXiv:1801.07698)