moonbox-ai/BiRefNet_lite-MLX

BiRefNet-lite MLX

MLX builds of BiRefNet-lite, made by Moonbox for Apple Silicon Macs. Each precision is a folder; keep a folder's files together. Only the parts listed below are converted: the rest of the pipeline comes from the original repository unless the section below says otherwise.

Files

The last column is how far each build is from the original PyTorch model on the same inputs (see the section below for what was measured).

Precision	File	Size	sha256	Checked against the original
bf16	bf16/config.json	0.1 kB	7a4673b5400f	MAE <= 0.0002, mask IoU >= 0.9987
bf16	bf16/mlx_model.py	20.1 kB	bd3cb5681703	MAE <= 0.0002, mask IoU >= 0.9987
bf16	bf16/model.safetensors	88.6 MB	7370536fcead	MAE <= 0.0002, mask IoU >= 0.9987
fp16	fp16/config.json	0.1 kB	22466f2eeac0	MAE <= 0.00002, mask IoU >= 0.9998
fp16	fp16/mlx_model.py	20.1 kB	bd3cb5681703	MAE <= 0.00002, mask IoU >= 0.9998
fp16	fp16/model.safetensors	88.6 MB	dedef2f909b6	MAE <= 0.00002, mask IoU >= 0.9998
8bit	8bit/config.json	0.2 kB	f4f927311626	MAE <= 0.0002, mask IoU >= 0.9978
8bit	8bit/mlx_model.py	20.1 kB	bd3cb5681703	MAE <= 0.0002, mask IoU >= 0.9978
8bit	8bit/model.safetensors	63.1 MB	f9f7b19cdb93	MAE <= 0.0002, mask IoU >= 0.9978
6bit	6bit/config.json	0.2 kB	e4cb180964ec	MAE <= 0.0003, mask IoU >= 0.9986
6bit	6bit/mlx_model.py	20.1 kB	bd3cb5681703	MAE <= 0.0003, mask IoU >= 0.9986
6bit	6bit/model.safetensors	56.3 MB	c40decce0ca3	MAE <= 0.0003, mask IoU >= 0.9986
4bit	4bit/config.json	0.2 kB	d3e54b90201b	MAE <= 0.0010, mask IoU >= 0.9845
4bit	4bit/mlx_model.py	20.1 kB	bd3cb5681703	MAE <= 0.0010, mask IoU >= 0.9845
4bit	4bit/model.safetensors	49.4 MB	dde1b2de46ff	MAE <= 0.0010, mask IoU >= 0.9845

How to run

BiRefNet-lite ported to MLX. Each folder holds model.safetensors, config.json and mlx_model.py, the port (only mlx, numpy and pillow needed):

import numpy as np, mlx_model
from PIL import Image
model = mlx_model.load("fp16")
rgb = np.array(Image.open("in.jpg").convert("RGB"))
alpha = mlx_model.matte(model, rgb)                   # HxW uint8; the input is resized to 1024x1024 as in handler.py
Image.fromarray(np.dstack([rgb, alpha])).save("cutout.png")

The port runs the released configuration (Swin-T at full and half resolution, ASPP-deformable decoder) with the BatchNorms folded into the convolutions and torchvision's deform_conv2d written out in MLX. The quantized folders quantize the Swin linears whose input width is a multiple of 64 (group size 64); the convolutions, the first Swin stage and the norms stay bf16, so they save less than the bit width suggests.

Precision	Size	MAE vs PyTorch fp32 (worst image)	Mask IoU (worst)
fp16/	88.7 MB	0.00002	0.9998
bf16/	88.7 MB	0.0002	0.9987
8bit/	63.1 MB	0.0002	0.9978
6bit/	56.3 MB	0.0003	0.9986
4bit/	49.5 MB	0.0010	0.9845

About 0.75 s per image (1024x1024) on an Apple M5 at every precision.

How these were made

From BiRefNet-lite at revision aa62cd87ea: the original weights loaded in MLX and stored as bf16 or fp16, or quantized with MLX's affine quantization (8bit, 6bit, 4bit, group size 64; layers MLX can't quantize, such as convolutions and norms, stay bf16). Every build was checked against the original before upload; the sha256 column above is what the files must hash to.

License

MIT. These files are a derivative of the original model and carry its license: it is copied in this repository, and it governs your use of these files. The original model card follows unchanged.

---

Original model card

Bilateral Reference for High-Resolution Dichotomous Image Segmentation

Peng Zheng ^1,4,5,6,  Dehong Gao ²,  Deng-Ping Fan ^1*,  Li Liu ³,  Jorma Laaksonen ⁴,  Wanli Ouyang ⁵,  Nicu Sebe ⁶

¹ Nankai University  ² Northwestern Polytechnical University  ³ National University of Defense Technology  ⁴ Aalto University  ⁵ Shanghai AI Laboratory  ⁶ University of Trento 

                 

DIS-Sample_1	DIS-Sample_2

This repo is the official implementation of "Bilateral Reference for High-Resolution Dichotomous Image Segmentation" (CAAI AIR 2024).

Visit our GitHub repo: https://github.com/ZhengPeng7/BiRefNet for more details -- codes, docs, and model zoo!

How to use (this tiny version)

0. Install Packages:

pip install -qr https://raw.githubusercontent.com/ZhengPeng7/BiRefNet/main/requirements.txt

1. Load BiRefNet:

Use codes + weights from HuggingFace

Only use the weights on HuggingFace -- Pro: No need to download BiRefNet codes manually; Con: Codes on HuggingFace might not be latest version (I'll try to keep them always latest).

# Load BiRefNet with weights
from transformers import AutoModelForImageSegmentation
birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet_lite', trust_remote_code=True)

Use codes from GitHub + weights from HuggingFace

Only use the weights on HuggingFace -- Pro: codes are always the latest; Con: Need to clone the BiRefNet repo from my GitHub.

# Download codes
git clone https://github.com/ZhengPeng7/BiRefNet.git
cd BiRefNet

# Use codes locally
from models.birefnet import BiRefNet

# Load weights from Hugging Face Models
### >>> Remember to set the `bb` in `config.py` as `swin_v1_t` to use this tiny version. <<< ###
birefnet = BiRefNet.from_pretrained('zhengpeng7/BiRefNet_lite')

Use codes from GitHub + weights from local space

Only use the weights and codes both locally.

# Use codes and weights locally
### >>> Remember to set the `bb` in `config.py` as `swin_v1_t` to use this tiny version. <<< ###
import torch
from utils import check_state_dict

birefnet = BiRefNet(bb_pretrained=False)
state_dict = torch.load(PATH_TO_WEIGHT, map_location='cpu')
state_dict = check_state_dict(state_dict)
birefnet.load_state_dict(state_dict)

Use the loaded BiRefNet for inference

# Imports
from PIL import Image
import matplotlib.pyplot as plt
import torch
from torchvision import transforms
from models.birefnet import BiRefNet

birefnet = ... # -- BiRefNet should be loaded with codes above, either way.
torch.set_float32_matmul_precision(['high', 'highest'][0])
birefnet.to('cuda')
birefnet.eval()
birefnet.half()

def extract_object(birefnet, imagepath):
    # Data settings
    image_size = (1024, 1024)
    transform_image = transforms.Compose([
        transforms.Resize(image_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    image = Image.open(imagepath)
    input_images = transform_image(image).unsqueeze(0).to('cuda').half()

    # Prediction
    with torch.no_grad():
        preds = birefnet(input_images)[-1].sigmoid().cpu()
    pred = preds[0].squeeze()
    pred_pil = transforms.ToPILImage()(pred)
    mask = pred_pil.resize(image.size)
    image.putalpha(mask)
    return image, mask

# Visualization
plt.axis("off")
plt.imshow(extract_object(birefnet, imagepath='PATH-TO-YOUR_IMAGE.jpg')[0])
plt.show()

This BiRefNet for standard dichotomous image segmentation (DIS) is trained on DIS-TR and validated on DIS-TEs and DIS-VD.

This repo holds the official model weights of "Bilateral Reference for High-Resolution Dichotomous Image Segmentation" (CAAI AIR 2024).

This repo contains the weights of BiRefNet proposed in our paper, which has achieved the SOTA performance on three tasks (DIS, HRSOD, and COD).

Go to my GitHub page for BiRefNet codes and the latest updates: https://github.com/ZhengPeng7/BiRefNet :)

Try our online demos for inference:

• Online Single Image Inference on Colab:
• Online Inference with GUI on Hugging Face with adjustable resolutions:
• Inference and evaluation of your given weights:

Acknowledgement:

• Many thanks to @fal for their generous support on GPU resources for training better BiRefNet models.
• Many thanks to @not-lain for his help on the better deployment of our BiRefNet model on HuggingFace.

Citation

@article{zheng2024birefnet,
  title={Bilateral Reference for High-Resolution Dichotomous Image Segmentation},
  author={Zheng, Peng and Gao, Dehong and Fan, Deng-Ping and Liu, Li and Laaksonen, Jorma and Ouyang, Wanli and Sebe, Nicu},
  journal={CAAI Artificial Intelligence Research},
  volume = {3},
  pages = {9150038},
  year={2024}
}