From ee252ee0419f616035db60b00b084588055602c3 Mon Sep 17 00:00:00 2001 From: ausboss <102309304+ausboss@users.noreply.github.com> Date: Thu, 13 Aug 2026 09:43:08 -0500 Subject: [PATCH] Add the Frame Chooser backend with a pausing selection flow --- __init__.py | 1 + nodes/_chooser_helpers.py | 274 ++++++++++++++++++++++++++++++++++ nodes/node_frame_chooser.py | 118 +++++++++++++++ scripts/validate_nodes.py | 1 + tests/test_chooser_helpers.py | 102 +++++++++++++ 5 files changed, 496 insertions(+) create mode 100644 nodes/_chooser_helpers.py create mode 100644 nodes/node_frame_chooser.py create mode 100644 tests/test_chooser_helpers.py diff --git a/__init__.py b/__init__.py index c2e173d..531b808 100644 --- a/__init__.py +++ b/__init__.py @@ -17,6 +17,7 @@ import traceback # One entry per file in nodes/. Keep alphabetical-ish; order only affects # nothing functional. NODE_MODULES = [ + "node_frame_chooser", "node_image_crop_rotate_pad", "node_lama_inpaint", "node_load_video", diff --git a/nodes/_chooser_helpers.py b/nodes/_chooser_helpers.py new file mode 100644 index 0000000..101d998 --- /dev/null +++ b/nodes/_chooser_helpers.py @@ -0,0 +1,274 @@ +"""Frame Chooser helpers: selection math, thumbnails, and pause plumbing. + +The pure functions at the top are covered by offline unit tests. Everything +that talks to ComfyUI (thumbnail files, the websocket announcement, the +answer route, the blocking wait) imports server modules lazily so this file +stays importable without ComfyUI installed. +""" + +from __future__ import annotations + +import threading +from pathlib import Path + +import torch + +try: + import folder_paths +except ImportError: # Offline tests import this module without ComfyUI. + folder_paths = None + + +POLL_SECONDS = 0.1 +THUMBNAIL_SUBFOLDER = "ausboss_chooser" +EVENT_NAME = "ausboss-frame-choose" + + +# --- pure selection logic ---------------------------------------------------- + +def normalize_selection(values, frame_count: int) -> list[int]: + """Validate a browser-supplied selection of one-based frame indices. + + Returns the indices deduplicated and ascending, so the kept frames always + stay in source order. An empty list is valid and means "keep every frame". + Anything else invalid raises, so a malformed or hostile request can never + release a paused graph with a selection the node did not offer.""" + if not isinstance(values, (list, tuple)): + raise ValueError("Frame selection must be a list of one-based integers.") + count = int(frame_count) + kept: set[int] = set() + for value in values: + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError("Frame selection entries must be plain integers.") + if value < 1 or value > count: + raise ValueError( + f"Frame {value} is outside this batch (1 through {count})." + ) + kept.add(value) + return sorted(kept) + + +def usable_remembered(values, frame_count: int) -> list[int] | None: + """Adapt a remembered selection to the current batch size. + + Indices past the end of the new batch are dropped. Returns None when the + memory named frames only beyond this batch (or is not a list at all): a + useless memory means the node should pause and ask again.""" + if not isinstance(values, (list, tuple)): + return None + count = int(frame_count) + valid = sorted( + { + int(value) + for value in values + if isinstance(value, int) + and not isinstance(value, bool) + and 1 <= value <= count + } + ) + if values and not valid: + return None + return valid + + +def keep_frames(frames: torch.Tensor, one_based: list[int]) -> torch.Tensor: + """Slice a BHWC batch down to the chosen frames, keeping source order. + + An empty selection keeps the whole batch unchanged.""" + if not isinstance(frames, torch.Tensor) or frames.ndim != 4: + raise ValueError("Frame Chooser expected a BHWC IMAGE batch.") + if not one_based: + return frames + index = torch.tensor( + [int(value) - 1 for value in one_based], dtype=torch.long, device=frames.device + ) + return frames.index_select(0, index) + + +def effective_indices(one_based: list[int], frame_count: int) -> list[int]: + """Resolve "empty means keep all" into the concrete kept indices.""" + if one_based: + return [int(value) for value in one_based] + return list(range(1, int(frame_count) + 1)) + + +def indices_string(one_based: list[int]) -> str: + """One-based indices joined for the STRING output, e.g. '1,4,9'.""" + return ",".join(str(value) for value in one_based) + + +# --- thumbnails --------------------------------------------------------------- + +def write_thumbnails(frames: torch.Tensor, run_token: str, max_size: int) -> list[dict]: + """Write one downscaled JPEG per frame into ComfyUI's temp folder. + + Returns /view descriptors ({filename, subfolder, type}) the frontend turns + into URLs. The temp folder is ComfyUI-managed, so the files disappear with + every normal temp cleanup.""" + from PIL import Image + + from ._media_helpers import encode_preview + + if folder_paths is None: + raise RuntimeError("Frame Chooser thumbnails require ComfyUI's folder_paths.") + root = Path(folder_paths.get_temp_directory()) / THUMBNAIL_SUBFOLDER + root.mkdir(parents=True, exist_ok=True) + files: list[dict] = [] + for index in range(int(frames.shape[0])): + array = ( + (frames[index].detach().clamp(0.0, 1.0) * 255.0) + .round() + .to(torch.uint8) + .cpu() + .numpy() + ) + if array.ndim == 3 and array.shape[-1] >= 3: + image = Image.fromarray(array[..., :3], "RGB") + else: + image = Image.fromarray(array.reshape(array.shape[0], array.shape[1]), "L") + name = f"{run_token}_{index + 1:04d}.jpg" + (root / name).write_bytes(encode_preview(image, int(max_size), int(max_size))) + files.append({"filename": name, "subfolder": THUMBNAIL_SUBFOLDER, "type": "temp"}) + return files + + +# --- pause plumbing ----------------------------------------------------------- + +class _PendingChoice: + """One paused Frame Chooser execution awaiting a browser answer.""" + + __slots__ = ("event", "frame_count", "selection", "cancelled") + + def __init__(self, frame_count: int): + self.event = threading.Event() + self.frame_count = int(frame_count) + self.selection: list[int] | None = None + self.cancelled = False + + +def _store() -> dict: + """Pending and remembered state, attached to the live PromptServer so a + module reload (ComfyUI-Manager, dev restarts) keeps sharing one copy.""" + from server import PromptServer + + server = PromptServer.instance + store = getattr(server, "_ausboss_frame_chooser", None) + if store is None: + store = {"pending": {}, "remembered": {}, "lock": threading.Lock()} + server._ausboss_frame_chooser = store + return store + + +def recall_selection(node_id: str) -> list[int] | None: + """Previous answer for this node id, or None when it never answered.""" + store = _store() + with store["lock"]: + remembered = store["remembered"].get(str(node_id)) + return list(remembered) if remembered is not None else None + + +def remember_selection(node_id: str, selection: list[int]) -> None: + store = _store() + with store["lock"]: + store["remembered"][str(node_id)] = list(selection) + + +def await_selection( + node_id: str, frame_count: int, files: list[dict], previous: list[int] | None +) -> list[int]: + """Announce the pause to the browser and block until it answers. + + Returns the validated one-based selection ([] = keep all). Raises + InterruptProcessingException when the browser cancels or the queue is + interrupted, so ComfyUI unwinds the run exactly like pressing stop.""" + import comfy.model_management as model_management + from server import PromptServer + + store = _store() + key = str(node_id) + pending = _PendingChoice(frame_count) + with store["lock"]: + store["pending"][key] = pending + try: + PromptServer.instance.send_sync( + EVENT_NAME, + { + "node_id": key, + "urls": files, + "count": int(frame_count), + "previous": list(previous) if previous else [], + }, + ) + while not pending.event.wait(POLL_SECONDS): + if model_management.processing_interrupted(): + raise model_management.InterruptProcessingException() + if pending.cancelled: + raise model_management.InterruptProcessingException() + return list(pending.selection or []) + finally: + with store["lock"]: + store["pending"].pop(key, None) + + +def register_chooser_route() -> None: + """POST /ausboss/frame_chooser answers or cancels a paused chooser.""" + try: + from aiohttp import web + from server import PromptServer + except ImportError: + return + server = getattr(PromptServer, "instance", None) + if server is None or getattr(server, "_ausboss_chooser_route", False): + return + server._ausboss_chooser_route = True + + @server.routes.post("/ausboss/frame_chooser") + async def ausboss_frame_chooser(request): + try: + data = await request.json() + except Exception: + return web.json_response({"error": "The request body must be JSON."}, status=400) + node_id = str(data.get("node_id", "")) + action = data.get("action", "") + store = _store() + with store["lock"]: + pending = store["pending"].get(node_id) + if pending is None: + return web.json_response( + {"error": "No Frame Chooser is paused under that id."}, status=404 + ) + if action == "cancel": + pending.cancelled = True + pending.event.set() + return web.json_response({"status": "cancelled"}) + if action == "continue": + try: + selection = normalize_selection(data.get("selected", []), pending.frame_count) + except ValueError as exc: + return web.json_response({"error": str(exc)}, status=400) + pending.selection = selection + pending.event.set() + return web.json_response( + { + "status": "continued", + "kept": len(effective_indices(selection, pending.frame_count)), + } + ) + return web.json_response({"error": "Action must be 'continue' or 'cancel'."}, status=400) + + +__all__ = [ + "EVENT_NAME", + "POLL_SECONDS", + "THUMBNAIL_SUBFOLDER", + "await_selection", + "effective_indices", + "indices_string", + "keep_frames", + "normalize_selection", + "recall_selection", + "register_chooser_route", + "remember_selection", + "usable_remembered", + "write_thumbnails", +] diff --git a/nodes/node_frame_chooser.py b/nodes/node_frame_chooser.py new file mode 100644 index 0000000..49774f9 --- /dev/null +++ b/nodes/node_frame_chooser.py @@ -0,0 +1,118 @@ +"""Frame Chooser (AusBoss).""" + +from __future__ import annotations + +import uuid + +from ._chooser_helpers import ( + await_selection, + effective_indices, + indices_string, + keep_frames, + normalize_selection, + recall_selection, + register_chooser_route, + remember_selection, + usable_remembered, + write_thumbnails, +) + + +NODE_ID = "AUSBOSS_NODES_FrameChooser" + +BEHAVIOR_PAUSE = "always pause" +BEHAVIOR_REMEMBER = "keep last selection" + +# IS_CHANGED tokens: one counter per node id, bumped once per queue (ComfyUI +# caches the value in the prompt), so a chooser re-runs and re-pauses even +# when none of its widgets changed between queues. +_RUN_COUNTERS: dict[str, int] = {} + + +class AusBossFrameChooser: + CATEGORY = "🆎 AusBoss/Video" + DESCRIPTION = ( + "Pauses the running graph and shows every incoming frame as a clickable " + "filmstrip on the node. Click the frames to keep and press Keep selected; " + "the graph resumes with just those frames in their original order. An " + "empty selection keeps every frame, and 'keep last selection' re-applies " + "the previous answer without pausing." + ) + SEARCH_ALIASES = ["pick frames", "pause", "choose frames", "filmstrip", "ausboss"] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "frames": ( + "IMAGE", + {"tooltip": "Video frames or any other BHWC IMAGE batch to choose from."}, + ), + "behavior": ( + [BEHAVIOR_PAUSE, BEHAVIOR_REMEMBER], + { + "default": BEHAVIOR_PAUSE, + "tooltip": ( + "'always pause' stops every run until you answer; 'keep last " + "selection' re-applies this node's previous answer without " + "pausing, and only pauses when no previous answer fits." + ), + }, + ), + "preview_max_size": ( + "INT", + { + "default": 256, + "min": 64, + "max": 1024, + "step": 32, + "tooltip": ( + "Longest edge of the filmstrip thumbnails. Smaller encodes and " + "loads faster; the full-resolution frames are untouched." + ), + }, + ), + }, + "hidden": {"unique_id": "UNIQUE_ID"}, + } + + RETURN_TYPES = ("IMAGE", "INT", "STRING") + RETURN_NAMES = ("frames", "count", "indices") + OUTPUT_TOOLTIPS = ( + "The kept frames in their original order; an empty selection keeps all.", + "How many frames were kept.", + "One-based indices of the kept frames, comma-joined, e.g. '1,4,9'.", + ) + FUNCTION = "choose" + + def choose(self, frames, behavior, preview_max_size, unique_id=None): + node_id = str(unique_id) + frame_count = int(frames.shape[0]) + previous = recall_selection(node_id) + + selection = None + if behavior == BEHAVIOR_REMEMBER and previous is not None: + selection = usable_remembered(previous, frame_count) + if selection is None: + files = write_thumbnails(frames, uuid.uuid4().hex[:12], int(preview_max_size)) + preselect = usable_remembered(previous or [], frame_count) + selection = await_selection(node_id, frame_count, files, preselect) + selection = normalize_selection(selection, frame_count) + remember_selection(node_id, selection) + + kept = effective_indices(selection, frame_count) + return (keep_frames(frames, selection), len(kept), indices_string(kept)) + + @classmethod + def IS_CHANGED(cls, unique_id=None, **_inputs): + key = str(unique_id) + _RUN_COUNTERS[key] = _RUN_COUNTERS.get(key, 0) + 1 + return f"{key}:{_RUN_COUNTERS[key]}" + + +register_chooser_route() + +NODE_CLASS_MAPPINGS = {NODE_ID: AusBossFrameChooser} +NODE_DISPLAY_NAME_MAPPINGS = {NODE_ID: "Frame Chooser (AusBoss)"} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/scripts/validate_nodes.py b/scripts/validate_nodes.py index 5a71d27..9333a70 100644 --- a/scripts/validate_nodes.py +++ b/scripts/validate_nodes.py @@ -26,6 +26,7 @@ errors = [] warnings = [] PUBLIC_NODE_IDS = { + "AUSBOSS_NODES_FrameChooser", "AUSBOSS_NODES_ImageCropRotatePad", "AUSBOSS_NODES_LaMaInpaint", "AUSBOSS_NODES_LoadVideo", diff --git a/tests/test_chooser_helpers.py b/tests/test_chooser_helpers.py new file mode 100644 index 0000000..4b9eb92 --- /dev/null +++ b/tests/test_chooser_helpers.py @@ -0,0 +1,102 @@ +"""Offline tests for the Frame Chooser's pure selection helpers. + +The blocking pause, the websocket event, and the answer route need a live +PromptServer and are exercised in the browser acceptance pass instead. +""" + +from __future__ import annotations + +import sys +import unittest +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) +if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"): + del sys.modules["nodes"] + +from nodes._chooser_helpers import ( + effective_indices, + indices_string, + keep_frames, + normalize_selection, + usable_remembered, +) + + +def batch(count: int) -> torch.Tensor: + """A BHWC batch where frame i is filled with the value i (zero-based).""" + frames = torch.zeros((count, 2, 3, 3), dtype=torch.float32) + for index in range(count): + frames[index] = float(index) + return frames + + +class NormalizeSelectionTests(unittest.TestCase): + def test_sorts_and_dedupes_into_source_order(self): + self.assertEqual(normalize_selection([9, 1, 4, 1, 9], 10), [1, 4, 9]) + + def test_empty_selection_is_valid_and_means_keep_all(self): + self.assertEqual(normalize_selection([], 10), []) + + def test_rejects_non_list_payloads(self): + with self.assertRaises(ValueError): + normalize_selection("1,2,3", 10) + + def test_rejects_non_integer_entries(self): + for bad in ([1.5], ["2"], [True], [None]): + with self.assertRaises(ValueError): + normalize_selection(bad, 10) + + def test_rejects_out_of_range_entries(self): + with self.assertRaises(ValueError): + normalize_selection([0], 10) + with self.assertRaises(ValueError): + normalize_selection([11], 10) + + +class UsableRememberedTests(unittest.TestCase): + def test_trims_to_the_current_batch(self): + self.assertEqual(usable_remembered([2, 5, 40], 10), [2, 5]) + + def test_empty_memory_means_keep_all(self): + self.assertEqual(usable_remembered([], 10), []) + + def test_entirely_stale_memory_is_useless(self): + self.assertIsNone(usable_remembered([40, 41], 10)) + + def test_garbage_memory_is_useless(self): + self.assertIsNone(usable_remembered("not a list", 10)) + self.assertIsNone(usable_remembered([True, "3"], 10)) + + +class KeepFramesTests(unittest.TestCase): + def test_keeps_the_chosen_frames_in_source_order(self): + kept = keep_frames(batch(6), [2, 5]) + self.assertEqual(kept.shape[0], 2) + self.assertEqual(float(kept[0, 0, 0, 0]), 1.0) + self.assertEqual(float(kept[1, 0, 0, 0]), 4.0) + + def test_empty_selection_keeps_the_whole_batch(self): + frames = batch(4) + self.assertIs(keep_frames(frames, []), frames) + + def test_rejects_non_bhwc_input(self): + with self.assertRaises(ValueError): + keep_frames(torch.zeros((3, 3)), [1]) + + +class ReportTests(unittest.TestCase): + def test_effective_indices_expand_keep_all(self): + self.assertEqual(effective_indices([], 4), [1, 2, 3, 4]) + self.assertEqual(effective_indices([2, 4], 4), [2, 4]) + + def test_indices_string_is_one_based_and_comma_joined(self): + self.assertEqual(indices_string([1, 4, 9]), "1,4,9") + self.assertEqual(indices_string([]), "") + + +if __name__ == "__main__": + unittest.main()