280 lines
8.6 KiB
Python
280 lines
8.6 KiB
Python
"""
|
|
Thread-safe JSON database for storing prompts with categories.
|
|
Prompts can optionally have LoRA attachments.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import uuid
|
|
import threading
|
|
|
|
DB_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "prompts.json")
|
|
_lock = threading.Lock()
|
|
|
|
DEFAULT_DB = {
|
|
"categories": ["Location", "Character", "Style", "Mood"],
|
|
"prompts": []
|
|
}
|
|
|
|
|
|
def load_db():
|
|
"""Load the database from disk. Returns default if file missing or corrupted."""
|
|
with _lock:
|
|
if not os.path.exists(DB_FILE):
|
|
save_db_unsafe(DEFAULT_DB)
|
|
return DEFAULT_DB.copy()
|
|
try:
|
|
with open(DB_FILE, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
if "categories" not in data or "prompts" not in data:
|
|
data = DEFAULT_DB.copy()
|
|
save_db_unsafe(data)
|
|
# Ensure all prompts have a loras field
|
|
for p in data.get("prompts", []):
|
|
if "loras" not in p:
|
|
p["loras"] = []
|
|
return data
|
|
except (json.JSONDecodeError, IOError):
|
|
save_db_unsafe(DEFAULT_DB)
|
|
return DEFAULT_DB.copy()
|
|
|
|
|
|
def save_db_unsafe(data):
|
|
"""Save database without acquiring lock (caller must hold lock)."""
|
|
with open(DB_FILE, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, indent=2, ensure_ascii=False)
|
|
|
|
|
|
def save_db(data):
|
|
"""Save database to disk (thread-safe)."""
|
|
with _lock:
|
|
save_db_unsafe(data)
|
|
|
|
|
|
def get_categories():
|
|
"""Return list of category names."""
|
|
data = load_db()
|
|
return data.get("categories", [])
|
|
|
|
|
|
def get_prompts(category=None):
|
|
"""Return prompts, optionally filtered by category."""
|
|
data = load_db()
|
|
prompts = data.get("prompts", [])
|
|
if category:
|
|
prompts = [p for p in prompts if p.get("category") == category]
|
|
return prompts
|
|
|
|
|
|
def get_prompt_by_id(prompt_id):
|
|
"""Return a single prompt by its ID, or None."""
|
|
data = load_db()
|
|
for p in data.get("prompts", []):
|
|
if p.get("id") == prompt_id:
|
|
return p
|
|
return None
|
|
|
|
|
|
def get_prompt_by_name(name, category=None):
|
|
"""Return a single prompt by its name and optionally category, or None."""
|
|
data = load_db()
|
|
for p in data.get("prompts", []):
|
|
if p.get("name") == name:
|
|
if category is None or p.get("category") == category:
|
|
return p
|
|
return None
|
|
|
|
|
|
def add_prompt(category, name, text, loras=None):
|
|
"""Add a new prompt. Returns the new prompt dict."""
|
|
with _lock:
|
|
data = _load_unsafe()
|
|
prompt = {
|
|
"id": str(uuid.uuid4()),
|
|
"category": category,
|
|
"name": name,
|
|
"text": text,
|
|
"loras": loras or [],
|
|
}
|
|
data["prompts"].append(prompt)
|
|
save_db_unsafe(data)
|
|
return prompt
|
|
|
|
|
|
def update_prompt(prompt_id, category=None, name=None, text=None, loras=None):
|
|
"""Update an existing prompt by ID. Returns updated prompt or None."""
|
|
with _lock:
|
|
data = _load_unsafe()
|
|
for p in data["prompts"]:
|
|
if p["id"] == prompt_id:
|
|
if category is not None:
|
|
p["category"] = category
|
|
if name is not None:
|
|
p["name"] = name
|
|
if text is not None:
|
|
p["text"] = text
|
|
if loras is not None:
|
|
p["loras"] = loras
|
|
save_db_unsafe(data)
|
|
return p
|
|
return None
|
|
|
|
|
|
def delete_prompt(prompt_id):
|
|
"""Delete a prompt by ID. Returns True if found and deleted."""
|
|
with _lock:
|
|
data = _load_unsafe()
|
|
original_len = len(data["prompts"])
|
|
data["prompts"] = [p for p in data["prompts"] if p["id"] != prompt_id]
|
|
if len(data["prompts"]) < original_len:
|
|
save_db_unsafe(data)
|
|
return True
|
|
return False
|
|
|
|
|
|
def add_category(name):
|
|
"""Add a new category. Returns True if added, False if already exists."""
|
|
with _lock:
|
|
data = _load_unsafe()
|
|
if name in data["categories"]:
|
|
return False
|
|
data["categories"].append(name)
|
|
save_db_unsafe(data)
|
|
return True
|
|
|
|
|
|
def delete_category(name):
|
|
"""Delete a category and all its prompts. Returns True if found and deleted."""
|
|
with _lock:
|
|
data = _load_unsafe()
|
|
if name not in data["categories"]:
|
|
return False
|
|
data["categories"].remove(name)
|
|
data["prompts"] = [p for p in data["prompts"] if p["category"] != name]
|
|
save_db_unsafe(data)
|
|
return True
|
|
|
|
|
|
def reorder_categories(ordered_names):
|
|
"""Reorder categories to match the given list. Returns True if successful."""
|
|
with _lock:
|
|
data = _load_unsafe()
|
|
existing = set(data["categories"])
|
|
given = set(ordered_names)
|
|
if existing != given:
|
|
return False
|
|
data["categories"] = list(ordered_names)
|
|
save_db_unsafe(data)
|
|
return True
|
|
|
|
|
|
def reorder_prompts(ordered_ids):
|
|
"""Reorder prompts to match the given ID list. Unmentioned prompts keep their relative order at the end."""
|
|
with _lock:
|
|
data = _load_unsafe()
|
|
by_id = {p["id"]: p for p in data["prompts"]}
|
|
reordered = []
|
|
for pid in ordered_ids:
|
|
if pid in by_id:
|
|
reordered.append(by_id.pop(pid))
|
|
# Append any remaining prompts not in the ordered list
|
|
for p in data["prompts"]:
|
|
if p["id"] in by_id:
|
|
reordered.append(p)
|
|
data["prompts"] = reordered
|
|
save_db_unsafe(data)
|
|
return True
|
|
|
|
|
|
def get_available_loras():
|
|
"""Return list of available LoRA filenames from ComfyUI's models folder."""
|
|
try:
|
|
import folder_paths
|
|
return folder_paths.get_filename_list("loras")
|
|
except Exception:
|
|
return []
|
|
|
|
|
|
def export_db():
|
|
"""Return full database dict for export."""
|
|
return load_db()
|
|
|
|
|
|
def import_db(incoming, mode="merge"):
|
|
"""Import prompts from an incoming dict.
|
|
mode='overwrite': replace entire DB.
|
|
mode='merge': add new categories and prompts, skip duplicates (by name+category).
|
|
Returns dict with stats: {categories_added, prompts_added, prompts_skipped}.
|
|
"""
|
|
with _lock:
|
|
if mode == "overwrite":
|
|
# Ensure all prompts have IDs and loras
|
|
for p in incoming.get("prompts", []):
|
|
if "id" not in p:
|
|
p["id"] = str(uuid.uuid4())
|
|
if "loras" not in p:
|
|
p["loras"] = []
|
|
save_db_unsafe(incoming)
|
|
return {
|
|
"categories_added": len(incoming.get("categories", [])),
|
|
"prompts_added": len(incoming.get("prompts", [])),
|
|
"prompts_skipped": 0,
|
|
}
|
|
|
|
# Merge mode
|
|
data = _load_unsafe()
|
|
cats_added = 0
|
|
prompts_added = 0
|
|
prompts_skipped = 0
|
|
|
|
# Add missing categories
|
|
existing_cats = set(data["categories"])
|
|
for cat in incoming.get("categories", []):
|
|
if cat not in existing_cats:
|
|
data["categories"].append(cat)
|
|
existing_cats.add(cat)
|
|
cats_added += 1
|
|
|
|
# Add prompts that don't exist yet (by name + category)
|
|
existing_prompts = {(p["name"], p["category"]) for p in data["prompts"]}
|
|
for p in incoming.get("prompts", []):
|
|
key = (p.get("name", ""), p.get("category", ""))
|
|
if key in existing_prompts:
|
|
prompts_skipped += 1
|
|
continue
|
|
prompt = {
|
|
"id": p.get("id") or str(uuid.uuid4()),
|
|
"name": p.get("name", ""),
|
|
"category": p.get("category", ""),
|
|
"text": p.get("text", ""),
|
|
"loras": p.get("loras", []),
|
|
}
|
|
data["prompts"].append(prompt)
|
|
existing_prompts.add(key)
|
|
prompts_added += 1
|
|
|
|
save_db_unsafe(data)
|
|
return {
|
|
"categories_added": cats_added,
|
|
"prompts_added": prompts_added,
|
|
"prompts_skipped": prompts_skipped,
|
|
}
|
|
|
|
|
|
def _load_unsafe():
|
|
"""Load database without lock (internal use only)."""
|
|
if not os.path.exists(DB_FILE):
|
|
save_db_unsafe(DEFAULT_DB)
|
|
return DEFAULT_DB.copy()
|
|
try:
|
|
with open(DB_FILE, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
# Ensure all prompts have loras field
|
|
for p in data.get("prompts", []):
|
|
if "loras" not in p:
|
|
p["loras"] = []
|
|
return data
|
|
except (json.JSONDecodeError, IOError):
|
|
save_db_unsafe(DEFAULT_DB)
|
|
return DEFAULT_DB.copy()
|