Implement BiRefNet and HAT Sharper models with background removal capabilities
- Added validation documentation for AI model tiers. - Introduced `background.py` for background removal logic, including inference functions for alpha and foreground extraction. - Created `onnx_birefnet.py` to handle BiRefNet model sessions and alpha inference. - Developed tests for background models, ensuring proper dispatch and session reuse. - Added tests for BiRefNet ONNX model to validate alpha output and error handling for invalid values. - Implemented model export scripts for BiRefNet and HAT Sharper, ensuring proper weight compression and licensing. - Created `deform_sampling.py` for deformable convolution operations used in BiRefNet. - Added utility functions for compressing ONNX weights without altering graph arithmetic. - Introduced a reference plane node for depth field adjustments in geometry processing.
This commit is contained in:
@@ -63,18 +63,13 @@ MODEL_BUILDS = {
|
||||
),
|
||||
exporter="realesrgan_wdn_x4",
|
||||
),
|
||||
"SAFMN_REAL_X4": BuildSpec(
|
||||
source_url=(
|
||||
"https://github.com/sunny2109/SAFMN/releases/download/"
|
||||
"v0.1.0/SAFMN_L_Real_LSDIR_x4-v2.pth"
|
||||
),
|
||||
source_filename="SAFMN_L_Real_LSDIR_x4-v2.pth",
|
||||
source_size=45002097,
|
||||
source_sha256=(
|
||||
"f1ac0ee3ee143fbbc49aff6584cb7b48b71ba8d1961dfb8ad076b98ca4799ae2"
|
||||
),
|
||||
exporter="safmn_real_x4",
|
||||
),
|
||||
}
|
||||
|
||||
ISOLATED_EXPORTERS = {
|
||||
"MOGE3_VITL": ("moge3_export.py", ()),
|
||||
"HAT_GAN_X4_SHARPER": ("hat_export.py", ()),
|
||||
"BIREFNET_LITE": ("birefnet_export.py", ("--variant", "lite")),
|
||||
"BIREFNET_HR_MATTING": ("birefnet_export.py", ("--variant", "hr-matting")),
|
||||
}
|
||||
|
||||
|
||||
@@ -194,21 +189,24 @@ def prepare_model(model):
|
||||
if model.build:
|
||||
build_model(model)
|
||||
return
|
||||
if model.key in {"PISA_SR", "MOGE3_VITL"}:
|
||||
if model.key in ISOLATED_EXPORTERS:
|
||||
model.destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
exporter = "pisa_sr_export.py" if model.key == "PISA_SR" else "moge3_export.py"
|
||||
exporter, arguments = ISOLATED_EXPORTERS[model.key]
|
||||
with TemporaryDirectory(prefix=f".{model.folder}-", dir=MODELS_DIR) as directory:
|
||||
subprocess.run(
|
||||
["uv", "run", "--script", str(Path(__file__).with_name(exporter)),
|
||||
"--destination", directory],
|
||||
*arguments, "--destination", directory],
|
||||
cwd=PROJECT_ROOT, check=True,
|
||||
)
|
||||
models = [item for item in MODEL_FILES if item.key == model.key]
|
||||
for item in models:
|
||||
if not verify_file(Path(directory) / item.filename, item.size, item.sha256):
|
||||
raise RuntimeError(f"{model.key} export differs from the published catalog")
|
||||
for path in Path(directory).iterdir():
|
||||
path.replace(model.destination.parent / path.name)
|
||||
for item in models:
|
||||
path = Path(directory) / item.filename
|
||||
destination = MODELS_DIR / item.folder / item.filename
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.replace(destination)
|
||||
return
|
||||
if not model.source_url:
|
||||
raise RuntimeError(f"No upstream source or builder for {model.key}")
|
||||
|
||||
Reference in New Issue
Block a user