342 lines
18 KiB
Python
342 lines
18 KiB
Python
"""Regression tests for typed dynamic image-list normalization and graph scheduling."""
|
|
import asyncio
|
|
import unittest
|
|
|
|
import torch
|
|
|
|
from image_list.runtime import (
|
|
END_CLASS,
|
|
IMAGE_OUT,
|
|
CONTEXT_OUT,
|
|
START_CLASS,
|
|
VALUE_OUT_BASE,
|
|
ImageListCancelledError,
|
|
build_worker_prompt,
|
|
create_master_job,
|
|
decode_item,
|
|
decode_result,
|
|
dynamic_schedule,
|
|
encode_item,
|
|
encode_result,
|
|
normalize_image_list_inputs,
|
|
normalize_worker_result,
|
|
remove_master_job,
|
|
set_master_result,
|
|
validate_image_list_regions,
|
|
)
|
|
|
|
|
|
def start_node(image_link):
|
|
return {"class_type": START_CLASS, "inputs": {"images": image_link}}
|
|
|
|
|
|
def graph_one_region():
|
|
return {
|
|
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
|
|
"2": {"class_type": "EmptyImage", "inputs": {"height": 8, "width": 8, "batch_size": 2}},
|
|
"3": start_node(["2", 0]),
|
|
"4": {"class_type": "GenericFrameOperation", "inputs": {
|
|
"image": ["3", IMAGE_OUT], "model": ["1", 0], "seed": ["3", VALUE_OUT_BASE],
|
|
}},
|
|
"5": {"class_type": END_CLASS, "inputs": {
|
|
"images": ["4", 0], "context": ["3", CONTEXT_OUT],
|
|
}},
|
|
"6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
|
|
}
|
|
|
|
|
|
def graph_two_sequential_regions():
|
|
return {
|
|
"1": {"class_type": "EmptyImage", "inputs": {"height": 4, "width": 4, "batch_size": 2}},
|
|
"2": start_node(["1", 0]),
|
|
"3": {"class_type": "FirstFrameOperation", "inputs": {"image": ["2", IMAGE_OUT]}},
|
|
"4": {"class_type": END_CLASS, "inputs": {"images": ["3", 0], "context": ["2", CONTEXT_OUT]}},
|
|
"5": start_node(["4", 0]),
|
|
"6": {"class_type": "SecondFrameOperation", "inputs": {"image": ["5", IMAGE_OUT]}},
|
|
"7": {"class_type": END_CLASS, "inputs": {"images": ["6", 0], "context": ["5", CONTEXT_OUT]}},
|
|
"8": {"class_type": "SaveImage", "inputs": {"images": ["7", 0]}},
|
|
}
|
|
|
|
|
|
class ImageListNormalizationTests(unittest.TestCase):
|
|
def test_batch_frame_normalization_preserves_sizes_channels_and_float_values(self):
|
|
images = [
|
|
torch.tensor([[[[0.125, 0.25, 0.5]]]], dtype=torch.float64),
|
|
torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(1, 2, 3, 4) / 17,
|
|
]
|
|
refs = torch.ones((2, 5, 7, 3), dtype=torch.float16) * 0.333
|
|
masks = [torch.full((1, 2, 3), 0.125, dtype=torch.float64),
|
|
torch.full((1, 8, 9), 0.875, dtype=torch.float32)]
|
|
rows = normalize_image_list_inputs(images, masks, {
|
|
"value_1": [refs[:1], refs[1:]], "value_2": ["first", "second"],
|
|
"value_3": [0.2], "value_4": [37],
|
|
})
|
|
self.assertEqual(len(rows), 2)
|
|
self.assertEqual([tuple(row["image"].shape) for row in rows], [(1, 1, 1, 3), (1, 2, 3, 4)])
|
|
self.assertEqual([row["image"].dtype for row in rows], [torch.float64, torch.float32])
|
|
self.assertTrue(torch.equal(rows[0]["values"]["value_1"], refs[:1]))
|
|
self.assertEqual(rows[1]["mask"].dtype, torch.float32)
|
|
self.assertEqual([row["index"] for row in rows], [0, 1])
|
|
self.assertEqual([row["values"]["value_2"] for row in rows], ["first", "second"])
|
|
self.assertEqual([row["values"]["value_3"] for row in rows], [0.2, 0.2])
|
|
self.assertEqual([row["values"]["value_4"] for row in rows], [37, 37])
|
|
|
|
def test_one_value_is_shared_and_n_values_pair_in_order(self):
|
|
image = torch.zeros((3, 2, 4, 3), dtype=torch.float32)
|
|
rows = normalize_image_list_inputs(image, values={
|
|
"value_1": ["same"], "value_2": [0.4], "value_3": [9], "value_4": [True]})
|
|
self.assertEqual([row["values"]["value_1"] for row in rows], ["same"] * 3)
|
|
self.assertEqual([row["values"]["value_2"] for row in rows], [0.4] * 3)
|
|
self.assertEqual([row["values"]["value_3"] for row in rows], [9] * 3)
|
|
self.assertEqual([row["values"]["value_4"] for row in rows], [True] * 3)
|
|
rows = normalize_image_list_inputs(image, values={"value_2": [1, 2, 3]})
|
|
self.assertEqual([row["values"] for row in rows], [{"value_2": 1}, {"value_2": 2}, {"value_2": 3}])
|
|
with self.assertRaisesRegex(ValueError, "value_1 has 2 values; connect exactly 1.*or 3"):
|
|
normalize_image_list_inputs(image, values={"value_1": ["one", "two"]})
|
|
with self.assertRaisesRegex(ValueError, "mask has 2 values"):
|
|
normalize_image_list_inputs(image, mask=torch.zeros((2, 2, 4)))
|
|
shared_mask = normalize_image_list_inputs(image, mask=torch.ones((1, 2, 4)))
|
|
self.assertTrue(all(torch.equal(row["mask"], torch.ones((1, 2, 4))) for row in shared_mask))
|
|
|
|
def test_value_batches_stay_whole_and_only_images_split(self):
|
|
images = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
|
|
pair = torch.rand((2, 4, 4, 3)) # e.g. two references for one item
|
|
rows = normalize_image_list_inputs(images, values={"value_1": [pair]})
|
|
self.assertTrue(all(torch.equal(row["values"]["value_1"], pair) for row in rows))
|
|
per_item = [torch.rand((2, 4, 4, 3)), torch.rand((2, 4, 4, 3))]
|
|
rows = normalize_image_list_inputs(images, values={"value_1": per_item})
|
|
self.assertTrue(torch.equal(rows[1]["values"]["value_1"], per_item[1]))
|
|
|
|
def test_rejects_unsupported_values_and_counts(self):
|
|
images = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
|
|
with self.assertRaisesRegex(ValueError, "value_1 has 3 values"):
|
|
normalize_image_list_inputs(images, values={"value_1": [1, 2, 3]})
|
|
with self.assertRaisesRegex(TypeError, "unsupported type dict"):
|
|
normalize_image_list_inputs(images, values={"value_1": [{"samples": 1}]})
|
|
with self.assertRaisesRegex(TypeError, "floating-point"):
|
|
normalize_image_list_inputs(images, values={"value_1": [torch.zeros(2, dtype=torch.int64)]})
|
|
with self.assertRaisesRegex(TypeError, "finite"):
|
|
normalize_image_list_inputs(images, values={"value_1": [float("nan")]})
|
|
with self.assertRaisesRegex(ValueError, "Unknown Start value slot"):
|
|
normalize_image_list_inputs(images, values={"value_11": [1]})
|
|
with self.assertRaisesRegex(ValueError, "at least one image"):
|
|
normalize_image_list_inputs(None)
|
|
# Unconnected Autogrow slots arrive as None and are ignored.
|
|
rows = normalize_image_list_inputs(images, values={"value_1": [None], "value_2": None})
|
|
self.assertEqual([row["values"] for row in rows], [{}, {}])
|
|
|
|
def test_typed_transport_round_trips_float_images_masks_and_scalar_values(self):
|
|
item = {"image": torch.arange(24, dtype=torch.float64).reshape(1, 2, 4, 3) / 7,
|
|
"mask": torch.full((1, 5, 6), 0.375, dtype=torch.float32),
|
|
"values": {"value_1": torch.ones((2, 1, 3, 4), dtype=torch.float16),
|
|
"value_2": "safe text", "value_3": 0.625, "value_4": 42,
|
|
"value_5": False},
|
|
"index": 3}
|
|
restored = decode_item(encode_item(item))
|
|
self.assertTrue(torch.equal(restored["image"], item["image"]))
|
|
self.assertEqual(restored["image"].dtype, torch.float64)
|
|
self.assertTrue(torch.equal(restored["values"]["value_1"], item["values"]["value_1"]))
|
|
self.assertEqual(restored["values"]["value_1"].dtype, torch.float16)
|
|
self.assertTrue(torch.equal(restored["mask"], item["mask"]))
|
|
self.assertEqual(restored["values"]["value_2"], "safe text")
|
|
self.assertEqual(restored["values"]["value_3"], 0.625)
|
|
self.assertIs(restored["values"]["value_4"], 42)
|
|
self.assertIs(restored["values"]["value_5"], False)
|
|
with self.assertRaisesRegex(ValueError, "unknown slot"):
|
|
decode_item({**encode_item(item), "values": {"strength": {"type": "FLOAT", "value": 1.0}}})
|
|
with self.assertRaisesRegex(ValueError, "unsupported type"):
|
|
decode_item({**encode_item(item), "values": {"value_1": {"type": "PATH", "value": "/x"}}})
|
|
with self.assertRaisesRegex(TypeError, "floating-point tensor"):
|
|
encode_item({**item, "image": torch.ones((1, 2, 4, 3), dtype=torch.int32)})
|
|
with self.assertRaisesRegex(ValueError, "Unsupported image-list transport type"):
|
|
from image_list.runtime import encode_value
|
|
encode_value("/untrusted/path", "PATH")
|
|
|
|
def test_result_transport_preserves_ragged_result_shapes(self):
|
|
result = normalize_worker_result(
|
|
[torch.ones((1, 3, 5, 4), dtype=torch.float32)],
|
|
[torch.ones((1, 3, 5), dtype=torch.float16)],
|
|
["done"],
|
|
)
|
|
encoded = encode_result(result)
|
|
decoded = decode_result(encoded)
|
|
self.assertEqual(tuple(decoded["image"].shape), (1, 3, 5, 4))
|
|
self.assertEqual(decoded["image"].dtype, torch.float32)
|
|
self.assertEqual(decoded["mask"].dtype, torch.float16)
|
|
self.assertEqual(decoded["text"], "done")
|
|
with self.assertRaisesRegex(ValueError, "exactly one IMAGE"):
|
|
normalize_worker_result(torch.zeros((2, 3, 4, 3)))
|
|
|
|
|
|
class ImageListGraphBoundaryTests(unittest.TestCase):
|
|
def test_prunes_upstream_item_source_but_keeps_external_model_loaders(self):
|
|
prompt = graph_one_region()
|
|
region = validate_image_list_regions(prompt)[0]
|
|
worker = build_worker_prompt(prompt, region, "token_12345678901234567890123456789012",
|
|
"ctx", 1, "http://master.invalid")
|
|
self.assertEqual(set(worker), {"1", "3", "4", "5"})
|
|
self.assertEqual(worker["3"]["inputs"]["worker_item_token"], "token_12345678901234567890123456789012")
|
|
self.assertNotIn("images", worker["3"]["inputs"])
|
|
|
|
def test_worker_prompt_drops_per_item_value_links(self):
|
|
prompt = graph_one_region()
|
|
prompt["8"] = {"class_type": "StringList", "inputs": {}}
|
|
prompt["3"]["inputs"].update({"values.value_1": ["8", 0], "values.value_2": ["8", 0]})
|
|
import sys, types
|
|
from unittest.mock import patch
|
|
registry = {"CheckpointLoaderSimple": type("Loader", (), {"RETURN_TYPES": ("MODEL",)})}
|
|
with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}):
|
|
region = validate_image_list_regions(prompt)[0]
|
|
worker = build_worker_prompt(prompt, region, "token_12345678901234567890123456789012", "ctx", 0, "http://m")
|
|
self.assertFalse([name for name in worker["3"]["inputs"] if name.startswith("values.")])
|
|
self.assertNotIn("8", worker)
|
|
self.assertEqual(worker["4"]["inputs"]["model"], ["1", 0])
|
|
self.assertNotIn("6", worker)
|
|
|
|
def test_rejects_upstream_processing_and_leaf_image_bypasses(self):
|
|
# Both a leaf source and a generated batch must cross Start exactly once.
|
|
for generated in (False, True):
|
|
with self.subTest(generated=generated):
|
|
prompt = graph_one_region()
|
|
if generated:
|
|
prompt["2"] = {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0]}}
|
|
prompt["7"] = {"class_type": "KSampler", "inputs": {"model": ["1", 0]}}
|
|
prompt["4"]["inputs"]["reference"] = ["2", 0]
|
|
for explicit_targets in (False, True):
|
|
with self.subTest(explicit_targets=explicit_targets):
|
|
with self.assertRaisesRegex(ValueError, "bypasses Start.*2"):
|
|
validate_image_list_regions(prompt, explicit_targets=explicit_targets)
|
|
|
|
def test_shared_typed_loader_and_leaf_constant_remain_valid(self):
|
|
import sys
|
|
import types
|
|
from unittest.mock import patch
|
|
prompt = graph_one_region()
|
|
prompt["7"] = {"class_type": "SeedConstant", "inputs": {"value": 9}}
|
|
prompt["2"]["inputs"].update({"model": ["1", 0], "seed": ["7", 0]})
|
|
prompt["4"]["inputs"]["shared_seed"] = ["7", 0]
|
|
registry = {
|
|
"CheckpointLoaderSimple": type("Loader", (), {"RETURN_TYPES": ("MODEL", "CLIP", "VAE")}),
|
|
"SeedConstant": type("Constant", (), {"RETURN_TYPES": ("INT",)}),
|
|
}
|
|
with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}):
|
|
region = validate_image_list_regions(prompt)[0]
|
|
worker = build_worker_prompt(prompt, region, "token", "ctx", 0, "http://master")
|
|
self.assertNotIn("2", worker)
|
|
self.assertTrue({"1", "7"}.issubset(worker))
|
|
|
|
def test_unknown_or_list_producing_shared_dependency_is_not_exempt(self):
|
|
import sys
|
|
import types
|
|
from unittest.mock import patch
|
|
for node_class in (None, type("ListLoader", (), {
|
|
"RETURN_TYPES": ("MODEL",), "OUTPUT_IS_LIST": (True,)})):
|
|
prompt = graph_one_region()
|
|
prompt["2"]["inputs"]["model"] = ["1", 0]
|
|
registry = {} if node_class is None else {"CheckpointLoaderSimple": node_class}
|
|
with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}):
|
|
with self.assertRaisesRegex(ValueError, "bypasses Start.*1"):
|
|
validate_image_list_regions(prompt)
|
|
|
|
def test_two_sequential_regions_are_valid_and_isolated(self):
|
|
regions = validate_image_list_regions(graph_two_sequential_regions())
|
|
self.assertEqual([(r["start_id"], r["end_id"]) for r in regions], [("2", "4"), ("5", "7")])
|
|
p = graph_two_sequential_regions()
|
|
second = build_worker_prompt(p, regions[1], "item_token_123456789012345678901234", "ctx", 0,
|
|
"http://master.invalid")
|
|
self.assertEqual(set(second), {"5", "6", "7"})
|
|
|
|
def test_rejects_missing_context_pair_nested_region_and_cross_boundary_branch(self):
|
|
prompt = graph_one_region()
|
|
prompt["5"]["inputs"]["context"] = ["2", 0]
|
|
for explicit_targets in (False, True):
|
|
with self.subTest(explicit_targets=explicit_targets):
|
|
with self.assertRaisesRegex(ValueError, "typed context output"):
|
|
validate_image_list_regions(prompt, explicit_targets=explicit_targets)
|
|
nested = graph_one_region()
|
|
nested["7"] = start_node(["4", 0])
|
|
nested["8"] = {"class_type": "NestedOperation", "inputs": {"image": ["7", IMAGE_OUT]}}
|
|
nested["5"]["inputs"]["images"] = ["8", 0]
|
|
nested["9"] = {"class_type": END_CLASS, "inputs": {"images": ["8", 0], "context": ["7", CONTEXT_OUT]}}
|
|
for explicit_targets in (False, True):
|
|
with self.subTest(explicit_targets=explicit_targets):
|
|
with self.assertRaisesRegex(ValueError, "Nested/ambiguous"):
|
|
validate_image_list_regions(nested, explicit_targets=explicit_targets)
|
|
outside = graph_one_region()
|
|
outside["7"] = {"class_type": "PreviewImage", "inputs": {"images": ["3", IMAGE_OUT]}}
|
|
for explicit_targets in (False, True):
|
|
with self.subTest(explicit_targets=explicit_targets):
|
|
with self.assertRaisesRegex(ValueError, "cross-boundary"):
|
|
validate_image_list_regions(outside, explicit_targets=explicit_targets)
|
|
|
|
def test_rejects_duplicate_or_missing_results(self):
|
|
prompt = graph_one_region()
|
|
prompt["5"]["inputs"].pop("images")
|
|
with self.assertRaisesRegex(ValueError, "images must be connected"):
|
|
validate_image_list_regions(prompt)
|
|
|
|
|
|
class DynamicImageSchedulingTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_fast_workers_get_more_items_and_order_is_stable(self):
|
|
assignments = {"fast": [], "slow": []}
|
|
async def dispatch(worker, index):
|
|
assignments[worker].append(index)
|
|
delay = 0.002 if worker == "fast" else 0.025
|
|
task = asyncio.create_task(asyncio.sleep(delay, result=(index, f"result-{index}")))
|
|
return {"index": index, "task": task}
|
|
async def wait_next(inflight):
|
|
done, _ = await asyncio.wait([h["task"] for h in inflight.values()], return_when=asyncio.FIRST_COMPLETED)
|
|
task = next(iter(done))
|
|
worker = next(w for w, handle in inflight.items() if handle["task"] is task)
|
|
index, result = task.result()
|
|
return worker, index, result
|
|
async def cleanup(_worker, handle):
|
|
handle["task"].cancel()
|
|
ordered = await dynamic_schedule(12, ["fast", "slow"], dispatch, wait_next, cleanup)
|
|
self.assertEqual(ordered, [f"result-{i}" for i in range(12)])
|
|
self.assertGreater(len(assignments["fast"]), len(assignments["slow"]))
|
|
self.assertEqual(sorted(assignments["fast"] + assignments["slow"]), list(range(12)))
|
|
|
|
async def test_failure_timeout_and_cancel_clean_only_current_assignments(self):
|
|
async def dispatch(worker, index):
|
|
return {"index": index, "prompt_id": f"{worker}-{index}"}
|
|
cancelled = []
|
|
async def cleanup(worker, handle):
|
|
cancelled.append((worker, handle["prompt_id"]))
|
|
async def timeout(_inflight):
|
|
raise TimeoutError("worker timed out")
|
|
with self.assertRaisesRegex(TimeoutError, "timed out"):
|
|
await dynamic_schedule(4, ["one", "two"], dispatch, timeout, cleanup)
|
|
self.assertCountEqual(cancelled, [("one", "one-0"), ("two", "two-1")])
|
|
|
|
cancelled.clear()
|
|
async def wrong_index(inflight):
|
|
worker, handle = next(iter(inflight.items()))
|
|
return worker, handle["index"] + 99, "bad"
|
|
with self.assertRaisesRegex(RuntimeError, "unexpected image index"):
|
|
await dynamic_schedule(2, ["one"], dispatch, wrong_index, cleanup)
|
|
self.assertEqual(len(cancelled), 1)
|
|
|
|
cancelled.clear()
|
|
with self.assertRaises(ImageListCancelledError):
|
|
await dynamic_schedule(2, ["one", "two"], dispatch, timeout, cleanup,
|
|
interrupted=lambda: True)
|
|
self.assertEqual(len(cancelled), 0)
|
|
|
|
async def test_no_workers_and_duplicate_results_fail_closed(self):
|
|
async def unused(*_args):
|
|
self.fail("must not dispatch without workers")
|
|
with self.assertRaisesRegex(RuntimeError, "No free"):
|
|
await dynamic_schedule(2, [], unused, unused, unused)
|
|
async def dispatch(worker, index):
|
|
return {"index": index}
|
|
async def duplicate(inflight):
|
|
worker, _handle = next(iter(inflight.items()))
|
|
return worker, 0, "same"
|
|
with self.assertRaisesRegex(RuntimeError, "unexpected image index"):
|
|
await dynamic_schedule(2, ["one"], dispatch, duplicate, lambda *_: asyncio.sleep(0))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|