Add batch image chooser gate
This commit is contained in:
+83
-1
@@ -1,5 +1,7 @@
|
||||
"""Blocking choice bus for the Image Gate node. Stdlib only — no comfy/torch."""
|
||||
"""Blocking coordination for manual gate nodes. Stdlib only — no comfy/torch."""
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
|
||||
|
||||
class GateCancelled(Exception):
|
||||
@@ -11,6 +13,11 @@ class GateBus:
|
||||
masks = {} # node_id(str) -> PNG bytes
|
||||
payloads = {} # node_id(str) -> arbitrary payload (e.g., edited text)
|
||||
cancelled = False
|
||||
active_tokens = {} # node_id(str) -> per-run token for scoped waiters
|
||||
token_payloads = {} # (node_id, token) -> arbitrary payload
|
||||
token_cancelled = set()
|
||||
token_contexts = {} # (node_id, token) -> waiter-specific validation data
|
||||
token_lock = threading.Lock()
|
||||
|
||||
@classmethod
|
||||
def arm(cls, node_id):
|
||||
@@ -57,3 +64,78 @@ class GateBus:
|
||||
raise GateCancelled()
|
||||
time.sleep(period)
|
||||
return cls.payloads.pop(sid)
|
||||
|
||||
@classmethod
|
||||
def arm_token(cls, node_id, context=None):
|
||||
"""Open a run-scoped channel and invalidate any older run for the node."""
|
||||
sid = str(node_id)
|
||||
with cls.token_lock:
|
||||
old_token = cls.active_tokens.get(sid)
|
||||
if old_token is not None:
|
||||
old_key = (sid, old_token)
|
||||
cls.token_payloads.pop(old_key, None)
|
||||
cls.token_cancelled.discard(old_key)
|
||||
cls.token_contexts.pop(old_key, None)
|
||||
|
||||
token = uuid.uuid4().hex
|
||||
key = (sid, token)
|
||||
cls.active_tokens[sid] = token
|
||||
cls.token_contexts[key] = context
|
||||
return token
|
||||
|
||||
@classmethod
|
||||
def token_context(cls, node_id, token):
|
||||
sid = str(node_id)
|
||||
with cls.token_lock:
|
||||
if cls.active_tokens.get(sid) != token:
|
||||
return None
|
||||
return cls.token_contexts.get((sid, token))
|
||||
|
||||
@classmethod
|
||||
def put_token_payload(cls, node_id, token, value):
|
||||
sid = str(node_id)
|
||||
with cls.token_lock:
|
||||
if cls.active_tokens.get(sid) != token:
|
||||
return False
|
||||
cls.token_payloads[(sid, token)] = value
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def cancel_token(cls, node_id, token):
|
||||
sid = str(node_id)
|
||||
with cls.token_lock:
|
||||
if cls.active_tokens.get(sid) != token:
|
||||
return False
|
||||
cls.token_cancelled.add((sid, token))
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def wait_token_payload(cls, node_id, token, period=0.1, should_cancel=None):
|
||||
sid = str(node_id)
|
||||
key = (sid, token)
|
||||
while True:
|
||||
with cls.token_lock:
|
||||
superseded = cls.active_tokens.get(sid) != token
|
||||
cancelled = key in cls.token_cancelled
|
||||
if superseded or cancelled:
|
||||
cls.token_cancelled.discard(key)
|
||||
raise GateCancelled()
|
||||
if key in cls.token_payloads:
|
||||
return cls.token_payloads.pop(key)
|
||||
|
||||
if should_cancel is not None and should_cancel():
|
||||
with cls.token_lock:
|
||||
cls.token_cancelled.discard(key)
|
||||
raise GateCancelled()
|
||||
time.sleep(period)
|
||||
|
||||
@classmethod
|
||||
def disarm_token(cls, node_id, token):
|
||||
sid = str(node_id)
|
||||
key = (sid, token)
|
||||
with cls.token_lock:
|
||||
if cls.active_tokens.get(sid) == token:
|
||||
cls.active_tokens.pop(sid, None)
|
||||
cls.token_payloads.pop(key, None)
|
||||
cls.token_cancelled.discard(key)
|
||||
cls.token_contexts.pop(key, None)
|
||||
|
||||
@@ -8,6 +8,7 @@ from PIL import Image
|
||||
from server import PromptServer
|
||||
|
||||
from .gate_bus import GateBus
|
||||
from .image_chooser import encode_previews, normalize_selection
|
||||
|
||||
routes = PromptServer.instance.routes
|
||||
|
||||
@@ -23,6 +24,22 @@ def send_preview(node_id, image, n_routes):
|
||||
)
|
||||
|
||||
|
||||
def send_image_choices(node_id, token, images):
|
||||
"""Show a lightweight preview of every image to the queuing client."""
|
||||
server = PromptServer.instance
|
||||
server.send_sync(
|
||||
"datasete-image-chooser-show",
|
||||
{
|
||||
"id": str(node_id),
|
||||
"display_id": str(getattr(server, "last_node_id", None) or node_id),
|
||||
"token": token,
|
||||
"images": encode_previews(images),
|
||||
"count": int(images.shape[0]),
|
||||
},
|
||||
getattr(server, "client_id", None),
|
||||
)
|
||||
|
||||
|
||||
@routes.post("/datasete_gate/choice")
|
||||
async def _choice(request):
|
||||
post = await request.post()
|
||||
@@ -44,6 +61,35 @@ async def _mask(request):
|
||||
return web.json_response({})
|
||||
|
||||
|
||||
@routes.post("/datasete_image_chooser/select")
|
||||
async def _image_chooser_select(request):
|
||||
post = await request.post()
|
||||
node_id = post.get("id")
|
||||
token = post.get("token")
|
||||
if node_id is None or token is None:
|
||||
return web.json_response({"error": "missing node id or token"}, status=400)
|
||||
|
||||
batch_size = GateBus.token_context(node_id, token)
|
||||
if batch_size is None:
|
||||
return web.json_response({"error": "chooser run is no longer active"}, status=409)
|
||||
|
||||
if post.get("action") == "cancel":
|
||||
accepted = GateBus.cancel_token(node_id, token)
|
||||
else:
|
||||
selection = post.get("selection")
|
||||
if selection is None:
|
||||
return web.json_response({"error": "missing selection"}, status=400)
|
||||
try:
|
||||
selection = normalize_selection(selection, batch_size)
|
||||
except ValueError as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
accepted = GateBus.put_token_payload(node_id, token, selection)
|
||||
|
||||
if not accepted:
|
||||
return web.json_response({"error": "chooser run is no longer active"}, status=409)
|
||||
return web.json_response({})
|
||||
|
||||
|
||||
def send_text(node_id, text):
|
||||
PromptServer.instance.send_sync(
|
||||
"datasete-textgate-show", {"id": str(node_id), "text": text or ""}
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
"""Manual gate that selects a subset from an IMAGE batch.
|
||||
|
||||
The selection helpers and thumbnail encoder intentionally avoid importing
|
||||
ComfyUI so this module remains straightforward to unit-test.
|
||||
"""
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from . import gate_bus
|
||||
|
||||
|
||||
PREVIEW_MAX_SIDE = 256
|
||||
PREVIEW_JPEG_QUALITY = 82
|
||||
|
||||
|
||||
def normalize_selection(selection, batch_size):
|
||||
"""Return validated, unique indices in their original batch order."""
|
||||
if isinstance(selection, str):
|
||||
try:
|
||||
selection = json.loads(selection)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError("Selection must be a JSON array of image indices") from exc
|
||||
|
||||
if not isinstance(selection, (list, tuple)):
|
||||
raise ValueError("Selection must be a list of image indices")
|
||||
if not selection:
|
||||
raise ValueError("Select at least one image")
|
||||
|
||||
unique = set()
|
||||
for index in selection:
|
||||
if isinstance(index, bool) or not isinstance(index, int):
|
||||
raise ValueError("Every selected image index must be an integer")
|
||||
if index < 0 or index >= batch_size:
|
||||
raise ValueError(
|
||||
f"Selected image index {index} is outside batch size {batch_size}"
|
||||
)
|
||||
unique.add(index)
|
||||
|
||||
# Batch order is deterministic and does not depend on click order.
|
||||
return tuple(sorted(unique))
|
||||
|
||||
|
||||
def select_batch(images, selection):
|
||||
"""Select one or more images while preserving the IMAGE batch dimension."""
|
||||
batch_size = int(images.shape[0])
|
||||
indices = normalize_selection(selection, batch_size)
|
||||
return images[list(indices)]
|
||||
|
||||
|
||||
def encode_previews(images, max_side=PREVIEW_MAX_SIDE,
|
||||
jpeg_quality=PREVIEW_JPEG_QUALITY):
|
||||
"""Encode small JPEG previews without modifying the original tensor."""
|
||||
previews = []
|
||||
for index, image in enumerate(images):
|
||||
array = (image.detach().cpu().float().numpy() * 255.0).clip(0, 255).astype(
|
||||
np.uint8
|
||||
)
|
||||
pil = Image.fromarray(array)
|
||||
if pil.mode != "RGB":
|
||||
pil = pil.convert("RGB")
|
||||
source_width, source_height = pil.size
|
||||
pil.thumbnail((max_side, max_side), Image.Resampling.LANCZOS)
|
||||
|
||||
buffer = io.BytesIO()
|
||||
pil.save(buffer, "JPEG", quality=jpeg_quality)
|
||||
previews.append({
|
||||
"index": index,
|
||||
"image": base64.b64encode(buffer.getvalue()).decode("ascii"),
|
||||
"width": source_width,
|
||||
"height": source_height,
|
||||
})
|
||||
return previews
|
||||
|
||||
|
||||
class ImageChooserGate:
|
||||
CATEGORY = "Dataset Gates"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("images",)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {"images": ("IMAGE",)},
|
||||
"hidden": {"unique_id": "UNIQUE_ID"},
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
return float("nan")
|
||||
|
||||
def run(self, images, unique_id):
|
||||
batch_size = int(images.shape[0])
|
||||
if batch_size < 1:
|
||||
raise ValueError("Image Chooser Gate requires a non-empty image batch")
|
||||
|
||||
from . import gate_server
|
||||
import comfy.model_management as mm
|
||||
|
||||
token = gate_bus.GateBus.arm_token(unique_id, context=batch_size)
|
||||
try:
|
||||
gate_server.send_image_choices(unique_id, token, images)
|
||||
selection = gate_bus.GateBus.wait_token_payload(
|
||||
unique_id, token, should_cancel=mm.processing_interrupted
|
||||
)
|
||||
except gate_bus.GateCancelled:
|
||||
raise mm.InterruptProcessingException()
|
||||
finally:
|
||||
gate_bus.GateBus.disarm_token(unique_id, token)
|
||||
|
||||
return (select_batch(images, selection),)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {"ImageChooserGate": ImageChooserGate}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageChooserGate": "Image Chooser Gate (Batch)",
|
||||
}
|
||||
Reference in New Issue
Block a user