chrmc40
9fd6498ba6
- Tag-builder: port ~1,477 v1 characters into the v2 chip format so the character picker shows 1,486 characters (was 9); 1,390 carry rich per-chip descriptions restructured from each character's own bio. Adds scripts/port_v1_characters_to_v2.py. - Ideogram: rewrite caption to the real Ideogram-4 schema + broaden refusal detection; sampler updates. - Pose-studio, regional-detailer, attention-couple: assorted fixes. - Add validate_templates / test_regions / test_crop_windows scripts.
325 lines
13 KiB
Python
325 lines
13 KiB
Python
# Pose Studio capture-file maintenance + imported-mesh storage.
|
|
#
|
|
# The 3D Poser viewport uploads content-addressed capture sets to
|
|
# input/promptchain_pose (promptchain_pose_<sha12>.png + _mask<i>.png). Names
|
|
# pin content, so sets are never overwritten — superseded ones are deleted here
|
|
# instead, and anything a crashed/closed tab missed is collected by an age
|
|
# sweep at server start. Legacy node-id-named files (promptchain_pose_<id>) are
|
|
# shared across workflows and never touched.
|
|
#
|
|
# Imported meshes (user .glb/.obj props) live in input/promptchain_meshes as
|
|
# <sha16>.<ext> — content-addressed so re-imports dedupe, and referenced from
|
|
# pose_state long-term, so they are deliberately OUTSIDE the capture age sweep.
|
|
|
|
import glob
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
import time
|
|
|
|
from aiohttp import web
|
|
|
|
import folder_paths
|
|
import server
|
|
|
|
from .api_utils import error_response, ok_response, parse_json
|
|
|
|
routes = server.PromptServer.instance.routes
|
|
|
|
_HASHED_BASE = re.compile(r"^promptchain_pose_[0-9a-f]{12}$")
|
|
_HASHED_FILE = re.compile(r"^promptchain_pose_([0-9a-f]{12})(?:_mask\d+)?\.png$")
|
|
_SWEEP_MAX_AGE_DAYS = 30
|
|
|
|
|
|
def _pose_dir() -> str:
|
|
return os.path.join(folder_paths.get_input_directory(), "promptchain_pose")
|
|
|
|
|
|
def _is_hashed(base_or_name: str) -> bool:
|
|
# A 12-digit decimal could in principle be a legacy node id; rule it out so
|
|
# legacy files stay deletable only by hand.
|
|
m = _HASHED_BASE.match(base_or_name)
|
|
return bool(m) and not base_or_name.rsplit("_", 1)[-1].isdigit()
|
|
|
|
|
|
@routes.post("/promptchain/pose-files/delete")
|
|
async def delete_pose_files(request):
|
|
data, err = await parse_json(request)
|
|
if err:
|
|
return err
|
|
base = data.get("base", "")
|
|
if not _is_hashed(base):
|
|
return error_response("not a content-addressed pose base")
|
|
if os.path.exists(os.path.join(_pose_dir(), base + ".keep")):
|
|
# A render consumed this set (the PoseStudio node pinned it). The
|
|
# viewport fires this delete on every re-pose to clear the SUPERSEDED
|
|
# set, but a set already baked into a generated image must survive —
|
|
# its per-figure masks are the only copy. Refuse, don't orphan.
|
|
return ok_response({"deleted": [], "kept": True})
|
|
deleted = []
|
|
for path in glob.glob(os.path.join(glob.escape(_pose_dir()), base + "*.png")):
|
|
name = os.path.basename(path)
|
|
if name == base + ".png" or re.match(re.escape(base) + r"_mask\d+\.png$", name):
|
|
try:
|
|
os.remove(path)
|
|
deleted.append(name)
|
|
except OSError:
|
|
pass
|
|
return ok_response({"deleted": deleted})
|
|
|
|
|
|
@routes.post("/promptchain/pose-files/pin")
|
|
async def pin_pose_files(request):
|
|
"""Mark a content-addressed capture set permanent (write <base>.keep).
|
|
|
|
Called when a prompt referencing the set is ENQUEUED, closing the window
|
|
between queue and the node's execute() (which also pins): during a busy
|
|
queue, a re-pose would otherwise fire the superseded-delete on a set whose
|
|
render is still waiting, losing its map + the only copy of its region masks.
|
|
The delete route already refuses any base carrying a .keep.
|
|
"""
|
|
data, err = await parse_json(request)
|
|
if err:
|
|
return err
|
|
base = data.get("base", "")
|
|
if not _is_hashed(base):
|
|
return error_response("not a content-addressed pose base")
|
|
pose_dir = _pose_dir()
|
|
# Only pin a set whose map is actually on disk — don't litter markers.
|
|
if not os.path.exists(os.path.join(pose_dir, base + ".png")):
|
|
return ok_response({"pinned": False})
|
|
marker = os.path.join(pose_dir, base + ".keep")
|
|
if not os.path.exists(marker):
|
|
try:
|
|
open(marker, "w").close()
|
|
except OSError as e:
|
|
return error_response(f"could not pin: {e}")
|
|
return ok_response({"pinned": True})
|
|
|
|
|
|
# ── imported meshes (.glb/.obj props) ────────────────────────────────────────
|
|
|
|
_MESH_EXTS = {".glb", ".gltf", ".obj"}
|
|
_MESH_FILE = re.compile(r"^[0-9a-f]{16}\.(glb|gltf|obj)$")
|
|
# 96MB: GLB weight is usually embedded TEXTURES, which the importer discards
|
|
# (matte prop material) — a 62MB rigged mannequin carried only 84k tris. The
|
|
# real viewport budget is triangles, not file bytes.
|
|
_MESH_MAX_BYTES = 96 * 1024 * 1024
|
|
|
|
|
|
def _mesh_dir() -> str:
|
|
return os.path.join(folder_paths.get_input_directory(), "promptchain_meshes")
|
|
|
|
|
|
@routes.post("/promptchain/pose-mesh/upload")
|
|
async def upload_pose_mesh(request):
|
|
"""Store an imported mesh content-addressed: <sha256[:16]>.<ext>.
|
|
|
|
Multipart field "mesh" (filename supplies the extension). Same content =
|
|
same name, so duplicate imports are a no-op and pose_state references stay
|
|
valid across workflows. Writes via a temp file + rename so a crashed upload
|
|
never leaves a half-written mesh under a valid hash name.
|
|
"""
|
|
reader = await request.multipart()
|
|
part = await reader.next()
|
|
while part is not None and part.name != "mesh":
|
|
part = await reader.next()
|
|
if part is None:
|
|
return error_response("multipart field 'mesh' missing")
|
|
ext = os.path.splitext(part.filename or "")[1].lower()
|
|
if ext not in _MESH_EXTS:
|
|
return error_response(f"unsupported mesh type '{ext}' — expected .glb/.gltf/.obj")
|
|
data = bytearray()
|
|
while True:
|
|
chunk = await part.read_chunk()
|
|
if not chunk:
|
|
break
|
|
data.extend(chunk)
|
|
if len(data) > _MESH_MAX_BYTES:
|
|
return error_response(f"mesh exceeds {_MESH_MAX_BYTES // (1024 * 1024)}MB")
|
|
if not data:
|
|
return error_response("empty upload")
|
|
name = hashlib.sha256(data).hexdigest()[:16] + ext # hash the bytearray directly — no full copy
|
|
mesh_dir = _mesh_dir()
|
|
os.makedirs(mesh_dir, exist_ok=True)
|
|
path = os.path.join(mesh_dir, name)
|
|
if not os.path.exists(path):
|
|
# Unique temp so two concurrent uploads of the same content (same path)
|
|
# don't write/replace each other's .part mid-flight.
|
|
tmp = f"{path}.part-{os.getpid()}-{os.urandom(4).hex()}"
|
|
with open(tmp, "wb") as f:
|
|
f.write(data)
|
|
os.replace(tmp, path)
|
|
return ok_response({"file": name})
|
|
|
|
|
|
def _heal_mesh_name(name: str) -> str | None:
|
|
"""Recover a referenced mesh whose file was renamed by hand.
|
|
|
|
The canonical filename IS the content hash, so the lookup key is derivable
|
|
from bytes alone: re-hash candidates and rename the match back to its
|
|
canonical name (using the reference's extension — it was minted from the
|
|
original upload, so it's right even if the rename mangled it). Only runs on
|
|
a 404, and only hashes STRAYS: a file still wearing a canonical hash name
|
|
is already addressable (the direct lookup would have found it), so a folder
|
|
of 10,000 well-named meshes costs nothing here — just the handful someone
|
|
renamed.
|
|
"""
|
|
want_stem = os.path.splitext(name)[0]
|
|
mesh_dir = _mesh_dir()
|
|
if not os.path.isdir(mesh_dir):
|
|
return None
|
|
for path in glob.glob(os.path.join(glob.escape(mesh_dir), "*")):
|
|
if not os.path.isfile(path) or os.path.splitext(path)[1].lower() not in _MESH_EXTS:
|
|
continue
|
|
if _MESH_FILE.match(os.path.basename(path)):
|
|
continue # canonical-named = already addressable, can't be the stray we want
|
|
try:
|
|
h = hashlib.sha256()
|
|
with open(path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(1 << 20), b""):
|
|
h.update(chunk)
|
|
if h.hexdigest()[:16] != want_stem:
|
|
continue
|
|
canonical = os.path.join(mesh_dir, name)
|
|
os.replace(path, canonical)
|
|
print(f"[PoseStudio] healed renamed mesh {os.path.basename(path)} -> {name}")
|
|
return canonical
|
|
except OSError:
|
|
continue
|
|
return None
|
|
|
|
|
|
@routes.get("/promptchain/pose-mesh/{file}")
|
|
async def get_pose_mesh(request):
|
|
name = request.match_info["file"]
|
|
if not _MESH_FILE.match(name):
|
|
return error_response("not a content-addressed mesh name")
|
|
path = os.path.join(_mesh_dir(), name)
|
|
if not os.path.isfile(path):
|
|
path = _heal_mesh_name(name) # renamed by hand? content still finds it
|
|
if not path:
|
|
return error_response("mesh not found", status=404)
|
|
# Content-addressed = immutable: cache hard so restores never refetch.
|
|
return web.FileResponse(path, headers={"Cache-Control": "public, max-age=31536000, immutable"})
|
|
|
|
|
|
# ── user-saved body poses (the 🤸 picker's user layer) ───────────────────────
|
|
# Bundled poses ship in code (BODY_POSES, pose-studio.js); user saves land in
|
|
# {user}/PromptChain/poses/<slug>.json and the frontend overlays them by NAME —
|
|
# a user pose named like a bundled one shadows it (the templates/prompts delta
|
|
# model). One file per pose: {"name", "rig", "pose": {joints, spine, digits}}.
|
|
|
|
_POSE_PRESET_MAX_BYTES = 256 * 1024
|
|
|
|
|
|
def _pose_presets_dir() -> str:
|
|
return os.path.join(folder_paths.get_user_directory(), "PromptChain", "poses")
|
|
|
|
|
|
def _pose_preset_slug(name: str) -> str:
|
|
slug = re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")
|
|
return slug or hashlib.sha256(name.encode()).hexdigest()[:12]
|
|
|
|
|
|
@routes.get("/promptchain/pose-presets/list")
|
|
async def list_pose_presets(request):
|
|
out = []
|
|
presets_dir = _pose_presets_dir()
|
|
if os.path.isdir(presets_dir):
|
|
for path in sorted(glob.glob(os.path.join(glob.escape(presets_dir), "*.json"))):
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
except (OSError, ValueError):
|
|
continue
|
|
if isinstance(data, dict) and data.get("name") and isinstance(data.get("pose"), dict):
|
|
out.append({"name": data["name"], "rig": data.get("rig") or "human", "pose": data["pose"]})
|
|
return web.json_response({"poses": out})
|
|
|
|
|
|
@routes.post("/promptchain/pose-presets")
|
|
async def save_pose_preset(request):
|
|
data, err = await parse_json(request)
|
|
if err:
|
|
return err
|
|
name = (data.get("name") or "").strip()
|
|
pose = data.get("pose")
|
|
if not name or len(name) > 64:
|
|
return error_response("pose name must be 1-64 characters")
|
|
if not isinstance(pose, dict) or not isinstance(pose.get("joints"), dict):
|
|
return error_response("pose must carry a joints map")
|
|
payload = {"name": name, "rig": (data.get("rig") or "human"), "pose": pose}
|
|
blob = json.dumps(payload, indent=2)
|
|
if len(blob) > _POSE_PRESET_MAX_BYTES:
|
|
return error_response("pose too large")
|
|
presets_dir = _pose_presets_dir()
|
|
os.makedirs(presets_dir, exist_ok=True)
|
|
path = os.path.join(presets_dir, _pose_preset_slug(name) + ".json")
|
|
if os.path.exists(path):
|
|
# Unique-name-on-write: refuse to clobber an existing saved pose. The client
|
|
# pre-checks too, but enforce it here so the write site is authoritative
|
|
# (a second tab, or two names that slugify alike, can't silently overwrite).
|
|
return error_response(f'A pose named "{name}" already exists', status=409)
|
|
tmp = path + ".part"
|
|
with open(tmp, "w", encoding="utf-8") as f:
|
|
f.write(blob)
|
|
os.replace(tmp, path) # atomic: a crashed save never leaves a half-written pose
|
|
return ok_response({"name": name})
|
|
|
|
|
|
@routes.delete("/promptchain/pose-presets/{name}")
|
|
async def delete_pose_preset(request):
|
|
name = request.match_info["name"]
|
|
path = os.path.join(_pose_presets_dir(), _pose_preset_slug(name) + ".json")
|
|
try:
|
|
os.remove(path)
|
|
except FileNotFoundError:
|
|
# Already gone (e.g. a second tab deleted it) — the user's intent is
|
|
# satisfied; report 404 rather than 500ing on the lost race.
|
|
return error_response("not found", status=404)
|
|
except OSError as e:
|
|
return error_response(f"could not delete: {e}")
|
|
return ok_response()
|
|
|
|
|
|
def _sweep_aged_captures() -> None:
|
|
pose_dir = _pose_dir()
|
|
if not os.path.isdir(pose_dir):
|
|
return
|
|
cutoff = time.time() - _SWEEP_MAX_AGE_DAYS * 86400
|
|
# Pinned bases were consumed by a render (the PoseStudio node drops a
|
|
# `<base>.keep`); a generated image needs them indefinitely, so age never
|
|
# collects them.
|
|
pinned = {
|
|
os.path.basename(p)[:-5]
|
|
for p in glob.glob(os.path.join(glob.escape(pose_dir), "promptchain_pose_*.keep"))
|
|
}
|
|
swept = 0
|
|
for path in glob.glob(os.path.join(glob.escape(pose_dir), "promptchain_pose_*.png")):
|
|
m = _HASHED_FILE.match(os.path.basename(path))
|
|
if not m or m.group(1).isdigit():
|
|
continue
|
|
if f"promptchain_pose_{m.group(1)}" in pinned:
|
|
continue
|
|
try:
|
|
if os.path.getmtime(path) < cutoff:
|
|
os.remove(path)
|
|
swept += 1
|
|
except OSError:
|
|
pass
|
|
# Drop orphaned pins whose capture set is already gone (hand-deleted, etc.)
|
|
# so markers don't leak forever.
|
|
for marker in glob.glob(os.path.join(glob.escape(pose_dir), "promptchain_pose_*.keep")):
|
|
if not os.path.exists(marker[:-5] + ".png"):
|
|
try:
|
|
os.remove(marker)
|
|
except OSError:
|
|
pass
|
|
if swept:
|
|
print(f"[PoseStudio] swept {swept} pose capture file(s) older than {_SWEEP_MAX_AGE_DAYS} days")
|
|
|
|
|
|
_sweep_aged_captures()
|