234 lines
10 KiB
Python
234 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from contextlib import redirect_stdout
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
from PIL import Image
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT))
|
|
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
|
del sys.modules["nodes"]
|
|
|
|
from nodes import _preview_helpers, node_lama_inpaint, node_refine_mask, node_select_frame
|
|
from nodes._preview_helpers import (
|
|
first_frame_to_pil,
|
|
mask_to_preview_batch,
|
|
preview_payload,
|
|
save_temp_preview,
|
|
temp_prefix,
|
|
)
|
|
|
|
|
|
def fake_folder_paths(temp_dir: str):
|
|
class FakeFolderPaths:
|
|
@staticmethod
|
|
def get_temp_directory():
|
|
return temp_dir
|
|
|
|
@staticmethod
|
|
def get_save_image_path(prefix, output_dir, _width, _height):
|
|
folder = Path(output_dir) / "ausboss"
|
|
folder.mkdir(parents=True, exist_ok=True)
|
|
return str(folder), prefix, 1, "ausboss", prefix
|
|
|
|
return FakeFolderPaths
|
|
|
|
|
|
class FirstFrameTests(unittest.TestCase):
|
|
def test_the_first_frame_of_the_batch_is_the_one_saved(self):
|
|
batch = torch.zeros(3, 4, 6, 3)
|
|
batch[0, :, :, 1] = 1.0 # first frame green, the rest black
|
|
preview = first_frame_to_pil(batch)
|
|
self.assertEqual((preview.width, preview.height), (6, 4))
|
|
self.assertEqual(preview.getpixel((0, 0)), (0, 255, 0))
|
|
|
|
def test_a_one_channel_frame_becomes_grayscale_not_a_broken_rgb(self):
|
|
mask = torch.full((1, 3, 5), 0.5)
|
|
preview = first_frame_to_pil(mask_to_preview_batch(mask))
|
|
self.assertEqual(preview.mode, "L")
|
|
self.assertEqual(preview.size, (5, 3))
|
|
self.assertEqual(preview.getpixel((0, 0)), 128)
|
|
|
|
def test_alpha_is_dropped_so_the_panel_shows_coverage_not_a_hole(self):
|
|
rgba = torch.zeros(1, 2, 2, 4)
|
|
rgba[..., 0] = 1.0 # opaque red pixels...
|
|
rgba[..., 3] = 0.0 # ...with a fully transparent alpha channel
|
|
preview = first_frame_to_pil(rgba)
|
|
self.assertEqual(preview.mode, "RGB")
|
|
self.assertEqual(preview.getpixel((0, 0)), (255, 0, 0))
|
|
|
|
def test_values_outside_the_range_are_clamped_rather_than_wrapped(self):
|
|
batch = torch.tensor([[[[-1.0, 0.5, 2.0]]]])
|
|
self.assertEqual(first_frame_to_pil(batch).getpixel((0, 0)), (0, 128, 255))
|
|
|
|
def test_the_shape_complaints_name_the_node_that_asked(self):
|
|
with self.assertRaisesRegex(ValueError, "Mask Refine expected a BHWC"):
|
|
first_frame_to_pil(torch.zeros(4, 6, 3), "Mask Refine")
|
|
with self.assertRaisesRegex(ValueError, "empty"):
|
|
first_frame_to_pil(torch.zeros(0, 4, 6, 3))
|
|
with self.assertRaisesRegex(ValueError, "BHW MASK"):
|
|
mask_to_preview_batch(torch.zeros(2, 3, 4, 1))
|
|
|
|
|
|
class SaveTempPreviewTests(unittest.TestCase):
|
|
def test_a_png_lands_in_the_temp_folder_with_a_matching_ui_entry(self):
|
|
image = torch.rand(2, 8, 12, 3)
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
with patch.object(_preview_helpers, "folder_paths", fake_folder_paths(tmp)):
|
|
entry = save_temp_preview(image, "ausboss_test", "Select Frame")
|
|
self.assertEqual(entry["type"], "temp")
|
|
self.assertEqual(entry["subfolder"], "ausboss")
|
|
self.assertEqual((entry["width"], entry["height"]), (12, 8))
|
|
saved = Path(tmp) / "ausboss" / entry["filename"]
|
|
with Image.open(saved) as png:
|
|
self.assertEqual(png.size, (12, 8))
|
|
|
|
def test_two_instances_of_a_node_never_share_a_prefix(self):
|
|
self.assertNotEqual(temp_prefix("mask_refine"), temp_prefix("mask_refine"))
|
|
self.assertTrue(temp_prefix("mask_refine").startswith("ausboss_mask_refine_"))
|
|
|
|
|
|
class PreviewPayloadTests(unittest.TestCase):
|
|
def test_the_payload_carries_both_the_preview_and_the_outputs(self):
|
|
image = torch.rand(1, 4, 4, 3)
|
|
outputs = (image, "second")
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
with patch.object(_preview_helpers, "folder_paths", fake_folder_paths(tmp)):
|
|
payload = preview_payload(image, "ausboss_test", "Select Frame", outputs)
|
|
self.assertEqual(payload["result"], outputs)
|
|
self.assertEqual(len(payload["ui"]["images"]), 1)
|
|
|
|
def test_a_preview_that_cannot_be_written_never_costs_the_outputs(self):
|
|
# A full disk or a swept temp folder must not turn a finished inpaint
|
|
# into a failed node; the panel is the only thing that goes quiet.
|
|
image = torch.rand(1, 4, 4, 3)
|
|
outputs = (image,)
|
|
with patch.object(_preview_helpers, "folder_paths", None):
|
|
printed = io.StringIO()
|
|
with redirect_stdout(printed):
|
|
payload = preview_payload(image, "ausboss_test", "LaMa Inpaint", outputs)
|
|
self.assertEqual(payload["result"], outputs)
|
|
self.assertNotIn("ui", payload)
|
|
note = printed.getvalue()
|
|
self.assertIn("LaMa Inpaint", note)
|
|
note.encode("ascii") # Windows consoles are cp1252
|
|
|
|
|
|
class NodePreviewContractTests(unittest.TestCase):
|
|
"""The three nodes that now show their own result on their own face.
|
|
|
|
Each returns a dict rather than a bare tuple, which ComfyUI unpacks in
|
|
execution.py: 'ui' is forwarded to the frontend and 'result' becomes the
|
|
node's outputs. The outputs must still match RETURN_TYPES exactly.
|
|
"""
|
|
|
|
def run_with_temp(self, call):
|
|
# mkdtemp with a cleanup hook, not a with-block: the returned paths
|
|
# are opened by the test body, which runs after a with-block would
|
|
# already have deleted the directory out from under them.
|
|
tmp = tempfile.mkdtemp()
|
|
self.addCleanup(shutil.rmtree, tmp, ignore_errors=True)
|
|
with patch.object(_preview_helpers, "folder_paths", fake_folder_paths(tmp)):
|
|
payload = call()
|
|
return payload, sorted((Path(tmp) / "ausboss").glob("*.png"))
|
|
|
|
def test_mask_refine_previews_the_refined_mask_and_still_returns_both_masks(self):
|
|
mask = torch.zeros(1, 16, 24)
|
|
mask[:, 4:12, 6:18] = 1.0
|
|
node = node_refine_mask.AusBossRefineMask()
|
|
payload, saved = self.run_with_temp(lambda: node.refine(
|
|
mask, expand=2, blur=1.0, fill_holes=False, smooth=0,
|
|
black_point=0.0, white_point=1.0, edge_refine="off",
|
|
))
|
|
refined, inverted = payload["result"]
|
|
self.assertEqual(len(payload["result"]), len(node.RETURN_TYPES))
|
|
self.assertEqual(tuple(refined.shape), (1, 16, 24))
|
|
torch.testing.assert_close(inverted, 1.0 - refined)
|
|
# The preview is the mask this node produced, at the mask's own size.
|
|
self.assertEqual(len(saved), 1)
|
|
with Image.open(saved[0]) as png:
|
|
self.assertEqual(png.size, (24, 16))
|
|
self.assertEqual(png.mode, "L")
|
|
|
|
def test_select_frame_previews_the_frame_it_selected(self):
|
|
frames = torch.zeros(4, 6, 8, 3)
|
|
frames[2, :, :, 0] = 1.0 # third frame is the red one
|
|
node = node_select_frame.AusBossSelectFrame()
|
|
payload, saved = self.run_with_temp(lambda: node.select_frame(frames, 3))
|
|
(selected,) = payload["result"]
|
|
self.assertEqual(len(payload["result"]), len(node.RETURN_TYPES))
|
|
torch.testing.assert_close(selected[0], frames[2])
|
|
with Image.open(saved[0]) as png:
|
|
self.assertEqual(png.getpixel((0, 0)), (255, 0, 0))
|
|
|
|
def test_lama_inpaint_previews_its_result_even_for_a_single_image(self):
|
|
# The bug this replaces: previews were attached to the progress bar
|
|
# only when the batch had more than one frame, so a one-image run
|
|
# finished with an empty panel and nothing to look at.
|
|
image = torch.rand(1, 8, 8, 3)
|
|
inpainted = torch.rand(1, 8, 8, 3)
|
|
node = node_lama_inpaint.AusBossLaMaInpaint()
|
|
with patch.object(node_lama_inpaint, "run_lama_inpaint", return_value=inpainted):
|
|
payload, saved = self.run_with_temp(
|
|
lambda: node.inpaint(image, torch.ones(1, 8, 8), "big-lama.pt")
|
|
)
|
|
self.assertIs(payload["result"][0], inpainted)
|
|
self.assertEqual(len(payload["result"]), len(node.RETURN_TYPES))
|
|
self.assertEqual(len(saved), 1)
|
|
|
|
def test_each_node_instance_writes_to_its_own_prefix(self):
|
|
# Two of the same node run in one graph; without distinct prefixes the
|
|
# second would overwrite the first and both panels show one picture.
|
|
for factory in (
|
|
node_refine_mask.AusBossRefineMask,
|
|
node_select_frame.AusBossSelectFrame,
|
|
node_lama_inpaint.AusBossLaMaInpaint,
|
|
):
|
|
with self.subTest(node=factory.__name__):
|
|
self.assertNotEqual(factory()._prefix, factory()._prefix)
|
|
|
|
|
|
|
|
|
|
class PreviewSwitchTests(unittest.TestCase):
|
|
"""preview=False returns the plain outputs and writes nothing."""
|
|
|
|
def test_every_preview_node_offers_the_switch_off_by_widget(self):
|
|
for module, cls in (
|
|
(node_select_frame, "AusBossSelectFrame"),
|
|
(node_refine_mask, "AusBossRefineMask"),
|
|
(node_lama_inpaint, "AusBossLaMaInpaint"),
|
|
):
|
|
optional = getattr(module, cls).INPUT_TYPES()["optional"]
|
|
self.assertIn("preview", optional, cls)
|
|
self.assertEqual(optional["preview"][0], "BOOLEAN", cls)
|
|
self.assertTrue(optional["preview"][1]["default"], cls)
|
|
|
|
def test_select_frame_off_returns_a_bare_tuple(self):
|
|
with patch.object(_preview_helpers, "save_temp_preview") as save:
|
|
result = node_select_frame.AusBossSelectFrame().select_frame(torch.zeros((3, 8, 8, 3)), 2, preview=False)
|
|
self.assertIsInstance(result, tuple)
|
|
self.assertEqual(tuple(result[0].shape), (1, 8, 8, 3))
|
|
save.assert_not_called()
|
|
|
|
def test_mask_refine_off_returns_both_masks_without_a_preview(self):
|
|
with patch.object(_preview_helpers, "save_temp_preview") as save:
|
|
result = node_refine_mask.AusBossRefineMask().refine(
|
|
torch.ones((1, 8, 8)), 0, 0.0, False, 0, 0.0, 1.0, "off", preview=False,
|
|
)
|
|
self.assertIsInstance(result, tuple)
|
|
self.assertEqual(len(result), 2)
|
|
save.assert_not_called()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|