Ethanfel
6d433ba371
Publish to Comfy registry / Publish Custom Node to registry (push) Waiting to run
Disable/enable no longer require ComfyUI Manager: - New pack_fs.py moves packs in/out of custom_nodes/.disabled/ (no import, delete, or re-clone). Fallback for hand-cloned packs, loose single-file nodes, or when Manager is absent. enable strips the @version suffix so packs restore as clean, importable dir names. - Routes: native-disable, native-enable, disabled-packs. - Frontend routes each disable per-pack (Manager queue vs native move), and shows an Enable button on recoverable packs in the Uninstalled tier. The restart banner degrades to a manual-restart notice when no Manager exists. Whitelist (packages-only): a star toggle protects a pack — pulled into its own pinned group, no Disable button, skipped by the 7-day trial auto-disable. - New whitelist_packages table; whitelisted flag on package stats. - Routes: whitelist, whitelist/add, whitelist/remove. Tests: test_pack_fs.py, test_whitelist.py (60 passing). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
277 lines
9.9 KiB
Python
277 lines
9.9 KiB
Python
import asyncio
|
|
import logging
|
|
import threading
|
|
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
|
|
from .mapper import NodePackageMapper, ModelMapper
|
|
from .node_introspect import find_disabled_pack_path, get_node_schema
|
|
from .pack_fs import disable_pack_native, enable_pack_native, list_disabled_packs
|
|
from .tracker import UsageTracker
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
NODE_CLASS_MAPPINGS = {}
|
|
WEB_DIRECTORY = "./js"
|
|
|
|
mapper = NodePackageMapper()
|
|
tracker = UsageTracker()
|
|
model_mapper = ModelMapper()
|
|
|
|
|
|
def on_prompt_handler(json_data):
|
|
"""Called on every prompt submission. Extracts class_types and queues recording."""
|
|
try:
|
|
prompt = json_data.get("prompt", {})
|
|
class_types = set()
|
|
for node_id, node_data in prompt.items():
|
|
ct = node_data.get("class_type")
|
|
if ct:
|
|
class_types.add(ct)
|
|
if class_types:
|
|
# Pass the full prompt to the thread — model extraction (which calls
|
|
# INPUT_TYPES() on every node) happens off the main request thread.
|
|
threading.Thread(
|
|
target=_record_prompt,
|
|
args=(class_types, prompt),
|
|
daemon=True,
|
|
).start()
|
|
except Exception:
|
|
logger.warning("nodes-stats: error recording usage", exc_info=True)
|
|
return json_data
|
|
|
|
|
|
def _record_prompt(class_types, prompt):
|
|
try:
|
|
tracker.record_usage(class_types, mapper)
|
|
except Exception:
|
|
logger.warning("nodes-stats: error recording node usage", exc_info=True)
|
|
try:
|
|
packages = {mapper.get_package(ct) for ct in class_types}
|
|
packages.discard("__builtin__")
|
|
packages.discard("__unknown__")
|
|
tracker.reset_trials_for(packages)
|
|
except Exception:
|
|
logger.warning("nodes-stats: error resetting trials", exc_info=True)
|
|
try:
|
|
models = model_mapper.extract_models_from_prompt(prompt)
|
|
if models:
|
|
tracker.record_model_usage(models)
|
|
except Exception:
|
|
logger.warning("nodes-stats: error recording model usage", exc_info=True)
|
|
|
|
|
|
# Age temporary-enable trials once per process start (one "boot").
|
|
try:
|
|
tracker.tick_boot_days()
|
|
except Exception:
|
|
logger.warning("nodes-stats: error ticking trial boot days", exc_info=True)
|
|
|
|
|
|
PromptServer.instance.add_on_prompt_handler(on_prompt_handler)
|
|
|
|
|
|
routes = PromptServer.instance.routes
|
|
|
|
|
|
@routes.get("/nodes-stats/packages")
|
|
async def get_package_stats(request):
|
|
try:
|
|
stats = tracker.get_package_stats(mapper)
|
|
return web.json_response(stats)
|
|
except Exception:
|
|
logger.error("nodes-stats: error getting package stats", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.get("/nodes-stats/usage")
|
|
async def get_node_stats(request):
|
|
try:
|
|
stats = tracker.get_node_stats()
|
|
return web.json_response(stats)
|
|
except Exception:
|
|
logger.error("nodes-stats: error getting node stats", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.get("/nodes-stats/models")
|
|
async def get_model_stats(request):
|
|
try:
|
|
installed_by_type = model_mapper.get_all_models()
|
|
stats = tracker.get_model_stats(installed_by_type)
|
|
return web.json_response(stats)
|
|
except Exception:
|
|
logger.error("nodes-stats: error getting model stats", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.post("/nodes-stats/reset")
|
|
async def reset_stats(request):
|
|
try:
|
|
tracker.reset()
|
|
mapper.invalidate()
|
|
model_mapper.invalidate()
|
|
return web.json_response({"status": "ok"})
|
|
except Exception:
|
|
logger.error("nodes-stats: error resetting stats", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.get("/nodes-stats/node-schema")
|
|
async def get_node_schema_route(request):
|
|
"""Parse a disabled pack's source and return one node's input/output schema.
|
|
|
|
Read-only and never imports/executes the pack. Used by the mirror-search
|
|
palette to draw a faithful node box for a node that isn't loaded.
|
|
"""
|
|
try:
|
|
class_type = request.query.get("class_type")
|
|
pack = request.query.get("pack")
|
|
if not class_type or not pack:
|
|
return web.json_response({"error": "class_type and pack required"}, status=400)
|
|
|
|
def _resolve():
|
|
path = find_disabled_pack_path(pack)
|
|
if not path:
|
|
return {"parseable": False, "reason": "source_not_found"}
|
|
return get_node_schema(class_type, path)
|
|
|
|
# AST-parsing a whole pack can take tens of ms; keep it off the event loop.
|
|
schema = await asyncio.get_event_loop().run_in_executor(None, _resolve)
|
|
return web.json_response(schema)
|
|
except Exception:
|
|
logger.error("nodes-stats: error parsing node schema", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.post("/nodes-stats/native-disable")
|
|
async def native_disable(request):
|
|
"""Disable a pack by moving it into custom_nodes/.disabled/ (no Manager).
|
|
|
|
Fallback for packs ComfyUI Manager doesn't manage. A restart is required for
|
|
ComfyUI to unload the pack.
|
|
"""
|
|
try:
|
|
data = await request.json()
|
|
package = data.get("package")
|
|
if not package:
|
|
return web.json_response({"error": "package required"}, status=400)
|
|
|
|
ok, message = await asyncio.get_event_loop().run_in_executor(
|
|
None, disable_pack_native, package
|
|
)
|
|
if not ok:
|
|
return web.json_response({"status": "error", "message": message}, status=409)
|
|
mapper.invalidate()
|
|
return web.json_response({"status": "ok", "message": message})
|
|
except Exception:
|
|
logger.error("nodes-stats: error in native disable", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.post("/nodes-stats/native-enable")
|
|
async def native_enable(request):
|
|
"""Re-enable a pack by moving it out of custom_nodes/.disabled/ (no Manager)."""
|
|
try:
|
|
data = await request.json()
|
|
package = data.get("package")
|
|
if not package:
|
|
return web.json_response({"error": "package required"}, status=400)
|
|
|
|
ok, message = await asyncio.get_event_loop().run_in_executor(
|
|
None, enable_pack_native, package
|
|
)
|
|
if not ok:
|
|
return web.json_response({"status": "error", "message": message}, status=409)
|
|
mapper.invalidate()
|
|
return web.json_response({"status": "ok", "message": message})
|
|
except Exception:
|
|
logger.error("nodes-stats: error in native enable", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.get("/nodes-stats/disabled-packs")
|
|
async def get_disabled_packs(request):
|
|
"""List packs present in custom_nodes/.disabled/ (re-enable candidates)."""
|
|
try:
|
|
names = await asyncio.get_event_loop().run_in_executor(None, list_disabled_packs)
|
|
return web.json_response(sorted(names))
|
|
except Exception:
|
|
logger.error("nodes-stats: error listing disabled packs", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.get("/nodes-stats/whitelist")
|
|
async def get_whitelist(request):
|
|
try:
|
|
return web.json_response(sorted(tracker.get_whitelist()))
|
|
except Exception:
|
|
logger.error("nodes-stats: error getting whitelist", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.post("/nodes-stats/whitelist/add")
|
|
async def whitelist_add(request):
|
|
try:
|
|
data = await request.json()
|
|
package = data.get("package")
|
|
if not package:
|
|
return web.json_response({"error": "package required"}, status=400)
|
|
tracker.add_to_whitelist(package)
|
|
return web.json_response({"status": "ok"})
|
|
except Exception:
|
|
logger.error("nodes-stats: error adding to whitelist", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.post("/nodes-stats/whitelist/remove")
|
|
async def whitelist_remove(request):
|
|
try:
|
|
data = await request.json()
|
|
package = data.get("package")
|
|
if not package:
|
|
return web.json_response({"error": "package required"}, status=400)
|
|
tracker.remove_from_whitelist(package)
|
|
return web.json_response({"status": "ok"})
|
|
except Exception:
|
|
logger.error("nodes-stats: error removing from whitelist", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.get("/nodes-stats/trials")
|
|
async def get_trials(request):
|
|
try:
|
|
return web.json_response(tracker.get_trials())
|
|
except Exception:
|
|
logger.error("nodes-stats: error getting trials", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.post("/nodes-stats/trials/start")
|
|
async def start_trial(request):
|
|
try:
|
|
data = await request.json()
|
|
package = data.get("package")
|
|
if not package:
|
|
return web.json_response({"error": "package required"}, status=400)
|
|
tracker.start_trial(package)
|
|
return web.json_response({"status": "ok"})
|
|
except Exception:
|
|
logger.error("nodes-stats: error starting trial", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.post("/nodes-stats/trials/stop")
|
|
async def stop_trial(request):
|
|
try:
|
|
data = await request.json()
|
|
package = data.get("package")
|
|
if not package:
|
|
return web.json_response({"error": "package required"}, status=400)
|
|
tracker.stop_trial(package)
|
|
return web.json_response({"status": "ok"})
|
|
except Exception:
|
|
logger.error("nodes-stats: error stopping trial", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|