915 lines
33 KiB
Python
915 lines
33 KiB
Python
import importlib
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import types
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
|
|
class FakeImageArray:
|
|
def __init__(self, width=8, height=6):
|
|
self.shape = (height, width, 3)
|
|
|
|
def cpu(self):
|
|
return self
|
|
|
|
def numpy(self):
|
|
return self
|
|
|
|
def astype(self, _dtype):
|
|
return self
|
|
|
|
def __rmul__(self, _value):
|
|
return self
|
|
|
|
|
|
class FakeTensorBatch:
|
|
def __init__(self, width=8, height=6, count=1):
|
|
self.image = FakeImageArray(width=width, height=height)
|
|
self.shape = (count, height, width, 3)
|
|
self.count = count
|
|
|
|
def __getitem__(self, index):
|
|
if isinstance(index, slice):
|
|
return FakeTensorBatch(
|
|
width=self.shape[2],
|
|
height=self.shape[1],
|
|
count=len(range(*index.indices(self.count))),
|
|
)
|
|
if index != 0:
|
|
raise IndexError(index)
|
|
return self.image
|
|
|
|
|
|
class FakePILImage:
|
|
saved_paths = []
|
|
|
|
def save(self, path, *args, **kwargs):
|
|
self.saved_paths.append((path, args, kwargs))
|
|
|
|
def convert(self, _mode):
|
|
return self
|
|
|
|
|
|
class DumasImageNodeTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.temp_dir = tempfile.mkdtemp(prefix="dumas-image-node-")
|
|
fake_numpy = types.SimpleNamespace(
|
|
clip=lambda array, _low, _high: array,
|
|
uint8="uint8",
|
|
)
|
|
fake_pil_image_module = types.SimpleNamespace(fromarray=lambda _array: FakePILImage())
|
|
fake_pil_module = types.SimpleNamespace(Image=fake_pil_image_module)
|
|
fake_folder_paths = types.SimpleNamespace(
|
|
get_temp_directory=lambda: cls.temp_dir,
|
|
get_output_directory=lambda: cls.temp_dir,
|
|
get_save_image_path=lambda prefix, _out, _width, _height: (
|
|
cls.temp_dir,
|
|
prefix,
|
|
1,
|
|
"",
|
|
prefix,
|
|
),
|
|
)
|
|
|
|
cls._saved_modules = {
|
|
name: sys.modules.get(name)
|
|
for name in ("numpy", "PIL", "PIL.Image", "folder_paths")
|
|
}
|
|
sys.modules["numpy"] = fake_numpy
|
|
sys.modules["PIL"] = fake_pil_module
|
|
sys.modules["PIL.Image"] = fake_pil_image_module
|
|
sys.modules["folder_paths"] = fake_folder_paths
|
|
cls.image_nodes = importlib.import_module("dumas_image_nodes")
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
for name, module in cls._saved_modules.items():
|
|
if module is None:
|
|
sys.modules.pop(name, None)
|
|
else:
|
|
sys.modules[name] = module
|
|
|
|
def setUp(self):
|
|
FakePILImage.saved_paths = []
|
|
|
|
def test_compare_images_returns_second_input_as_new_image(self):
|
|
node = self.image_nodes.DumasImageCompareNode()
|
|
image1 = FakeTensorBatch()
|
|
image2 = FakeTensorBatch()
|
|
|
|
result = node.compare_images(image1=image1, image2=image2)
|
|
|
|
self.assertIs(result["result"][0], image2)
|
|
self.assertEqual([item["slot"] for item in result["ui"]["images"]], [1, 2])
|
|
self.assertEqual(len(FakePILImage.saved_paths), 2)
|
|
|
|
def test_compare_images_falls_back_to_first_image(self):
|
|
node = self.image_nodes.DumasImageCompareNode()
|
|
image1 = FakeTensorBatch()
|
|
|
|
result = node.compare_images(image1=image1)
|
|
|
|
self.assertIs(result["result"][0], image1)
|
|
self.assertEqual([item["slot"] for item in result["ui"]["images"]], [1])
|
|
|
|
def test_compare_images_handles_missing_inputs(self):
|
|
node = self.image_nodes.DumasImageCompareNode()
|
|
|
|
result = node.compare_images()
|
|
|
|
self.assertIsNone(result["result"][0])
|
|
self.assertEqual(result["ui"]["images"], [])
|
|
|
|
def test_saved_filenames_use_dumas_prefix(self):
|
|
node = self.image_nodes.DumasImageCompareNode()
|
|
image1 = FakeTensorBatch()
|
|
|
|
result = node.compare_images(image1=image1)
|
|
|
|
self.assertTrue(result["ui"]["images"][0]["filename"].startswith("dumas_compare"))
|
|
self.assertTrue(os.path.basename(FakePILImage.saved_paths[0][0]).startswith("dumas_compare"))
|
|
|
|
def test_save_image_uses_second_input_token(self):
|
|
node = self.image_nodes.DumasSaveImageNode()
|
|
image = FakeTensorBatch(width=10, height=12)
|
|
|
|
node.save_images(
|
|
images=image,
|
|
folder=self.temp_dir,
|
|
pattern="shot_%input%_%input2%_%counter%",
|
|
format="png",
|
|
quality=100,
|
|
embed_workflow=False,
|
|
save_on_run=True,
|
|
name="alpha.png",
|
|
name_2="beta/final",
|
|
)
|
|
|
|
saved_name = os.path.basename(FakePILImage.saved_paths[0][0])
|
|
self.assertEqual(saved_name, "shot_alpha_beta_final_001.png")
|
|
|
|
def test_save_image_resolves_date_size_and_batch_tokens(self):
|
|
node = self.image_nodes.DumasSaveImageNode()
|
|
image = FakeTensorBatch(width=10, height=12, count=2)
|
|
fake_now = time.struct_time((2026, 8, 5, 13, 7, 9, 2, 217, -1))
|
|
|
|
with mock.patch.object(self.image_nodes.time, "localtime", return_value=fake_now):
|
|
node.save_images(
|
|
images=image,
|
|
folder=self.temp_dir,
|
|
pattern="asset_%date:yyyy-MM-dd%_%date:hh-mm-ss%_%width%x%height%_%batch_num%_%counter%",
|
|
format="png",
|
|
quality=100,
|
|
embed_workflow=False,
|
|
save_on_run=True,
|
|
)
|
|
|
|
saved_names = [os.path.basename(path) for path, _args, _kwargs in FakePILImage.saved_paths]
|
|
self.assertEqual(
|
|
saved_names,
|
|
[
|
|
"asset_2026-08-05_13-07-09_10x12_0_001.png",
|
|
"asset_2026-08-05_13-07-09_10x12_1_001.png",
|
|
],
|
|
)
|
|
|
|
def test_save_image_counter_increments_for_existing_files(self):
|
|
node = self.image_nodes.DumasSaveImageNode()
|
|
image = FakeTensorBatch()
|
|
|
|
node.save_images(
|
|
images=image,
|
|
folder=self.temp_dir,
|
|
pattern="counter_%counter%",
|
|
format="png",
|
|
quality=100,
|
|
embed_workflow=False,
|
|
save_on_run=True,
|
|
)
|
|
open(FakePILImage.saved_paths[0][0], "a", encoding="utf-8").close()
|
|
node.save_images(
|
|
images=image,
|
|
folder=self.temp_dir,
|
|
pattern="counter_%counter%",
|
|
format="png",
|
|
quality=100,
|
|
embed_workflow=False,
|
|
save_on_run=True,
|
|
)
|
|
|
|
saved_names = [os.path.basename(path) for path, _args, _kwargs in FakePILImage.saved_paths]
|
|
self.assertEqual(saved_names, ["counter_001.png", "counter_002.png"])
|
|
|
|
def test_save_image_uses_same_counter_for_folder_and_filename(self):
|
|
node = self.image_nodes.DumasSaveImageNode()
|
|
image = FakeTensorBatch()
|
|
|
|
node.save_images(
|
|
images=image,
|
|
folder=self.temp_dir,
|
|
pattern="Char%counter%/Char%counter%",
|
|
format="png",
|
|
quality=100,
|
|
embed_workflow=False,
|
|
save_on_run=True,
|
|
)
|
|
|
|
saved_path = FakePILImage.saved_paths[0][0].replace("\\", "/")
|
|
self.assertTrue(saved_path.endswith("/Char001/Char001.png"))
|
|
|
|
def test_save_image_creates_nested_directories_before_saving(self):
|
|
node = self.image_nodes.DumasSaveImageNode()
|
|
image = FakeTensorBatch()
|
|
|
|
node.save_images(
|
|
images=image,
|
|
folder=self.temp_dir,
|
|
pattern="Char_%counter%/Base",
|
|
format="png",
|
|
quality=100,
|
|
embed_workflow=False,
|
|
save_on_run=True,
|
|
)
|
|
|
|
saved_path = FakePILImage.saved_paths[0][0]
|
|
self.assertTrue(os.path.isdir(os.path.dirname(saved_path)))
|
|
|
|
def test_character_reference_builds_structured_reference(self):
|
|
node = self.image_nodes.DumasCharacterReferenceNode()
|
|
image = FakeTensorBatch()
|
|
|
|
result = node.build_reference(
|
|
image=image,
|
|
character_id="char_dave",
|
|
name="Dave",
|
|
alias="The Locksmith",
|
|
gender="male",
|
|
age="41",
|
|
nationality="English",
|
|
occupation="a detective",
|
|
height_feet="6",
|
|
height_inches="2",
|
|
accent="English",
|
|
description="Square jaw, tired eyes, cropped brown hair.",
|
|
general="wears a long grey coat",
|
|
wardrobe="weathered red flight jacket, grey cargo shorts, black boots",
|
|
)
|
|
|
|
reference = result[0]
|
|
self.assertIs(reference["image"], image)
|
|
self.assertEqual(
|
|
reference,
|
|
{
|
|
"kind": "character",
|
|
"id": "char-dave",
|
|
"name": "Dave",
|
|
"aliases": ["The Locksmith"],
|
|
"picture_id": None,
|
|
"picture_label": "",
|
|
"image": image,
|
|
"summary": "Dave reference.",
|
|
"description": "Square jaw, tired eyes, cropped brown hair.",
|
|
"wardrobe": "weathered red flight jacket, grey cargo shorts, black boots",
|
|
"general": "wears a long grey coat",
|
|
"facts": {
|
|
"gender": "male",
|
|
"age": "41",
|
|
"nationality": "English",
|
|
"occupation": "a detective",
|
|
"height_feet": "6",
|
|
"height_inches": "2",
|
|
"accent": "English",
|
|
},
|
|
},
|
|
)
|
|
|
|
def test_character_reference_input_types_do_not_expose_picture_id(self):
|
|
required = self.image_nodes.DumasCharacterReferenceNode.INPUT_TYPES()["required"]
|
|
self.assertNotIn("picture_id", required)
|
|
|
|
def test_character_reference_handles_missing_optional_fields(self):
|
|
node = self.image_nodes.DumasCharacterReferenceNode()
|
|
image = FakeTensorBatch()
|
|
|
|
result = node.build_reference(
|
|
image=image,
|
|
character_id="",
|
|
name="",
|
|
alias="",
|
|
gender="",
|
|
age="unknown",
|
|
nationality="",
|
|
occupation="",
|
|
height_feet="",
|
|
height_inches="",
|
|
accent="",
|
|
description="",
|
|
general="",
|
|
wardrobe="",
|
|
)
|
|
|
|
reference = result[0]
|
|
self.assertEqual(reference["kind"], "character")
|
|
self.assertIsNone(reference["picture_id"])
|
|
self.assertEqual(reference["picture_label"], "")
|
|
self.assertEqual(reference["wardrobe"], "")
|
|
self.assertEqual(reference["general"], "")
|
|
self.assertEqual(reference["facts"]["age"], "")
|
|
|
|
def test_location_reference_builds_structured_reference(self):
|
|
node = self.image_nodes.DumasLocationReferenceNode()
|
|
image = FakeTensorBatch()
|
|
|
|
result = node.build_reference(
|
|
image=image,
|
|
location_id="coffee-shop-01",
|
|
name="Coffee Shop",
|
|
alias="Cafe Interior",
|
|
description="Warm tungsten lighting, narrow counter, rainy front window.",
|
|
general="Evening ambience, cramped but cozy.",
|
|
)
|
|
|
|
self.assertEqual(
|
|
result[0],
|
|
{
|
|
"kind": "location",
|
|
"id": "coffee-shop-01",
|
|
"name": "Coffee Shop",
|
|
"aliases": ["Cafe Interior"],
|
|
"picture_id": None,
|
|
"picture_label": "",
|
|
"image": image,
|
|
"summary": "Coffee Shop reference.",
|
|
"description": "Warm tungsten lighting, narrow counter, rainy front window.",
|
|
"wardrobe": "",
|
|
"general": "Evening ambience, cramped but cozy.",
|
|
"facts": {},
|
|
},
|
|
)
|
|
|
|
def test_location_reference_input_types_do_not_expose_picture_id(self):
|
|
required = self.image_nodes.DumasLocationReferenceNode.INPUT_TYPES()["required"]
|
|
self.assertNotIn("picture_id", required)
|
|
|
|
def test_character_helper_restores_image_and_text_outputs(self):
|
|
node = self.image_nodes.DumasCharacterHelperNode()
|
|
image1 = FakeTensorBatch()
|
|
image2 = FakeTensorBatch()
|
|
|
|
result = node.build_character_text(
|
|
image1=image1,
|
|
image2=image2,
|
|
image1_picture_id="1",
|
|
image2_picture_id="2",
|
|
character_id="char_dave",
|
|
name="Dave",
|
|
alias="The Locksmith",
|
|
gender="male",
|
|
age="41",
|
|
nationality="English",
|
|
occupation="a detective",
|
|
height_feet="6",
|
|
height_inches="2",
|
|
accent="English",
|
|
general="Moves carefully and notices every exit",
|
|
wardrobe="weathered red flight jacket, grey cargo shorts, black boots",
|
|
)
|
|
|
|
self.assertIs(result[0], image1)
|
|
self.assertIs(result[1], image2)
|
|
self.assertIn("<Picture 1> and <Picture 2> reference the same character", result[2])
|
|
self.assertIn("Dave is also known as The Locksmith", result[2])
|
|
self.assertIn("is 41 years old", result[2])
|
|
self.assertEqual(result[3], "Dave = weathered red flight jacket, grey cargo shorts, black boots")
|
|
|
|
def test_location_helper_matches_character_helper_shape_without_wardrobe(self):
|
|
node = self.image_nodes.DumasLocationHelperNode()
|
|
image1 = FakeTensorBatch()
|
|
image2 = FakeTensorBatch()
|
|
|
|
result = node.build_location_text(
|
|
image1=image1,
|
|
image2=image2,
|
|
image1_picture_id="3",
|
|
image2_picture_id="4",
|
|
location_id="coffee-shop-01",
|
|
name="Coffee Shop",
|
|
alias="Cafe Interior",
|
|
description="Warm tungsten lighting, narrow counter, rainy front window",
|
|
general="Evening ambience, cramped but cozy",
|
|
)
|
|
|
|
self.assertIs(result[0], image1)
|
|
self.assertIs(result[1], image2)
|
|
self.assertIn("<Picture 3> and <Picture 4> reference the same location", result[2])
|
|
self.assertIn("Coffee Shop is also known as Cafe Interior", result[2])
|
|
self.assertIn("Warm tungsten lighting, narrow counter, rainy front window.", result[2])
|
|
self.assertIn("Evening ambience, cramped but cozy.", result[2])
|
|
self.assertEqual(len(result), 3)
|
|
|
|
def test_soundscape_helper_defaults_to_selected_preset_description(self):
|
|
node = self.image_nodes.DumasSoundscapeHelperNode()
|
|
|
|
result = node.build_soundscape("rainy street", "")
|
|
|
|
self.assertEqual(result[0], "steady rain, wet pavement, distant traffic hum")
|
|
|
|
def test_h3_prompt_curator_compacts_named_references(self):
|
|
node = self.image_nodes.DumasH3PromptCuratorNode()
|
|
dave_image = FakeTensorBatch()
|
|
cafe_image = FakeTensorBatch()
|
|
van_image = FakeTensorBatch()
|
|
dave = self.image_nodes.make_reference(
|
|
kind="character",
|
|
image=dave_image,
|
|
name="Dave",
|
|
aliases="The Locksmith",
|
|
description="tired eyes, cropped brown hair",
|
|
wardrobe="red flight jacket",
|
|
)
|
|
cafe = self.image_nodes.make_reference(
|
|
kind="location",
|
|
image=cafe_image,
|
|
name="Coffee Shop",
|
|
description="warm tungsten lighting and rainy windows",
|
|
)
|
|
van = self.image_nodes.make_reference(
|
|
kind="location",
|
|
image=van_image,
|
|
name="Blue Van",
|
|
description="scuffed blue delivery van",
|
|
)
|
|
|
|
result = node.curate_prompt(
|
|
action_prompt="Dave runs from the Coffee Shop into the rain.",
|
|
anatomy_guard="auto",
|
|
subject_count_guard="auto",
|
|
anchor="grounded handheld thriller",
|
|
soundscape="steady rain",
|
|
ref_1=dave,
|
|
ref_2=van,
|
|
ref_3=cafe,
|
|
)
|
|
|
|
prompt = result[0]
|
|
self.assertIn("<Picture 1> Dave", prompt)
|
|
self.assertIn("<Picture 2> Coffee Shop", prompt)
|
|
self.assertIn("Action: Dave runs from the Coffee Shop into the rain.", prompt)
|
|
self.assertIn("Anatomy guard:", prompt)
|
|
self.assertIn("Subject count guard:", prompt)
|
|
self.assertIn("exactly one named character: <Picture 1> Dave", prompt)
|
|
self.assertIs(result[1], dave_image)
|
|
self.assertIs(result[2], cafe_image)
|
|
self.assertIsNone(result[3])
|
|
self.assertEqual(result[10], 2)
|
|
self.assertIn("input 3-><Picture 2> Coffee Shop", result[11])
|
|
|
|
def test_h3_prompt_curator_renumbers_explicit_reference_tags(self):
|
|
node = self.image_nodes.DumasH3PromptCuratorNode()
|
|
image1 = FakeTensorBatch()
|
|
image3 = FakeTensorBatch()
|
|
unused = FakeTensorBatch()
|
|
first = self.image_nodes.make_reference(kind="character", image=image1, name="Maya")
|
|
second = self.image_nodes.make_reference(kind="location", image=unused, name="Lobby")
|
|
third = self.image_nodes.make_reference(kind="location", image=image3, name="Rooftop")
|
|
|
|
result = node.curate_prompt(
|
|
action_prompt="<Picture 1> Maya crosses to <ref3> as the wind rises.",
|
|
anatomy_guard="off",
|
|
subject_count_guard="off",
|
|
ref_1=first,
|
|
ref_2=second,
|
|
ref_3=third,
|
|
)
|
|
|
|
prompt = result[0]
|
|
self.assertIn("<Picture 1> Maya crosses to <Picture 2>", prompt)
|
|
self.assertNotIn("<Picture 3>", prompt)
|
|
self.assertIs(result[1], image1)
|
|
self.assertIs(result[2], image3)
|
|
self.assertIsNone(result[3])
|
|
self.assertEqual(result[10], 2)
|
|
|
|
def test_h3_prompt_curator_can_force_subject_count_without_character_refs(self):
|
|
node = self.image_nodes.DumasH3PromptCuratorNode()
|
|
|
|
result = node.curate_prompt(
|
|
action_prompt="A locked-off shot of the empty corridor.",
|
|
anatomy_guard="off",
|
|
subject_count_guard="on",
|
|
)
|
|
|
|
self.assertIn("Subject count guard:", result[0])
|
|
self.assertIn("Only include the people explicitly described", result[0])
|
|
self.assertEqual(result[10], 0)
|
|
|
|
def test_h3_prompt_curator_defaults_anatomy_guard_to_on(self):
|
|
required = self.image_nodes.DumasH3PromptCuratorNode.INPUT_TYPES()["required"]
|
|
|
|
self.assertEqual(required["anatomy_guard"][1]["default"], "on")
|
|
|
|
def test_helper_node_mappings_use_general_purpose_helpers(self):
|
|
mappings = self.image_nodes.NODE_CLASS_MAPPINGS
|
|
display = self.image_nodes.NODE_DISPLAY_NAME_MAPPINGS
|
|
|
|
self.assertIs(mappings["DumasCharacterHelper"], self.image_nodes.DumasCharacterHelperNode)
|
|
self.assertIs(mappings["DumasLocationHelper"], self.image_nodes.DumasLocationHelperNode)
|
|
self.assertIs(mappings["DumasSoundscapeHelper"], self.image_nodes.DumasSoundscapeHelperNode)
|
|
self.assertIs(mappings["DumasH3PromptCurator"], self.image_nodes.DumasH3PromptCuratorNode)
|
|
self.assertEqual(display["DumasCharacterHelper"], "Dumas Character Helper")
|
|
self.assertEqual(display["DumasLocationHelper"], "Dumas Location Helper")
|
|
self.assertEqual(display["DumasSoundscapeHelper"], "Dumas Soundscape Helper")
|
|
self.assertEqual(display["DumasH3PromptCurator"], "Dumas H3 Prompt Curator")
|
|
|
|
def test_h3_prompt_curator_uses_documented_reference_limits(self):
|
|
node = self.image_nodes.DumasH3PromptCuratorNode()
|
|
self.assertEqual(len(node.RETURN_TYPES), 12)
|
|
self.assertEqual(node.RETURN_NAMES[1:10], tuple(f"ref_image_{i}" for i in range(1, 10)))
|
|
|
|
def test_normalize_reference_upgrades_generic_summary_with_socket_picture_id(self):
|
|
image = FakeTensorBatch()
|
|
|
|
reference = self.image_nodes.normalize_reference(
|
|
{
|
|
"kind": "character",
|
|
"name": "Dave",
|
|
"image": image,
|
|
"summary": "Dave reference.",
|
|
},
|
|
picture_id=3,
|
|
allow_image_fallback=False,
|
|
)
|
|
|
|
self.assertEqual(reference["picture_id"], 3)
|
|
self.assertEqual(reference["picture_label"], "<Picture 3>")
|
|
self.assertEqual(reference["summary"], "Dave shown in <Picture 3>.")
|
|
|
|
def test_normalize_reference_keeps_custom_summary_when_socket_picture_id_is_added(self):
|
|
image = FakeTensorBatch()
|
|
|
|
reference = self.image_nodes.normalize_reference(
|
|
{
|
|
"kind": "character",
|
|
"name": "Dave",
|
|
"image": image,
|
|
"summary": "Primary hero look for the opening close-up.",
|
|
},
|
|
picture_id=3,
|
|
allow_image_fallback=False,
|
|
)
|
|
|
|
self.assertEqual(reference["picture_id"], 3)
|
|
self.assertEqual(reference["summary"], "Primary hero look for the opening close-up.")
|
|
|
|
def test_anchor_style_node_exposes_requested_presets(self):
|
|
input_types = self.image_nodes.DumasAnchorStyleNode.INPUT_TYPES()
|
|
options = input_types["required"]["anchor_style"][0]
|
|
|
|
self.assertGreaterEqual(len(options), 40)
|
|
self.assertIn("cinematic action movie", options)
|
|
self.assertIn("comedy", options)
|
|
self.assertIn("found footage", options)
|
|
self.assertIn("90s sitcom", options)
|
|
self.assertIn("mobile/cell phone captured", options)
|
|
self.assertIn("news broadcast", options)
|
|
self.assertIn("mockumentary", options)
|
|
self.assertIn("heist thriller", options)
|
|
self.assertIn("cyberpunk neon", options)
|
|
self.assertIn("nature documentary", options)
|
|
self.assertIn("courtroom drama", options)
|
|
|
|
def test_anchor_style_node_defaults_to_selected_preset_description(self):
|
|
node = self.image_nodes.DumasAnchorStyleNode()
|
|
|
|
result = node.build_anchor("found footage", "")
|
|
|
|
self.assertIn("found-footage", result[0])
|
|
self.assertIn("real time", result[0])
|
|
self.assertNotIn("persistent camera language", result[0])
|
|
|
|
def test_anchor_style_node_prefers_manual_description_edits(self):
|
|
node = self.image_nodes.DumasAnchorStyleNode()
|
|
custom = "Lo-fi pirate broadcast with smeared highlights and anxious zoom corrections."
|
|
|
|
result = node.build_anchor("news broadcast", custom)
|
|
|
|
self.assertEqual(result[0], custom)
|
|
|
|
def test_save_image_returns_ui_entries_for_output_folder(self):
|
|
node = self.image_nodes.DumasSaveImageNode()
|
|
image = FakeTensorBatch()
|
|
|
|
result = node.save_images(
|
|
images=image,
|
|
folder="",
|
|
pattern="result_%counter%",
|
|
format="jpg",
|
|
quality=90,
|
|
embed_workflow=False,
|
|
save_on_run=True,
|
|
)
|
|
|
|
self.assertEqual(result["ui"]["images"][0]["type"], "output")
|
|
self.assertTrue(result["ui"]["images"][0]["filename"].endswith(".jpg"))
|
|
|
|
def test_save_image_returns_ui_entries_for_external_folder(self):
|
|
node = self.image_nodes.DumasSaveImageNode()
|
|
image = FakeTensorBatch()
|
|
external_dir = tempfile.mkdtemp(prefix="dumas-image-node-external-")
|
|
|
|
result = node.save_images(
|
|
images=image,
|
|
folder=external_dir,
|
|
pattern="external_%counter%",
|
|
format="png",
|
|
quality=100,
|
|
embed_workflow=False,
|
|
save_on_run=True,
|
|
)
|
|
|
|
self.assertEqual(result["ui"]["images"][0]["type"], "external")
|
|
self.assertEqual(result["ui"]["images"][0]["subfolder"], external_dir.replace("\\", "/"))
|
|
self.assertTrue(result["ui"]["images"][0]["token"])
|
|
|
|
def test_save_image_handles_windows_different_drive_paths(self):
|
|
node = self.image_nodes.DumasSaveImageNode()
|
|
image = FakeTensorBatch()
|
|
|
|
with mock.patch.object(
|
|
self.image_nodes.folder_paths,
|
|
"get_output_directory",
|
|
return_value="C:\\ComfyUI\\output",
|
|
), mock.patch.object(
|
|
self.image_nodes.os.path,
|
|
"abspath",
|
|
return_value="D:\\renders",
|
|
):
|
|
result = node.save_images(
|
|
images=image,
|
|
folder="D:\\renders",
|
|
pattern="drive_%counter%",
|
|
format="png",
|
|
quality=100,
|
|
embed_workflow=False,
|
|
save_on_run=True,
|
|
)
|
|
|
|
self.assertEqual(result["ui"]["images"][0]["type"], "external")
|
|
self.assertEqual(result["ui"]["images"][0]["subfolder"], "D:/renders")
|
|
self.assertTrue(result["ui"]["images"][0]["token"])
|
|
|
|
def test_save_image_skips_when_save_is_disabled(self):
|
|
node = self.image_nodes.DumasSaveImageNode()
|
|
image = FakeTensorBatch()
|
|
|
|
result = node.save_images(
|
|
images=image,
|
|
folder=self.temp_dir,
|
|
pattern="ignored_%counter%",
|
|
format="png",
|
|
quality=100,
|
|
embed_workflow=False,
|
|
save_on_run=False,
|
|
)
|
|
|
|
self.assertEqual(result["ui"]["images"], [])
|
|
self.assertEqual(FakePILImage.saved_paths, [])
|
|
|
|
def test_load_images_folder_uses_manual_selected_files(self):
|
|
node = self.image_nodes.DumasLoadImagesFolderNode()
|
|
folder = tempfile.mkdtemp(prefix="dumas-load-folder-manual-")
|
|
for name in ("b.png", "a.png", "notes.txt"):
|
|
open(os.path.join(folder, name), "a", encoding="utf-8").close()
|
|
state = json.dumps(
|
|
{
|
|
"folder": folder,
|
|
"recursive": False,
|
|
"sort": "name",
|
|
"sort_dir": "asc",
|
|
"selection_mode": "selected",
|
|
"selected": ["b.png", "a.png"],
|
|
}
|
|
)
|
|
|
|
with mock.patch.object(
|
|
self.image_nodes,
|
|
"_load_folder_image",
|
|
side_effect=lambda path: (f"image:{os.path.basename(path)}", f"mask:{os.path.basename(path)}", 32, 24),
|
|
):
|
|
result = node.load(state)
|
|
|
|
self.assertEqual(result[0], ["image:b.png", "image:a.png"])
|
|
self.assertEqual(result[1], ["mask:b.png", "mask:a.png"])
|
|
self.assertEqual(result[4], ["b", "a"])
|
|
self.assertEqual(result[5], [1, 2])
|
|
self.assertEqual(result[6], [2, 2])
|
|
|
|
def test_load_images_folder_first_n_uses_sorted_files(self):
|
|
node = self.image_nodes.DumasLoadImagesFolderNode()
|
|
folder = tempfile.mkdtemp(prefix="dumas-load-folder-first-")
|
|
for name in ("c.png", "a.png", "b.png"):
|
|
open(os.path.join(folder, name), "a", encoding="utf-8").close()
|
|
state = json.dumps(
|
|
{
|
|
"folder": folder,
|
|
"recursive": False,
|
|
"sort": "name",
|
|
"sort_dir": "asc",
|
|
"selection_mode": "first_n",
|
|
"first_n": 2,
|
|
}
|
|
)
|
|
|
|
with mock.patch.object(
|
|
self.image_nodes,
|
|
"_load_folder_image",
|
|
side_effect=lambda path: (f"image:{os.path.basename(path)}", f"mask:{os.path.basename(path)}", 64, 48),
|
|
):
|
|
result = node.load(state)
|
|
|
|
self.assertEqual(result[0], ["image:a.png", "image:b.png"])
|
|
self.assertEqual(result[4], ["a", "b"])
|
|
self.assertEqual(result[6], [2, 2])
|
|
|
|
def test_load_images_folder_random_selects_one_visible_image(self):
|
|
node = self.image_nodes.DumasLoadImagesFolderNode()
|
|
folder = tempfile.mkdtemp(prefix="dumas-load-folder-random-")
|
|
for name in ("a.png", "b.png", "c.png"):
|
|
open(os.path.join(folder, name), "a", encoding="utf-8").close()
|
|
state = json.dumps(
|
|
{
|
|
"folder": folder,
|
|
"recursive": False,
|
|
"sort": "name",
|
|
"sort_dir": "asc",
|
|
"selection_mode": "random",
|
|
}
|
|
)
|
|
|
|
with mock.patch.object(
|
|
self.image_nodes.random,
|
|
"choice",
|
|
side_effect=lambda files: files[1],
|
|
), mock.patch.object(
|
|
self.image_nodes,
|
|
"_load_folder_image",
|
|
side_effect=lambda path: (f"image:{os.path.basename(path)}", f"mask:{os.path.basename(path)}", 80, 60),
|
|
):
|
|
result = node.load(state)
|
|
|
|
self.assertEqual(result[0], ["image:b.png"])
|
|
self.assertEqual(result[4], ["b"])
|
|
self.assertEqual(result[5], [1])
|
|
self.assertEqual(result[6], [1])
|
|
|
|
def test_load_images_folder_random_is_changed_changes_each_run(self):
|
|
node = self.image_nodes.DumasLoadImagesFolderNode()
|
|
folder = tempfile.mkdtemp(prefix="dumas-load-folder-changed-")
|
|
open(os.path.join(folder, "a.png"), "a", encoding="utf-8").close()
|
|
state = json.dumps(
|
|
{
|
|
"folder": folder,
|
|
"recursive": False,
|
|
"sort": "name",
|
|
"sort_dir": "asc",
|
|
"selection_mode": "random",
|
|
}
|
|
)
|
|
|
|
with mock.patch.object(self.image_nodes.time, "time_ns", side_effect=[1, 2]):
|
|
changed_1 = node.IS_CHANGED(state)
|
|
changed_2 = node.IS_CHANGED(state)
|
|
|
|
self.assertNotEqual(changed_1, changed_2)
|
|
|
|
def test_h3_plan_scene_images_attach_and_extract(self):
|
|
attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode()
|
|
extract_node = self.image_nodes.DumasH3PlanExtractSceneImagesNode()
|
|
plan = {"shots": [{"id": "intro"}, {"id": "middle"}]}
|
|
image_a = FakeTensorBatch()
|
|
image_b = FakeTensorBatch(width=10, height=10)
|
|
|
|
attached_plan, connected = attach_node.attach(
|
|
plan=plan,
|
|
scene_index=2,
|
|
image1=image_a,
|
|
image3=image_b,
|
|
)
|
|
extracted = extract_node.extract(attached_plan, 2)
|
|
|
|
self.assertEqual(connected, 2)
|
|
self.assertEqual(attached_plan["_dumas_scene_image_bindings"]["scene_counts"], {"2": 2})
|
|
self.assertIs(extracted[1], image_a)
|
|
self.assertIsNone(extracted[2])
|
|
self.assertIs(extracted[3], image_b)
|
|
self.assertEqual(extracted[-1], 2)
|
|
self.assertNotIn("_dumas_scene_image_bindings", plan)
|
|
|
|
def test_h3_plan_scene_images_support_multiple_scenes(self):
|
|
attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode()
|
|
extract_node = self.image_nodes.DumasH3PlanExtractSceneImagesNode()
|
|
plan = {"shots": [{"id": "one"}, {"id": "two"}]}
|
|
image_1 = FakeTensorBatch()
|
|
image_2 = FakeTensorBatch(width=12, height=9)
|
|
image_3 = FakeTensorBatch(width=8, height=8)
|
|
image_4 = FakeTensorBatch(width=16, height=16)
|
|
image_5 = FakeTensorBatch(width=20, height=12)
|
|
|
|
plan_after_first, _connected = attach_node.attach(plan=plan, scene_index=1, image2=image_1)
|
|
plan_after_second, _connected = attach_node.attach(
|
|
plan=plan_after_first,
|
|
scene_index=2,
|
|
image6=image_2,
|
|
image7=image_3,
|
|
image8=image_4,
|
|
image9=image_5,
|
|
)
|
|
|
|
scene1 = extract_node.extract(plan_after_second, 1)
|
|
scene2 = extract_node.extract(plan_after_second, 2)
|
|
|
|
self.assertIs(scene1[2], image_1)
|
|
self.assertEqual(scene1[-1], 1)
|
|
self.assertIs(scene2[6], image_2)
|
|
self.assertIs(scene2[7], image_3)
|
|
self.assertIs(scene2[8], image_4)
|
|
self.assertIs(scene2[9], image_5)
|
|
self.assertEqual(scene2[-1], 4)
|
|
self.assertEqual(
|
|
plan_after_second["_dumas_scene_image_bindings"]["scene_counts"],
|
|
{"1": 1, "2": 4},
|
|
)
|
|
|
|
def test_h3_plan_scene_images_metadata_is_json_serializable(self):
|
|
attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode()
|
|
plan = {"shots": [{"id": "one"}]}
|
|
|
|
attached_plan, connected = attach_node.attach(
|
|
plan=plan,
|
|
scene_index=1,
|
|
image4=FakeTensorBatch(),
|
|
)
|
|
|
|
self.assertEqual(connected, 1)
|
|
json.dumps(attached_plan)
|
|
|
|
def test_h3_plan_scene_images_support_nine_slots(self):
|
|
attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode()
|
|
extract_node = self.image_nodes.DumasH3PlanExtractSceneImagesNode()
|
|
plan = {"shots": [{"id": "one"}]}
|
|
images = [FakeTensorBatch(width=8 + index, height=8 + index) for index in range(9)]
|
|
|
|
attached_plan, connected = attach_node.attach(
|
|
plan=plan,
|
|
scene_index=1,
|
|
image1=images[0],
|
|
image2=images[1],
|
|
image3=images[2],
|
|
image4=images[3],
|
|
image5=images[4],
|
|
image6=images[5],
|
|
image7=images[6],
|
|
image8=images[7],
|
|
image9=images[8],
|
|
)
|
|
extracted = extract_node.extract(attached_plan, 1)
|
|
|
|
self.assertEqual(connected, 9)
|
|
for index, image in enumerate(images, start=1):
|
|
self.assertIs(extracted[index], image)
|
|
self.assertEqual(extracted[-1], 9)
|
|
|
|
def test_h3_plan_scene_images_can_clear_a_scene_binding(self):
|
|
attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode()
|
|
extract_node = self.image_nodes.DumasH3PlanExtractSceneImagesNode()
|
|
plan = {"shots": [{"id": "one"}]}
|
|
|
|
attached_plan, connected = attach_node.attach(
|
|
plan=plan,
|
|
scene_index=1,
|
|
image1=FakeTensorBatch(),
|
|
)
|
|
cleared_plan, cleared = attach_node.attach(plan=attached_plan, scene_index=1)
|
|
extracted = extract_node.extract(cleared_plan, 1)
|
|
|
|
self.assertEqual(connected, 1)
|
|
self.assertEqual(cleared, 0)
|
|
self.assertNotIn("_dumas_scene_image_bindings", cleared_plan)
|
|
self.assertEqual(extracted[-1], 0)
|
|
|
|
def test_h3_plan_scene_images_reject_invalid_scene_index(self):
|
|
attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode()
|
|
plan = {"shots": [{"id": "one"}]}
|
|
|
|
with self.assertRaises(ValueError):
|
|
attach_node.attach(plan=plan, scene_index=2, image1=FakeTensorBatch())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|