Add H3 scene image plan helpers
This commit is contained in:
@@ -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",
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user