alisson-anjos-ComfyUI-BFSNodes/bfs_frame_loop.py

"""BFS Frame Loop: run any part of a workflow once per frame, one after another, and collect the results.

Wire a batch of frames (e.g. the Keyframe Planner's source_frames, or BFS Keyframe Marks' marked_frames) into
BFS Frame Loop Start; it gives one frame (plus its mask, index and count). Build the per-frame work after it (an image
edit: Qwen 2.1 body swap, Klein, ...), and wire its result into BFS Frame Loop End. The End node expands the body
between Start and End again for the next frame until all are done, then returns them as one batch (same order).
Models stay outside the loop body. Results are kept on disk between iterations; no global state.

The expansion follows the BFS H3 Native Shot Loop pattern (start/end graph expansion).
"""
import json
import os
import uuid
from pathlib import Path

import torch
import torch.nn.functional as F

try:
    import folder_paths
except Exception:  # tests run without ComfyUI
    folder_paths = None

CATEGORY = "BFS/Loop"


class BFSFrameLoopStart:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"images": ("IMAGE", {"tooltip": "The frames to process one by one."})},
                "optional": {"masks": ("MASK", {"tooltip": "One mask per frame (optional), handed out with its frame."}),
                             "loop_state": ("BFS_FRAME_LOOP", {"tooltip": "Set by Frame Loop End; leave unconnected."})}}

    RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "BFS_FRAME_LOOP")
    RETURN_NAMES = ("image", "mask", "index", "count", "loop_state")
    OUTPUT_TOOLTIPS = ("The current frame.", "Its mask (zeros when no masks are given).", "Index of the current frame (0-based).",
                       "How many frames the loop runs.", "Connect to BFS Frame Loop End.")
    FUNCTION = "start"
    CATEGORY = CATEGORY
    DESCRIPTION = "Hands out one frame per loop iteration; pair it with BFS Frame Loop End."

    def start(self, images, masks=None, loop_state=None):
        n = int(images.shape[0])
        if loop_state is None:
            temp = folder_paths.get_temp_directory() if folder_paths else "."
            folder = os.path.join(temp, "bfs_frame_loop", uuid.uuid4().hex[:12])
            os.makedirs(folder, exist_ok=True)
            loop_state = {"folder": folder, "index": 0, "count": n, "done": []}
        state = dict(loop_state)
        if state["count"] != n:
            raise ValueError(f"The loop started with {state['count']} frames but now gets {n}.")
        i = int(state["index"])
        h, w = images.shape[1:3]
        if masks is not None:
            m = masks if masks.ndim == 3 else masks[None]
            if m.shape[0] != n:
                raise ValueError(f"{m.shape[0]} masks for {n} frames.")
            mask = m[i:i + 1]
            if mask.shape[1:] != (h, w):
                mask = F.interpolate(mask[:, None].float(), size=(h, w), mode="bilinear", align_corners=False)[:, 0]
        else:
            mask = torch.zeros(1, h, w)
        print(f"[BFS Frame Loop] frame {i + 1}/{n}", flush=True)
        return images[i:i + 1], mask, i, n, state


class BFSFrameLoopEnd:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"image": ("IMAGE", {"tooltip": "The processed frame of this iteration."}),
                             "loop_state": ("BFS_FRAME_LOOP", {"tooltip": "From BFS Frame Loop Start."})},
                "optional": {"resize_to_first": ("BOOLEAN", {"default": True,
                                                             "tooltip": "Resize every result to the first result's size so they form one batch."})},
                "hidden": {"dynprompt": "DYNPROMPT", "unique_id": "UNIQUE_ID"}}

    RETURN_TYPES = ("IMAGE", "INT")
    RETURN_NAMES = ("images", "count")
    OUTPUT_TOOLTIPS = ("All processed frames, in order, as one batch.", "How many frames.")
    FUNCTION = "finish"
    CATEGORY = CATEGORY
    DESCRIPTION = "Collects one processed frame per iteration and expands the loop body until all frames are done."

    def finish(self, image, loop_state, resize_to_first=True, dynprompt=None, unique_id=None):
        state = dict(loop_state)
        i = int(state["index"])
        path = Path(state["folder"]) / f"{i:05d}.pt"
        tmp = path.with_suffix(".tmp"); torch.save(image[:1].detach().cpu().contiguous(), tmp); tmp.replace(path)
        state["done"] = list(state["done"]) + [str(path)]
        state["index"] = i + 1
        if state["index"] < state["count"]:
            if dynprompt is None:
                raise ValueError("BFS Frame Loop needs ComfyUI's dynamic graph execution")
            return self.expand(dynprompt, str(unique_id), state)
        parts = [torch.load(p, map_location="cpu", weights_only=True) for p in state["done"]]
        if resize_to_first:
            h, w = parts[0].shape[1:3]
            parts = [p if p.shape[1:3] == (h, w) else
                     F.interpolate(p.movedim(-1, 1), size=(h, w), mode="bilinear", align_corners=False).movedim(1, -1) for p in parts]
        if len({tuple(p.shape[1:]) for p in parts}) != 1:
            raise ValueError("The processed frames have different sizes; turn on resize_to_first.")
        (Path(state["folder"]) / "manifest.json").write_text(json.dumps(state, indent=1))
        return torch.cat(parts), len(parts)

    @staticmethod
    def expand(dynprompt, end, state):
        from comfy_execution.graph_utils import GraphBuilder, is_link
        feeds, seen = {}, set()

        def visit(nid):
            if nid in seen:
                return
            seen.add(nid)
            for value in dynprompt.get_node(nid).get("inputs", {}).values():
                if is_link(value):
                    parent = str(value[0]); feeds.setdefault(parent, []).append(nid); visit(parent)
        visit(end)
        starts = [n for n in seen if dynprompt.get_node(n)["class_type"] == "BFSFrameLoopStart"]
        if len(starts) != 1:
            raise ValueError("Connect exactly one BFS Frame Loop Start upstream of this Frame Loop End")
        start = starts[0]; body = {start}; pending = [start]
        while pending:
            for child in feeds.get(pending.pop(), []):
                if child not in body:
                    body.add(child); pending.append(child)
        graph = GraphBuilder()
        clones = {n: graph.node(dynprompt.get_node(n)["class_type"], "Next" if n == end else n) for n in body}
        for nid, node in clones.items():
            node.set_override_display_id(dynprompt.get_display_node_id(nid))
            for key, value in dynprompt.get_node(nid).get("inputs", {}).items():
                node.set_input(key, clones[str(value[0])].out(value[1]) if is_link(value) and str(value[0]) in body else value)
        clones[start].set_input("loop_state", state)
        return {"result": tuple(clones[end].out(i) for i in range(2)), "expand": graph.finalize()}


NODE_CLASS_MAPPINGS = {"BFSFrameLoopStart": BFSFrameLoopStart, "BFSFrameLoopEnd": BFSFrameLoopEnd}
NODE_DISPLAY_NAME_MAPPINGS = {"BFSFrameLoopStart": "BFS Frame Loop Start", "BFSFrameLoopEnd": "BFS Frame Loop End"}