301 lines
10 KiB
Python
301 lines
10 KiB
Python
import os
|
|
import time
|
|
import uuid
|
|
import threading
|
|
from typing import Any, Dict, Optional, Tuple
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from aiohttp import web
|
|
|
|
import comfy.utils
|
|
import folder_paths
|
|
from server import PromptServer
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Thread-safe, per-node state store
|
|
# -----------------------------------------------------------------------------
|
|
_STATE_LOCK = threading.Lock()
|
|
_STATE: Dict[str, Dict[str, Any]] = {}
|
|
# _STATE[node_id] = {"event": Event, "token": str, "data": dict|None, "created": float}
|
|
|
|
WAIT_TIMEOUT_SECONDS = 600 # 10 minutes
|
|
|
|
|
|
def _get_or_create_state(node_id: str) -> Dict[str, Any]:
|
|
with _STATE_LOCK:
|
|
st = _STATE.get(node_id)
|
|
if st is None:
|
|
st = {
|
|
"event": threading.Event(),
|
|
"token": "",
|
|
"data": None,
|
|
"created": time.time(),
|
|
}
|
|
_STATE[node_id] = st
|
|
return st
|
|
|
|
|
|
def _reset_state_for_run(node_id: str) -> str:
|
|
st = _get_or_create_state(node_id)
|
|
token = uuid.uuid4().hex
|
|
with _STATE_LOCK:
|
|
st["token"] = token
|
|
st["data"] = None
|
|
st["created"] = time.time()
|
|
st["event"].clear()
|
|
return token
|
|
|
|
|
|
def _set_payload(node_id: str, token: str, payload: dict) -> Tuple[bool, str]:
|
|
st = _get_or_create_state(node_id)
|
|
with _STATE_LOCK:
|
|
expected = st.get("token") or ""
|
|
if expected and token != expected:
|
|
return False, f"Token mismatch for node_id={node_id}. expected={expected} got={token}"
|
|
st["data"] = payload
|
|
st["event"].set()
|
|
return True, "ok"
|
|
|
|
|
|
def _pop_payload(node_id: str) -> Optional[dict]:
|
|
with _STATE_LOCK:
|
|
st = _STATE.get(node_id)
|
|
if not st:
|
|
return None
|
|
data = st.get("data")
|
|
_STATE.pop(node_id, None)
|
|
return data
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# API handler
|
|
# -----------------------------------------------------------------------------
|
|
async def submit_crop_handler(request: web.Request) -> web.Response:
|
|
try:
|
|
json_data = await request.json()
|
|
|
|
node_id = str(json_data.get("node_id", "")).strip()
|
|
token = str(json_data.get("token", "")).strip()
|
|
action = str(json_data.get("action", "submit")).strip().lower()
|
|
crop_data = json_data.get("crop_data")
|
|
|
|
print(f"[VisualMarquee] API Received: node_id={node_id} token={token} action={action}")
|
|
|
|
if not node_id:
|
|
return web.json_response({"status": "error", "message": "Missing node_id"}, status=400)
|
|
if not token:
|
|
return web.json_response({"status": "error", "message": "Missing token"}, status=400)
|
|
|
|
if action not in ("submit", "cancel"):
|
|
return web.json_response({"status": "error", "message": "Invalid action"}, status=400)
|
|
|
|
if action == "cancel":
|
|
ok, msg = _set_payload(node_id, token, {"__action__": "cancel"})
|
|
if not ok:
|
|
return web.json_response({"status": "error", "message": msg}, status=409)
|
|
return web.json_response({"status": "success", "action": "cancel"})
|
|
|
|
# submit
|
|
if not isinstance(crop_data, dict):
|
|
return web.json_response({"status": "error", "message": "crop_data must be an object"}, status=400)
|
|
|
|
ok, msg = _set_payload(node_id, token, {"__action__": "submit", "crop_data": crop_data})
|
|
if not ok:
|
|
return web.json_response({"status": "error", "message": msg}, status=409)
|
|
|
|
return web.json_response({"status": "success", "action": "submit"})
|
|
except Exception as e:
|
|
print(f"[VisualMarquee] API Error: {e}")
|
|
return web.json_response({"status": "error", "message": str(e)}, status=500)
|
|
|
|
|
|
def _route_exists(app: web.Application, path: str, method: str = "POST") -> bool:
|
|
try:
|
|
for r in app.router.routes():
|
|
if r.method != method:
|
|
continue
|
|
res = getattr(r, "resource", None)
|
|
if res is None:
|
|
continue
|
|
if getattr(res, "canonical", None) == path:
|
|
return True
|
|
except Exception:
|
|
pass
|
|
return False
|
|
|
|
|
|
def _register_routes() -> None:
|
|
if not hasattr(PromptServer.instance, "app"):
|
|
print("[VisualMarquee] PromptServer has no app yet. Routes not registered.")
|
|
return
|
|
|
|
app = PromptServer.instance.app
|
|
|
|
# Register BOTH paths (ComfyUI frontend variations)
|
|
paths = ["/flow_assistor/submit_crop", "/api/flow_assistor/submit_crop"]
|
|
for p in paths:
|
|
if not _route_exists(app, p, "POST"):
|
|
app.router.add_post(p, submit_crop_handler)
|
|
print(f"[VisualMarquee] API route registered: POST {p}")
|
|
|
|
|
|
_register_routes()
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Node
|
|
# -----------------------------------------------------------------------------
|
|
class VisualMarqueeSelection:
|
|
"""
|
|
Visual Marquee (Interactive)
|
|
|
|
- Shows popup to select area
|
|
- Pauses until confirm or exit/cancel
|
|
- "Original size" (default ON) returns exact selection w/h
|
|
- If Original size OFF, upscale to max_resolution (longest side)
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_temp_directory()
|
|
self.type = "temp"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"max_resolution": ("INT", {"default": 1024, "min": 512, "max": 4096}),
|
|
# Requested toggle (default enabled)
|
|
"original_size": ("BOOLEAN", {"default": True}),
|
|
# keep old snapping behavior only when upscaling
|
|
"force_multiple_of_8": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "TILE_DATA")
|
|
RETURN_NAMES = ("cropped_image", "mask", "tile_data")
|
|
FUNCTION = "process"
|
|
CATEGORY = "flow-assistor/interactive"
|
|
OUTPUT_NODE = False
|
|
|
|
def IS_CHANGED(self, **kwargs):
|
|
return float("nan")
|
|
|
|
def process(
|
|
self,
|
|
image,
|
|
max_resolution,
|
|
original_size=True,
|
|
force_multiple_of_8=True,
|
|
unique_id=None,
|
|
extra_pnginfo=None
|
|
):
|
|
node_id = str(unique_id).strip() or uuid.uuid4().hex
|
|
|
|
# 1) Save preview
|
|
img_tensor = image[0] # [H, W, C]
|
|
arr = (255.0 * img_tensor.detach().cpu().numpy())
|
|
arr = np.clip(arr, 0, 255).astype(np.uint8)
|
|
img = Image.fromarray(arr)
|
|
|
|
filename = f"marquee_preview_{node_id}.png"
|
|
full_path = os.path.join(self.output_dir, filename)
|
|
img.save(full_path)
|
|
|
|
# 2) Reset state and create token for this run
|
|
token = _reset_state_for_run(node_id)
|
|
|
|
# 3) Notify browser (also tell it current mode for display text)
|
|
PromptServer.instance.send_sync("flow_assistor_marquee_show", {
|
|
"node_id": node_id,
|
|
"token": token,
|
|
"filename": filename,
|
|
"max_resolution": int(max_resolution),
|
|
"original_size": bool(original_size),
|
|
})
|
|
|
|
print(f"[VisualMarquee] Node {node_id} paused. Waiting... token={token}")
|
|
|
|
# 4) Wait
|
|
st = _get_or_create_state(node_id)
|
|
ok = st["event"].wait(timeout=WAIT_TIMEOUT_SECONDS)
|
|
if not ok:
|
|
with _STATE_LOCK:
|
|
_STATE.pop(node_id, None)
|
|
raise TimeoutError("[VisualMarquee] Timed out waiting for selection.")
|
|
|
|
payload = _pop_payload(node_id)
|
|
if not isinstance(payload, dict):
|
|
raise RuntimeError("[VisualMarquee] Resume signaled but payload missing/invalid.")
|
|
|
|
if payload.get("__action__") == "cancel":
|
|
# Stop whole process
|
|
raise RuntimeError("[VisualMarquee] Cancelled by user.")
|
|
|
|
crop_data = payload.get("crop_data")
|
|
if not isinstance(crop_data, dict):
|
|
raise RuntimeError("[VisualMarquee] crop_data missing/invalid.")
|
|
|
|
# 5) Parse crop
|
|
img_h, img_w = int(image.shape[1]), int(image.shape[2])
|
|
|
|
x = int(round(float(crop_data.get("x", 0))))
|
|
y = int(round(float(crop_data.get("y", 0))))
|
|
w = int(round(float(crop_data.get("w", 512))))
|
|
h = int(round(float(crop_data.get("h", 512))))
|
|
|
|
# Clamp
|
|
x = max(0, min(x, img_w - 1))
|
|
y = max(0, min(y, img_h - 1))
|
|
w = max(1, min(w, img_w - x))
|
|
h = max(1, min(h, img_h - y))
|
|
|
|
cropped = image[:, y:y + h, x:x + w, :] # [B, h, w, C]
|
|
|
|
# 6) Output size logic
|
|
if original_size:
|
|
final_image = cropped
|
|
out_w, out_h = w, h
|
|
else:
|
|
scale = float(max_resolution) / float(max(w, h))
|
|
out_w = max(1, int(round(w * scale)))
|
|
out_h = max(1, int(round(h * scale)))
|
|
|
|
if force_multiple_of_8:
|
|
out_w = max(8, (out_w // 8) * 8)
|
|
out_h = max(8, (out_h // 8) * 8)
|
|
|
|
samples = cropped.movedim(-1, 1) # [B, C, h, w]
|
|
resized = comfy.utils.common_upscale(samples, out_w, out_h, "lanczos", "disabled")
|
|
final_image = resized.movedim(1, -1) # [B, out_h, out_w, C]
|
|
|
|
mask = torch.ones(
|
|
(final_image.shape[0], final_image.shape[1], final_image.shape[2]),
|
|
dtype=torch.float32,
|
|
device=final_image.device
|
|
)
|
|
|
|
tile_data = {
|
|
"original_bbox": (x, y, w, h),
|
|
"target_size": (int(out_w), int(out_h)),
|
|
"original_image_shape": (img_h, img_w),
|
|
"original_size_mode": bool(original_size),
|
|
}
|
|
|
|
print(f"[VisualMarquee] Resumed: selected={w}x{h} output={out_w}x{out_h} original_size={original_size}")
|
|
return (final_image, mask, tile_data)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"VisualMarqueeSelection": VisualMarqueeSelection,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"VisualMarqueeSelection": "Visual Marquee (Interactive)",
|
|
} |