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:
@@ -93,7 +93,7 @@ class ModelInferenceTest(unittest.TestCase):
|
||||
|
||||
upscale = load_server_module("models.upscale")
|
||||
|
||||
@pytest.mark.parametrize("key", ["REALESRGAN_GENERAL_WDN_X4V3", "SAFMN_REAL_X4", "REALESRGAN_X4PLUS"])
|
||||
@pytest.mark.parametrize("key", ["REALESRGAN_GENERAL_WDN_X4V3", "HAT_GAN_X4_SHARPER", "REALESRGAN_X4PLUS"])
|
||||
def test_upscale_models_resize_prediction_then_restore_alpha(tmp_path, key):
|
||||
rng = np.random.default_rng(1)
|
||||
source = Image.fromarray(rng.integers(0,256,(5,7,4),dtype=np.uint8))
|
||||
@@ -101,7 +101,8 @@ def test_upscale_models_resize_prediction_then_restore_alpha(tmp_path, key):
|
||||
source.save(path)
|
||||
prediction = Image.fromarray(rng.integers(0,256,(20,28,3),dtype=np.uint8))
|
||||
manager = SimpleNamespace(directory=lambda _: tmp_path, get_upscale=lambda *args: object())
|
||||
with patch("server.models.onnx_upscale.infer", return_value=prediction):
|
||||
with patch("server.models.onnx_upscale.infer", return_value=prediction) as infer:
|
||||
output = upscale.infer_one(manager, {"model":key,"device":"cpu"}, path)
|
||||
assert infer.call_args.kwargs["tile_border"] == (48 if key == "HAT_GAN_X4_SHARPER" else 16)
|
||||
np.testing.assert_array_equal(output.convert("RGB"), prediction.resize((14,10), Image.Resampling.LANCZOS))
|
||||
np.testing.assert_array_equal(output.getchannel("A"), source.getchannel("A").resize((14,10), Image.Resampling.LANCZOS))
|
||||
|
||||
Reference in New Issue
Block a user