icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement BiRefNet and HAT Sharper models with background removal capabilities

Browse Source 
- 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:
马南
2026-09-10 20:50:44 +08:00
parent 8ae3c1d5b7
commit e8129fe495
70 changed files with 1890 additions and 1208 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
tests/tools/test_pack.py
View File

@@ -47,7 +47,7 @@ class ExtensionPackagingTest(unittest.TestCase):
records = {}
with tempfile.TemporaryDirectory() as directory:
directory_path = Path(directory)
wheel_path = directory_path / "blendjob-0.1.16-py3-none-any.whl"
wheel_path = directory_path / "blendjob-0.1.17-py3-none-any.whl"
wheel_path.write_bytes(b"wheel")
original_dist = pack.DIST_DIR
pack.DIST_DIR = directory_path / "dist"
@@ -141,7 +141,7 @@ class ExtensionPackagingTest(unittest.TestCase):
self.assertNotIn("server/selection/polygon_mask.py", names)
self.assertNotIn("server/jobs/prepare_geometry_data.py", names)
self.assertNotIn("server/geometry/cutout.py", names)
self.assertIn("server/models/ben2.py", names)
self.assertIn("server/models/background.py", names)
self.assertNotIn("server/models/da3.py", names)
self.assertIn("server/models/moge.py", names)
self.assertNotIn("server/jobs/debug_predictions.py", names)
@@ -150,6 +150,8 @@ class ExtensionPackagingTest(unittest.TestCase):
self.assertIn("server/model_catalog.py", names)
self.assertIn("server/model_download.py", names)
self.assertIn("server/models/onnx_ben2.py", names)
self.assertIn("server/models/onnx_birefnet.py", names)
self.assertNotIn("server/models/onnx_pisa_sr.py", names)
self.assertNotIn("server/models/onnx_da3.py", names)
self.assertIn("server/models/onnx_moge2.py", names)
self.assertIn("server/models/onnx_moge3.py", names)
@@ -159,7 +161,7 @@ class ExtensionPackagingTest(unittest.TestCase):
self.assertIn("operators/cutout_tool/polygon.py", names)
self.assertIn("licenses/scikit-image.txt", names)
self.assertIn(
"wheels/blendjob-0.1.16-py3-none-any.whl",
"wheels/blendjob-0.1.17-py3-none-any.whl",
names,
)
self.assertFalse(any(name.startswith("blendjob/") for name in names))
@@ -177,7 +179,7 @@ class ExtensionPackagingTest(unittest.TestCase):
self.assertEqual(
pack.project_dependencies(),
[
"blendjob==0.1.16",
"blendjob==0.1.17",
"numpy>=1.26,<2.0",
"scipy==1.15.3",
],
@@ -188,7 +190,7 @@ class ExtensionPackagingTest(unittest.TestCase):
wheels = [
Path(name)
for name in (
"blendjob-0.1.16-py3-none-any.whl",
"blendjob-0.1.17-py3-none-any.whl",
"numpy-1.26.4-cp311-cp311-win_amd64.whl",
"scipy-1.15.3-cp311-cp311-win_amd64.whl",
)
@@ -217,12 +219,12 @@ class ExtensionPackagingTest(unittest.TestCase):
if target.name == "3.13":
constraints = Path(command[command.index("--constraint") + 1])
self.assertEqual(
constraints.read_text(), "blendjob==0.1.16"
constraints.read_text(), "blendjob==0.1.17"
)
else:
self.assertNotIn("--constraint", command)
for name in (
"blendjob-0.1.16-py3-none-any.whl",
"blendjob-0.1.17-py3-none-any.whl",
f"scipy-1.15.3-{tag}-{tag}-{config.pip_platform}.whl",
f"numpy-2.1.0-{tag}-{tag}-{config.pip_platform}.whl",
):
@@ -238,7 +240,7 @@ class ExtensionPackagingTest(unittest.TestCase):
("numpy>=1.26,<2.0", "numpy>=2.1,<3.0"),
):
command = call.args[0]
self.assertIn("blendjob==0.1.16", command)
self.assertIn("blendjob==0.1.17", command)
self.assertIn(str(pack.SOURCE_DIR / "wheels"), command)
self.assertIn(f"--python-version={version}", command)
self.assertIn("--implementation=cp", command)
@@ -272,7 +274,7 @@ class ExtensionPackagingTest(unittest.TestCase):
def fake_run(command, check):
target = Path(command[command.index("--dest") + 1])
if target.name == "3.11":
(target / "blendjob-0.1.16-py3-none-any.whl").write_bytes(b"wheel")
(target / "blendjob-0.1.17-py3-none-any.whl").write_bytes(b"wheel")
with patch.object(pack.subprocess, "run", side_effect=fake_run):
with self.assertRaisesRegex(RuntimeError, "Python 3.13"):
@@ -341,7 +343,7 @@ class ExtensionPackagingTest(unittest.TestCase):
target = Path(command[command.index("--dest") + 1])
if target.name == "3.13":
raise RuntimeError("download failed")
(target / "blendjob-0.1.16-py3-none-any.whl").write_bytes(b"wheel")
(target / "blendjob-0.1.17-py3-none-any.whl").write_bytes(b"wheel")
with (
patch.object(pack, "SOURCE_DIR", source),