nan
e8129fe495
- 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.
126 lines
4.3 KiB
Python
126 lines
4.3 KiB
Python
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock, patch
|
|
from tests.support.blender import BlenderTestCase
|
|
|
|
|
|
class AiImageOperatorsTest(BlenderTestCase):
|
|
def test_upscale_operator_defaults_to_wdn_and_sends_selected_model(self):
|
|
operator_type = self.upscale.UpscaleImage
|
|
self.assertIn(operator_type, self.anyimage.CLASSES)
|
|
self.assertEqual(
|
|
operator_type.__annotations__["model"]["default"],
|
|
"REALESRGAN_GENERAL_WDN_X4V3",
|
|
)
|
|
self.assertEqual(
|
|
operator_type.__annotations__["max_input_size"]["default"],
|
|
2048,
|
|
)
|
|
|
|
self.assertIn(self.runtime.JobOperatorBase, operator_type.__mro__)
|
|
self.assertEqual(operator_type.bl_options, {"UNDO"})
|
|
operator = operator_type()
|
|
operator.input_path = "input.png"
|
|
operator._image_target = SimpleNamespace(start_frame=4)
|
|
operator.model = "HAT_GAN_X4_SHARPER"
|
|
operator.max_input_size = 1024
|
|
context = SimpleNamespace(
|
|
scene=SimpleNamespace(anyimage_settings=SimpleNamespace(video_frames=7))
|
|
)
|
|
with (
|
|
patch.object(self.upscale, "require_environment"),
|
|
patch.object(self.upscale, "require_model") as require_model,
|
|
patch.object(
|
|
self.upscale, "require_input_path", return_value=Path("input.png")
|
|
),
|
|
patch.object(self.upscale, "production_device", return_value="CUDA"),
|
|
):
|
|
request = operator.request(context)
|
|
|
|
require_model.assert_called_once_with("HAT_GAN_X4_SHARPER")
|
|
self.assertEqual(
|
|
request,
|
|
{
|
|
"input": "input.png",
|
|
"model": "HAT_GAN_X4_SHARPER",
|
|
"device": "cuda",
|
|
"video_frames": 7,
|
|
"start_frame": 4,
|
|
"max_input_size": 1024,
|
|
},
|
|
)
|
|
|
|
|
|
def test_upscale_operator_opens_setup_when_required_ai_is_missing(self):
|
|
expected = {"CANCELLED"}
|
|
operator = self.upscale.UpscaleImage()
|
|
with patch.object(
|
|
self.upscale,
|
|
"invoke_ai_setup_if_needed",
|
|
return_value=expected,
|
|
) as setup:
|
|
result = operator.invoke(None, None)
|
|
|
|
setup.assert_called_once_with(operator)
|
|
self.assertIs(result, expected)
|
|
|
|
|
|
def test_upscale_operator_runs_immediately_when_ai_is_ready(self):
|
|
operator = self.upscale.UpscaleImage()
|
|
context = object()
|
|
expected = {"FINISHED"}
|
|
with (
|
|
patch.object(
|
|
self.upscale,
|
|
"invoke_ai_setup_if_needed",
|
|
return_value=None,
|
|
),
|
|
patch.object(
|
|
operator,
|
|
"execute",
|
|
return_value=expected,
|
|
) as execute,
|
|
):
|
|
result = operator.invoke(context, None)
|
|
|
|
execute.assert_called_once_with(context)
|
|
self.assertIs(result, expected)
|
|
self.assertNotIn("REGISTER", operator.bl_options)
|
|
|
|
|
|
def test_remove_background_loads_a_plain_image_edit_result(self):
|
|
result = SimpleNamespace(directory=Path("result"))
|
|
operator = self.remove_background.RunBackgroundRemoval()
|
|
operator._image_target = Mock()
|
|
ops = SimpleNamespace(
|
|
ed=SimpleNamespace(undo_push=Mock(return_value={"FINISHED"}))
|
|
)
|
|
with (
|
|
patch.object(self.remove_background.bpy, "ops", ops, create=True),
|
|
):
|
|
operator.response(None, result)
|
|
|
|
operator._image_target.apply.assert_called_once_with(Path("result"))
|
|
ops.ed.undo_push.assert_called_once_with(message="Remove Background")
|
|
|
|
|
|
def test_upscale_loads_a_plain_image_edit_result(self):
|
|
result = SimpleNamespace(directory=Path("result"))
|
|
operator = self.upscale.UpscaleImage()
|
|
operator._image_target = Mock()
|
|
operator.model = "HAT_GAN_X4_SHARPER"
|
|
with (
|
|
patch.object(
|
|
self.upscale,
|
|
"upscale_model_label",
|
|
return_value="HAT Sharper",
|
|
),
|
|
):
|
|
message = operator.response(None, result)
|
|
|
|
operator._image_target.apply.assert_called_once_with(
|
|
Path("result"),
|
|
frame_prefix="upscale.",
|
|
)
|
|
self.assertEqual(message, "Image is upscaled with HAT Sharper")
|