185 lines
7.4 KiB
Python
185 lines
7.4 KiB
Python
import os
|
|
import unittest
|
|
import torch
|
|
from tests.workflow_runner import (
|
|
WorkflowRunner,
|
|
validate_workflow_run,
|
|
validate_crop_outputs,
|
|
validate_stitch_outputs,
|
|
MockLoadImage,
|
|
MockMaskToImage,
|
|
MockImageInvert,
|
|
MockImpactMakeImageBatch,
|
|
MockImpactMakeMaskBatch,
|
|
MockImageCompositeMasked,
|
|
)
|
|
|
|
|
|
class TestMockNodes(unittest.TestCase):
|
|
"""Unit tests for the self-contained mock nodes used in workflow execution."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
cls.testimgs_dir = os.path.join(cls.repo_root, "testimgs")
|
|
|
|
def test_mock_load_image_rgb(self):
|
|
loader = MockLoadImage(self.testimgs_dir)
|
|
img, mask = loader.load_image("example.png")
|
|
self.assertEqual(img.ndim, 4)
|
|
self.assertEqual(img.shape[0], 1)
|
|
self.assertEqual(img.shape[-1], 3)
|
|
self.assertEqual(mask.ndim, 3)
|
|
self.assertEqual(mask.shape, (1, 64, 64)) # No alpha -> 64x64 zero mask
|
|
self.assertEqual(mask.sum().item(), 0.0)
|
|
|
|
def test_mock_load_image_clipspace_rgba(self):
|
|
loader = MockLoadImage(self.testimgs_dir)
|
|
img, mask = loader.load_image("clipspace/clipspace-mask-105444.59999999404.png [input]")
|
|
self.assertEqual(img.ndim, 4)
|
|
self.assertEqual(img.shape[0], 1)
|
|
self.assertEqual(img.shape[-1], 3)
|
|
self.assertEqual(mask.ndim, 3)
|
|
self.assertEqual(mask.shape[1:], img.shape[1:3])
|
|
self.assertTrue((mask >= 0.0).all() and (mask <= 1.0).all())
|
|
|
|
def test_mock_load_image_not_found(self):
|
|
loader = MockLoadImage(self.testimgs_dir)
|
|
with self.assertRaises(FileNotFoundError):
|
|
loader.load_image("non_existent_image_12345.png")
|
|
|
|
def test_mock_mask_to_image(self):
|
|
node = MockMaskToImage()
|
|
mask_3d = torch.rand(2, 32, 48)
|
|
(img_4d,) = node.mask_to_image(mask_3d)
|
|
self.assertEqual(img_4d.shape, (2, 32, 48, 3))
|
|
# Channels should be identical grayscale replicated
|
|
self.assertTrue(torch.equal(img_4d[..., 0], img_4d[..., 1]))
|
|
self.assertTrue(torch.equal(img_4d[..., 1], img_4d[..., 2]))
|
|
|
|
mask_2d = torch.rand(32, 48)
|
|
(img_from_2d,) = node.mask_to_image(mask_2d)
|
|
self.assertEqual(img_from_2d.shape, (1, 32, 48, 3))
|
|
|
|
def test_mock_image_invert(self):
|
|
node = MockImageInvert()
|
|
img = torch.tensor([[[[0.0, 0.25], [0.75, 1.0]]]])
|
|
(inv,) = node.invert(img)
|
|
self.assertTrue(torch.allclose(inv, 1.0 - img))
|
|
|
|
def test_mock_impact_image_batch(self):
|
|
node = MockImpactMakeImageBatch()
|
|
img1 = torch.rand(1, 16, 16, 3)
|
|
img2 = torch.rand(2, 16, 16, 3)
|
|
(batched,) = node.make_batch(image1=img1, image2=img2, image3=None)
|
|
self.assertEqual(batched.shape, (3, 16, 16, 3))
|
|
|
|
def test_mock_impact_mask_batch(self):
|
|
node = MockImpactMakeMaskBatch()
|
|
m1 = torch.rand(1, 16, 16)
|
|
m2 = torch.rand(16, 16) # 2D mask
|
|
(batched,) = node.make_batch(mask1=m1, mask2=m2, mask3=None)
|
|
self.assertEqual(batched.shape, (2, 16, 16))
|
|
|
|
def test_mock_image_composite_masked(self):
|
|
node = MockImageCompositeMasked()
|
|
dest = torch.zeros(1, 40, 40, 3)
|
|
src = torch.ones(1, 20, 20, 3)
|
|
mask = torch.ones(1, 20, 20)
|
|
(res,) = node.composite(dest, src, x=10, y=10, mask=mask)
|
|
self.assertEqual(res.shape, (1, 40, 40, 3))
|
|
# Inner region (10:30, 10:30) should be 1.0, outer should be 0.0
|
|
self.assertAlmostEqual(res[0, 15, 15, 0].item(), 1.0)
|
|
self.assertAlmostEqual(res[0, 0, 0, 0].item(), 0.0)
|
|
|
|
|
|
class TestWorkflowExecution(unittest.TestCase):
|
|
"""Executes testscpu.json and testsgpu.json end-to-end and validates outputs."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
cls.cpu_path = os.path.join(cls.repo_root, "testscpu.json")
|
|
cls.gpu_path = os.path.join(cls.repo_root, "testsgpu.json")
|
|
cls.testimgs_dir = os.path.join(cls.repo_root, "testimgs")
|
|
cls.runner = WorkflowRunner(testimgs_dir=cls.testimgs_dir, verbose=False)
|
|
# Run each workflow once and cache for all assertion methods
|
|
cls.cpu_result = cls.runner.run_file(cls.cpu_path)
|
|
cls.gpu_result = cls.runner.run_file(cls.gpu_path)
|
|
|
|
def test_testscpu_workflow_execution_and_outputs(self):
|
|
"""Execute all 793 nodes in testscpu.json and validate outputs."""
|
|
result = self.cpu_result
|
|
self.assertEqual(len(result.nodes), 793)
|
|
self.assertEqual(len(result.outputs), 793)
|
|
self.assertEqual(len(result.crop_node_ids), 105)
|
|
self.assertEqual(len(result.stitch_node_ids), 34)
|
|
self.assertEqual(len(result.preview_node_ids), 349)
|
|
self.assertEqual(len(result.load_node_ids), 78)
|
|
|
|
# Full validation across all crop, stitch, and preview nodes
|
|
validate_workflow_run(result)
|
|
|
|
# Specific sample checks:
|
|
# Check node 15 (InpaintCropImproved)
|
|
crop_15_outs = result.outputs[15]
|
|
self.assertEqual(len(crop_15_outs), 26)
|
|
stitcher_15 = crop_15_outs[0]
|
|
self.assertEqual(stitcher_15["device_mode"], "cpu (compatible)")
|
|
self.assertIn("canvas_image", stitcher_15)
|
|
cropped_img_15 = crop_15_outs[1]
|
|
self.assertEqual(cropped_img_15.ndim, 4)
|
|
self.assertEqual(cropped_img_15.shape[-1], 3)
|
|
|
|
# Check node 478 (InpaintStitchImproved)
|
|
stitch_478_out = result.outputs[478][0]
|
|
self.assertEqual(stitch_478_out.ndim, 4)
|
|
self.assertEqual(stitch_478_out.shape[-1], 3)
|
|
self.assertTrue((stitch_478_out >= 0.0).all() and (stitch_478_out <= 1.0).all())
|
|
|
|
def test_testsgpu_workflow_execution_and_outputs(self):
|
|
"""Execute all 793 nodes in testsgpu.json and validate outputs."""
|
|
result = self.gpu_result
|
|
self.assertEqual(len(result.nodes), 793)
|
|
self.assertEqual(len(result.outputs), 793)
|
|
self.assertEqual(len(result.crop_node_ids), 105)
|
|
self.assertEqual(len(result.stitch_node_ids), 34)
|
|
self.assertEqual(len(result.preview_node_ids), 349)
|
|
|
|
# Full validation across all crop, stitch, and preview nodes
|
|
validate_workflow_run(result)
|
|
|
|
# Verify device mode was passed through stitcher
|
|
sample_crop_id = result.crop_node_ids[0]
|
|
stitcher = result.outputs[sample_crop_id][0]
|
|
self.assertEqual(stitcher["device_mode"], "gpu (much faster)")
|
|
|
|
def test_cpu_gpu_workflow_consistency(self):
|
|
"""Compare CPU vs GPU workflows: output shapes and coordinate consistency."""
|
|
cpu_res = self.cpu_result
|
|
gpu_res = self.gpu_result
|
|
|
|
# Verify all PreviewImage nodes have identical output shapes
|
|
for nid in cpu_res.preview_node_ids:
|
|
cpu_tensor = cpu_res.outputs[nid][0]
|
|
gpu_tensor = gpu_res.outputs[nid][0]
|
|
self.assertEqual(
|
|
cpu_tensor.shape,
|
|
gpu_tensor.shape,
|
|
f"Preview node {nid} shape mismatch between CPU and GPU",
|
|
)
|
|
|
|
# Verify all InpaintStitchImproved nodes have identical shapes
|
|
for nid in cpu_res.stitch_node_ids:
|
|
cpu_stitched = cpu_res.outputs[nid][0]
|
|
gpu_stitched = gpu_res.outputs[nid][0]
|
|
self.assertEqual(
|
|
cpu_stitched.shape,
|
|
gpu_stitched.shape,
|
|
f"Stitch node {nid} shape mismatch between CPU and GPU",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|