115 lines
3.9 KiB
Python
115 lines
3.9 KiB
Python
import unittest
|
|
|
|
try:
|
|
from .test_pro_edit_helpers import load_plugin_package, make_torch_image, make_torch_mask, torch
|
|
except ImportError:
|
|
from test_pro_edit_helpers import load_plugin_package, make_torch_image, make_torch_mask, torch
|
|
|
|
|
|
class TestProEditFinish(unittest.TestCase):
|
|
def setUp(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSProImageEditFinish"]
|
|
self.node = node_cls()
|
|
|
|
def test_region_composite_only_replaces_masked_pixels(self):
|
|
if torch is None:
|
|
self.skipTest("torch is required for finish compositing tests")
|
|
|
|
original = make_torch_image(4, 4, 0.0)
|
|
generated = make_torch_image(4, 4, 1.0)
|
|
mask = make_torch_mask(4, 4, (1, 1, 3, 3))
|
|
|
|
final_image, preview_image = self.node.finish(
|
|
original_image=original,
|
|
generated_image=generated,
|
|
edit_info={
|
|
"backend_name": "sdxl",
|
|
"edit_scope": "region",
|
|
"original_size": [4, 4],
|
|
"work_size": [4, 4],
|
|
},
|
|
feather=0.0,
|
|
color_match="disabled",
|
|
brightness_match="disabled",
|
|
blend_strength=1.0,
|
|
restore_unmasked_area=True,
|
|
edge_fix="none",
|
|
preview_mode="final",
|
|
work_mask=mask,
|
|
sample_info=None,
|
|
)
|
|
|
|
self.assertEqual(float(final_image[0, 0, 0, 0].item()), 0.0)
|
|
self.assertEqual(float(final_image[0, 1, 1, 0].item()), 1.0)
|
|
self.assertEqual(tuple(preview_image.shape), tuple(final_image.shape))
|
|
|
|
def test_crop_composite_pastes_generated_crop_back_into_original(self):
|
|
if torch is None:
|
|
self.skipTest("torch is required for finish compositing tests")
|
|
|
|
original = make_torch_image(4, 4, 0.0)
|
|
generated = make_torch_image(2, 2, 0.75)
|
|
mask = make_torch_mask(2, 2, (0, 0, 2, 2))
|
|
|
|
final_image, preview_image = self.node.finish(
|
|
original_image=original,
|
|
generated_image=generated,
|
|
edit_info={
|
|
"backend_name": "sdxl",
|
|
"edit_scope": "crop",
|
|
"original_size": [4, 4],
|
|
"work_size": [2, 2],
|
|
"crop_box": [1, 1, 3, 3],
|
|
},
|
|
feather=0.0,
|
|
color_match="disabled",
|
|
brightness_match="disabled",
|
|
blend_strength=1.0,
|
|
restore_unmasked_area=True,
|
|
edge_fix="none",
|
|
preview_mode="mask",
|
|
work_mask=mask,
|
|
sample_info=None,
|
|
)
|
|
|
|
self.assertEqual(float(final_image[0, 0, 0, 0].item()), 0.0)
|
|
self.assertEqual(float(final_image[0, 1, 1, 0].item()), 0.75)
|
|
self.assertEqual(tuple(preview_image.shape), tuple(original.shape))
|
|
|
|
def test_canvas_output_keeps_expanded_canvas_size(self):
|
|
if torch is None:
|
|
self.skipTest("torch is required for finish compositing tests")
|
|
|
|
original = make_torch_image(4, 4, 0.0)
|
|
generated = make_torch_image(6, 6, 0.5)
|
|
mask = make_torch_mask(6, 6, (0, 0, 6, 6))
|
|
|
|
final_image, preview_image = self.node.finish(
|
|
original_image=original,
|
|
generated_image=generated,
|
|
edit_info={
|
|
"backend_name": "flux",
|
|
"edit_scope": "canvas",
|
|
"original_size": [4, 4],
|
|
"work_size": [6, 6],
|
|
"original_box_in_canvas": [1, 1, 5, 5],
|
|
},
|
|
feather=0.0,
|
|
color_match="disabled",
|
|
brightness_match="disabled",
|
|
blend_strength=1.0,
|
|
restore_unmasked_area=True,
|
|
edge_fix="none",
|
|
preview_mode="compare",
|
|
work_mask=mask,
|
|
sample_info=None,
|
|
)
|
|
|
|
self.assertEqual(tuple(final_image.shape[1:3]), (6, 6))
|
|
self.assertEqual(tuple(preview_image.shape[1:3]), (6, 12))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|