76 lines
3.0 KiB
Python
76 lines
3.0 KiB
Python
"""GridImagePool — the Image Pool (Grid) ComfyUI node."""
|
|
import os
|
|
from .gates_compat import grid_pool_base as _grid_pool_base
|
|
from . import pool, imaging
|
|
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
|
|
class GridImagePool:
|
|
CATEGORY = "Datasete Gates"
|
|
FUNCTION = "run"
|
|
RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "STRING")
|
|
RETURN_NAMES = ("image", "mask", "index", "count", "label")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# pool_id is a per-node UUID owned by the JS extension. It must be a
|
|
# NORMAL input, not a "hidden" one: ComfyUI only populates hidden inputs
|
|
# for built-in types (UNIQUE_ID, PROMPT, ...), so a custom hidden
|
|
# "POOL_ID" would always arrive as None. The frontend, however, forwards
|
|
# every serialized widget value by name, so a declared STRING input
|
|
# backed by a (hidden-rendered) widget reliably reaches run().
|
|
return {
|
|
"required": {
|
|
"index": ("INT", {"default": -1, "min": -1, "max": 9999}),
|
|
"pool_id": ("STRING", {"default": "default"}),
|
|
},
|
|
# optional companion input: a Pool Profile node feeds the profile id
|
|
# here; when connected it overrides pool_id (see `effective` below).
|
|
"optional": {
|
|
"profile": ("POOL_PROFILE",),
|
|
},
|
|
}
|
|
|
|
@staticmethod
|
|
def _resolve(index, pool_id):
|
|
base = _grid_pool_base()
|
|
m = pool.read_manifest(base, pool_id)
|
|
idx = pool.resolve_slot(m, index)
|
|
return base, m, idx
|
|
|
|
def run(self, index, pool_id="default", profile=None):
|
|
effective = profile or pool_id
|
|
base, m, idx = self._resolve(index, effective)
|
|
if idx < 0:
|
|
img, mask = imaging.empty_outputs()
|
|
return (img, mask, 0, 0, "")
|
|
slot = m["slots"][idx]
|
|
d = pool.pool_dir(base, effective)
|
|
img = imaging.load_image_tensor(str(d / slot["image"]))
|
|
h, w = int(img.shape[1]), int(img.shape[2])
|
|
mask_name = slot.get("mask")
|
|
mask = imaging.load_mask_tensor(str(d / mask_name) if mask_name else None, h, w)
|
|
return (img, mask, idx, len(m["slots"]), slot.get("label", ""))
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, index, pool_id="default", profile=None, **kwargs):
|
|
effective = profile or pool_id
|
|
base, m, idx = cls._resolve(index, effective)
|
|
if idx < 0:
|
|
return imaging.change_hash(effective, -1, [])
|
|
slot = m["slots"][idx]
|
|
d = pool.pool_dir(base, effective)
|
|
mtimes = []
|
|
for key in ("image", "mask"):
|
|
name = slot.get(key)
|
|
p = d / name if name else None
|
|
mtimes.append(os.path.getmtime(p) if p and p.exists() else 0.0)
|
|
# include active so manual selection changes invalidate cache
|
|
return imaging.change_hash(effective, f"{idx}:{m.get('active')}", mtimes)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"GridImagePool": GridImagePool}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"GridImagePool": "Image Pool (Grid)"}
|