shootthesound
9260d9fada
Folder, route base (/cleanfreak/), storage path (<ComfyUI>/user/cleanfreak/role_overrides.json), JS extension name (CleanFreak.TidyByRole), README title, and all internal log prefixes all renamed to match.
114 lines
3.7 KiB
Python
114 lines
3.7 KiB
Python
"""comfyui-CleanFreak — one-click workflow layout by node role.
|
|
|
|
Adds a "Tidy by Role" action set to the canvas right-click menu. Every node
|
|
in the graph is classified into a role bucket (loaders, encoders,
|
|
conditioning, samplers, decoders, outputs, image-input, prompts, post,
|
|
misc) and laid out in width-aware columns. Connections are never touched —
|
|
ComfyUI's links are by node id, so moving a node never breaks a wire.
|
|
|
|
Also exposes a small HTTP API for the user-overrides JSON file so the
|
|
frontend's classification editor can persist per-class role assignments
|
|
across sessions and workflows. The file gradually accumulates the user's
|
|
preferred role for every node class they've explicitly re-assigned, so
|
|
future tidies start with a richer classifier than the built-in heuristics.
|
|
|
|
Stored at::
|
|
|
|
<ComfyUI>/user/cleanfreak/role_overrides.json
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import threading
|
|
|
|
import folder_paths
|
|
|
|
try:
|
|
from server import PromptServer
|
|
_HAS_SERVER = True
|
|
except ImportError:
|
|
_HAS_SERVER = False
|
|
|
|
|
|
WEB_DIRECTORY = "./web"
|
|
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
|
|
|
|
|
# --- Persistent user overrides --------------------------------------
|
|
|
|
_LOCK = threading.Lock()
|
|
|
|
|
|
def _overrides_path() -> str:
|
|
base = os.path.join(folder_paths.get_user_directory(), "cleanfreak")
|
|
os.makedirs(base, exist_ok=True)
|
|
return os.path.join(base, "role_overrides.json")
|
|
|
|
|
|
def _load_overrides() -> dict:
|
|
path = _overrides_path()
|
|
if not os.path.exists(path):
|
|
return {}
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
except (json.JSONDecodeError, OSError):
|
|
return {}
|
|
if isinstance(data, dict) and isinstance(data.get("overrides"), dict):
|
|
return {str(k): str(v) for k, v in data["overrides"].items()}
|
|
if isinstance(data, dict):
|
|
# Tolerate a flat {"ClassName": "role"} shape too.
|
|
return {str(k): str(v) for k, v in data.items() if isinstance(v, str)}
|
|
return {}
|
|
|
|
|
|
def _save_overrides(overrides: dict) -> None:
|
|
path = _overrides_path()
|
|
payload = {"version": 1, "overrides": overrides}
|
|
tmp = path + ".tmp"
|
|
with open(tmp, "w", encoding="utf-8") as f:
|
|
json.dump(payload, f, indent=2, sort_keys=True)
|
|
os.replace(tmp, path)
|
|
|
|
|
|
# --- HTTP API ------------------------------------------------------
|
|
|
|
if _HAS_SERVER:
|
|
from aiohttp import web
|
|
|
|
routes = PromptServer.instance.routes
|
|
|
|
@routes.get("/cleanfreak/overrides")
|
|
async def _get_overrides_route(request):
|
|
with _LOCK:
|
|
overrides = _load_overrides()
|
|
return web.json_response({"overrides": overrides})
|
|
|
|
@routes.post("/cleanfreak/overrides")
|
|
async def _save_overrides_route(request):
|
|
try:
|
|
data = await request.json()
|
|
except Exception:
|
|
return web.json_response({"error": "invalid JSON"}, status=400)
|
|
incoming = (data or {}).get("overrides")
|
|
if not isinstance(incoming, dict):
|
|
return web.json_response({"error": "overrides must be an object"}, status=400)
|
|
# Whitelist the value type — anything non-string is dropped silently
|
|
# so a malformed frontend can't poison the file.
|
|
cleaned = {str(k): str(v) for k, v in incoming.items() if isinstance(v, str) and v}
|
|
with _LOCK:
|
|
_save_overrides(cleaned)
|
|
return web.json_response({"ok": True, "overrides": cleaned})
|
|
|
|
@routes.post("/cleanfreak/overrides/clear")
|
|
async def _clear_overrides_route(request):
|
|
with _LOCK:
|
|
_save_overrides({})
|
|
return web.json_response({"ok": True, "overrides": {}})
|