272 lines
8.3 KiB
Python
272 lines
8.3 KiB
Python
"""
|
|
ComfyUI custom nodes for Prompt Database.
|
|
"""
|
|
|
|
import json as _json
|
|
import re as _re
|
|
import folder_paths
|
|
|
|
|
|
def _clean_name(name):
|
|
"""Strip the 📖 LoRA marker suffix added by the JS frontend."""
|
|
return _re.sub(r'\s*📖$', '', name).strip()
|
|
import comfy.sd
|
|
import comfy.utils
|
|
from . import db
|
|
|
|
|
|
class PromptDBManager:
|
|
"""
|
|
Management node — provides a UI button to open the prompt database manager.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "manage"
|
|
CATEGORY = "Prompt Database"
|
|
OUTPUT_NODE = True
|
|
|
|
def manage(self, unique_id=None):
|
|
return {}
|
|
|
|
|
|
class PromptSelector:
|
|
"""
|
|
Selector node — pick a category and prompt.
|
|
Outputs:
|
|
- prompt_text (STRING): the selected prompt text (no LoRA tags)
|
|
- db_selection (DB_SELECTION): the category name (for PromptCombiner)
|
|
- lora_stack (LORA_STACK): LoRAs attached to the selected prompt
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
categories = db.get_categories()
|
|
if not categories:
|
|
categories = ["(no categories)"]
|
|
|
|
all_prompts = db.get_prompts()
|
|
prompt_names = list({p["name"] for p in all_prompts})
|
|
if not prompt_names:
|
|
prompt_names = ["(no prompts)"]
|
|
|
|
return {
|
|
"required": {
|
|
"category": (categories, {"default": categories[0]}),
|
|
"prompt_name": (prompt_names, {"default": prompt_names[0]}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "DB_SELECTION", "LORA_STACK")
|
|
RETURN_NAMES = ("prompt_text", "db_selection", "lora_stack")
|
|
FUNCTION = "select_prompt"
|
|
CATEGORY = "Prompt Database"
|
|
|
|
def select_prompt(self, category, prompt_name):
|
|
prompt_name = _clean_name(prompt_name)
|
|
prompt = db.get_prompt_by_name(prompt_name, category)
|
|
if not prompt:
|
|
prompt = db.get_prompt_by_name(prompt_name)
|
|
if not prompt:
|
|
return ("", category, [])
|
|
|
|
text = prompt["text"]
|
|
loras = prompt.get("loras", [])
|
|
# Normalize lora entries
|
|
lora_stack = []
|
|
for l in loras:
|
|
lora_stack.append({
|
|
"name": l.get("name", ""),
|
|
"strength": float(l.get("strength", 1.0)),
|
|
})
|
|
|
|
return (text, category, lora_stack)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, category, prompt_name):
|
|
prompt_name = _clean_name(prompt_name)
|
|
prompt = db.get_prompt_by_name(prompt_name, category)
|
|
if prompt:
|
|
return prompt["text"] + str(prompt.get("loras", []))
|
|
return ""
|
|
|
|
|
|
class PromptCombiner:
|
|
"""
|
|
Combiner node — accepts multiple DB_SELECTION inputs (category names),
|
|
shows a prompt dropdown per category inside the node,
|
|
and combines selected prompts into one STRING + merged LORA_STACK.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"separator": (
|
|
["comma_space", "newline", "space", "none"],
|
|
{"default": "comma_space"},
|
|
),
|
|
"prompt_selections": (
|
|
"STRING",
|
|
{"default": "{}", "multiline": False},
|
|
),
|
|
},
|
|
"optional": {
|
|
"selection_1": ("DB_SELECTION", {"forceInput": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "LORA_STACK")
|
|
RETURN_NAMES = ("combined_prompt", "lora_stack")
|
|
FUNCTION = "combine"
|
|
CATEGORY = "Prompt Database"
|
|
OUTPUT_NODE = True
|
|
|
|
def combine(self, separator="comma_space", prompt_selections="{}", **kwargs):
|
|
sep_map = {
|
|
"comma_space": ", ",
|
|
"newline": "\n",
|
|
"space": " ",
|
|
"none": "",
|
|
}
|
|
sep = sep_map.get(separator, ", ")
|
|
|
|
try:
|
|
selections = _json.loads(prompt_selections)
|
|
except _json.JSONDecodeError:
|
|
selections = {}
|
|
|
|
texts = []
|
|
merged_loras = []
|
|
seen_loras = set()
|
|
|
|
selection_keys = [k for k in kwargs.keys() if k.startswith("selection_")]
|
|
selection_keys.sort(key=lambda k: int(k.split("_")[1]))
|
|
|
|
for key in selection_keys:
|
|
category = kwargs[key]
|
|
if not category:
|
|
continue
|
|
prompt_name = _clean_name(selections.get(key, ""))
|
|
if not prompt_name:
|
|
continue
|
|
prompt = db.get_prompt_by_name(prompt_name, category)
|
|
if prompt:
|
|
texts.append(prompt["text"])
|
|
# Collect LoRAs (deduplicate by name)
|
|
for l in prompt.get("loras", []):
|
|
lname = l.get("name", "")
|
|
if lname and lname not in seen_loras:
|
|
seen_loras.add(lname)
|
|
merged_loras.append({
|
|
"name": lname,
|
|
"strength": float(l.get("strength", 1.0)),
|
|
})
|
|
|
|
combined = sep.join(texts)
|
|
return (combined, merged_loras)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, separator, prompt_selections="{}", **kwargs):
|
|
# Build a fingerprint that includes actual prompt content,
|
|
# so edits to prompt text invalidate the cache.
|
|
try:
|
|
selections = _json.loads(prompt_selections)
|
|
except _json.JSONDecodeError:
|
|
selections = {}
|
|
|
|
parts = [separator, prompt_selections]
|
|
for key in sorted(kwargs.keys()):
|
|
if not key.startswith("selection_"):
|
|
continue
|
|
category = kwargs[key]
|
|
prompt_name = _clean_name(selections.get(key, ""))
|
|
if category and prompt_name:
|
|
prompt = db.get_prompt_by_name(prompt_name, category)
|
|
if prompt:
|
|
parts.append(prompt["text"])
|
|
parts.append(str(prompt.get("loras", [])))
|
|
return "|".join(parts)
|
|
|
|
|
|
class ApplyLoraStack:
|
|
"""
|
|
Applies LoRAs from a LORA_STACK to a model and clip.
|
|
The lora_config hidden widget stores per-LoRA overrides
|
|
(enable/disable + custom strength).
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"clip": ("CLIP",),
|
|
"lora_config": (
|
|
"STRING",
|
|
{"default": "{}", "multiline": False},
|
|
),
|
|
},
|
|
"optional": {
|
|
"lora_stack": ("LORA_STACK",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP")
|
|
RETURN_NAMES = ("model", "clip")
|
|
FUNCTION = "apply_loras"
|
|
CATEGORY = "Prompt Database"
|
|
|
|
def apply_loras(self, model, clip, lora_config="{}", lora_stack=None):
|
|
if not lora_stack:
|
|
return (model, clip)
|
|
|
|
# Parse overrides
|
|
try:
|
|
config = _json.loads(lora_config)
|
|
except _json.JSONDecodeError:
|
|
config = {}
|
|
|
|
for lora_entry in lora_stack:
|
|
name = lora_entry.get("name", "")
|
|
default_strength = lora_entry.get("strength", 1.0)
|
|
if not name:
|
|
continue
|
|
|
|
# Apply per-LoRA overrides from config
|
|
override = config.get(name, {})
|
|
if not override.get("enabled", True):
|
|
print(f"[PromptDB] LoRA skipped (disabled): {name}")
|
|
continue
|
|
|
|
strength = float(override.get("strength", default_strength))
|
|
|
|
# Load LoRA file
|
|
lora_path = folder_paths.get_full_path("loras", name)
|
|
if not lora_path:
|
|
print(f"[PromptDB] LoRA file not found: {name}")
|
|
continue
|
|
|
|
try:
|
|
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
model, clip = comfy.sd.load_lora_for_models(
|
|
model, clip, lora, strength, strength
|
|
)
|
|
print(f"[PromptDB] LoRA applied: {name} (strength={strength:.2f})")
|
|
except Exception as e:
|
|
print(f"[PromptDB] Error loading LoRA {name}: {e}")
|
|
|
|
return (model, clip)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, model, clip, lora_config="{}", lora_stack=None):
|
|
return str(lora_stack) + lora_config
|