diff --git a/README.md b/README.md index cffdf6c..cd23f57 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,17 @@ - Saves images to any folder, not just ComfyUI's output directory. - Supports filename tokens including `%input%`, `%input2%`, `%date:yyyy-MM-dd%`, `%counter%`, `%width%`, `%height%`, and `%batch_num%`. +- `Dumas H3 Plan Attach Scene Images` + - Inputs: `plan`, `scene_index`, optional `image1`..`image6` + - Outputs: `plan`, `connected_images` + - Attaches up to six optional image sockets to one MiniMax H3 plan scene while keeping the plan JSON-serializable for the upstream archive system. + - Chain one node per scene that needs its own six-image bundle. + +- `Dumas H3 Plan Extract Scene Images` + - Inputs: `plan`, `scene_index` + - Outputs: `plan`, `image1`..`image6`, `connected_images` + - Reads back the six optional images for a selected MiniMax H3 plan scene, for example by connecting the current `clip_index`. + - `Dumas JSON String to Object` - Input: `json_string` - Output: parsed `JSON` @@ -170,6 +181,8 @@ decr -> use index - 1 `Save Image Dumas` leaves `folder` empty to use ComfyUI's output directory. `name` feeds `%input%` and `name_2` feeds `%input2%`, so a pattern like `project/%input%_%input2%_%counter%` can combine two upstream strings into the saved filename. +`Dumas H3 Plan Attach Scene Images` and `Dumas H3 Plan Extract Scene Images` are a companion pair for `ComfyUI-MiniMaxH3-Contex-Loop`. The upstream H3 plan node cannot dynamically grow six new image sockets for every JSON-defined scene, so Dumas stores scene image bindings beside the plan using a lightweight token and an in-memory registry. That keeps `plan.json` archiving intact while still letting you wire up six IMAGE sockets per scene through chained helper nodes. + `Dumas Strip Iteration Suffix` keeps the part before the first underscore and drops the rest. Names like `char123_pose_final.png` become `char123.png`, while names with no underscore such as `char123.png` are left untouched. `Dumas Slugify String` lowercases text, strips accents, replaces non-alphanumeric runs with `-`, and trims leading or trailing dashes. diff --git a/dumas_image_nodes.py b/dumas_image_nodes.py index c7c706d..8eb9358 100644 --- a/dumas_image_nodes.py +++ b/dumas_image_nodes.py @@ -18,6 +18,10 @@ _MEDIA_EXT_RE = re.compile( _DATE_TOKEN_RE = re.compile(r"%date:([^%]+)%") _SERVE_TOKENS = OrderedDict() _SERVE_CAP = 256 +_H3_PLAN_TYPE = "H3_CHAIN_PLAN" +_H3_PLAN_IMAGE_BINDINGS_KEY = "_dumas_scene_image_bindings" +_H3_PLAN_IMAGE_BINDINGS = OrderedDict() +_H3_PLAN_IMAGE_BINDINGS_CAP = 128 def _clean_input_token_value(value): @@ -111,6 +115,71 @@ def resolve_serve_token(token): return _SERVE_TOKENS.get(str(token or "")) +def _touch_plan_image_binding(token): + token = str(token or "") + if not token or token not in _H3_PLAN_IMAGE_BINDINGS: + return + binding = _H3_PLAN_IMAGE_BINDINGS.pop(token) + _H3_PLAN_IMAGE_BINDINGS[token] = binding + + +def _prune_plan_image_bindings(): + while len(_H3_PLAN_IMAGE_BINDINGS) > _H3_PLAN_IMAGE_BINDINGS_CAP: + _H3_PLAN_IMAGE_BINDINGS.popitem(last=False) + + +def _clone_h3_plan(plan): + if not isinstance(plan, dict): + raise ValueError("Dumas H3 plan helpers require a plan dictionary.") + shots = plan.get("shots") + if not isinstance(shots, list): + raise ValueError("Dumas H3 plan helpers require a plan with a shots list.") + + cloned = dict(plan) + cloned["shots"] = [dict(shot) if isinstance(shot, dict) else shot for shot in shots] + bindings = plan.get(_H3_PLAN_IMAGE_BINDINGS_KEY) + if isinstance(bindings, dict): + cloned[_H3_PLAN_IMAGE_BINDINGS_KEY] = { + "token": str(bindings.get("token") or ""), + "scene_counts": { + str(key): int(value) + for key, value in dict(bindings.get("scene_counts") or {}).items() + }, + } + return cloned + + +def _normalize_h3_scene_index(plan, scene_index): + shots = plan.get("shots") + total = len(shots) if isinstance(shots, list) else 0 + index = int(scene_index) + if index < 1 or index > total: + raise ValueError( + f"Dumas H3 scene index {index} is outside the plan's {total} scenes." + ) + return index + + +def _h3_plan_binding_entry(plan): + bindings = plan.get(_H3_PLAN_IMAGE_BINDINGS_KEY) + if not isinstance(bindings, dict): + return "", {} + token = str(bindings.get("token") or "") + counts = { + str(key): int(value) + for key, value in dict(bindings.get("scene_counts") or {}).items() + } + return token, counts + + +def _scene_images_tuple(image1=None, image2=None, image3=None, image4=None, image5=None, image6=None): + return (image1, image2, image3, image4, image5, image6) + + +def _connected_image_count(images): + return sum(1 for image in images if image is not None) + + def _is_within_directory(parent_path, child_path): try: return os.path.commonpath([parent_path, child_path]) == parent_path @@ -418,12 +487,170 @@ class DumasSaveImageNode: return {"ui": {"images": ui_images}} +class DumasH3PlanAttachSceneImagesNode: + DESCRIPTION = ( + "Attach up to six optional IMAGE sockets to one H3 Chain Plan scene " + "without breaking the upstream plan archive format. Chain multiple " + "copies of this node to bind different scene indexes." + ) + RETURN_TYPES = (_H3_PLAN_TYPE, "INT") + RETURN_NAMES = ("plan", "connected_images") + FUNCTION = "attach" + CATEGORY = "Dumas/MiniMax" + + @classmethod + def INPUT_TYPES(cls): + optional = {} + for slot in range(1, 7): + optional[f"image{slot}"] = ( + "IMAGE", + { + "tooltip": ( + f"Optional image for slot {slot} on the selected H3 plan scene." + ) + }, + ) + return { + "required": { + "plan": ( + _H3_PLAN_TYPE, + { + "tooltip": ( + "Validated MiniMax H3 chain plan to enrich with scene-level " + "image bindings." + ) + }, + ), + "scene_index": ( + "INT", + { + "default": 1, + "min": 1, + "max": 9999, + "step": 1, + "tooltip": ( + "1-based scene index inside the H3 plan. Use one node per " + "scene that needs up to six image sockets." + ), + }, + ), + }, + "optional": optional, + } + + def attach( + self, + plan, + scene_index, + image1=None, + image2=None, + image3=None, + image4=None, + image5=None, + image6=None, + ): + updated_plan = _clone_h3_plan(plan) + scene_index = _normalize_h3_scene_index(updated_plan, scene_index) + images = _scene_images_tuple(image1, image2, image3, image4, image5, image6) + connected_count = _connected_image_count(images) + + token, scene_counts = _h3_plan_binding_entry(updated_plan) + if not token: + token = uuid.uuid4().hex + registry = _H3_PLAN_IMAGE_BINDINGS.setdefault(token, {}) + _touch_plan_image_binding(token) + + if connected_count: + registry[int(scene_index)] = images + scene_counts[str(scene_index)] = connected_count + else: + registry.pop(int(scene_index), None) + scene_counts.pop(str(scene_index), None) + + if registry: + updated_plan[_H3_PLAN_IMAGE_BINDINGS_KEY] = { + "token": token, + "scene_counts": scene_counts, + } + else: + _H3_PLAN_IMAGE_BINDINGS.pop(token, None) + updated_plan.pop(_H3_PLAN_IMAGE_BINDINGS_KEY, None) + + _prune_plan_image_bindings() + return (updated_plan, connected_count) + + +class DumasH3PlanExtractSceneImagesNode: + DESCRIPTION = ( + "Read back the six optional image bindings for one H3 Chain Plan scene. " + "Connect clip_index or another scene selector to recover the matching " + "scene images downstream." + ) + RETURN_TYPES = (_H3_PLAN_TYPE, "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "INT") + RETURN_NAMES = ( + "plan", + "image1", + "image2", + "image3", + "image4", + "image5", + "image6", + "connected_images", + ) + FUNCTION = "extract" + CATEGORY = "Dumas/MiniMax" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "plan": ( + _H3_PLAN_TYPE, + { + "tooltip": ( + "H3 plan previously enriched by Dumas H3 Plan Attach Scene Images." + ) + }, + ), + "scene_index": ( + "INT", + { + "default": 1, + "min": 1, + "max": 9999, + "step": 1, + "tooltip": ( + "1-based scene index to retrieve. Connect Current Shot " + "clip_index to get the active scene's images." + ), + }, + ), + } + } + + def extract(self, plan, scene_index): + passthrough_plan = _clone_h3_plan(plan) + scene_index = _normalize_h3_scene_index(passthrough_plan, scene_index) + token, _scene_counts = _h3_plan_binding_entry(passthrough_plan) + if not token: + return (passthrough_plan, None, None, None, None, None, None, 0) + + registry = _H3_PLAN_IMAGE_BINDINGS.get(token) or {} + _touch_plan_image_binding(token) + images = registry.get(int(scene_index)) or (None, None, None, None, None, None) + return (passthrough_plan, *images, _connected_image_count(images)) + + NODE_CLASS_MAPPINGS = { "DumasImageCompare": DumasImageCompareNode, "DumasSaveImage": DumasSaveImageNode, + "DumasH3PlanAttachSceneImages": DumasH3PlanAttachSceneImagesNode, + "DumasH3PlanExtractSceneImages": DumasH3PlanExtractSceneImagesNode, } NODE_DISPLAY_NAME_MAPPINGS = { "DumasImageCompare": "Dumas Image Compare", "DumasSaveImage": "Save Image Dumas", + "DumasH3PlanAttachSceneImages": "Dumas H3 Plan Attach Scene Images", + "DumasH3PlanExtractSceneImages": "Dumas H3 Plan Extract Scene Images", } diff --git a/tests/test_dumas_image_nodes.py b/tests/test_dumas_image_nodes.py index a3d335a..f459b6e 100644 --- a/tests/test_dumas_image_nodes.py +++ b/tests/test_dumas_image_nodes.py @@ -1,4 +1,5 @@ import importlib +import json import os import sys import tempfile @@ -318,6 +319,93 @@ class DumasImageNodeTests(unittest.TestCase): self.assertEqual(result["ui"]["images"], []) self.assertEqual(FakePILImage.saved_paths, []) + 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) + + 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, + ) + + 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.assertEqual(scene2[-1], 1) + self.assertEqual( + plan_after_second["_dumas_scene_image_bindings"]["scene_counts"], + {"1": 1, "2": 1}, + ) + + 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_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()