nodes/kijai-ComfyUI-NativeLooping_testing
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Some fixes

Browse Source
This commit is contained in:
kijai
2026-06-11 18:51:32 +03:00
parent 514fcb5645
commit 31f07f9def
1 changed files with 53 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+53 -32
View File
@@ -18,8 +18,9 @@ def _temporal_frame_count(samples):
return samples.shape[2] # (B,C,T,H,W)
return samples.shape[0] # (B,C,H,W) — batch count
def _accum_count(accum):
"""Count items in an accumulation, handling tensors (Image/Mask) and dicts (Latent/Video Latent)."""
def _accum_count(accum, blend_overlap=0):
"""Count items in an accumulation, handling tensors (Image/Mask) and dicts (Latent/Video Latent).
blend_overlap: frames consumed per seam by post-loop crossfade blending (fade modes)."""
if not isinstance(accum, dict) or "accum" not in accum:
return 0
total = 0
@@ -28,6 +29,8 @@ def _accum_count(accum):
total += _temporal_frame_count(item["samples"])
else:
total += item.shape[0] # IMAGE/MASK: count batch items
if blend_overlap > 0 and len(accum["accum"]) > 1:
total -= (len(accum["accum"]) - 1) * blend_overlap
return total
@@ -85,6 +88,7 @@ class TensorLoopOpen(io.ComfyNode):
accum = state["accum"]
previous_value = state["previous_value"]
open_node_id = state["open_node_id"]
blend_overlap = state.get("blend_overlap", 0)
else:
selected_mode = mode.get("mode", "iterations")
count = mode.get("iterations", 4) if selected_mode == "iterations" else 0
@@ -93,8 +97,9 @@ class TensorLoopOpen(io.ComfyNode):
accum = None
previous_value = initial_value
open_node_id = unique_id
blend_overlap = 0
accumulated_count = _accum_count(accum)
accumulated_count = _accum_count(accum, blend_overlap)
# In total_frames mode, count=0 and remaining goes negative each iteration.
# The math still produces correct 1-based iteration: 0-0+1=1, 0-(-1)+1=2, etc.
current_iteration = count - remaining + 1
@@ -102,10 +107,17 @@ class TensorLoopOpen(io.ComfyNode):
comfy.utils.ProgressBar(total_frames_val, node_id=open_node_id).update_absolute(accumulated_count)
elif count > 0:
comfy.utils.ProgressBar(count, node_id=open_node_id).update_absolute(count - remaining)
loop_state = {"remaining": remaining, "accum": accum, "previous_value": previous_value, "count": count, "open_node_id": open_node_id, "total_frames": total_frames_val}
loop_state = {"remaining": remaining, "accum": accum, "previous_value": previous_value, "count": count, "open_node_id": open_node_id, "total_frames": total_frames_val, "blend_overlap": blend_overlap}
return io.NodeOutput(loop_state, previous_value, accumulated_count, current_iteration)
def _overlap_option(name, tooltip):
"""DynamicCombo option with an overlap_frames count input."""
return io.DynamicCombo.Option(name, [
io.Int.Input("overlap_frames", default=8, min=1, tooltip=tooltip),
])
class TensorLoopClose(io.ComfyNode):
"""
Closes the loop started by TensorLoopOpen.
@@ -127,24 +139,15 @@ class TensorLoopClose(io.ComfyNode):
io.MatchType.Input("processed", template=cls.MATCHTYPE, raw_link=True, tooltip="Output generated this iteration."),
io.Boolean.Input("accumulate", default=True,
tooltip="When enabled, collects all iterations into a batch. When disabled, only outputs the final iteration's result."),
io.DynamicCombo.Input("overlap", tooltip="Remove or blend duplicate frames where consecutive iterations overlap.", options=[
io.DynamicCombo.Input("overlap", tooltip="How to handle duplicate frames where consecutive iterations overlap (e.g. context frames a video model re-generates).\n"
"- disabled: keep every frame as-is\n"
"- start/end: cut the duplicates from the start or end of each iteration's output\n"
"- fade_linear/fade_smooth: crossfade the end of each iteration into the start of the next", options=[
io.DynamicCombo.Option("disabled", []),
io.DynamicCombo.Option("start", [
io.Int.Input("overlap_frames", default=8, min=1,
tooltip="Number of frames to trim. Use when the model re-generates context frames at the start of its output (most common for video continuation)."),
]),
io.DynamicCombo.Option("end", [
io.Int.Input("overlap_frames", default=8, min=1,
tooltip="Number of frames to trim. Use when the model generates look-ahead frames at the end of its output."),
]),
io.DynamicCombo.Option("fade_linear", [
io.Int.Input("overlap_frames", default=8, min=1,
tooltip="Number of frames to crossfade with a linear blend between consecutive iterations."),
]),
io.DynamicCombo.Option("fade_smooth", [
io.Int.Input("overlap_frames", default=8, min=1,
tooltip="Number of frames to crossfade with a smoothstep (ease in/out) blend between consecutive iterations."),
]),
_overlap_option("start", "Number of frames to cut from the START of each iteration's output. Use when the model re-generates its context frames at the start of its output (most common for video continuation)."),
_overlap_option("end", "Number of frames to cut from the END of each iteration's output. Use when the model generates look-ahead frames at the end of its output that the next iteration re-generates."),
_overlap_option("fade_linear", "Number of overlapping frames to crossfade: the end of each iteration is blended into the start of the next with a linear ramp."),
_overlap_option("fade_smooth", "Number of overlapping frames to crossfade: the end of each iteration is blended into the start of the next with a smoothstep (ease in/out) ramp."),
]),
io.Boolean.Input("stop", optional=True, default=False, raw_link=True, force_input=True,
tooltip="Optional early stop signal from inside the loop body. When True, the loop stops after the current iteration regardless of remaining iterations or total_frames target."),
@@ -198,6 +201,7 @@ class TensorLoopClose(io.ComfyNode):
accum_out = graph.node("_AccumulateNode", to_add=to_accum, accumulation=accum_out).out(0)
# Disable total_frames when not accumulating to avoid infinite loops
blend_overlap = overlap_frames if overlap_mode in ("fade_linear", "fade_smooth") else 0
pack = graph.node("_ImageAccumStatePack",
remaining=sub.out(0),
accum=accum_out,
@@ -206,6 +210,7 @@ class TensorLoopClose(io.ComfyNode):
open_node_id=unpack.out(5),
total_frames=unpack.out(6) if accumulate else 0,
prev_accumulated_count=unpack.out(3),
blend_overlap=blend_overlap if accumulate else 0,
)
# Optional early stop from loop body
@@ -480,13 +485,19 @@ def _concat_tensor(a, b, dim):
if _is_nested(sa):
result["samples"] = NestedTensor([torch.cat([ta, tb], dim=dim) for ta, tb in zip(sa.tensors, sb.tensors)])
ma, mb = a.get("noise_mask"), b.get("noise_mask")
if ma is not None and mb is not None and _is_nested(ma):
if ma is not None and mb is not None and _is_nested(ma) and _is_nested(mb):
result["noise_mask"] = NestedTensor([torch.cat([ta, tb], dim=dim) for ta, tb in zip(ma.tensors, mb.tensors)])
elif ma is not None and _is_nested(ma):
# Per-frame mask can't cover the concatenated frames — drop it
result.pop("noise_mask", None)
else:
result["samples"] = torch.cat([sa, sb], dim=dim)
ma, mb = a.get("noise_mask"), b.get("noise_mask")
if ma is not None and mb is not None and ma.ndim == sa.ndim:
if ma is not None and mb is not None and ma.ndim == sa.ndim and mb.ndim == sb.ndim:
result["noise_mask"] = torch.cat([ma, mb], dim=dim)
elif ma is not None and ma.ndim == sa.ndim:
# Per-frame mask can't cover the concatenated frames — drop it
result.pop("noise_mask", None)
return result
else:
return torch.cat([a, b], dim=dim)
@@ -501,7 +512,10 @@ def _blend_overlap(items, overlap_frames, mode):
overlap_frames = min(overlap_frames, min_frames - 1) if min_frames > 1 else 0
if overlap_frames <= 0:
# Nothing to blend — fall through to simple concat
return _concat_tensor(items[0], items[1], dim) if len(items) == 2 else items[0]
result = items[0]
for item in items[1:]:
result = _concat_tensor(result, item, dim)
return result
t = torch.linspace(0, 1, overlap_frames)
if mode == "fade_smooth":
t = t * t * (3 - 2 * t)
@@ -572,8 +586,9 @@ class _AccumulationToImageBatch(io.ComfyNode):
result = items[0].copy()
result["samples"] = NestedTensor(catted)
result.pop("noise_mask", None)
masks = [item.get("noise_mask") for item in items if item.get("noise_mask") is not None]
if masks and _is_nested(masks[0]):
# Only keep masks when every item has one, otherwise the mask wouldn't cover all frames
masks = [item.get("noise_mask") for item in items]
if all(m is not None and _is_nested(m) for m in masks):
mask_catted = []
for i in range(len(masks[0].tensors)):
mask_catted.append(torch.cat([m.tensors[i] for m in masks], dim=2))
@@ -583,14 +598,18 @@ class _AccumulationToImageBatch(io.ComfyNode):
result = items[0].copy()
result["samples"] = torch.cat([item["samples"] for item in items], dim=2)
result.pop("noise_mask", None)
masks = [item.get("noise_mask") for item in items if item.get("noise_mask") is not None]
if masks and masks[0].ndim == 5:
masks = [item.get("noise_mask") for item in items]
if all(m is not None and m.ndim == 5 for m in masks):
result["noise_mask"] = torch.cat(masks, dim=2)
return io.NodeOutput(result)
else:
# Image latent — batch along dim 0
result = items[0].copy()
result["samples"] = torch.cat([item["samples"] for item in items], dim=0)
result.pop("noise_mask", None)
masks = [item.get("noise_mask") for item in items]
if all(m is not None and m.ndim == samples.ndim for m in masks):
result["noise_mask"] = torch.cat(masks, dim=0)
return io.NodeOutput(result)
else:
return io.NodeOutput(torch.cat(items, dim=0))
@@ -646,6 +665,7 @@ class _ImageAccumStatePack(io.ComfyNode):
io.AnyType.Input("open_node_id"),
io.AnyType.Input("total_frames"),
io.Int.Input("prev_accumulated_count", default=0),
io.Int.Input("blend_overlap", default=0),
],
outputs=[
io.AnyType.Output("loop_state"),
@@ -654,8 +674,9 @@ class _ImageAccumStatePack(io.ComfyNode):
)
@classmethod
def execute(cls, remaining, accum, previous_value, count, open_node_id, total_frames, prev_accumulated_count=0) -> io.NodeOutput:
accumulated_count = _accum_count(accum)
def execute(cls, remaining, accum, previous_value, count, open_node_id, total_frames, prev_accumulated_count=0, blend_overlap=0) -> io.NodeOutput:
# Effective count: fade modes consume blend_overlap frames per seam in the post-loop blend
accumulated_count = _accum_count(accum, blend_overlap)
if total_frames > 0:
should_continue = accumulated_count < total_frames
@@ -669,7 +690,7 @@ class _ImageAccumStatePack(io.ComfyNode):
comfy.utils.ProgressBar(count, node_id=open_node_id).update_absolute(current_iteration)
return io.NodeOutput(
{"remaining": remaining, "accum": accum, "previous_value": previous_value, "count": count, "open_node_id": open_node_id, "total_frames": total_frames},
{"remaining": remaining, "accum": accum, "previous_value": previous_value, "count": count, "open_node_id": open_node_id, "total_frames": total_frames, "blend_overlap": blend_overlap},
should_continue,
)
@@ -703,7 +724,7 @@ class _ImageAccumStateUnpack(io.ComfyNode):
count = loop_state.get("count", 0)
open_node_id = loop_state.get("open_node_id")
total_frames = loop_state.get("total_frames", 0)
accumulated_count = _accum_count(accum)
accumulated_count = _accum_count(accum, loop_state.get("blend_overlap", 0))
return io.NodeOutput(remaining, accum, previous_value, accumulated_count, count, open_node_id, total_frames)