ausboss
2b34005150
The node gets a second output, image_b: the B batch, untouched, the way image_a already passes through. image_a stays the first output, so saved workflows keep their wires. The slider, its resting side, Erase and Save are unchanged. Help page, tooltips, README line and CHANGELOG say so; the node API snapshot gains the one new output.
105 lines
4.6 KiB
Python
105 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
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_compare
|
|
from nodes._compare_helpers import pair_pictures
|
|
from nodes.node_compare import AusBossCompare, first_frame_to_pil
|
|
|
|
|
|
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 CompareNodeTests(unittest.TestCase):
|
|
def test_first_frame_conversion_and_validation(self):
|
|
batch = torch.zeros(3, 4, 6, 3)
|
|
batch[0, :, :, 0] = 1.0 # first frame is pure red
|
|
preview = first_frame_to_pil(batch)
|
|
self.assertEqual((preview.width, preview.height), (6, 4))
|
|
self.assertEqual(preview.getpixel((0, 0)), (255, 0, 0))
|
|
with self.assertRaisesRegex(ValueError, "BHWC"):
|
|
first_frame_to_pil(torch.zeros(4, 6, 3))
|
|
with self.assertRaisesRegex(ValueError, "empty"):
|
|
first_frame_to_pil(torch.zeros(0, 4, 6, 3))
|
|
|
|
def test_compare_saves_temp_previews_and_passes_both_through(self):
|
|
image_a = torch.rand(2, 8, 12, 3)
|
|
image_b = torch.rand(1, 8, 12, 3)
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
# folder_paths is imported by the shared preview helper now, which
|
|
# is where every node's temp preview is written.
|
|
with patch.object(_preview_helpers, "folder_paths", fake_folder_paths(tmp)), \
|
|
patch.object(node_compare, "folder_paths", fake_folder_paths(tmp)):
|
|
result = AusBossCompare().compare(image_a, image_b)
|
|
|
|
# image_a stays the first output, so saved links keep their slot;
|
|
# image_b was appended after it.
|
|
self.assertEqual(AusBossCompare.RETURN_NAMES, ("image_a", "image_b"))
|
|
self.assertEqual(len(AusBossCompare.OUTPUT_TOOLTIPS), 2)
|
|
self.assertIs(result["result"][0], image_a)
|
|
self.assertIs(result["result"][1], image_b)
|
|
for key in ("a_images", "b_images"):
|
|
(ref,) = result["ui"][key]
|
|
self.assertEqual(ref["type"], "temp")
|
|
self.assertEqual(ref["subfolder"], "ausboss")
|
|
self.assertEqual((ref["width"], ref["height"]), (12, 8))
|
|
saved = Path(tmp) / "ausboss" / ref["filename"]
|
|
with Image.open(saved) as png:
|
|
self.assertEqual(png.size, (12, 8))
|
|
# The run's key stands for exactly these two files; Erase's Save
|
|
# sends the key back and never a file name.
|
|
(key,) = result["ui"]["pair"]
|
|
self.assertEqual(
|
|
pair_pictures(key)[:2],
|
|
tuple(Path(tmp) / "ausboss" / result["ui"][name][0]["filename"] for name in ("a_images", "b_images")),
|
|
)
|
|
# No name wired in: the one typed on the node will be used.
|
|
self.assertEqual((result["ui"]["save_name"], pair_pictures(key)[2]), ([""], ""))
|
|
|
|
def test_a_wired_filename_names_the_saved_files(self):
|
|
image = torch.rand(1, 8, 12, 3)
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
with patch.object(_preview_helpers, "folder_paths", fake_folder_paths(tmp)), \
|
|
patch.object(node_compare, "folder_paths", fake_folder_paths(tmp)):
|
|
result = AusBossCompare().compare(image, image, filename="sets/a/photo 12.jpg")
|
|
self.assertEqual(result["ui"]["save_name"], ["photo 12"])
|
|
self.assertEqual(pair_pictures(result["ui"]["pair"][0])[2], "photo 12")
|
|
self.assertIs(result["result"][0], image)
|
|
optional = AusBossCompare.INPUT_TYPES()["optional"]["filename"]
|
|
self.assertEqual(optional[0], "STRING")
|
|
self.assertTrue(optional[1]["forceInput"])
|
|
self.assertEqual(list(AusBossCompare.INPUT_TYPES()["required"]), ["image_a", "image_b"])
|
|
|
|
def test_two_nodes_use_distinct_temp_prefixes(self):
|
|
# Prefixes are randomized per instance so parallel Compare nodes
|
|
# never overwrite each other's previews.
|
|
self.assertNotEqual(AusBossCompare()._prefix, AusBossCompare()._prefix)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|