248 lines
9.8 KiB
Python
248 lines
9.8 KiB
Python
import asyncio
|
|
import os
|
|
import json
|
|
import tempfile
|
|
import time
|
|
import uuid
|
|
import re
|
|
from aiohttp import web
|
|
import folder_paths
|
|
from .utils import require_filename, resolve_within
|
|
from .recipes import MAX_RECIPE_BYTES, _parameter_signature, _parameter_gallery_images
|
|
|
|
def get_parameters_dir():
|
|
# Store parameter notebooks in the user directory
|
|
# ComfyUI/user/default/workflows/anomalous_parameters
|
|
user_dir = (
|
|
folder_paths.get_user_directory()
|
|
if hasattr(folder_paths, "get_user_directory")
|
|
else None
|
|
)
|
|
if not user_dir:
|
|
# Fallback if user_dir is not supported in this ComfyUI version
|
|
base_dir = folder_paths.base_path
|
|
user_dir = os.path.join(base_dir, "user", "default")
|
|
|
|
workflows_dir = os.path.join(user_dir, "workflows")
|
|
parameters_dir = os.path.join(workflows_dir, "anomalous_parameters")
|
|
|
|
if not os.path.exists(parameters_dir):
|
|
os.makedirs(parameters_dir, exist_ok=True)
|
|
return parameters_dir
|
|
|
|
|
|
|
|
async def api_get_parameters(request):
|
|
try:
|
|
recipe_filename = request.query.get("recipe_filename")
|
|
if recipe_filename:
|
|
recipe_filename = require_filename(recipe_filename)
|
|
if not recipe_filename.endswith(".json"):
|
|
raise ValueError("Invalid recipe filename")
|
|
except (AttributeError, ValueError):
|
|
return web.json_response({"status": "error", "message": "Invalid recipe filename"}, status=400)
|
|
parameters_dir = get_parameters_dir()
|
|
notebooks = await asyncio.to_thread(_read_parameter_notebooks, parameters_dir, recipe_filename)
|
|
return web.json_response({"notebooks": notebooks})
|
|
|
|
|
|
_parameter_notebooks_cache = None
|
|
_parameter_notebooks_cache_time = 0
|
|
|
|
def _get_all_notebooks_cached(parameters_dir):
|
|
global _parameter_notebooks_cache, _parameter_notebooks_cache_time
|
|
# Cache for 2 seconds to prevent rapid disk I/O on UI interactions
|
|
if _parameter_notebooks_cache is not None and time.time() - _parameter_notebooks_cache_time < 2.0:
|
|
return _parameter_notebooks_cache
|
|
notebooks = _read_parameter_notebooks(parameters_dir)
|
|
_parameter_notebooks_cache = notebooks
|
|
_parameter_notebooks_cache_time = time.time()
|
|
return notebooks
|
|
|
|
async def api_get_parameters_by_type(request):
|
|
try:
|
|
node_type = request.query.get("type")
|
|
if not node_type:
|
|
raise ValueError("Missing node type")
|
|
except ValueError:
|
|
return web.json_response({"status": "error", "message": "Missing node type"}, status=400)
|
|
|
|
parameters_dir = get_parameters_dir()
|
|
notebooks = await asyncio.to_thread(_get_all_notebooks_cached, parameters_dir)
|
|
|
|
# Filter and extract node data
|
|
# Structure: [ { "recipe_filename": ..., "notebooks": [ { "name": ..., "nodes": [...] } ] } ]
|
|
|
|
grouped = {}
|
|
for nb in notebooks:
|
|
workflow = nb.get("data", {}).get("workflow")
|
|
if not isinstance(workflow, dict):
|
|
continue
|
|
|
|
nodes = workflow.get("nodes", [])
|
|
if not isinstance(nodes, list):
|
|
continue
|
|
|
|
matched_nodes = []
|
|
for n in nodes:
|
|
if not isinstance(n, dict):
|
|
continue
|
|
if n.get("type") == node_type:
|
|
matched_nodes.append({
|
|
"id": n.get("id"),
|
|
"title": n.get("title") or n.get("type"),
|
|
"widgets_values": n.get("widgets_values", [])
|
|
})
|
|
|
|
if not matched_nodes:
|
|
continue
|
|
|
|
recipe_fn = nb.get("data", {}).get("recipe_filename") or "unbound"
|
|
if recipe_fn not in grouped:
|
|
grouped[recipe_fn] = []
|
|
|
|
grouped[recipe_fn].append({
|
|
"filename": nb.get("filename"),
|
|
"name": nb.get("name"),
|
|
"timestamp": nb.get("timestamp"),
|
|
"nodes": matched_nodes
|
|
})
|
|
|
|
result = []
|
|
for r_fn, nbs in grouped.items():
|
|
result.append({
|
|
"recipe_filename": r_fn,
|
|
"notebooks": nbs
|
|
})
|
|
|
|
return web.json_response({"groups": result})
|
|
|
|
def _read_parameter_notebooks(parameters_dir, recipe_filename=None):
|
|
notebooks = []
|
|
try:
|
|
filenames = os.listdir(parameters_dir)
|
|
except OSError:
|
|
return notebooks
|
|
for filename in filenames:
|
|
if not filename.endswith(".json"):
|
|
continue
|
|
file_path = resolve_within(parameters_dir, filename)
|
|
try:
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
if not isinstance(data, dict):
|
|
continue
|
|
if recipe_filename and data.get("recipe_filename") != recipe_filename:
|
|
continue
|
|
notebooks.append({
|
|
"filename": filename,
|
|
"name": data.get("name", "Untitled Parameter Notebook"),
|
|
"data": data,
|
|
"timestamp": data.get("timestamp", 0),
|
|
})
|
|
except (OSError, ValueError, json.JSONDecodeError):
|
|
continue
|
|
notebooks.sort(key=lambda item: item.get("timestamp", 0), reverse=True)
|
|
return notebooks
|
|
|
|
async def api_save_parameter(request):
|
|
try:
|
|
data = await request.json()
|
|
except (ValueError, json.JSONDecodeError):
|
|
return web.json_response({"status": "error", "message": "Invalid parameter data"}, status=400)
|
|
|
|
if not isinstance(data, dict) or not isinstance(data.get("workflow"), dict):
|
|
return web.json_response({"status": "error", "message": "Invalid parameter workflow"}, status=400)
|
|
|
|
recipe_filename = data.get("recipe_filename")
|
|
if recipe_filename is not None:
|
|
try:
|
|
recipe_filename = require_filename(recipe_filename)
|
|
if not recipe_filename.endswith(".json"):
|
|
raise ValueError("Invalid recipe filename")
|
|
except (AttributeError, ValueError):
|
|
return web.json_response({"status": "error", "message": "Invalid recipe filename"}, status=400)
|
|
data["recipe_filename"] = recipe_filename
|
|
|
|
filename_stem = os.path.splitext(recipe_filename)[0] if recipe_filename else "unbound"
|
|
filename_stem = re.sub(r"[^A-Za-z0-9_-]+", "_", filename_stem)[:64]
|
|
filename = f"params_{filename_stem}_{int(time.time())}_{uuid.uuid4().hex[:8]}.json"
|
|
data["parameter_signature"] = _parameter_signature(data["workflow"])
|
|
data["timestamp"] = int(time.time() * 1000)
|
|
encoded = json.dumps(data, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
|
if len(encoded) > MAX_RECIPE_BYTES:
|
|
return web.json_response({"status": "error", "message": "Parameter notebook is too large"}, status=413)
|
|
|
|
try:
|
|
parameters_dir = get_parameters_dir()
|
|
file_path = resolve_within(parameters_dir, filename)
|
|
fd, temp_path = tempfile.mkstemp(prefix=".parameter-", suffix=".tmp", dir=parameters_dir)
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, ensure_ascii=False, separators=(",", ":"))
|
|
f.flush()
|
|
os.fsync(f.fileno())
|
|
os.replace(temp_path, file_path)
|
|
finally:
|
|
if os.path.exists(temp_path):
|
|
os.remove(temp_path)
|
|
except OSError:
|
|
return web.json_response({"status": "error", "message": "Could not save parameter notebook"}, status=500)
|
|
|
|
return web.json_response({"status": "success", "filename": filename})
|
|
|
|
async def api_delete_parameter(request):
|
|
try:
|
|
data = await request.json()
|
|
filename = require_filename(data.get("filename", ""))
|
|
if not filename.endswith(".json"):
|
|
raise ValueError("Invalid filename")
|
|
except (ValueError, AttributeError):
|
|
return web.json_response({"status": "error", "message": "Invalid filename"}, status=400)
|
|
|
|
try:
|
|
parameters_dir = get_parameters_dir()
|
|
file_path = resolve_within(parameters_dir, filename)
|
|
if os.path.exists(file_path):
|
|
os.remove(file_path)
|
|
except OSError:
|
|
return web.json_response({"status": "error", "message": "Could not delete parameter notebook"}, status=500)
|
|
|
|
return web.json_response({"status": "success"})
|
|
|
|
async def api_get_parameter_gallery(request):
|
|
"""Find recent output PNGs whose embedded parameters match one parameter notebook."""
|
|
try:
|
|
filename = request.query.get("filename")
|
|
fingerprint = request.query.get("fingerprint")
|
|
if filename:
|
|
filename = require_filename(filename)
|
|
if not filename.endswith(".json"):
|
|
raise ValueError("Invalid parameter notebook")
|
|
parameters_dir = get_parameters_dir()
|
|
file_path = resolve_within(parameters_dir, filename)
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
|
|
fingerprint = data.get("parameter_signature", {}).get("value")
|
|
if not fingerprint:
|
|
fingerprint = _parameter_signature(data.get("workflow"))["value"]
|
|
elif not isinstance(fingerprint, str) or not re.fullmatch(r"[0-9a-f]{64}", fingerprint, re.IGNORECASE):
|
|
raise ValueError("Invalid fingerprint")
|
|
|
|
images, scanned = await asyncio.to_thread(_parameter_gallery_images, fingerprint.lower())
|
|
except (AttributeError, ValueError, json.JSONDecodeError):
|
|
return web.json_response({"status": "error", "message": "Invalid parameter gallery request"}, status=400)
|
|
except FileNotFoundError:
|
|
return web.json_response({"status": "error", "message": "File not found"}, status=404)
|
|
except OSError:
|
|
return web.json_response({"status": "error", "message": "Could not read parameter gallery"}, status=500)
|
|
|
|
return web.json_response({
|
|
"status": "success",
|
|
"fingerprint": fingerprint.lower(),
|
|
"match_mode": "params",
|
|
"images": images,
|
|
"scanned": scanned,
|
|
})
|