274 lines
8.5 KiB
Python
274 lines
8.5 KiB
Python
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from urllib.parse import urlencode
|
|
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
|
|
try:
|
|
import folder_paths
|
|
except Exception:
|
|
folder_paths = None
|
|
|
|
|
|
NODE_DIR = Path(__file__).resolve().parent
|
|
DEFAULT_LIBRARY_PATH = NODE_DIR / "data" / "expression_parts.json"
|
|
LIBRARIES_DIR = NODE_DIR / "data" / "libraries"
|
|
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
|
|
|
|
|
|
def load_default_library():
|
|
try:
|
|
return json.loads(DEFAULT_LIBRARY_PATH.read_text(encoding="utf-8"))
|
|
except Exception:
|
|
return {"categories": []}
|
|
|
|
|
|
DEFAULT_LIBRARY = load_default_library()
|
|
|
|
|
|
def safe_library_name(name):
|
|
cleaned = "".join(c for c in str(name or "default") if c.isalnum() or c in ("-", "_")).strip()
|
|
return cleaned or "default"
|
|
|
|
|
|
def library_path(name):
|
|
LIBRARIES_DIR.mkdir(parents=True, exist_ok=True)
|
|
return LIBRARIES_DIR / f"{safe_library_name(name)}.json"
|
|
|
|
|
|
def get_output_directory():
|
|
if folder_paths is not None:
|
|
try:
|
|
return Path(folder_paths.get_output_directory()).resolve()
|
|
except Exception:
|
|
pass
|
|
return (NODE_DIR.parent.parent / "output").resolve()
|
|
|
|
|
|
def ensure_default_library_file():
|
|
path = library_path("default")
|
|
if not path.exists():
|
|
path.write_text(json.dumps(DEFAULT_LIBRARY, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
|
|
|
|
def read_library_file(name):
|
|
path = library_path(name)
|
|
if not path.exists():
|
|
return None
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def write_library_file(name, payload):
|
|
categories = payload.get("categories") if isinstance(payload, dict) else None
|
|
if not isinstance(categories, list):
|
|
raise ValueError("library JSON must contain a categories array")
|
|
|
|
normalized = {
|
|
"version": payload.get("version", 1),
|
|
"meta": payload.get("meta", {}),
|
|
"categories": categories,
|
|
}
|
|
path = library_path(name)
|
|
path.write_text(json.dumps(normalized, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
return normalized
|
|
|
|
|
|
def list_library_files():
|
|
ensure_default_library_file()
|
|
items = []
|
|
for path in sorted(LIBRARIES_DIR.glob("*.json")):
|
|
try:
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
except Exception:
|
|
data = {}
|
|
categories = data.get("categories") if isinstance(data, dict) else []
|
|
parts_count = 0
|
|
if isinstance(categories, list):
|
|
for category in categories:
|
|
if isinstance(category, dict) and isinstance(category.get("parts"), list):
|
|
parts_count += len(category["parts"])
|
|
items.append(
|
|
{
|
|
"name": path.stem,
|
|
"label": data.get("meta", {}).get("title") or path.stem,
|
|
"categories": len(categories) if isinstance(categories, list) else 0,
|
|
"parts": parts_count,
|
|
}
|
|
)
|
|
return items
|
|
|
|
|
|
def list_recent_output_images(limit=80):
|
|
output_dir = get_output_directory()
|
|
if not output_dir.exists():
|
|
return []
|
|
|
|
files = []
|
|
for path in output_dir.rglob("*"):
|
|
if not path.is_file() or path.suffix.lower() not in IMAGE_EXTENSIONS:
|
|
continue
|
|
try:
|
|
stat = path.stat()
|
|
rel = path.relative_to(output_dir)
|
|
except Exception:
|
|
continue
|
|
|
|
subfolder = "" if rel.parent == Path(".") else rel.parent.as_posix()
|
|
query = {
|
|
"filename": rel.name,
|
|
"type": "output",
|
|
}
|
|
if subfolder:
|
|
query["subfolder"] = subfolder
|
|
|
|
files.append(
|
|
{
|
|
"filename": rel.name,
|
|
"subfolder": subfolder,
|
|
"mtime": stat.st_mtime,
|
|
"size": stat.st_size,
|
|
"url": f"/view?{urlencode(query)}",
|
|
}
|
|
)
|
|
|
|
files.sort(key=lambda item: item["mtime"], reverse=True)
|
|
return files[: max(1, min(int(limit or 80), 200))]
|
|
|
|
|
|
@PromptServer.instance.routes.get("/visual_prompt_library/libraries")
|
|
async def vpl_list_libraries(_request):
|
|
return web.json_response({"libraries": list_library_files()})
|
|
|
|
|
|
@PromptServer.instance.routes.get("/visual_prompt_library/recent_outputs")
|
|
async def vpl_recent_outputs(request):
|
|
limit = request.query.get("limit", "80")
|
|
try:
|
|
images = list_recent_output_images(limit)
|
|
except Exception as error:
|
|
return web.json_response({"error": str(error)}, status=400)
|
|
return web.json_response({"images": images, "output_dir": str(get_output_directory())})
|
|
|
|
|
|
@PromptServer.instance.routes.post("/visual_prompt_library/open_output_folder")
|
|
async def vpl_open_output_folder(_request):
|
|
output_dir = get_output_directory()
|
|
try:
|
|
if os.name == "nt":
|
|
os.startfile(str(output_dir))
|
|
else:
|
|
import subprocess
|
|
|
|
subprocess.Popen(["xdg-open", str(output_dir)])
|
|
except Exception as error:
|
|
return web.json_response({"error": str(error)}, status=400)
|
|
return web.json_response({"opened": True, "output_dir": str(output_dir)})
|
|
|
|
|
|
@PromptServer.instance.routes.get("/visual_prompt_library/libraries/{name}")
|
|
async def vpl_get_library(request):
|
|
name = safe_library_name(request.match_info.get("name"))
|
|
data = read_library_file(name)
|
|
if data is None:
|
|
return web.json_response({"error": "library not found"}, status=404)
|
|
return web.json_response({"name": name, "library": data})
|
|
|
|
|
|
@PromptServer.instance.routes.post("/visual_prompt_library/libraries/{name}")
|
|
async def vpl_save_library(request):
|
|
name = safe_library_name(request.match_info.get("name"))
|
|
try:
|
|
payload = await request.json()
|
|
library = payload.get("library", payload) if isinstance(payload, dict) else payload
|
|
saved = write_library_file(name, library)
|
|
except Exception as error:
|
|
return web.json_response({"error": str(error)}, status=400)
|
|
return web.json_response({"name": name, "library": saved})
|
|
|
|
|
|
@PromptServer.instance.routes.delete("/visual_prompt_library/libraries/{name}")
|
|
async def vpl_delete_library(request):
|
|
name = safe_library_name(request.match_info.get("name"))
|
|
if name == "default":
|
|
return web.json_response({"error": "default library cannot be deleted"}, status=400)
|
|
path = library_path(name)
|
|
if not path.exists():
|
|
return web.json_response({"error": "library not found"}, status=404)
|
|
try:
|
|
path.unlink()
|
|
except Exception as error:
|
|
return web.json_response({"error": str(error)}, status=400)
|
|
return web.json_response({"name": name, "deleted": True})
|
|
|
|
|
|
class VisualPromptPicker:
|
|
DESCRIPTION = "Thumbnail-driven prompt part picker. Click visual cards to emit selected prompt fragments."
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"delimiter": ("STRING", {"default": ", ", "multiline": False}),
|
|
"preview_prompt": ("STRING", {"default": "", "multiline": True}),
|
|
},
|
|
"hidden": {
|
|
"selected_parts": ("STRING", {"default": ""}),
|
|
"library_json": ("STRING", {"default": json.dumps(DEFAULT_LIBRARY, ensure_ascii=False)}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("prompt",)
|
|
FUNCTION = "build_prompt"
|
|
CATEGORY = "prompt/visual"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, delimiter, preview_prompt="", **kwargs):
|
|
payload = json.dumps(
|
|
{
|
|
"delimiter": delimiter,
|
|
"selected_parts": kwargs.get("selected_parts", ""),
|
|
"library_json": kwargs.get("library_json", ""),
|
|
},
|
|
ensure_ascii=False,
|
|
sort_keys=True,
|
|
)
|
|
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
|
|
|
def build_prompt(self, delimiter, preview_prompt="", **kwargs):
|
|
selected_raw = kwargs.get("selected_parts") or ""
|
|
|
|
try:
|
|
state = json.loads(selected_raw) if selected_raw else {}
|
|
except json.JSONDecodeError:
|
|
state = {}
|
|
|
|
selected = state.get("selected", [])
|
|
if not isinstance(selected, list):
|
|
selected = []
|
|
|
|
fragments = []
|
|
|
|
for item in selected:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
|
|
prompt = str(item.get("prompt", "")).strip()
|
|
if prompt:
|
|
fragments.append(prompt)
|
|
|
|
prompt = delimiter.join([p for p in fragments if p])
|
|
return (prompt,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"VisualPromptPicker": VisualPromptPicker,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"VisualPromptPicker": "Visual Prompt Picker",
|
|
}
|