You can deploy this model and start earning money today!

Rohanraje10/
unet-mobile-bald

Name: unet-mobile-bald
Availability: InStock

Views

No views yet

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:

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

Usage

Training

bash

# 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)

python

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

bash

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

iOS Integration

swift

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)

Rohanraje10/unet-mobile-bald

UNet-Mobile: Face → Bald-Face Translation

Model Architecture

Specs

Training

Loss Functions

Optimizer & Schedule

Dataset

Usage

Training

Inference (PyTorch)

CoreML Export

iOS Integration

ANE Compatibility Notes

Files

Citation

Rohanraje10/
unet-mobile-bald