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

Add the Frame Chooser backend with a pausing selection flow

Browse Source
This commit is contained in:
ausboss
2026-08-13 09:43:08 -05:00
parent 3d5f0502bf
commit ee252ee041
5 changed files with 496 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+1
View File
@@ -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",
nodes/_chooser_helpers.py
+274
View File
@@ -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",
]
nodes/node_frame_chooser.py
+118
View File
@@ -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"]
scripts/validate_nodes.py
+1
View File
@@ -26,6 +26,7 @@ errors = []
warnings = []
PUBLIC_NODE_IDS = {
"AUSBOSS_NODES_FrameChooser",
"AUSBOSS_NODES_ImageCropRotatePad",
"AUSBOSS_NODES_LaMaInpaint",
"AUSBOSS_NODES_LoadVideo",
tests/test_chooser_helpers.py
+102
View File
@@ -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()