Merge feat/pool-profiles: Pool Profile companion node + portable profiles
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
+7
-2
@@ -16,11 +16,16 @@ if __package__:
|
|||||||
NODE_DISPLAY_NAME_MAPPINGS as _GATE_NAMES
|
NODE_DISPLAY_NAME_MAPPINGS as _GATE_NAMES
|
||||||
from .gates.textgate import NODE_CLASS_MAPPINGS as _TEXT_NODES, \
|
from .gates.textgate import NODE_CLASS_MAPPINGS as _TEXT_NODES, \
|
||||||
NODE_DISPLAY_NAME_MAPPINGS as _TEXT_NAMES
|
NODE_DISPLAY_NAME_MAPPINGS as _TEXT_NAMES
|
||||||
|
from .gates.profile_node import NODE_CLASS_MAPPINGS as _PROF_NODES, \
|
||||||
|
NODE_DISPLAY_NAME_MAPPINGS as _PROF_NAMES
|
||||||
from .gates import routes # noqa: F401 (registers aiohttp routes on import)
|
from .gates import routes # noqa: F401 (registers aiohttp routes on import)
|
||||||
from .gates import gate_server # noqa: F401 (registers /datasete_gate/* + text routes)
|
from .gates import gate_server # noqa: F401 (registers /datasete_gate/* + text routes)
|
||||||
|
from .gates import profiles_routes # noqa: F401 (registers /grid_pool/profiles/*)
|
||||||
|
|
||||||
NODE_CLASS_MAPPINGS = {**_POOL_NODES, **_LOADER_NODES, **_GATE_NODES, **_TEXT_NODES}
|
NODE_CLASS_MAPPINGS = {**_POOL_NODES, **_LOADER_NODES, **_GATE_NODES,
|
||||||
NODE_DISPLAY_NAME_MAPPINGS = {**_POOL_NAMES, **_LOADER_NAMES, **_GATE_NAMES, **_TEXT_NAMES}
|
**_TEXT_NODES, **_PROF_NODES}
|
||||||
|
NODE_DISPLAY_NAME_MAPPINGS = {**_POOL_NAMES, **_LOADER_NAMES, **_GATE_NAMES,
|
||||||
|
**_TEXT_NAMES, **_PROF_NAMES}
|
||||||
else: # pragma: no cover - exercised only under pytest collection
|
else: # pragma: no cover - exercised only under pytest collection
|
||||||
NODE_CLASS_MAPPINGS = {}
|
NODE_CLASS_MAPPINGS = {}
|
||||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||||
|
|||||||
+15
-8
@@ -26,6 +26,11 @@ class GridImagePool:
|
|||||||
"index": ("INT", {"default": -1, "min": -1, "max": 9999}),
|
"index": ("INT", {"default": -1, "min": -1, "max": 9999}),
|
||||||
"pool_id": ("STRING", {"default": "default"}),
|
"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
|
@staticmethod
|
||||||
@@ -35,13 +40,14 @@ class GridImagePool:
|
|||||||
idx = pool.resolve_slot(m, index)
|
idx = pool.resolve_slot(m, index)
|
||||||
return base, m, idx
|
return base, m, idx
|
||||||
|
|
||||||
def run(self, index, pool_id="default"):
|
def run(self, index, pool_id="default", profile=None):
|
||||||
base, m, idx = self._resolve(index, pool_id)
|
effective = profile or pool_id
|
||||||
|
base, m, idx = self._resolve(index, effective)
|
||||||
if idx < 0:
|
if idx < 0:
|
||||||
img, mask = imaging.empty_outputs()
|
img, mask = imaging.empty_outputs()
|
||||||
return (img, mask, 0, 0, "")
|
return (img, mask, 0, 0, "")
|
||||||
slot = m["slots"][idx]
|
slot = m["slots"][idx]
|
||||||
d = pool.pool_dir(base, pool_id)
|
d = pool.pool_dir(base, effective)
|
||||||
img = imaging.load_image_tensor(str(d / slot["image"]))
|
img = imaging.load_image_tensor(str(d / slot["image"]))
|
||||||
h, w = int(img.shape[1]), int(img.shape[2])
|
h, w = int(img.shape[1]), int(img.shape[2])
|
||||||
mask_name = slot.get("mask")
|
mask_name = slot.get("mask")
|
||||||
@@ -49,19 +55,20 @@ class GridImagePool:
|
|||||||
return (img, mask, idx, len(m["slots"]), slot.get("label", ""))
|
return (img, mask, idx, len(m["slots"]), slot.get("label", ""))
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def IS_CHANGED(cls, index, pool_id="default", **kwargs):
|
def IS_CHANGED(cls, index, pool_id="default", profile=None, **kwargs):
|
||||||
base, m, idx = cls._resolve(index, pool_id)
|
effective = profile or pool_id
|
||||||
|
base, m, idx = cls._resolve(index, effective)
|
||||||
if idx < 0:
|
if idx < 0:
|
||||||
return imaging.change_hash(pool_id, -1, [])
|
return imaging.change_hash(effective, -1, [])
|
||||||
slot = m["slots"][idx]
|
slot = m["slots"][idx]
|
||||||
d = pool.pool_dir(base, pool_id)
|
d = pool.pool_dir(base, effective)
|
||||||
mtimes = []
|
mtimes = []
|
||||||
for key in ("image", "mask"):
|
for key in ("image", "mask"):
|
||||||
name = slot.get(key)
|
name = slot.get(key)
|
||||||
p = d / name if name else None
|
p = d / name if name else None
|
||||||
mtimes.append(os.path.getmtime(p) if p and p.exists() else 0.0)
|
mtimes.append(os.path.getmtime(p) if p and p.exists() else 0.0)
|
||||||
# include active so manual selection changes invalidate cache
|
# include active so manual selection changes invalidate cache
|
||||||
return imaging.change_hash(pool_id, f"{idx}:{m.get('active')}", mtimes)
|
return imaging.change_hash(effective, f"{idx}:{m.get('active')}", mtimes)
|
||||||
|
|
||||||
|
|
||||||
NODE_CLASS_MAPPINGS = {"GridImagePool": GridImagePool}
|
NODE_CLASS_MAPPINGS = {"GridImagePool": GridImagePool}
|
||||||
|
|||||||
@@ -0,0 +1,30 @@
|
|||||||
|
# gates/profile_node.py
|
||||||
|
NODE_CLASS_MAPPINGS = {}
|
||||||
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||||
|
|
||||||
|
|
||||||
|
class PoolProfile:
|
||||||
|
CATEGORY = "Datasete Gates"
|
||||||
|
FUNCTION = "run"
|
||||||
|
RETURN_TYPES = ("POOL_PROFILE",)
|
||||||
|
RETURN_NAMES = ("profile",)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(cls):
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"profile": ("STRING", {"default": ""}), # name; JS renders a dropdown
|
||||||
|
"profile_id": ("STRING", {"default": ""}), # hidden, JS-owned id
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
def run(self, profile, profile_id=""):
|
||||||
|
return (profile_id or "default",)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def IS_CHANGED(cls, profile, profile_id="", **kwargs):
|
||||||
|
return profile_id
|
||||||
|
|
||||||
|
|
||||||
|
NODE_CLASS_MAPPINGS = {"PoolProfile": PoolProfile}
|
||||||
|
NODE_DISPLAY_NAME_MAPPINGS = {"PoolProfile": "Pool Profile"}
|
||||||
@@ -0,0 +1,160 @@
|
|||||||
|
"""Named-profile registry + dir ops for the Image Pool. Stdlib only."""
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import shutil
|
||||||
|
import zipfile
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
REGISTRY_NAME = "profiles.json"
|
||||||
|
|
||||||
|
|
||||||
|
def registry_path(base):
|
||||||
|
return Path(base) / REGISTRY_NAME
|
||||||
|
|
||||||
|
|
||||||
|
def empty_registry():
|
||||||
|
return {"profiles": []}
|
||||||
|
|
||||||
|
|
||||||
|
def read_registry(base):
|
||||||
|
p = registry_path(base)
|
||||||
|
if not p.exists():
|
||||||
|
return empty_registry()
|
||||||
|
try:
|
||||||
|
with open(p, "r", encoding="utf-8") as f:
|
||||||
|
reg = json.load(f)
|
||||||
|
if not isinstance(reg, dict) or "profiles" not in reg:
|
||||||
|
raise ValueError("bad registry")
|
||||||
|
return reg
|
||||||
|
except (ValueError, json.JSONDecodeError):
|
||||||
|
return empty_registry()
|
||||||
|
|
||||||
|
|
||||||
|
def write_registry(base, reg):
|
||||||
|
Path(base).mkdir(parents=True, exist_ok=True)
|
||||||
|
final = registry_path(base)
|
||||||
|
tmp = final.with_name(REGISTRY_NAME + ".tmp")
|
||||||
|
with open(tmp, "w", encoding="utf-8") as f:
|
||||||
|
json.dump(reg, f, indent=2)
|
||||||
|
os.replace(tmp, final)
|
||||||
|
return reg
|
||||||
|
|
||||||
|
|
||||||
|
def find_by_id(reg, pid):
|
||||||
|
return next((p for p in reg["profiles"] if p["id"] == pid), None)
|
||||||
|
|
||||||
|
|
||||||
|
def find_by_name(reg, name):
|
||||||
|
return next((p for p in reg["profiles"] if p["name"] == name), None)
|
||||||
|
|
||||||
|
|
||||||
|
def create_profile(base, name, pid, ts=0):
|
||||||
|
reg = read_registry(base)
|
||||||
|
if find_by_name(reg, name):
|
||||||
|
raise ValueError(f"profile name already exists: {name}")
|
||||||
|
(Path(base) / pid).mkdir(parents=True, exist_ok=True)
|
||||||
|
entry = {"id": pid, "name": name, "created": ts}
|
||||||
|
reg["profiles"].append(entry)
|
||||||
|
write_registry(base, reg)
|
||||||
|
return entry
|
||||||
|
|
||||||
|
|
||||||
|
def rename_profile(base, pid, name):
|
||||||
|
reg = read_registry(base)
|
||||||
|
entry = find_by_id(reg, pid)
|
||||||
|
if not entry:
|
||||||
|
raise KeyError(pid)
|
||||||
|
other = find_by_name(reg, name)
|
||||||
|
if other and other["id"] != pid:
|
||||||
|
raise ValueError(f"profile name already exists: {name}")
|
||||||
|
entry["name"] = name
|
||||||
|
write_registry(base, reg)
|
||||||
|
return entry
|
||||||
|
|
||||||
|
|
||||||
|
def delete_profile(base, pid):
|
||||||
|
reg = read_registry(base)
|
||||||
|
reg["profiles"] = [p for p in reg["profiles"] if p["id"] != pid]
|
||||||
|
write_registry(base, reg)
|
||||||
|
d = Path(base) / pid
|
||||||
|
if d.exists():
|
||||||
|
shutil.rmtree(d)
|
||||||
|
return reg
|
||||||
|
|
||||||
|
|
||||||
|
def duplicate_profile(base, src_id, name, new_id, ts=0):
|
||||||
|
reg = read_registry(base)
|
||||||
|
if not find_by_id(reg, src_id):
|
||||||
|
raise KeyError(src_id)
|
||||||
|
if find_by_name(reg, name):
|
||||||
|
raise ValueError(f"profile name already exists: {name}")
|
||||||
|
src = Path(base) / src_id
|
||||||
|
dst = Path(base) / new_id
|
||||||
|
if src.exists():
|
||||||
|
shutil.copytree(src, dst)
|
||||||
|
else:
|
||||||
|
dst.mkdir(parents=True, exist_ok=True)
|
||||||
|
entry = {"id": new_id, "name": name, "created": ts}
|
||||||
|
reg["profiles"].append(entry)
|
||||||
|
write_registry(base, reg)
|
||||||
|
return entry
|
||||||
|
|
||||||
|
|
||||||
|
def export_profile(base, pid, dest_zip):
|
||||||
|
src = Path(base) / pid
|
||||||
|
if not src.exists():
|
||||||
|
raise KeyError(pid)
|
||||||
|
entry = find_by_id(read_registry(base), pid)
|
||||||
|
name = entry["name"] if entry else pid
|
||||||
|
with zipfile.ZipFile(dest_zip, "w", zipfile.ZIP_DEFLATED) as z:
|
||||||
|
z.writestr("profile_meta.json", json.dumps({"name": name}))
|
||||||
|
for f in src.rglob("*"):
|
||||||
|
if f.is_file():
|
||||||
|
z.write(f, arcname=str(Path("pool") / f.relative_to(src)))
|
||||||
|
return dest_zip
|
||||||
|
|
||||||
|
|
||||||
|
def seed_profile(base, from_id, profile_id):
|
||||||
|
"""Copy a pool dir's files (images/masks/manifest) into a profile dir.
|
||||||
|
|
||||||
|
Used to save an Image Pool's current contents into a freshly-selected empty
|
||||||
|
profile. Copies top-level files only (the pool layout is flat); returns the
|
||||||
|
number of files copied. No-op (0) if the source dir is missing.
|
||||||
|
"""
|
||||||
|
src = Path(base) / from_id
|
||||||
|
dst = Path(base) / profile_id
|
||||||
|
if not src.exists():
|
||||||
|
return 0
|
||||||
|
dst.mkdir(parents=True, exist_ok=True)
|
||||||
|
n = 0
|
||||||
|
for f in src.iterdir():
|
||||||
|
if f.is_file():
|
||||||
|
shutil.copy2(f, dst / f.name)
|
||||||
|
n += 1
|
||||||
|
return n
|
||||||
|
|
||||||
|
|
||||||
|
def import_profile(base, src_zip, new_id, name=None, ts=0):
|
||||||
|
reg = read_registry(base)
|
||||||
|
meta_name = None
|
||||||
|
dst = Path(base) / new_id
|
||||||
|
dst.mkdir(parents=True, exist_ok=True)
|
||||||
|
with zipfile.ZipFile(src_zip) as z:
|
||||||
|
names = z.namelist()
|
||||||
|
if "profile_meta.json" in names:
|
||||||
|
meta_name = json.loads(z.read("profile_meta.json")).get("name")
|
||||||
|
for n in names:
|
||||||
|
if n.startswith("pool/") and not n.endswith("/"):
|
||||||
|
target = dst / n[len("pool/"):]
|
||||||
|
target.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
with z.open(n) as srcf, open(target, "wb") as out:
|
||||||
|
shutil.copyfileobj(srcf, out)
|
||||||
|
final = name or meta_name or new_id
|
||||||
|
candidate, i = final, 2
|
||||||
|
while find_by_name(reg, candidate):
|
||||||
|
candidate = f"{final} ({i})"
|
||||||
|
i += 1
|
||||||
|
entry = {"id": new_id, "name": candidate, "created": ts}
|
||||||
|
reg["profiles"].append(entry)
|
||||||
|
write_registry(base, reg)
|
||||||
|
return entry
|
||||||
@@ -0,0 +1,81 @@
|
|||||||
|
# gates/profiles_routes.py
|
||||||
|
import os
|
||||||
|
import tempfile
|
||||||
|
import uuid
|
||||||
|
|
||||||
|
from aiohttp import web
|
||||||
|
from server import PromptServer
|
||||||
|
|
||||||
|
from . import profiles
|
||||||
|
from .gates_compat import grid_pool_base
|
||||||
|
|
||||||
|
routes = PromptServer.instance.routes
|
||||||
|
|
||||||
|
|
||||||
|
def _base():
|
||||||
|
return grid_pool_base()
|
||||||
|
|
||||||
|
|
||||||
|
@routes.get("/grid_pool/profiles/list")
|
||||||
|
async def _list(request):
|
||||||
|
return web.json_response(profiles.read_registry(_base()))
|
||||||
|
|
||||||
|
|
||||||
|
@routes.post("/grid_pool/profiles/create")
|
||||||
|
async def _create(request):
|
||||||
|
body = await request.json()
|
||||||
|
e = profiles.create_profile(_base(), body["name"], uuid.uuid4().hex)
|
||||||
|
return web.json_response(e)
|
||||||
|
|
||||||
|
|
||||||
|
@routes.post("/grid_pool/profiles/rename")
|
||||||
|
async def _rename(request):
|
||||||
|
body = await request.json()
|
||||||
|
return web.json_response(profiles.rename_profile(_base(), body["id"], body["name"]))
|
||||||
|
|
||||||
|
|
||||||
|
@routes.post("/grid_pool/profiles/delete")
|
||||||
|
async def _delete(request):
|
||||||
|
body = await request.json()
|
||||||
|
return web.json_response(profiles.delete_profile(_base(), body["id"]))
|
||||||
|
|
||||||
|
|
||||||
|
@routes.post("/grid_pool/profiles/duplicate")
|
||||||
|
async def _duplicate(request):
|
||||||
|
body = await request.json()
|
||||||
|
e = profiles.duplicate_profile(_base(), body["id"], body["name"], uuid.uuid4().hex)
|
||||||
|
return web.json_response(e)
|
||||||
|
|
||||||
|
|
||||||
|
@routes.post("/grid_pool/profiles/seed")
|
||||||
|
async def _seed(request):
|
||||||
|
body = await request.json()
|
||||||
|
n = profiles.seed_profile(_base(), body["from"], body["id"])
|
||||||
|
return web.json_response({"copied": n})
|
||||||
|
|
||||||
|
|
||||||
|
@routes.get("/grid_pool/profiles/export")
|
||||||
|
async def _export(request):
|
||||||
|
pid = request.query["id"]
|
||||||
|
reg = profiles.read_registry(_base())
|
||||||
|
entry = profiles.find_by_id(reg, pid)
|
||||||
|
fname = (entry["name"] if entry else pid) + ".zip"
|
||||||
|
tmp = os.path.join(tempfile.gettempdir(), f"profile_{pid}.zip")
|
||||||
|
profiles.export_profile(_base(), pid, tmp)
|
||||||
|
return web.FileResponse(tmp, headers={"Content-Disposition": f'attachment; filename="{fname}"'})
|
||||||
|
|
||||||
|
|
||||||
|
@routes.post("/grid_pool/profiles/import")
|
||||||
|
async def _import(request):
|
||||||
|
reader = await request.multipart()
|
||||||
|
tmp = os.path.join(tempfile.gettempdir(), f"import_{uuid.uuid4().hex}.zip")
|
||||||
|
async for part in reader:
|
||||||
|
if part.name == "file":
|
||||||
|
with open(tmp, "wb") as f:
|
||||||
|
while True:
|
||||||
|
chunk = await part.read_chunk()
|
||||||
|
if not chunk:
|
||||||
|
break
|
||||||
|
f.write(chunk)
|
||||||
|
e = profiles.import_profile(_base(), tmp, uuid.uuid4().hex)
|
||||||
|
return web.json_response(e)
|
||||||
@@ -54,3 +54,17 @@ def test_is_changed_differs_after_active_change(tmp_path, monkeypatch):
|
|||||||
pool.set_active(base, "p1", 1)
|
pool.set_active(base, "p1", 1)
|
||||||
h2 = node.GridImagePool.IS_CHANGED(index=-1, pool_id="p1")
|
h2 = node.GridImagePool.IS_CHANGED(index=-1, pool_id="p1")
|
||||||
assert h1 != h2
|
assert h1 != h2
|
||||||
|
|
||||||
|
|
||||||
|
def test_profile_input_overrides_pool_id(tmp_path, monkeypatch):
|
||||||
|
base = str(tmp_path / "grid_pool")
|
||||||
|
monkeypatch.setattr(node, "_grid_pool_base", lambda: base)
|
||||||
|
import io
|
||||||
|
from PIL import Image
|
||||||
|
from gates import pool
|
||||||
|
buf = io.BytesIO(); Image.new("RGB", (4, 6), (255, 0, 0)).save(buf, "PNG")
|
||||||
|
pool.add_image(base, "prof1", buf.getvalue(), ts=1) # images under the PROFILE id
|
||||||
|
n = node.GridImagePool()
|
||||||
|
# pool_id is "default" (empty) but profile points at prof1
|
||||||
|
img, mask, idx, count, label = n.run(index=-1, pool_id="default", profile="prof1")
|
||||||
|
assert count == 1 and idx == 0
|
||||||
|
|||||||
@@ -0,0 +1,13 @@
|
|||||||
|
# tests/test_profile_node.py
|
||||||
|
from gates import profile_node as pn
|
||||||
|
|
||||||
|
def test_io():
|
||||||
|
assert pn.PoolProfile.RETURN_TYPES == ("POOL_PROFILE",)
|
||||||
|
assert pn.PoolProfile.RETURN_NAMES == ("profile",)
|
||||||
|
|
||||||
|
def test_run_returns_id_or_default():
|
||||||
|
assert pn.PoolProfile().run(profile="setA", profile_id="id1") == ("id1",)
|
||||||
|
assert pn.PoolProfile().run(profile="", profile_id="") == ("default",)
|
||||||
|
|
||||||
|
def test_is_changed_tracks_id():
|
||||||
|
assert pn.PoolProfile.IS_CHANGED(profile="x", profile_id="id1") == "id1"
|
||||||
@@ -0,0 +1,108 @@
|
|||||||
|
# tests/test_profiles.py
|
||||||
|
from gates import profiles as pr
|
||||||
|
|
||||||
|
def test_empty_registry():
|
||||||
|
assert pr.empty_registry() == {"profiles": []}
|
||||||
|
|
||||||
|
def test_read_missing_is_empty(tmp_path):
|
||||||
|
assert pr.read_registry(str(tmp_path)) == {"profiles": []}
|
||||||
|
|
||||||
|
def test_write_then_read(tmp_path):
|
||||||
|
reg = {"profiles": [{"id": "a", "name": "n", "created": 1}]}
|
||||||
|
pr.write_registry(str(tmp_path), reg)
|
||||||
|
assert (tmp_path / "profiles.json").exists()
|
||||||
|
assert pr.read_registry(str(tmp_path)) == reg
|
||||||
|
|
||||||
|
def test_read_corrupt_is_empty(tmp_path):
|
||||||
|
(tmp_path / "profiles.json").write_text("{ not json")
|
||||||
|
assert pr.read_registry(str(tmp_path)) == {"profiles": []}
|
||||||
|
|
||||||
|
def test_find_helpers():
|
||||||
|
reg = {"profiles": [{"id": "a", "name": "x"}, {"id": "b", "name": "y"}]}
|
||||||
|
assert pr.find_by_id(reg, "b")["name"] == "y"
|
||||||
|
assert pr.find_by_name(reg, "x")["id"] == "a"
|
||||||
|
assert pr.find_by_id(reg, "z") is None
|
||||||
|
|
||||||
|
def test_create_profile(tmp_path):
|
||||||
|
e = pr.create_profile(str(tmp_path), "setA", "id1", ts=10)
|
||||||
|
assert e == {"id": "id1", "name": "setA", "created": 10}
|
||||||
|
assert (tmp_path / "id1").is_dir()
|
||||||
|
assert pr.find_by_name(pr.read_registry(str(tmp_path)), "setA")["id"] == "id1"
|
||||||
|
|
||||||
|
def test_create_duplicate_name_raises(tmp_path):
|
||||||
|
import pytest
|
||||||
|
pr.create_profile(str(tmp_path), "setA", "id1")
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
pr.create_profile(str(tmp_path), "setA", "id2")
|
||||||
|
|
||||||
|
def test_rename_profile(tmp_path):
|
||||||
|
pr.create_profile(str(tmp_path), "old", "id1")
|
||||||
|
e = pr.rename_profile(str(tmp_path), "id1", "new")
|
||||||
|
assert e["name"] == "new"
|
||||||
|
assert pr.find_by_name(pr.read_registry(str(tmp_path)), "new")["id"] == "id1"
|
||||||
|
|
||||||
|
def test_rename_to_existing_name_raises(tmp_path):
|
||||||
|
import pytest
|
||||||
|
pr.create_profile(str(tmp_path), "a", "id1")
|
||||||
|
pr.create_profile(str(tmp_path), "b", "id2")
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
pr.rename_profile(str(tmp_path), "id2", "a")
|
||||||
|
|
||||||
|
def test_delete_profile_removes_dir_and_entry(tmp_path):
|
||||||
|
pr.create_profile(str(tmp_path), "a", "id1")
|
||||||
|
(tmp_path / "id1" / "img_0001.png").write_bytes(b"x")
|
||||||
|
pr.delete_profile(str(tmp_path), "id1")
|
||||||
|
assert not (tmp_path / "id1").exists()
|
||||||
|
assert pr.find_by_id(pr.read_registry(str(tmp_path)), "id1") is None
|
||||||
|
|
||||||
|
def test_duplicate_copies_images(tmp_path):
|
||||||
|
pr.create_profile(str(tmp_path), "src", "id1")
|
||||||
|
(tmp_path / "id1" / "img_0001.png").write_bytes(b"abc")
|
||||||
|
e = pr.duplicate_profile(str(tmp_path), "id1", "copy", "id2", ts=5)
|
||||||
|
assert e == {"id": "id2", "name": "copy", "created": 5}
|
||||||
|
assert (tmp_path / "id2" / "img_0001.png").read_bytes() == b"abc"
|
||||||
|
|
||||||
|
def test_duplicate_duplicate_name_raises(tmp_path):
|
||||||
|
import pytest
|
||||||
|
pr.create_profile(str(tmp_path), "src", "id1")
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
pr.duplicate_profile(str(tmp_path), "id1", "src", "id2")
|
||||||
|
|
||||||
|
def test_export_import_roundtrip(tmp_path):
|
||||||
|
src_base = str(tmp_path / "a"); dst_base = str(tmp_path / "b")
|
||||||
|
pr.create_profile(src_base, "setA", "id1", ts=1)
|
||||||
|
from pathlib import Path
|
||||||
|
(Path(src_base) / "id1" / "img_0001.png").write_bytes(b"hello")
|
||||||
|
zpath = str(tmp_path / "setA.zip")
|
||||||
|
pr.export_profile(src_base, "id1", zpath)
|
||||||
|
assert (tmp_path / "setA.zip").exists()
|
||||||
|
# import into a different base, fresh id
|
||||||
|
e = pr.import_profile(dst_base, zpath, "id99", ts=2)
|
||||||
|
assert e["id"] == "id99"
|
||||||
|
assert e["name"] == "setA" # name carried in zip meta
|
||||||
|
assert (Path(dst_base) / "id99" / "img_0001.png").read_bytes() == b"hello"
|
||||||
|
|
||||||
|
def test_import_name_collision_suffixes(tmp_path):
|
||||||
|
base = str(tmp_path)
|
||||||
|
pr.create_profile(base, "setA", "id1")
|
||||||
|
from pathlib import Path
|
||||||
|
(Path(base) / "id1" / "f.png").write_bytes(b"x")
|
||||||
|
z = str(tmp_path / "e.zip"); pr.export_profile(base, "id1", z)
|
||||||
|
e = pr.import_profile(base, z, "id2")
|
||||||
|
assert e["name"] == "setA (2)"
|
||||||
|
|
||||||
|
def test_seed_profile_copies_pool_into_empty(tmp_path):
|
||||||
|
from pathlib import Path
|
||||||
|
base = str(tmp_path)
|
||||||
|
pr.create_profile(base, "A", "id1") # empty profile dir
|
||||||
|
(Path(base) / "srcpool").mkdir() # a pool's own-UUID dir
|
||||||
|
(Path(base) / "srcpool" / "img_0001.png").write_bytes(b"img")
|
||||||
|
(Path(base) / "srcpool" / "manifest.json").write_text("{}")
|
||||||
|
n = pr.seed_profile(base, "srcpool", "id1")
|
||||||
|
assert n == 2 # image + manifest copied
|
||||||
|
assert (Path(base) / "id1" / "img_0001.png").read_bytes() == b"img"
|
||||||
|
|
||||||
|
def test_seed_profile_missing_source_is_noop(tmp_path):
|
||||||
|
base = str(tmp_path)
|
||||||
|
pr.create_profile(base, "A", "id1")
|
||||||
|
assert pr.seed_profile(base, "nope", "id1") == 0
|
||||||
@@ -593,6 +593,10 @@ function setupGridNode(node) {
|
|||||||
if (node._countEl) node._countEl.textContent = `${n} image${n === 1 ? "" : "s"}`;
|
if (node._countEl) node._countEl.textContent = `${n} image${n === 1 ? "" : "s"}`;
|
||||||
};
|
};
|
||||||
|
|
||||||
|
// public hook for the Pool Profile companion: after it sets our pool_id widget
|
||||||
|
// to the selected profile id, it calls this to repaint the grid from that dir.
|
||||||
|
node._datasetePoolRefresh = () => node._gridRefresh();
|
||||||
|
|
||||||
// initial width (a sensible wide default) + content-driven height; the node
|
// initial width (a sensible wide default) + content-driven height; the node
|
||||||
// stays freely resizable (no width floor) and the grid flex-wraps to fit.
|
// stays freely resizable (no width floor) and the grid flex-wraps to fit.
|
||||||
node._lastCount = 0;
|
node._lastCount = 0;
|
||||||
|
|||||||
@@ -0,0 +1,267 @@
|
|||||||
|
import { app } from "../../scripts/app.js";
|
||||||
|
import { api } from "../../scripts/api.js";
|
||||||
|
|
||||||
|
// Pool Profile — companion to the Image Pool. A dropdown of named profiles
|
||||||
|
// (registry under input/grid_pool/profiles.json) plus create/rename/delete/
|
||||||
|
// duplicate/export/import actions. The pool is switched ONLY when the user
|
||||||
|
// actively picks a profile in the dropdown (or creates/duplicates/imports one) —
|
||||||
|
// connecting the node never changes the pool. Selecting an *empty* profile while
|
||||||
|
// a pool with images is connected offers to seed it from those images, so the
|
||||||
|
// current pool is never silently lost. (Modeled on JSON-Manager/project_key.)
|
||||||
|
|
||||||
|
const NODE = "PoolProfile";
|
||||||
|
const POOL_NODE = "GridImagePool";
|
||||||
|
const R = "/grid_pool/profiles";
|
||||||
|
|
||||||
|
// ---- server calls -----------------------------------------------------------
|
||||||
|
|
||||||
|
async function listProfiles() {
|
||||||
|
const r = await api.fetchApi(`${R}/list`);
|
||||||
|
return (await r.json()).profiles || [];
|
||||||
|
}
|
||||||
|
|
||||||
|
async function listPoolSlots(poolId) {
|
||||||
|
try {
|
||||||
|
const r = await api.fetchApi(`/grid_pool/list?pool_id=${encodeURIComponent(poolId)}`);
|
||||||
|
return (await r.json()).slots || [];
|
||||||
|
} catch (e) { return []; }
|
||||||
|
}
|
||||||
|
|
||||||
|
async function postJson(path, body) {
|
||||||
|
const r = await api.fetchApi(`${R}/${path}`, {
|
||||||
|
method: "POST",
|
||||||
|
headers: { "Content-Type": "application/json" },
|
||||||
|
body: JSON.stringify(body),
|
||||||
|
});
|
||||||
|
if (!r.ok) throw new Error(await r.text());
|
||||||
|
return await r.json();
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- widget helpers ---------------------------------------------------------
|
||||||
|
|
||||||
|
function hideWidget(w) {
|
||||||
|
if (!w) return;
|
||||||
|
if (w.origType === undefined) w.origType = w.type;
|
||||||
|
w.type = "hidden";
|
||||||
|
w.hidden = true;
|
||||||
|
w.computeSize = () => [0, -4];
|
||||||
|
}
|
||||||
|
|
||||||
|
function profileWidget(node) {
|
||||||
|
return node.widgets?.find((w) => w.name === "profile");
|
||||||
|
}
|
||||||
|
|
||||||
|
function idWidget(node) {
|
||||||
|
return node.widgets?.find((w) => w.name === "profile_id");
|
||||||
|
}
|
||||||
|
|
||||||
|
function currentEntry(node) {
|
||||||
|
const combo = profileWidget(node);
|
||||||
|
return (node._profiles || []).find((p) => p.name === combo?.value);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Replace a STRING widget with a real combo, preserving its serialized value.
|
||||||
|
function replaceWithCombo(node, name, values, callback) {
|
||||||
|
const idx = node.widgets?.findIndex((w) => w.name === name);
|
||||||
|
if (idx === undefined || idx === -1) return null;
|
||||||
|
const old = node.widgets[idx];
|
||||||
|
const saved = old.value || "";
|
||||||
|
const vals = values.length ? values.slice() : [""];
|
||||||
|
if (saved && !vals.includes(saved)) vals.unshift(saved);
|
||||||
|
node.widgets.splice(idx, 1);
|
||||||
|
const combo = node.addWidget("combo", name, saved || vals[0], callback, { values: vals });
|
||||||
|
node.widgets.splice(node.widgets.length - 1, 1);
|
||||||
|
node.widgets.splice(idx, 0, combo);
|
||||||
|
return combo;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- connected pools + switching --------------------------------------------
|
||||||
|
|
||||||
|
function connectedPools(node) {
|
||||||
|
const res = [];
|
||||||
|
const out = node.outputs?.[0];
|
||||||
|
if (!out?.links) return res;
|
||||||
|
for (const linkId of out.links) {
|
||||||
|
const link = node.graph?.links?.[linkId];
|
||||||
|
if (!link) continue;
|
||||||
|
const t = node.graph?.getNodeById?.(link.target_id);
|
||||||
|
if (t && t.type === POOL_NODE) res.push(t);
|
||||||
|
}
|
||||||
|
return res;
|
||||||
|
}
|
||||||
|
|
||||||
|
function setIdFromCombo(node) {
|
||||||
|
const entry = currentEntry(node);
|
||||||
|
const idw = idWidget(node);
|
||||||
|
if (idw) idw.value = entry?.id || "";
|
||||||
|
}
|
||||||
|
|
||||||
|
// Push the current profile id into every connected pool's pool_id widget (the
|
||||||
|
// grid keys off getPoolId) and repaint. Only ever called from user actions.
|
||||||
|
function switchPools(node) {
|
||||||
|
const id = idWidget(node)?.value || "default";
|
||||||
|
for (const pool of connectedPools(node)) {
|
||||||
|
const pw = pool.widgets?.find((w) => w.name === "pool_id");
|
||||||
|
if (pw) pw.value = id;
|
||||||
|
pool._datasetePoolRefresh?.();
|
||||||
|
pool.setDirtyCanvas?.(true, true);
|
||||||
|
}
|
||||||
|
node.setDirtyCanvas?.(true, true);
|
||||||
|
}
|
||||||
|
|
||||||
|
// If the selected profile is empty and a connected pool has images, offer to
|
||||||
|
// copy those images into the profile (so switching never loses the current pool).
|
||||||
|
async function maybeSeed(node, entry) {
|
||||||
|
const profSlots = await listPoolSlots(entry.id);
|
||||||
|
if (profSlots.length > 0) return; // profile already has images
|
||||||
|
for (const pool of connectedPools(node)) {
|
||||||
|
const curId = pool.widgets?.find((w) => w.name === "pool_id")?.value;
|
||||||
|
if (!curId || curId === entry.id) continue;
|
||||||
|
const curSlots = await listPoolSlots(curId);
|
||||||
|
if (curSlots.length === 0) continue;
|
||||||
|
if (confirm(`Profile "${entry.name}" is empty. Copy the ${curSlots.length} current pool image(s) into it?`)) {
|
||||||
|
try { await postJson("seed", { from: curId, id: entry.id }); }
|
||||||
|
catch (err) { alert("Seed failed: " + err); }
|
||||||
|
}
|
||||||
|
return; // seed from the first match only
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// user-initiated: set id from the dropdown, optionally offer to seed, then switch
|
||||||
|
async function selectProfile(node) {
|
||||||
|
setIdFromCombo(node);
|
||||||
|
const entry = currentEntry(node);
|
||||||
|
if (entry) await maybeSeed(node, entry);
|
||||||
|
switchPools(node);
|
||||||
|
}
|
||||||
|
|
||||||
|
// programmatic: refresh the dropdown options + hidden id only — never switches
|
||||||
|
async function refreshList(node, selectName) {
|
||||||
|
const profs = await listProfiles();
|
||||||
|
node._profiles = profs;
|
||||||
|
const names = profs.map((p) => p.name);
|
||||||
|
const combo = profileWidget(node);
|
||||||
|
if (combo) {
|
||||||
|
combo.options = combo.options || {};
|
||||||
|
combo.options.values = names.length ? names : [""];
|
||||||
|
if (selectName !== undefined) combo.value = selectName;
|
||||||
|
else if (!names.includes(combo.value)) combo.value = names[0] || "";
|
||||||
|
}
|
||||||
|
setIdFromCombo(node);
|
||||||
|
node.setDirtyCanvas?.(true, true);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- actions ----------------------------------------------------------------
|
||||||
|
|
||||||
|
async function actionCreate(node) {
|
||||||
|
const name = prompt("New profile name:");
|
||||||
|
if (!name) return;
|
||||||
|
try {
|
||||||
|
const e = await postJson("create", { name });
|
||||||
|
await refreshList(node, e.name);
|
||||||
|
await selectProfile(node); // new profile is empty → offer to seed current pool
|
||||||
|
} catch (err) { alert("Create failed: " + err); }
|
||||||
|
}
|
||||||
|
|
||||||
|
async function actionRename(node) {
|
||||||
|
const e = currentEntry(node);
|
||||||
|
if (!e) return alert("Select a profile first");
|
||||||
|
const name = prompt("Rename profile:", e.name);
|
||||||
|
if (!name || name === e.name) return;
|
||||||
|
try {
|
||||||
|
await postJson("rename", { id: e.id, name });
|
||||||
|
await refreshList(node, name); // same id, no pool switch needed
|
||||||
|
} catch (err) { alert("Rename failed: " + err); }
|
||||||
|
}
|
||||||
|
|
||||||
|
async function actionDuplicate(node) {
|
||||||
|
const e = currentEntry(node);
|
||||||
|
if (!e) return alert("Select a profile first");
|
||||||
|
const name = prompt("Duplicate as:", e.name + " copy");
|
||||||
|
if (!name) return;
|
||||||
|
try {
|
||||||
|
const ne = await postJson("duplicate", { id: e.id, name });
|
||||||
|
await refreshList(node, ne.name);
|
||||||
|
await selectProfile(node); // already has images → maybeSeed no-ops, just switch
|
||||||
|
} catch (err) { alert("Duplicate failed: " + err); }
|
||||||
|
}
|
||||||
|
|
||||||
|
async function actionDelete(node) {
|
||||||
|
const e = currentEntry(node);
|
||||||
|
if (!e) return alert("Select a profile first");
|
||||||
|
if (!confirm(`Delete profile "${e.name}"? This removes its images.`)) return;
|
||||||
|
try {
|
||||||
|
await postJson("delete", { id: e.id });
|
||||||
|
await refreshList(node); // update dropdown; leave the pool as-is
|
||||||
|
} catch (err) { alert("Delete failed: " + err); }
|
||||||
|
}
|
||||||
|
|
||||||
|
function actionExport(node) {
|
||||||
|
const e = currentEntry(node);
|
||||||
|
if (!e) return alert("Select a profile first");
|
||||||
|
window.open(`${R}/export?id=${encodeURIComponent(e.id)}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
function actionImport(node) {
|
||||||
|
const input = document.createElement("input");
|
||||||
|
input.type = "file";
|
||||||
|
input.accept = ".zip";
|
||||||
|
input.onchange = async () => {
|
||||||
|
if (!input.files?.length) return;
|
||||||
|
const fd = new FormData();
|
||||||
|
fd.append("file", input.files[0], input.files[0].name);
|
||||||
|
try {
|
||||||
|
const r = await api.fetchApi(`${R}/import`, { method: "POST", body: fd });
|
||||||
|
if (!r.ok) throw new Error(await r.text());
|
||||||
|
const e = await r.json();
|
||||||
|
await refreshList(node, e.name);
|
||||||
|
await selectProfile(node); // imported profile has images → just switch
|
||||||
|
} catch (err) { alert("Import failed: " + err); }
|
||||||
|
};
|
||||||
|
input.click();
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- node setup -------------------------------------------------------------
|
||||||
|
|
||||||
|
function setupProfileNode(node) {
|
||||||
|
hideWidget(idWidget(node));
|
||||||
|
// combo callback = active user selection → switch (and maybe seed)
|
||||||
|
replaceWithCombo(node, "profile", [], () => { selectProfile(node); });
|
||||||
|
|
||||||
|
node.addWidget("button", "➕ Create", null, () => actionCreate(node));
|
||||||
|
node.addWidget("button", "✎ Rename", null, () => actionRename(node));
|
||||||
|
node.addWidget("button", "⧉ Duplicate", null, () => actionDuplicate(node));
|
||||||
|
node.addWidget("button", "🗑 Delete", null, () => actionDelete(node));
|
||||||
|
node.addWidget("button", "⬇ Export", null, () => actionExport(node));
|
||||||
|
node.addWidget("button", "⬆ Import", null, () => actionImport(node));
|
||||||
|
|
||||||
|
node.setSize(node.computeSize());
|
||||||
|
refreshList(node); // populate the dropdown; does NOT switch any pool
|
||||||
|
}
|
||||||
|
|
||||||
|
app.registerExtension({
|
||||||
|
name: "datasete.gates.poolprofile",
|
||||||
|
|
||||||
|
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||||
|
if (nodeData.name !== NODE) return;
|
||||||
|
|
||||||
|
const onNodeCreated = nodeType.prototype.onNodeCreated;
|
||||||
|
nodeType.prototype.onNodeCreated = function () {
|
||||||
|
const r = onNodeCreated?.apply(this, arguments);
|
||||||
|
setupProfileNode(this);
|
||||||
|
return r;
|
||||||
|
};
|
||||||
|
|
||||||
|
// on load the pool already has its saved pool_id, so just refresh the
|
||||||
|
// dropdown to show the saved name — no switching, no seeding.
|
||||||
|
const onConfigure = nodeType.prototype.onConfigure;
|
||||||
|
nodeType.prototype.onConfigure = function () {
|
||||||
|
const r = onConfigure?.apply(this, arguments);
|
||||||
|
const node = this;
|
||||||
|
queueMicrotask(() => refreshList(node, profileWidget(node)?.value));
|
||||||
|
return r;
|
||||||
|
};
|
||||||
|
// NOTE: intentionally no onConnectionsChange handler — connecting a profile
|
||||||
|
// must never change the pool (the user switches via the dropdown).
|
||||||
|
},
|
||||||
|
});
|
||||||
Reference in New Issue
Block a user