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

12
tests/anyimage/operators/test_ai_image_operators.py
View File

@@ -22,7 +22,7 @@ class AiImageOperatorsTest(BlenderTestCase):
operator = operator_type()
operator.input_path = "input.png"
operator._image_target = SimpleNamespace(start_frame=4)
operator.model = "SAFMN_REAL_X4"
operator.model = "HAT_GAN_X4_SHARPER"
operator.max_input_size = 1024
context = SimpleNamespace(
scene=SimpleNamespace(anyimage_settings=SimpleNamespace(video_frames=7))
@@ -37,12 +37,12 @@ class AiImageOperatorsTest(BlenderTestCase):
):
request = operator.request(context)
require_model.assert_called_once_with("SAFMN_REAL_X4")
require_model.assert_called_once_with("HAT_GAN_X4_SHARPER")
self.assertEqual(
request,
{
"input": "input.png",
"model": "SAFMN_REAL_X4",
"model": "HAT_GAN_X4_SHARPER",
"device": "cuda",
"video_frames": 7,
"start_frame": 4,
@@ -108,12 +108,12 @@ class AiImageOperatorsTest(BlenderTestCase):
result = SimpleNamespace(directory=Path("result"))
operator = self.upscale.UpscaleImage()
operator._image_target = Mock()
operator.model = "SAFMN_REAL_X4"
operator.model = "HAT_GAN_X4_SHARPER"
with (
patch.object(
self.upscale,
"upscale_model_label",
return_value="SAFMN Real x4",
return_value="HAT Sharper",
),
):
message = operator.response(None, result)
@@ -122,4 +122,4 @@ class AiImageOperatorsTest(BlenderTestCase):
Path("result"),
frame_prefix="upscale.",
)
self.assertEqual(message, "Image is upscaled with SAFMN Real x4")
self.assertEqual(message, "Image is upscaled with HAT Sharper")

2
tests/anyimage/operators/test_ai_input_failures.py
View File

@@ -30,7 +30,7 @@ class AiInputFailuresTest(BlenderTestCase):
stack.enter_context(patch.object(self.runtime.runtime, "active_job", None))
submit = stack.enter_context(patch.object(self.runtime.runtime, "begin_job"))
if module == self.upscale:
stack.enter_context(patch.object(module, "configured_upscale_model", return_value="SAFMN_REAL_X4"))
stack.enter_context(patch.object(module, "configured_upscale_model", return_value="HAT_GAN_X4_SHARPER"))
stack.enter_context(patch.object(module, "configured_max_ai_input_size", return_value=2048))
self.assertEqual(operator.execute(context), {"CANCELLED"})
submit.assert_not_called()

29
tests/anyimage/operators/test_ai_setup.py
View File

@@ -31,7 +31,7 @@ class AiSetupTest(BlenderTestCase):
calls,
[
("download-model", {"model": "MOGE2_VITS_NORMAL"}),
("download-model", {"model": "BEN2_BASE"}),
("download-model", {"model": "BIREFNET_LITE"}),
(
"download-model",
{"model": "REALESRGAN_GENERAL_WDN_X4V3"},
@@ -50,9 +50,9 @@ class AiSetupTest(BlenderTestCase):
with self.assertRaisesRegex(
RuntimeError,
(
"MoGe-2 ViT-S Normal: network unavailable; "
"BEN2 Base: network unavailable; "
"Real-ESRGAN General WDN x4v3: network unavailable"
"Fast \\| MoGe-2 ViT-S Normal: network unavailable; "
"Fast \\| BiRefNet Lite: network unavailable; "
"Fast \\| Real-ESRGAN General WDN x4v3: network unavailable"
),
):
self.runtime.post_install(SimpleNamespace(request=request))
@@ -61,7 +61,7 @@ class AiSetupTest(BlenderTestCase):
calls,
[
"MOGE2_VITS_NORMAL",
"BEN2_BASE",
"BIREFNET_LITE",
"REALESRGAN_GENERAL_WDN_X4V3",
],
)
@@ -77,7 +77,7 @@ class AiSetupTest(BlenderTestCase):
{
"models": [
"MOGE2_VITS_NORMAL",
"BEN2_BASE",
"BIREFNET_LITE",
"REALESRGAN_GENERAL_WDN_X4V3",
]
},
@@ -93,7 +93,7 @@ class AiSetupTest(BlenderTestCase):
"model_catalog",
return_value=(
{"key": "MOGE2_VITS_NORMAL", "ready": True},
{"key": "BEN2_BASE", "ready": False},
{"key": "BIREFNET_LITE", "ready": False},
{
"key": "REALESRGAN_GENERAL_WDN_X4V3",
"ready": True,
@@ -103,7 +103,7 @@ class AiSetupTest(BlenderTestCase):
):
status = properties.ai_status()
self.assertTrue(status["environment_ready"])
self.assertEqual(status["missing_models"], ("BEN2_BASE",))
self.assertEqual(status["missing_models"], ("BIREFNET_LITE",))
self.assertFalse(status["ready"])
self.assertEqual(properties.ai_setup_label(status), "Download Required Models…")
@@ -114,7 +114,7 @@ class AiSetupTest(BlenderTestCase):
"model_catalog",
return_value=tuple(
{"key": key, "ready": True}
for key in properties.REQUIRED_MODEL_KEYS
for key in properties.shared_model_catalog.DEFAULT_MODEL_KEYS
),
),
):
@@ -148,7 +148,7 @@ class AiSetupTest(BlenderTestCase):
"ai_status",
return_value={
"environment_ready": False,
"missing_models": ("MOGE2_VITS_NORMAL", "BEN2_BASE"),
"missing_models": ("MOGE2_VITS_NORMAL", "BIREFNET_LITE"),
"ready": False,
},
):
@@ -159,7 +159,7 @@ class AiSetupTest(BlenderTestCase):
"ai_status",
return_value={
"environment_ready": True,
"missing_models": ("BEN2_BASE",),
"missing_models": ("BIREFNET_LITE",),
"ready": False,
},
):
@@ -185,13 +185,14 @@ class AiSetupTest(BlenderTestCase):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
moge = root / "models" / "moge-2-vits-normal-onnx"
ben2 = root / "models" / "BEN2-ONNX" / "onnx"
lite = root / "models" / "birefnet-lite"
wdn = root / "models" / "realesr-general-wdn-x4v3"
moge.mkdir(parents=True)
ben2.mkdir(parents=True)
lite.mkdir(parents=True)
wdn.mkdir(parents=True)
(moge / "model.onnx").touch()
(ben2 / "model_fp16.onnx").touch()
(lite / "model.onnx").touch()
(lite / "LICENSE.txt").touch()
(wdn / "realesr-general-wdn-x4v3.onnx").touch()
with (

8
tests/anyimage/operators/test_cutout_operators.py
View File

@@ -249,7 +249,7 @@ class CutoutOperatorsTest(BlenderTestCase):
operator.color_path = __file__
operator.source_object_name = "Source"
operator.shape = "DEPTH_SOLID"
operator.model = self.anyimage.properties.DEFAULT_MOGE2_MODEL_KEY
operator.model = self.anyimage.preferences.DEFAULT_GEOMETRY_MODEL_KEY
operator.resolution_level = 5
operator.max_ai_input_size = 2048
operator.generate_normal = True
@@ -288,10 +288,8 @@ class CutoutOperatorsTest(BlenderTestCase):
geometry_model="MOGE2_VITB_NORMAL",
geometry_resolution_level=7,
)
with patch.object(
self.cutout_main,
"addon_preferences",
return_value=preferences,
with patch.object(self.cutout_main, "addon_preferences", return_value=preferences), patch.object(
self.anyimage.preferences, "addon_preferences", return_value=preferences
):
settings = self.cutout_main.cutout_geometry_settings(None)

3
tests/anyimage/operators/test_panorama_conversion.py
View File

@@ -43,11 +43,12 @@ def test_entry_submits_current_alpha_preferences_and_cleans_cancelled_input(sour
return {"CANCELLED"}
monkeypatch.setattr(operators, "invoke_ai_setup_if_needed", lambda _operator: None)
monkeypatch.setattr(operators, "addon_preferences", lambda _: SimpleNamespace(geometry_model="MOGE2_VITB", geometry_resolution_level=2))
monkeypatch.setattr(operators, "configured_geometry_model", lambda _: "MOGE2_VITB_NORMAL")
monkeypatch.setattr(operators, "configured_max_ai_input_size", lambda _: 1024)
monkeypatch.setattr(operators, "bpy", SimpleNamespace(ops=SimpleNamespace(anyimage=SimpleNamespace(generate_panorama=submit))))
operator = SimpleNamespace(mesh_detail=4, report=Mock())
assert operators.ConvertToPanorama.execute(operator, bpy.context) == {"CANCELLED"}
assert submitted["model"] == "MOGE2_VITB" and submitted["resolution_level"] == 2
assert submitted["model"] == "MOGE2_VITB_NORMAL" and submitted["resolution_level"] == 2
assert submitted["max_input_size"] == 1024 and submitted["mesh_detail"] == 4
assert submitted["color_path"] == submitted["input_path"]
from pathlib import Path

6
tests/anyimage/operators/test_texture_node_actions.py
View File

@@ -113,7 +113,7 @@ def test_texture_action_undo_redo_restores_image_and_binding(texture, tmp_path,
node = bpy.data.materials["Texture material"].node_tree.nodes["Source texture"]
operator._image_target = ImageEditTarget.capture(texture_context(node))
if operation == upscale.UpscaleImage:
operator.model = "PISA_SR"
operator.model = "HAT_GAN_X4_SHARPER"
operator.response(context, SimpleNamespace(directory=tmp_path))
return {"FINISHED"}
@@ -175,7 +175,7 @@ def test_texture_job_request_and_completion(texture, tmp_path, operation, module
stack.enter_context(patch.object(runtime, "update_ui"))
stack.enter_context(patch.object(runtime, "redraw_ui"))
if operation == upscale.UpscaleImage:
stack.enter_context(patch.object(upscale, "configured_upscale_model", return_value="SAFMN_REAL_X4"))
stack.enter_context(patch.object(upscale, "configured_upscale_model", return_value="HAT_GAN_X4_SHARPER"))
stack.enter_context(patch.object(upscale, "configured_max_ai_input_size", return_value=2048))
assert original_execute(operator, context) == {"RUNNING_MODAL"}
job = runtime.active_job
@@ -185,7 +185,7 @@ def test_texture_job_request_and_completion(texture, tmp_path, operation, module
assert kind == operation.job_type
assert parameters["start_frame"] == 1
assert parameters["device"] == "cpu"
assert parameters["model"] == ("SAFMN_REAL_X4" if operation == upscale.UpscaleImage else "BEN2_BASE")
assert parameters["model"] == ("HAT_GAN_X4_SHARPER" if operation == upscale.UpscaleImage else "BIREFNET_LITE")
input_path = operator.input_path
assert operator.delete_input
# Selection changes must not redirect completion.

2
tests/anyimage/operators/test_upscale_undo.py
View File

@@ -23,7 +23,7 @@ def test_upscale_result_replacement_undo_restores_source(tmp_path):
def execute(operator, context):
operator._image_target = ImageEditTarget.capture(SimpleNamespace(object=owner))
operator.model = "PISA_SR"
operator.model = "HAT_GAN_X4_SHARPER"
operator.response(context, SimpleNamespace(directory=tmp_path))
return {"FINISHED"}