Files

125 lines
4.7 KiB
Python

"""Catalogs and identifiers for LayerForge background-removal models."""
_BIREFNET_REPOSITORY = "ZhengPeng7/BiRefNet"
_BIREFNET_FILENAME = "model.safetensors"
_BIREFNET_DEFAULT_LOCAL_FILENAME = "BiRefNet-general.safetensors"
_BIREFNET_PROJECT_URL = "https://github.com/ZhengPeng7/BiRefNet"
_BIREFNET_REMOTE_PREFIX = "remote:"
_BIREFNET_REQUIRED_KEYS = {
"bb.layers.1.blocks.0.attn.relative_position_index",
"bb.layers.2.blocks.17.attn.qkv.weight",
}
_RMBG_REMOTE_PREFIX = "remote:"
_RMBG_MODEL_CATALOG = (
{
"id": "rmbg_2_0",
"label": "BRIA RMBG 2.0",
"description": (
"Local BRIA background removal with strong general-purpose segmentation. "
"Hugging Face access approval is required before the gated weights can be downloaded."
),
"repo_id": "briaai/RMBG-2.0",
"local_directory": "RMBG-2.0",
"url": "https://huggingface.co/briaai/RMBG-2.0",
"project_url": "https://github.com/Bria-AI/RMBG-2.0",
"backend": "rmbg",
},
)
_BIREFNET_MODEL_CATALOG = (
{
"id": "general",
"label": "BiRefNet — General",
"description": "The best starting point for everyday images and general background removal.",
"repo_id": "ZhengPeng7/BiRefNet",
"filename": "model.safetensors",
"local_filename": "BiRefNet-general.safetensors",
},
{
"id": "high_resolution",
"label": "BiRefNet — High Resolution",
"description": "High-resolution segmentation for detailed edges and larger source images; it uses more memory.",
"repo_id": "ZhengPeng7/BiRefNet_HR",
"filename": "model.safetensors",
"local_filename": "BiRefNet-HR.safetensors",
},
{
"id": "portrait",
"label": "BiRefNet — Portrait",
"description": "Portrait matting for people, hair, and portrait-focused cutouts.",
"repo_id": "ZhengPeng7/BiRefNet-portrait",
"filename": "model.safetensors",
"local_filename": "BiRefNet-portrait.safetensors",
},
{
"id": "matting",
"label": "BiRefNet — Matting",
"description": "General matting with a focus on soft alpha edges such as hair and semi-transparent details.",
"repo_id": "ZhengPeng7/BiRefNet-matting",
"filename": "model.safetensors",
"local_filename": "BiRefNet-matting.safetensors",
},
{
"id": "high_resolution_matting",
"label": "BiRefNet — High Resolution Matting",
"description": "High-resolution general matting for fine details; it requires more memory than the standard matting model.",
"repo_id": "ZhengPeng7/BiRefNet_HR-matting",
"filename": "model.safetensors",
"local_filename": "BiRefNet-HR-matting.safetensors",
},
{
"id": "dynamic",
"label": "BiRefNet — Dynamic",
"description": "Dynamic-shape segmentation for inputs with varying aspect ratios and resolutions.",
"repo_id": "ZhengPeng7/BiRefNet_dynamic",
"filename": "model.safetensors",
"local_filename": "BiRefNet_dynamic.safetensors",
},
{
"id": "dynamic_matting",
"label": "BiRefNet — Dynamic Matting",
"description": "Dynamic-shape matting for arbitrary input sizes, with a focus on soft alpha edges.",
"repo_id": "ZhengPeng7/BiRefNet_dynamic-matting",
"filename": "model.safetensors",
"local_filename": "BiRefNet_dynamic-matting.safetensors",
},
{
"id": "hrsod",
"label": "BiRefNet — HRSOD",
"description": "High-resolution salient-object detection; useful when the main subject should stand out from its surroundings.",
"repo_id": "ZhengPeng7/BiRefNet-HRSOD",
"filename": "model.safetensors",
"local_filename": "BiRefNet-HRSOD.safetensors",
},
{
"id": "dis5k",
"label": "BiRefNet — DIS5K",
"description": "Dichotomous image segmentation trained for clean foreground/background separation.",
"repo_id": "ZhengPeng7/BiRefNet-DIS5K",
"filename": "model.safetensors",
"local_filename": "BiRefNet-DIS5K.safetensors",
},
{
"id": "cod",
"label": "BiRefNet — COD",
"description": "Camouflaged-object detection; use it for subjects that blend into the background.",
"repo_id": "ZhengPeng7/BiRefNet-COD",
"filename": "model.safetensors",
"local_filename": "BiRefNet-COD.safetensors",
},
)
__all__ = [
"_BIREFNET_DEFAULT_LOCAL_FILENAME",
"_BIREFNET_FILENAME",
"_BIREFNET_MODEL_CATALOG",
"_BIREFNET_PROJECT_URL",
"_BIREFNET_REMOTE_PREFIX",
"_BIREFNET_REPOSITORY",
"_BIREFNET_REQUIRED_KEYS",
"_RMBG_MODEL_CATALOG",
"_RMBG_REMOTE_PREFIX",
]