ausboss
d41aeaa1d6
The four mode buttons become one switch, Slide | A | B | Erase, whose marker slides to the choice. In Erase the brush's two jobs and the picture on top are switches of the same kind, and PSD is a plain off | on. Save moves to the bottom right of the node, larger, with a disk icon, in a row of its own with the name of the files and the PSD choice. The files can be named. A Name field on the node takes letters, digits, spaces, dots, dashes and underscores; the server checks the same rule and refuses anything else, so a folder is never part of a name. A new optional filename input takes a wired name, such as the source picture's file name: its folder and extension are dropped. The save route still names no file to read. Size and Soft scrub from their label and from the unit inside the box as well as from the number (the scrub control gains handle() for this). Workflows saved before load as they were: 60 that use the node were opened (the pack's examples, published workflows, the oldest saved shapes), 72 nodes kept their size, mode and links. The node is one input socket taller.
100 lines
4.3 KiB
Python
100 lines
4.3 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_a_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)
|
|
|
|
self.assertIs(result["result"][0], image_a)
|
|
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()
|