366 lines
16 KiB
Python
366 lines
16 KiB
Python
import unittest
|
|
import torch
|
|
from inpaint_cropandstitch import (
|
|
InpaintCropImproved,
|
|
InpaintStitchImproved,
|
|
_sanitize_image_tensor,
|
|
_sanitize_mask_tensor,
|
|
)
|
|
|
|
|
|
class TestInputShielding(unittest.TestCase):
|
|
def setUp(self):
|
|
self.crop_node = InpaintCropImproved()
|
|
self.stitch_node = InpaintStitchImproved()
|
|
|
|
# =========================================================================
|
|
# 1. Tests for _sanitize_image_tensor
|
|
# =========================================================================
|
|
def test_sanitize_image_non_tensor_raises(self):
|
|
with self.assertRaises(TypeError):
|
|
_sanitize_image_tensor("not_a_tensor")
|
|
with self.assertRaises(TypeError):
|
|
_sanitize_image_tensor([1.0, 2.0, 3.0])
|
|
|
|
def test_sanitize_image_2d(self):
|
|
# [H, W] -> [1, H, W, 3] (allow_grayscale=False default)
|
|
t = torch.rand(40, 50, dtype=torch.float32)
|
|
sanitized = _sanitize_image_tensor(t)
|
|
self.assertEqual(sanitized.shape, (1, 40, 50, 3))
|
|
# With allow_grayscale=True -> [1, 40, 50, 1]
|
|
sanitized_gray = _sanitize_image_tensor(t, allow_grayscale=True)
|
|
self.assertEqual(sanitized_gray.shape, (1, 40, 50, 1))
|
|
|
|
def test_sanitize_image_3d_unbatched_hwc(self):
|
|
# [H, W, 3] -> [1, H, W, 3]
|
|
t = torch.rand(30, 40, 3, dtype=torch.float32)
|
|
sanitized = _sanitize_image_tensor(t)
|
|
self.assertEqual(sanitized.shape, (1, 30, 40, 3))
|
|
|
|
def test_sanitize_image_3d_channels_first_chw(self):
|
|
# [3, H, W] -> [1, H, W, 3]
|
|
t = torch.rand(3, 30, 40, dtype=torch.float32)
|
|
sanitized = _sanitize_image_tensor(t)
|
|
self.assertEqual(sanitized.shape, (1, 30, 40, 3))
|
|
|
|
def test_sanitize_image_4d_channels_first_bchw(self):
|
|
# [2, 3, 30, 40] -> [2, 30, 40, 3]
|
|
t = torch.rand(2, 3, 30, 40, dtype=torch.float32)
|
|
sanitized = _sanitize_image_tensor(t)
|
|
self.assertEqual(sanitized.shape, (2, 30, 40, 3))
|
|
|
|
def test_sanitize_image_5d_singleton(self):
|
|
# [1, 1, 30, 40, 3] -> [1, 30, 40, 3]
|
|
t = torch.rand(1, 1, 30, 40, 3, dtype=torch.float32)
|
|
sanitized = _sanitize_image_tensor(t)
|
|
self.assertEqual(sanitized.shape, (1, 30, 40, 3))
|
|
|
|
def test_sanitize_image_2channel_to_rgba(self):
|
|
# [1, 30, 40, 2] (grayscale + alpha) -> [1, 30, 40, 4]
|
|
t = torch.rand(1, 30, 40, 2, dtype=torch.float32)
|
|
sanitized = _sanitize_image_tensor(t)
|
|
self.assertEqual(sanitized.shape, (1, 30, 40, 4))
|
|
|
|
def test_sanitize_image_6channel_clamped(self):
|
|
# >4 channels -> clamped to first 4 channels
|
|
t = torch.rand(1, 30, 40, 6, dtype=torch.float32)
|
|
sanitized = _sanitize_image_tensor(t)
|
|
self.assertEqual(sanitized.shape, (1, 30, 40, 4))
|
|
|
|
def test_sanitize_image_uint8_conversion(self):
|
|
# uint8 0..255 -> float32 0..1
|
|
t = (torch.rand(1, 20, 20, 3) * 255).to(torch.uint8)
|
|
sanitized = _sanitize_image_tensor(t)
|
|
self.assertEqual(sanitized.dtype, torch.float32)
|
|
self.assertLessEqual(sanitized.max().item(), 1.0)
|
|
self.assertGreaterEqual(sanitized.min().item(), 0.0)
|
|
|
|
def test_sanitize_image_nans_and_infs(self):
|
|
t = torch.tensor([[[[float("nan"), float("inf"), float("-inf"), 0.5]]]])
|
|
sanitized = _sanitize_image_tensor(t)
|
|
self.assertFalse(torch.isnan(sanitized).any())
|
|
self.assertFalse(torch.isinf(sanitized).any())
|
|
|
|
# =========================================================================
|
|
# 2. Tests for _sanitize_mask_tensor
|
|
# =========================================================================
|
|
def test_sanitize_mask_none_and_empty(self):
|
|
self.assertIsNone(_sanitize_mask_tensor(None))
|
|
self.assertIsNone(_sanitize_mask_tensor(torch.tensor([])))
|
|
|
|
def test_sanitize_mask_non_tensor_raises(self):
|
|
with self.assertRaises(TypeError):
|
|
_sanitize_mask_tensor("not_a_mask")
|
|
|
|
def test_sanitize_mask_2d(self):
|
|
# [H, W] -> [1, H, W]
|
|
t = torch.rand(30, 40)
|
|
sanitized = _sanitize_mask_tensor(t)
|
|
self.assertEqual(sanitized.shape, (1, 30, 40))
|
|
|
|
def test_sanitize_mask_3d_standard(self):
|
|
# [B, H, W] stays [B, H, W]
|
|
t = torch.rand(2, 30, 40)
|
|
sanitized = _sanitize_mask_tensor(t)
|
|
self.assertEqual(sanitized.shape, (2, 30, 40))
|
|
|
|
def test_sanitize_mask_3d_trailing_singleton(self):
|
|
# [30, 40, 1] unbatched with channel -> [1, 30, 40]
|
|
t = torch.rand(30, 40, 1)
|
|
sanitized = _sanitize_mask_tensor(t)
|
|
self.assertEqual(sanitized.shape, (1, 30, 40))
|
|
|
|
def test_sanitize_mask_4d_rgb_image_input(self):
|
|
# RGB IMAGE [B, H, W, 3] wired to mask input -> converted to [B, H, W] luminance
|
|
t = torch.rand(2, 30, 40, 3)
|
|
sanitized = _sanitize_mask_tensor(t)
|
|
self.assertEqual(sanitized.shape, (2, 30, 40))
|
|
|
|
def test_sanitize_mask_4d_rgba_image_input(self):
|
|
# RGBA IMAGE [B, H, W, 4] wired to mask input -> converted to [B, H, W] luminance
|
|
t = torch.rand(1, 30, 40, 4)
|
|
sanitized = _sanitize_mask_tensor(t)
|
|
self.assertEqual(sanitized.shape, (1, 30, 40))
|
|
|
|
def test_sanitize_mask_4d_channels_first_image_input(self):
|
|
# Channels-first [B, 3, H, W] -> converted to [B, H, W]
|
|
t = torch.rand(2, 3, 30, 40)
|
|
sanitized = _sanitize_mask_tensor(t)
|
|
self.assertEqual(sanitized.shape, (2, 30, 40))
|
|
|
|
def test_sanitize_mask_uint8_and_clamping(self):
|
|
t = (torch.rand(1, 20, 20) * 255).to(torch.uint8)
|
|
sanitized = _sanitize_mask_tensor(t)
|
|
self.assertEqual(sanitized.dtype, torch.float32)
|
|
self.assertLessEqual(sanitized.max().item(), 1.0)
|
|
self.assertGreaterEqual(sanitized.min().item(), 0.0)
|
|
|
|
def test_sanitize_mask_float_slightly_above_one_is_not_rescaled(self):
|
|
# Issue #182: blur/resize float noise can push a 0..1 mask slightly above 1.0.
|
|
# That must not be mistaken for a 0..255 mask.
|
|
for peak in [1.0000006, 1.05, 1.2]:
|
|
with self.subTest(peak=peak):
|
|
t = torch.zeros(1, 20, 20)
|
|
t[0, 5:15, 5:15] = 0.5
|
|
t[0, 10, 10] = peak
|
|
sanitized = _sanitize_mask_tensor(t)
|
|
self.assertAlmostEqual(sanitized[0, 6, 6].item(), 0.5, places=6)
|
|
self.assertEqual(sanitized[0, 10, 10].item(), 1.0)
|
|
|
|
def test_sanitize_mask_0_255_still_normalized(self):
|
|
float_255 = torch.zeros(1, 20, 20)
|
|
float_255[0, 5:15, 5:15] = 127.5
|
|
float_255[0, 10, 10] = 255.0
|
|
sanitized = _sanitize_mask_tensor(float_255)
|
|
self.assertAlmostEqual(sanitized[0, 6, 6].item(), 0.5, places=6)
|
|
self.assertAlmostEqual(sanitized[0, 10, 10].item(), 1.0, places=6)
|
|
|
|
uint8_255 = torch.zeros(1, 20, 20, dtype=torch.uint8)
|
|
uint8_255[0, 5:15, 5:15] = 255
|
|
self.assertAlmostEqual(_sanitize_mask_tensor(uint8_255)[0, 6, 6].item(), 1.0, places=6)
|
|
|
|
# Integer 0/1 masks keep their meaning
|
|
uint8_01 = torch.zeros(1, 20, 20, dtype=torch.uint8)
|
|
uint8_01[0, 5:15, 5:15] = 1
|
|
self.assertEqual(_sanitize_mask_tensor(uint8_01)[0, 6, 6].item(), 1.0)
|
|
|
|
def test_sanitize_mask_nans_and_infs(self):
|
|
t = torch.tensor([[float("nan"), float("inf"), float("-inf"), 0.5]])
|
|
sanitized = _sanitize_mask_tensor(t)
|
|
self.assertFalse(torch.isnan(sanitized).any())
|
|
self.assertFalse(torch.isinf(sanitized).any())
|
|
self.assertLessEqual(sanitized.max().item(), 1.0)
|
|
self.assertGreaterEqual(sanitized.min().item(), 0.0)
|
|
|
|
# =========================================================================
|
|
# 3. Integration Tests for InpaintCropImproved Shielding
|
|
# =========================================================================
|
|
def _default_crop_args(self, image, mask=None, **kwargs):
|
|
args = {
|
|
"image": image,
|
|
"mask": mask,
|
|
"optional_context_mask": None,
|
|
"downscale_algorithm": "bilinear",
|
|
"upscale_algorithm": "bicubic",
|
|
"preresize": False,
|
|
"preresize_mode": "no preresize",
|
|
"preresize_min_width": 256,
|
|
"preresize_min_height": 256,
|
|
"preresize_max_width": 1024,
|
|
"preresize_max_height": 1024,
|
|
"extend_for_outpainting": False,
|
|
"extend_up_factor": 1.0,
|
|
"extend_down_factor": 1.0,
|
|
"extend_left_factor": 1.0,
|
|
"extend_right_factor": 1.0,
|
|
"mask_hipass_filter": 0.0,
|
|
"mask_fill_holes": False,
|
|
"mask_expand_pixels": 0,
|
|
"mask_invert": False,
|
|
"mask_blend_pixels": 0,
|
|
"context_from_mask_extend_factor": 1.2,
|
|
"output_resize_to_target_size": False,
|
|
"output_target_width": 256,
|
|
"output_target_height": 256,
|
|
"output_padding": 8,
|
|
"device_mode": "cpu (compatible)",
|
|
}
|
|
args.update(kwargs)
|
|
return args
|
|
|
|
def test_crop_with_rgb_image_as_mask(self):
|
|
# User mistakenly wired IMAGE [1, 64, 64, 3] to MASK input
|
|
image = torch.rand(1, 64, 64, 3, dtype=torch.float32)
|
|
image_as_mask = torch.zeros(1, 64, 64, 3, dtype=torch.float32)
|
|
image_as_mask[0, 20:40, 20:40, :] = 1.0 # White square
|
|
|
|
args = self._default_crop_args(image, mask=image_as_mask)
|
|
stitcher, crop_im, crop_m = self.crop_node.inpaint_crop(**args)[:3]
|
|
self.assertEqual(crop_im.ndim, 4)
|
|
self.assertEqual(crop_im.shape[-1], 3)
|
|
self.assertEqual(crop_m.ndim, 3)
|
|
|
|
def test_crop_with_rgba_image_as_mask(self):
|
|
# User mistakenly wired RGBA IMAGE [1, 64, 64, 4] to MASK input
|
|
image = torch.rand(1, 64, 64, 3, dtype=torch.float32)
|
|
rgba_as_mask = torch.zeros(1, 64, 64, 4, dtype=torch.float32)
|
|
rgba_as_mask[0, 20:40, 20:40, :] = 1.0
|
|
|
|
args = self._default_crop_args(image, mask=rgba_as_mask)
|
|
stitcher, crop_im, crop_m = self.crop_node.inpaint_crop(**args)[:3]
|
|
self.assertEqual(crop_im.ndim, 4)
|
|
self.assertEqual(crop_m.ndim, 3)
|
|
|
|
def test_crop_with_unbatched_inputs(self):
|
|
# 3D unbatched image [64, 64, 3] and 2D unbatched mask [64, 64]
|
|
image = torch.rand(64, 64, 3, dtype=torch.float32)
|
|
mask = torch.zeros(64, 64, dtype=torch.float32)
|
|
mask[20:40, 20:40] = 1.0
|
|
|
|
args = self._default_crop_args(image, mask=mask)
|
|
stitcher, crop_im, crop_m = self.crop_node.inpaint_crop(**args)[:3]
|
|
self.assertEqual(crop_im.shape[0], 1)
|
|
self.assertEqual(crop_m.shape[0], 1)
|
|
|
|
def test_crop_mismatched_batch_harmonization(self):
|
|
# Image batch = 1, Mask batch = 3
|
|
image = torch.rand(1, 64, 64, 3, dtype=torch.float32)
|
|
mask = torch.zeros(3, 64, 64, dtype=torch.float32)
|
|
mask[:, 20:40, 20:40] = 1.0
|
|
|
|
args = self._default_crop_args(image, mask=mask)
|
|
stitcher, crop_im, crop_m = self.crop_node.inpaint_crop(**args)[:3]
|
|
self.assertEqual(crop_im.shape[0], 3)
|
|
self.assertEqual(crop_m.shape[0], 3)
|
|
self.assertEqual(len(stitcher["canvas_image"]), 3)
|
|
|
|
def test_crop_with_extreme_and_negative_widget_values(self):
|
|
# Negative factors, high blend pixels, negative padding should not crash
|
|
image = torch.rand(1, 64, 64, 3, dtype=torch.float32)
|
|
mask = torch.zeros(1, 64, 64, dtype=torch.float32)
|
|
mask[0, 20:40, 20:40] = 1.0
|
|
|
|
args = self._default_crop_args(
|
|
image,
|
|
mask=mask,
|
|
extend_up_factor=-0.5,
|
|
extend_down_factor=-1.0,
|
|
mask_blend_pixels=-5,
|
|
mask_expand_pixels=-2,
|
|
context_from_mask_extend_factor=0.5, # should be bounded to >= 1.0
|
|
output_padding=-4,
|
|
mask_hipass_filter=2.5, # should be clamped to 1.0
|
|
)
|
|
stitcher, crop_im, crop_m = self.crop_node.inpaint_crop(**args)[:3]
|
|
self.assertIsNotNone(crop_im)
|
|
self.assertIsNotNone(stitcher)
|
|
|
|
# =========================================================================
|
|
# 4. Integration Tests for InpaintStitchImproved Shielding (Issue #179 & Edge Cases)
|
|
# =========================================================================
|
|
def test_stitch_qwen_rgba_inpaint_into_rgb_canvas_issue_179(self):
|
|
# Issue #179 scenario:
|
|
# Canvas was cropped from standard RGB [1, 100, 100, 3]
|
|
# Inpaint node (e.g. Qwen 2.1 / WanVAE) outputs RGBA [1, H, W, 4]
|
|
orig_image = torch.rand(1, 100, 100, 3, dtype=torch.float32)
|
|
mask = torch.zeros(1, 100, 100, dtype=torch.float32)
|
|
mask[0, 30:60, 30:60] = 1.0
|
|
|
|
args = self._default_crop_args(orig_image, mask=mask)
|
|
stitcher, crop_im, crop_m = self.crop_node.inpaint_crop(**args)[:3]
|
|
|
|
# Simulate Qwen 2.1 returning RGBA with 4 channels for a regular (opaque) edit
|
|
h, w = crop_im.shape[1], crop_im.shape[2]
|
|
rgba_inpaint = torch.rand(1, h, w, 4, dtype=torch.float32)
|
|
rgba_inpaint[..., 3] = 1.0
|
|
|
|
# Execute stitch
|
|
(stitched,) = self.stitch_node.inpaint_stitch(stitcher, rgba_inpaint)
|
|
self.assertEqual(stitched.shape, (1, 100, 100, 3))
|
|
self.assertFalse(torch.isnan(stitched).any())
|
|
|
|
def test_stitch_rgb_inpaint_into_rgba_canvas(self):
|
|
# Canvas was RGBA [1, 100, 100, 4]
|
|
# Inpainted image is RGB [1, H, W, 3]
|
|
orig_image = torch.rand(1, 100, 100, 4, dtype=torch.float32)
|
|
mask = torch.zeros(1, 100, 100, dtype=torch.float32)
|
|
mask[0, 30:60, 30:60] = 1.0
|
|
|
|
args = self._default_crop_args(orig_image, mask=mask)
|
|
stitcher, crop_im, crop_m = self.crop_node.inpaint_crop(**args)[:3]
|
|
|
|
# Standard 3-channel RGB inpaint
|
|
h, w = crop_im.shape[1], crop_im.shape[2]
|
|
rgb_inpaint = torch.rand(1, h, w, 3, dtype=torch.float32)
|
|
|
|
(stitched,) = self.stitch_node.inpaint_stitch(stitcher, rgb_inpaint)
|
|
# Stitched preserves the canvas RGBA channels
|
|
self.assertEqual(stitched.shape, (1, 100, 100, 4))
|
|
self.assertFalse(torch.isnan(stitched).any())
|
|
|
|
def test_stitch_channels_first_inpainted_image(self):
|
|
# Upstream node outputs [B, C, H, W] instead of [B, H, W, C]
|
|
orig_image = torch.rand(1, 80, 80, 3, dtype=torch.float32)
|
|
mask = torch.zeros(1, 80, 80, dtype=torch.float32)
|
|
mask[0, 20:50, 20:50] = 1.0
|
|
|
|
args = self._default_crop_args(orig_image, mask=mask)
|
|
stitcher, crop_im, crop_m = self.crop_node.inpaint_crop(**args)[:3]
|
|
|
|
h, w = crop_im.shape[1], crop_im.shape[2]
|
|
# Inpainted image passed as [1, 3, H, W]
|
|
bchw_inpaint = torch.rand(1, 3, h, w, dtype=torch.float32)
|
|
|
|
(stitched,) = self.stitch_node.inpaint_stitch(stitcher, bchw_inpaint)
|
|
self.assertEqual(stitched.shape, (1, 80, 80, 3))
|
|
|
|
def test_stitch_unbatched_inpainted_image(self):
|
|
# Upstream node outputs 3D unbatched [H, W, C]
|
|
orig_image = torch.rand(1, 80, 80, 3, dtype=torch.float32)
|
|
mask = torch.zeros(1, 80, 80, dtype=torch.float32)
|
|
mask[0, 20:50, 20:50] = 1.0
|
|
|
|
args = self._default_crop_args(orig_image, mask=mask)
|
|
stitcher, crop_im, crop_m = self.crop_node.inpaint_crop(**args)[:3]
|
|
|
|
h, w = crop_im.shape[1], crop_im.shape[2]
|
|
hwc_inpaint = torch.rand(h, w, 3, dtype=torch.float32)
|
|
|
|
(stitched,) = self.stitch_node.inpaint_stitch(stitcher, hwc_inpaint)
|
|
self.assertEqual(stitched.shape, (1, 80, 80, 3))
|
|
|
|
def test_stitch_invalid_stitcher_types_raise(self):
|
|
inpaint = torch.rand(1, 30, 30, 3)
|
|
with self.assertRaises(TypeError):
|
|
self.stitch_node.inpaint_stitch("not_a_dict", inpaint)
|
|
with self.assertRaises(TypeError):
|
|
self.stitch_node.inpaint_stitch(None, inpaint)
|
|
with self.assertRaises(ValueError):
|
|
self.stitch_node.inpaint_stitch({}, inpaint)
|
|
with self.assertRaises(ValueError):
|
|
self.stitch_node.inpaint_stitch({"cropped_to_canvas_x": []}, inpaint)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|