1052 lines
43 KiB
Python
1052 lines
43 KiB
Python
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
import folder_paths
|
|
import comfy.sd
|
|
import comfy.utils
|
|
import node_helpers
|
|
import numpy as np
|
|
import torch
|
|
from aiohttp import web
|
|
from PIL import Image, ImageOps
|
|
from server import PromptServer
|
|
|
|
|
|
EXTENSION_DIR = Path(__file__).resolve().parent
|
|
STORAGE_ROOT = Path(folder_paths.get_user_directory()) / "template_reference"
|
|
LEGACY_PROMPT_TEMPLATES_DIR = EXTENSION_DIR / "prompt_templates"
|
|
LEGACY_REFERENCE_TEMPLATES_DIR = EXTENSION_DIR / "Template Reference"
|
|
PROMPT_TEMPLATES_DIR = STORAGE_ROOT / "prompt_templates"
|
|
PROMPT_TEMPLATE_IMAGES_DIR = PROMPT_TEMPLATES_DIR / "images"
|
|
PROMPT_TEMPLATES_LORA_DIR = STORAGE_ROOT / "prompt_templates_lora"
|
|
PROMPT_TEMPLATE_LORA_IMAGES_DIR = PROMPT_TEMPLATES_LORA_DIR / "images"
|
|
REFERENCE_TEMPLATES_DIR = STORAGE_ROOT / "reference_templates"
|
|
REFERENCE_TEMPLATE_IMAGES_DIR = REFERENCE_TEMPLATES_DIR / "image"
|
|
LEGACY_PROMPT_TEMPLATE_IMAGES_DIR = LEGACY_PROMPT_TEMPLATES_DIR / "images"
|
|
LEGACY_REFERENCE_TEMPLATE_IMAGES_DIR = LEGACY_REFERENCE_TEMPLATES_DIR / "image"
|
|
LOGGER = logging.getLogger(__name__)
|
|
_LOADED_LORA_CACHE = None
|
|
|
|
|
|
def _ensure_prompt_templates_dir():
|
|
PROMPT_TEMPLATES_DIR.mkdir(parents=True, exist_ok=True)
|
|
PROMPT_TEMPLATE_IMAGES_DIR.mkdir(parents=True, exist_ok=True)
|
|
return PROMPT_TEMPLATES_DIR
|
|
|
|
|
|
def _ensure_prompt_templates_lora_dir():
|
|
PROMPT_TEMPLATES_LORA_DIR.mkdir(parents=True, exist_ok=True)
|
|
PROMPT_TEMPLATE_LORA_IMAGES_DIR.mkdir(parents=True, exist_ok=True)
|
|
return PROMPT_TEMPLATES_LORA_DIR
|
|
|
|
|
|
def _ensure_reference_templates_dir():
|
|
REFERENCE_TEMPLATES_DIR.mkdir(parents=True, exist_ok=True)
|
|
REFERENCE_TEMPLATE_IMAGES_DIR.mkdir(parents=True, exist_ok=True)
|
|
return REFERENCE_TEMPLATES_DIR
|
|
|
|
|
|
def _existing_template_dirs(primary, legacy):
|
|
dirs = [primary]
|
|
if legacy != primary and legacy.is_dir():
|
|
dirs.append(legacy)
|
|
return dirs
|
|
|
|
|
|
def _safe_prompt_template_name(name):
|
|
name = str(name or "").strip()
|
|
if name.lower().endswith(".json"):
|
|
name = name[:-5]
|
|
if not name or re.search(r"[\x00-\x1f]", name):
|
|
raise ValueError("Invalid template name.")
|
|
if any(separator in name for separator in ("/", "\\")) or ":" in name:
|
|
raise ValueError("Template name must not contain path separators.")
|
|
safe = re.sub(r"[^0-9A-Za-z._ \-\u3040-\u30ff\u3400-\u9fff\uff00-\uffef]", "_", name).strip(" .")
|
|
if not safe:
|
|
raise ValueError("Invalid template name after sanitization.")
|
|
return safe
|
|
|
|
|
|
def _prompt_template_path(name):
|
|
safe = _safe_prompt_template_name(name)
|
|
base = _ensure_prompt_templates_dir().resolve()
|
|
path = (base / f"{safe}.json").resolve()
|
|
if base != path.parent:
|
|
raise ValueError("Template path escapes storage directory.")
|
|
return safe, path
|
|
|
|
|
|
def _prompt_template_lora_path(name):
|
|
safe = _safe_prompt_template_name(name)
|
|
if not safe.lower().endswith("lora"):
|
|
safe = f"{safe}_lora"
|
|
base = _ensure_prompt_templates_lora_dir().resolve()
|
|
path = (base / f"{safe}.json").resolve()
|
|
if base != path.parent:
|
|
raise ValueError("Prompt Template LoRA path escapes storage directory.")
|
|
return safe, path
|
|
|
|
|
|
def _existing_prompt_template_path(name):
|
|
safe = _safe_prompt_template_name(name)
|
|
for base in _existing_template_dirs(_ensure_prompt_templates_dir(), LEGACY_PROMPT_TEMPLATES_DIR):
|
|
path = (base.resolve() / f"{safe}.json").resolve()
|
|
if base.resolve() == path.parent and path.is_file():
|
|
return safe, path
|
|
return safe, (_ensure_prompt_templates_dir().resolve() / f"{safe}.json").resolve()
|
|
|
|
|
|
def _existing_prompt_template_lora_path(name):
|
|
safe = _safe_prompt_template_name(name)
|
|
base = _ensure_prompt_templates_lora_dir().resolve()
|
|
path = (base / f"{safe}.json").resolve()
|
|
if base == path.parent and path.is_file():
|
|
return safe, path
|
|
return safe, path
|
|
|
|
|
|
def _reference_template_path(name):
|
|
safe = _safe_prompt_template_name(name)
|
|
base = _ensure_reference_templates_dir().resolve()
|
|
path = (base / f"{safe}.json").resolve()
|
|
if base != path.parent:
|
|
raise ValueError("Template Reference path escapes storage directory.")
|
|
return safe, path
|
|
|
|
|
|
def _existing_reference_template_path(name):
|
|
safe = _safe_prompt_template_name(name)
|
|
for base in _existing_template_dirs(_ensure_reference_templates_dir(), LEGACY_REFERENCE_TEMPLATES_DIR):
|
|
path = (base.resolve() / f"{safe}.json").resolve()
|
|
if base.resolve() == path.parent and path.is_file():
|
|
return safe, path
|
|
return safe, (_ensure_reference_templates_dir().resolve() / f"{safe}.json").resolve()
|
|
|
|
|
|
def _safe_image_filename(name):
|
|
name = Path(str(name or "")).name.strip()
|
|
if not name or re.search(r"[\x00-\x1f]", name):
|
|
raise ValueError("Invalid image file name.")
|
|
safe = re.sub(r"[^0-9A-Za-z._ \-\u3040-\u30ff\u3400-\u9fff\uff00-\uffef]", "_", name).strip(" .")
|
|
if not safe:
|
|
raise ValueError("Invalid image file name after sanitization.")
|
|
return safe
|
|
|
|
|
|
def _stored_image_dir(template_name, base_dir=PROMPT_TEMPLATE_IMAGES_DIR):
|
|
safe = _safe_prompt_template_name(template_name)
|
|
base = base_dir.resolve()
|
|
folder = (base / safe).resolve()
|
|
if base not in folder.parents and base != folder:
|
|
raise ValueError("Image directory escapes storage directory.")
|
|
folder.mkdir(parents=True, exist_ok=True)
|
|
return safe, folder
|
|
|
|
|
|
def _stored_image_path(template_name, image_name, base_dir=PROMPT_TEMPLATE_IMAGES_DIR):
|
|
safe_template, folder = _stored_image_dir(template_name, base_dir)
|
|
safe_image = _safe_image_filename(image_name)
|
|
path = (folder / safe_image).resolve()
|
|
if folder != path.parent:
|
|
raise ValueError("Image path escapes storage directory.")
|
|
return safe_template, safe_image, path
|
|
|
|
|
|
def _unique_image_path(template_name, image_name, base_dir=PROMPT_TEMPLATE_IMAGES_DIR):
|
|
safe_template, safe_image, path = _stored_image_path(template_name, image_name, base_dir)
|
|
stem = Path(safe_image).stem
|
|
suffix = Path(safe_image).suffix
|
|
index = 1
|
|
while path.exists():
|
|
safe_image = f"{stem}_{index}{suffix}"
|
|
_, _, path = _stored_image_path(safe_template, safe_image, base_dir)
|
|
index += 1
|
|
return safe_template, safe_image, path
|
|
|
|
|
|
def _source_image_path(image):
|
|
"""Resolve an image dict to a filesystem Path.
|
|
|
|
Resolution order:
|
|
1. Prompt-template storage (storage == 'prompt_template')
|
|
2. Reference-template storage (storage == 'template_reference_file')
|
|
3. ComfyUI annotated file path fallback
|
|
"""
|
|
if not isinstance(image, dict):
|
|
return None
|
|
local_path = _local_saved_image_path(image.get("path"))
|
|
if local_path is not None:
|
|
return local_path
|
|
if image.get("storage") in ("prompt_template", "prompt_template_lora") or image.get("type") in ("prompt_template", "prompt_template_lora"):
|
|
template_name = image.get("template") or image.get("template_name") or ""
|
|
name = image.get("name") or ""
|
|
if not template_name and image.get("path"):
|
|
parts = str(image.get("path") or "").replace("\\", "/").split("/")
|
|
if "images" in parts:
|
|
index = parts.index("images")
|
|
if len(parts) > index + 1:
|
|
template_name = parts[index + 1]
|
|
if template_name and name:
|
|
bases = [PROMPT_TEMPLATE_IMAGES_DIR, LEGACY_PROMPT_TEMPLATE_IMAGES_DIR]
|
|
if image.get("storage") == "prompt_template_lora" or image.get("type") == "prompt_template_lora":
|
|
bases = [PROMPT_TEMPLATE_LORA_IMAGES_DIR, PROMPT_TEMPLATE_IMAGES_DIR, LEGACY_PROMPT_TEMPLATE_IMAGES_DIR]
|
|
for base in bases:
|
|
_, _, path = _stored_image_path(template_name, name, base)
|
|
if path.is_file():
|
|
return path
|
|
if image.get("storage") == "template_reference_file" or image.get("type") == "template_reference_file":
|
|
template_name = image.get("template") or image.get("template_name") or ""
|
|
name = image.get("name") or ""
|
|
if not template_name and image.get("path"):
|
|
parts = str(image.get("path") or "").replace("\\", "/").split("/")
|
|
if "image" in parts:
|
|
index = parts.index("image")
|
|
if len(parts) > index + 1:
|
|
template_name = parts[index + 1]
|
|
if template_name and name:
|
|
for base in (REFERENCE_TEMPLATE_IMAGES_DIR, LEGACY_REFERENCE_TEMPLATE_IMAGES_DIR):
|
|
_, _, path = _stored_image_path(template_name, name, base)
|
|
if path.is_file():
|
|
return path
|
|
image_path = TemplateReference._image_path_from_item({"image": image})
|
|
if not image_path:
|
|
return None
|
|
try:
|
|
full_path = Path(folder_paths.get_annotated_filepath(image_path))
|
|
except Exception:
|
|
return None
|
|
return full_path if full_path.is_file() else None
|
|
|
|
|
|
def _local_saved_image_path(image_path):
|
|
normalized = str(image_path or "").replace("\\", "/").strip("/")
|
|
if not normalized:
|
|
return None
|
|
|
|
parts = normalized.split("/")
|
|
if len(parts) >= 4 and parts[0] == "prompt_templates_lora" and parts[1] == "images":
|
|
try:
|
|
_, _, path = _stored_image_path(parts[2], parts[-1], PROMPT_TEMPLATE_LORA_IMAGES_DIR)
|
|
if path.is_file():
|
|
return path
|
|
except ValueError:
|
|
pass
|
|
|
|
if len(parts) >= 4 and parts[0] == "prompt_templates" and parts[1] == "images":
|
|
for base in (PROMPT_TEMPLATE_IMAGES_DIR, LEGACY_PROMPT_TEMPLATE_IMAGES_DIR):
|
|
try:
|
|
_, _, path = _stored_image_path(parts[2], parts[-1], base)
|
|
if path.is_file():
|
|
return path
|
|
except ValueError:
|
|
continue
|
|
|
|
if len(parts) >= 4 and parts[0] == "Template Reference" and parts[1] == "image":
|
|
for base in (REFERENCE_TEMPLATE_IMAGES_DIR, LEGACY_REFERENCE_TEMPLATE_IMAGES_DIR):
|
|
try:
|
|
_, _, path = _stored_image_path(parts[2], parts[-1], base)
|
|
if path.is_file():
|
|
return path
|
|
except ValueError:
|
|
continue
|
|
|
|
return None
|
|
|
|
|
|
def _copy_prompt_template_images(template_name, data):
|
|
used_names = set()
|
|
for item in data.get("items", []):
|
|
image = item.get("image") if isinstance(item.get("image"), dict) else {}
|
|
source = _source_image_path(image)
|
|
if source is None:
|
|
continue
|
|
|
|
preferred = image.get("name") or source.name
|
|
safe_template, safe_image, destination = _stored_image_path(template_name, preferred)
|
|
if safe_image in used_names and source.resolve() != destination.resolve():
|
|
safe_template, safe_image, destination = _unique_image_path(template_name, preferred)
|
|
used_names.add(safe_image)
|
|
|
|
if source.resolve() != destination.resolve():
|
|
shutil.copy2(source, destination)
|
|
|
|
image["name"] = safe_image
|
|
image["subfolder"] = ""
|
|
image["type"] = "prompt_template"
|
|
image["path"] = f"prompt_templates/images/{safe_template}/{safe_image}"
|
|
image["storage"] = "prompt_template"
|
|
image["template"] = safe_template
|
|
image["template_name"] = safe_template
|
|
item["image"] = image
|
|
return data
|
|
|
|
|
|
def _copy_prompt_template_lora_images(template_name, data):
|
|
used_names = set()
|
|
for item in data.get("items", []):
|
|
image = item.get("image") if isinstance(item.get("image"), dict) else {}
|
|
source = _source_image_path(image)
|
|
if source is None:
|
|
continue
|
|
|
|
preferred = image.get("name") or source.name
|
|
safe_template, safe_image, destination = _stored_image_path(template_name, preferred, PROMPT_TEMPLATE_LORA_IMAGES_DIR)
|
|
if safe_image in used_names and source.resolve() != destination.resolve():
|
|
safe_template, safe_image, destination = _unique_image_path(template_name, preferred, PROMPT_TEMPLATE_LORA_IMAGES_DIR)
|
|
used_names.add(safe_image)
|
|
|
|
if source.resolve() != destination.resolve():
|
|
shutil.copy2(source, destination)
|
|
|
|
image["name"] = safe_image
|
|
image["subfolder"] = ""
|
|
image["type"] = "prompt_template_lora"
|
|
image["path"] = f"prompt_templates_lora/images/{safe_template}/{safe_image}"
|
|
image["storage"] = "prompt_template_lora"
|
|
image["template"] = safe_template
|
|
image["template_name"] = safe_template
|
|
item["image"] = image
|
|
return data
|
|
|
|
|
|
def _copy_reference_template_images(template_name, data):
|
|
used_names = set()
|
|
for item in data.get("items", []):
|
|
if item.get("type") != "image":
|
|
continue
|
|
image = item.get("image") if isinstance(item.get("image"), dict) else {}
|
|
source = _source_image_path(image)
|
|
if source is None:
|
|
continue
|
|
|
|
preferred = image.get("name") or source.name
|
|
safe_template, safe_image, destination = _stored_image_path(template_name, preferred, REFERENCE_TEMPLATE_IMAGES_DIR)
|
|
if safe_image in used_names and source.resolve() != destination.resolve():
|
|
safe_template, safe_image, destination = _unique_image_path(template_name, preferred, REFERENCE_TEMPLATE_IMAGES_DIR)
|
|
used_names.add(safe_image)
|
|
|
|
if source.resolve() != destination.resolve():
|
|
shutil.copy2(source, destination)
|
|
|
|
image["name"] = safe_image
|
|
image["subfolder"] = ""
|
|
image["type"] = "template_reference_file"
|
|
image["path"] = f"Template Reference/image/{safe_template}/{safe_image}"
|
|
image["storage"] = "template_reference_file"
|
|
image["template"] = safe_template
|
|
image["template_name"] = safe_template
|
|
item["image"] = image
|
|
return data
|
|
|
|
|
|
def _normalize_prompt_template_payload(data):
|
|
if isinstance(data, str):
|
|
payload = json.loads(data or "{}")
|
|
elif isinstance(data, dict):
|
|
payload = data
|
|
else:
|
|
raise ValueError("Template data must be a JSON object.")
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("Template data must be a JSON object.")
|
|
items = PromptTemplate._parse_items(json.dumps(payload, ensure_ascii=False))
|
|
selected_id = PromptTemplate._selected_id(json.dumps(payload, ensure_ascii=False))
|
|
return {
|
|
"version": 1,
|
|
"selected_id": selected_id,
|
|
"list_hidden": bool(payload.get("list_hidden")),
|
|
"items": items,
|
|
}
|
|
|
|
|
|
def _normalize_prompt_template_lora_payload(data):
|
|
if isinstance(data, str):
|
|
payload = json.loads(data or "{}")
|
|
elif isinstance(data, dict):
|
|
payload = data
|
|
else:
|
|
raise ValueError("Prompt Template LoRA data must be a JSON object.")
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("Prompt Template LoRA data must be a JSON object.")
|
|
items = PromptTemplateLoRA._parse_items(json.dumps(payload, ensure_ascii=False))
|
|
selected_id = PromptTemplateLoRA._selected_id(json.dumps(payload, ensure_ascii=False))
|
|
return {
|
|
"version": 1,
|
|
"selected_id": selected_id,
|
|
"list_hidden": bool(payload.get("list_hidden")),
|
|
"items": items,
|
|
}
|
|
|
|
|
|
@PromptServer.instance.routes.get("/template_reference/loras")
|
|
async def list_loras(request):
|
|
try:
|
|
return web.json_response({"success": True, "loras": folder_paths.get_filename_list("loras")})
|
|
except Exception as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
|
|
|
|
|
def _normalize_reference_template_payload(data):
|
|
if isinstance(data, str):
|
|
payload = json.loads(data or "{}")
|
|
elif isinstance(data, dict):
|
|
payload = data
|
|
else:
|
|
raise ValueError("Template Reference data must be a JSON object.")
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("Template Reference data must be a JSON object.")
|
|
items = TemplateReference._parse_items(json.dumps(payload, ensure_ascii=False))
|
|
selected_id = TemplateReference._selected_id(json.dumps(payload, ensure_ascii=False))
|
|
return {
|
|
"version": 1,
|
|
"selected_id": selected_id,
|
|
"list_hidden": bool(payload.get("list_hidden")),
|
|
"items": items,
|
|
}
|
|
|
|
|
|
@PromptServer.instance.routes.get("/template_reference/prompt_templates")
|
|
async def list_prompt_templates(request):
|
|
try:
|
|
folder = _ensure_prompt_templates_dir()
|
|
files = {
|
|
path.stem
|
|
for base in _existing_template_dirs(folder, LEGACY_PROMPT_TEMPLATES_DIR)
|
|
for path in base.glob("*.json")
|
|
if path.is_file()
|
|
}
|
|
files = sorted(files)
|
|
return web.json_response({"success": True, "templates": files})
|
|
except Exception as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
|
|
|
|
|
@PromptServer.instance.routes.get("/template_reference/prompt_templates/{name}")
|
|
async def open_prompt_template(request):
|
|
try:
|
|
safe, path = _existing_prompt_template_path(request.match_info.get("name", ""))
|
|
if not path.is_file():
|
|
return web.json_response({"success": False, "error": "Template file not found."}, status=404)
|
|
with path.open("r", encoding="utf-8") as file:
|
|
payload = json.load(file)
|
|
data = _normalize_prompt_template_payload(payload)
|
|
return web.json_response({"success": True, "name": safe, "data": data})
|
|
except ValueError as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=400)
|
|
except json.JSONDecodeError:
|
|
return web.json_response({"success": False, "error": "Template JSON is invalid."}, status=400)
|
|
except Exception as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
|
|
|
|
|
@PromptServer.instance.routes.get("/template_reference/prompt_templates/image/{template}/{image}")
|
|
async def view_prompt_template_image(request):
|
|
try:
|
|
_, _, path = _stored_image_path(request.match_info.get("template", ""), request.match_info.get("image", ""))
|
|
if not path.is_file():
|
|
return web.json_response({"success": False, "error": "Image file not found."}, status=404)
|
|
return web.FileResponse(path)
|
|
except ValueError as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=400)
|
|
except Exception as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
|
|
|
|
|
@PromptServer.instance.routes.post("/template_reference/prompt_templates/save")
|
|
async def save_prompt_template(request):
|
|
try:
|
|
body = await request.json()
|
|
safe, path = _prompt_template_path(body.get("name", ""))
|
|
data = _normalize_prompt_template_payload(body.get("data"))
|
|
data = _copy_prompt_template_images(safe, data)
|
|
with path.open("w", encoding="utf-8") as file:
|
|
json.dump(data, file, ensure_ascii=False, indent=2)
|
|
file.write("\n")
|
|
files = sorted(item.stem for item in _ensure_prompt_templates_dir().glob("*.json") if item.is_file())
|
|
return web.json_response({"success": True, "name": safe, "templates": files, "data": data})
|
|
except ValueError as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=400)
|
|
except json.JSONDecodeError:
|
|
return web.json_response({"success": False, "error": "Template JSON is invalid."}, status=400)
|
|
except Exception as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
|
|
|
|
|
@PromptServer.instance.routes.get("/template_reference/prompt_templates_lora")
|
|
async def list_prompt_templates_lora(request):
|
|
try:
|
|
folder = _ensure_prompt_templates_lora_dir()
|
|
files = {
|
|
path.stem
|
|
for path in folder.glob("*.json")
|
|
if path.is_file()
|
|
}
|
|
files.update(
|
|
path.stem
|
|
for base in _existing_template_dirs(_ensure_prompt_templates_dir(), LEGACY_PROMPT_TEMPLATES_DIR)
|
|
for path in base.glob("*.json")
|
|
if path.is_file()
|
|
)
|
|
files = sorted(files)
|
|
return web.json_response({"success": True, "templates": files})
|
|
except Exception as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
|
|
|
|
|
@PromptServer.instance.routes.get("/template_reference/prompt_templates_lora/{name}")
|
|
async def open_prompt_template_lora(request):
|
|
try:
|
|
safe, path = _existing_prompt_template_lora_path(request.match_info.get("name", ""))
|
|
if not path.is_file():
|
|
safe, path = _existing_prompt_template_path(request.match_info.get("name", ""))
|
|
if not path.is_file():
|
|
return web.json_response({"success": False, "error": "Prompt Template file not found."}, status=404)
|
|
with path.open("r", encoding="utf-8") as file:
|
|
payload = json.load(file)
|
|
data = _normalize_prompt_template_lora_payload(payload)
|
|
return web.json_response({"success": True, "name": safe, "data": data})
|
|
except ValueError as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=400)
|
|
except json.JSONDecodeError:
|
|
return web.json_response({"success": False, "error": "Prompt Template LoRA JSON is invalid."}, status=400)
|
|
except Exception as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
|
|
|
|
|
@PromptServer.instance.routes.get("/template_reference/prompt_templates_lora/image/{template}/{image}")
|
|
async def view_prompt_template_lora_image(request):
|
|
try:
|
|
_, _, path = _stored_image_path(
|
|
request.match_info.get("template", ""),
|
|
request.match_info.get("image", ""),
|
|
PROMPT_TEMPLATE_LORA_IMAGES_DIR,
|
|
)
|
|
if not path.is_file():
|
|
return web.json_response({"success": False, "error": "Image file not found."}, status=404)
|
|
return web.FileResponse(path)
|
|
except ValueError as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=400)
|
|
except Exception as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
|
|
|
|
|
@PromptServer.instance.routes.post("/template_reference/prompt_templates_lora/save")
|
|
async def save_prompt_template_lora(request):
|
|
try:
|
|
body = await request.json()
|
|
safe, path = _prompt_template_lora_path(body.get("name", ""))
|
|
data = _normalize_prompt_template_lora_payload(body.get("data"))
|
|
data = _copy_prompt_template_lora_images(safe, data)
|
|
with path.open("w", encoding="utf-8") as file:
|
|
json.dump(data, file, ensure_ascii=False, indent=2)
|
|
file.write("\n")
|
|
files = sorted(item.stem for item in _ensure_prompt_templates_lora_dir().glob("*.json") if item.is_file())
|
|
return web.json_response({"success": True, "name": safe, "templates": files, "data": data})
|
|
except ValueError as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=400)
|
|
except json.JSONDecodeError:
|
|
return web.json_response({"success": False, "error": "Prompt Template LoRA JSON is invalid."}, status=400)
|
|
except Exception as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
|
|
|
|
|
@PromptServer.instance.routes.get("/template_reference/reference_templates")
|
|
async def list_reference_templates(request):
|
|
try:
|
|
folder = _ensure_reference_templates_dir()
|
|
files = {
|
|
path.stem
|
|
for base in _existing_template_dirs(folder, LEGACY_REFERENCE_TEMPLATES_DIR)
|
|
for path in base.glob("*.json")
|
|
if path.is_file()
|
|
}
|
|
files = sorted(files)
|
|
return web.json_response({"success": True, "templates": files})
|
|
except Exception as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
|
|
|
|
|
@PromptServer.instance.routes.get("/template_reference/reference_templates/{name}")
|
|
async def open_reference_template(request):
|
|
try:
|
|
safe, path = _existing_reference_template_path(request.match_info.get("name", ""))
|
|
if not path.is_file():
|
|
return web.json_response({"success": False, "error": "Template Reference file not found."}, status=404)
|
|
with path.open("r", encoding="utf-8") as file:
|
|
payload = json.load(file)
|
|
data = _normalize_reference_template_payload(payload)
|
|
return web.json_response({"success": True, "name": safe, "data": data})
|
|
except ValueError as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=400)
|
|
except json.JSONDecodeError:
|
|
return web.json_response({"success": False, "error": "Template Reference JSON is invalid."}, status=400)
|
|
except Exception as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
|
|
|
|
|
@PromptServer.instance.routes.get("/template_reference/reference_templates/image/{template}/{image}")
|
|
async def view_reference_template_image(request):
|
|
try:
|
|
_, _, path = _stored_image_path(
|
|
request.match_info.get("template", ""),
|
|
request.match_info.get("image", ""),
|
|
REFERENCE_TEMPLATE_IMAGES_DIR,
|
|
)
|
|
if not path.is_file():
|
|
return web.json_response({"success": False, "error": "Image file not found."}, status=404)
|
|
return web.FileResponse(path)
|
|
except ValueError as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=400)
|
|
except Exception as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
|
|
|
|
|
@PromptServer.instance.routes.post("/template_reference/reference_templates/save")
|
|
async def save_reference_template(request):
|
|
try:
|
|
body = await request.json()
|
|
safe, path = _reference_template_path(body.get("name", ""))
|
|
data = _normalize_reference_template_payload(body.get("data"))
|
|
data = _copy_reference_template_images(safe, data)
|
|
with path.open("w", encoding="utf-8") as file:
|
|
json.dump(data, file, ensure_ascii=False, indent=2)
|
|
file.write("\n")
|
|
files = sorted(item.stem for item in _ensure_reference_templates_dir().glob("*.json") if item.is_file())
|
|
return web.json_response({"success": True, "name": safe, "templates": files, "data": data})
|
|
except ValueError as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=400)
|
|
except json.JSONDecodeError:
|
|
return web.json_response({"success": False, "error": "Template Reference JSON is invalid."}, status=400)
|
|
except Exception as exc:
|
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
|
|
|
|
|
def _parse_selected_id(json_str, fallback=""):
|
|
try:
|
|
payload = json.loads(json_str or "{}")
|
|
except Exception:
|
|
return str(fallback or "")
|
|
if not isinstance(payload, dict):
|
|
return str(fallback or "")
|
|
if "selected_id" in payload:
|
|
return str(payload.get("selected_id") or "")
|
|
return str(fallback or "")
|
|
|
|
|
|
def _parse_list_hidden(json_str):
|
|
try:
|
|
payload = json.loads(json_str or "{}")
|
|
except Exception:
|
|
return False
|
|
if not isinstance(payload, dict):
|
|
return False
|
|
return bool(payload.get("list_hidden"))
|
|
|
|
|
|
class TemplateReference:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"template_json": (
|
|
"STRING",
|
|
{
|
|
"default": '{"version":1,"items":[]}',
|
|
"multiline": True,
|
|
"tooltip": "Serialized text and image reference blocks edited by the custom UI.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "IMAGE", "STRING")
|
|
RETURN_NAMES = ("select_text", "select_image", "template_json")
|
|
FUNCTION = "build_reference"
|
|
CATEGORY = "template_reference"
|
|
DESCRIPTION = "Outputs text or image from the selected reference block."
|
|
|
|
def build_reference(self, template_json):
|
|
items = self._parse_items(template_json)
|
|
selected_id = self._selected_id(template_json)
|
|
list_hidden = self._list_hidden(template_json)
|
|
select_text = ""
|
|
select_image = self._placeholder_image()
|
|
|
|
if items and selected_id:
|
|
selected = next((item for item in items if str(item.get("id") or "") == str(selected_id)), None)
|
|
if selected and selected.get("type") == "text":
|
|
select_text = str(selected.get("text") or "")
|
|
elif selected and selected.get("type") == "image":
|
|
loaded = self._load_image(self._image_path_from_item(selected))
|
|
if loaded is not None:
|
|
select_image = loaded
|
|
|
|
normalized_json = json.dumps({"version": 1, "selected_id": selected_id, "list_hidden": list_hidden, "items": items}, ensure_ascii=False)
|
|
return (select_text, select_image, normalized_json)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, template_json):
|
|
hasher = hashlib.sha256()
|
|
hasher.update(str(template_json).encode("utf-8", errors="replace"))
|
|
|
|
for item in cls._parse_items(template_json):
|
|
image_path = cls._image_path_from_item(item)
|
|
if not image_path:
|
|
continue
|
|
try:
|
|
full_path = _local_saved_image_path(image_path) or Path(folder_paths.get_annotated_filepath(image_path))
|
|
stat = full_path.stat()
|
|
hasher.update(f"{full_path}:{stat.st_mtime}:{stat.st_size}".encode())
|
|
except Exception:
|
|
hasher.update(f"missing:{image_path}".encode("utf-8", errors="replace"))
|
|
|
|
return hasher.hexdigest()
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, template_json):
|
|
try:
|
|
payload = json.loads(template_json or "{}")
|
|
except (json.JSONDecodeError, TypeError):
|
|
return "template_json is not valid JSON."
|
|
if not isinstance(payload, dict):
|
|
return "template_json must be a JSON object."
|
|
return True
|
|
|
|
@staticmethod
|
|
def _selected_id(template_json):
|
|
return _parse_selected_id(template_json)
|
|
|
|
@staticmethod
|
|
def _list_hidden(template_json):
|
|
return _parse_list_hidden(template_json)
|
|
|
|
@staticmethod
|
|
def _parse_items(template_json):
|
|
try:
|
|
payload = json.loads(template_json or "{}")
|
|
except Exception:
|
|
return []
|
|
|
|
raw_items = payload.get("items", payload) if isinstance(payload, dict) else payload
|
|
if not isinstance(raw_items, list):
|
|
return []
|
|
|
|
items = []
|
|
for raw in raw_items:
|
|
if not isinstance(raw, dict):
|
|
continue
|
|
|
|
item_type = raw.get("type")
|
|
if item_type not in {"text", "image"}:
|
|
continue
|
|
|
|
item = {
|
|
"id": str(raw.get("id") or ""),
|
|
"type": item_type,
|
|
"title": str(raw.get("title") or ""),
|
|
"collapsed": bool(raw.get("collapsed")),
|
|
}
|
|
|
|
if item_type == "text":
|
|
item["text"] = str(raw.get("text") or "")
|
|
try:
|
|
item["text_height"] = int(raw.get("text_height") or 82)
|
|
except (TypeError, ValueError):
|
|
item["text_height"] = 82
|
|
else:
|
|
image = raw.get("image") if isinstance(raw.get("image"), dict) else {}
|
|
try:
|
|
preview_height = int(image.get("preview_height") or 138)
|
|
except (TypeError, ValueError):
|
|
preview_height = 138
|
|
|
|
item["image"] = {
|
|
"name": str(image.get("name") or ""),
|
|
"subfolder": str(image.get("subfolder") or ""),
|
|
"type": str(image.get("type") or "input"),
|
|
"path": str(image.get("path") or ""),
|
|
"preview_height": preview_height,
|
|
"storage": str(image.get("storage") or ""),
|
|
"template": str(image.get("template") or image.get("template_name") or ""),
|
|
"template_name": str(image.get("template_name") or image.get("template") or ""),
|
|
}
|
|
|
|
items.append(item)
|
|
|
|
return items
|
|
|
|
@staticmethod
|
|
def _image_path_from_item(item):
|
|
image = item.get("image") if isinstance(item.get("image"), dict) else {}
|
|
path = str(image.get("path") or "").replace("\\", "/").strip("/")
|
|
if path:
|
|
return path
|
|
|
|
name = str(image.get("name") or "").replace("\\", "/").strip("/")
|
|
subfolder = str(image.get("subfolder") or "").replace("\\", "/").strip("/")
|
|
if not name:
|
|
return ""
|
|
return f"{subfolder}/{name}" if subfolder else name
|
|
|
|
@staticmethod
|
|
def _load_image(image_path):
|
|
try:
|
|
full_path = _local_saved_image_path(image_path) or Path(folder_paths.get_annotated_filepath(image_path))
|
|
if not os.path.isfile(full_path):
|
|
LOGGER.warning("Template Reference image not found: %s", image_path)
|
|
return None
|
|
|
|
img = node_helpers.pillow(Image.open, full_path)
|
|
img = node_helpers.pillow(ImageOps.exif_transpose, img)
|
|
|
|
image = img.convert("RGB")
|
|
image_np = np.array(image).astype(np.float32) / 255.0
|
|
image_tensor = torch.from_numpy(image_np)[None,]
|
|
|
|
return image_tensor
|
|
except Exception as exc:
|
|
LOGGER.warning("Failed to load Template Reference image %r: %s", image_path, exc)
|
|
return None
|
|
|
|
@staticmethod
|
|
def _placeholder_image():
|
|
image = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
|
return image
|
|
|
|
|
|
class PromptTemplate:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt_template_json": (
|
|
"STRING",
|
|
{
|
|
"default": '{"version":1,"items":[]}',
|
|
"multiline": True,
|
|
"tooltip": "Serialized prompt template blocks edited by the custom UI.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("selected_prompt", "prompt", "template_json")
|
|
FUNCTION = "build_prompt"
|
|
CATEGORY = "template_reference"
|
|
DESCRIPTION = "Outputs the prompt text from the template selected in the custom UI."
|
|
|
|
def build_prompt(self, prompt_template_json):
|
|
items = self._parse_items(prompt_template_json)
|
|
selected_template_id = self._selected_id(prompt_template_json)
|
|
list_hidden = self._list_hidden(prompt_template_json)
|
|
selected_prompt = ""
|
|
|
|
if items and selected_template_id:
|
|
selected = next((item for item in items if str(item.get("id") or "") == str(selected_template_id or "")), None)
|
|
if selected is not None:
|
|
selected_prompt = str(selected.get("text") or "")
|
|
|
|
normalized_json = json.dumps(
|
|
{"version": 1, "selected_id": selected_template_id, "list_hidden": list_hidden, "items": items},
|
|
ensure_ascii=False,
|
|
)
|
|
return (selected_prompt, selected_prompt, normalized_json)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, prompt_template_json):
|
|
hasher = hashlib.sha256()
|
|
hasher.update(str(prompt_template_json).encode("utf-8", errors="replace"))
|
|
return hasher.hexdigest()
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, prompt_template_json):
|
|
try:
|
|
payload = json.loads(prompt_template_json or "{}")
|
|
except (json.JSONDecodeError, TypeError):
|
|
return "prompt_template_json is not valid JSON."
|
|
if not isinstance(payload, dict):
|
|
return "prompt_template_json must be a JSON object."
|
|
return True
|
|
|
|
@staticmethod
|
|
def _selected_id(prompt_template_json, fallback=""):
|
|
return _parse_selected_id(prompt_template_json, fallback)
|
|
|
|
@staticmethod
|
|
def _list_hidden(prompt_template_json):
|
|
return _parse_list_hidden(prompt_template_json)
|
|
|
|
@staticmethod
|
|
def _parse_items(prompt_template_json):
|
|
return PromptTemplate._parse_prompt_items(prompt_template_json, include_lora=False)
|
|
|
|
@staticmethod
|
|
def _parse_prompt_items(prompt_template_json, include_lora=False):
|
|
try:
|
|
payload = json.loads(prompt_template_json or "{}")
|
|
except Exception:
|
|
return []
|
|
|
|
raw_items = payload.get("items", payload) if isinstance(payload, dict) else payload
|
|
if not isinstance(raw_items, list):
|
|
return []
|
|
|
|
items = []
|
|
for raw in raw_items:
|
|
if not isinstance(raw, dict):
|
|
continue
|
|
|
|
try:
|
|
text_height = int(raw.get("text_height") or 82)
|
|
except (TypeError, ValueError):
|
|
text_height = 82
|
|
|
|
image = raw.get("image") if isinstance(raw.get("image"), dict) else {}
|
|
try:
|
|
preview_height = int(image.get("preview_height") or 138)
|
|
except (TypeError, ValueError):
|
|
preview_height = 138
|
|
item = {
|
|
"id": str(raw.get("id") or ""),
|
|
"type": "text",
|
|
"title": str(raw.get("title") or ""),
|
|
"text": str(raw.get("text") or ""),
|
|
"text_height": text_height,
|
|
"collapsed": bool(raw.get("collapsed")),
|
|
"image_collapsed": bool(raw.get("image_collapsed", True)),
|
|
"image": {
|
|
"name": str(image.get("name") or ""),
|
|
"subfolder": str(image.get("subfolder") or ""),
|
|
"type": str(image.get("type") or "input"),
|
|
"path": str(image.get("path") or ""),
|
|
"preview_height": preview_height,
|
|
"storage": str(image.get("storage") or ""),
|
|
"template": str(image.get("template") or image.get("template_name") or ""),
|
|
"template_name": str(image.get("template_name") or image.get("template") or ""),
|
|
},
|
|
}
|
|
if include_lora:
|
|
lora = raw.get("lora") if isinstance(raw.get("lora"), dict) else {}
|
|
try:
|
|
lora_strength = float(lora.get("strength", 1.0))
|
|
except (TypeError, ValueError):
|
|
lora_strength = 1.0
|
|
item["lora_collapsed"] = bool(raw.get("lora_collapsed", True))
|
|
item["lora"] = {
|
|
"name": str(lora.get("name") or ""),
|
|
"strength": lora_strength,
|
|
}
|
|
items.append(item)
|
|
|
|
return items
|
|
|
|
|
|
class PromptTemplateLoRA(PromptTemplate):
|
|
@staticmethod
|
|
def _parse_items(prompt_template_json):
|
|
return PromptTemplate._parse_prompt_items(prompt_template_json, include_lora=True)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL", {"tooltip": "Model to apply the selected prompt block LoRA to."}),
|
|
"clip": ("CLIP", {"tooltip": "CLIP to apply the selected prompt block LoRA to."}),
|
|
"prompt_template_json": (
|
|
"STRING",
|
|
{
|
|
"default": '{"version":1,"items":[]}',
|
|
"multiline": True,
|
|
"tooltip": "Serialized prompt template blocks edited by the custom UI.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "MODEL", "CLIP", "STRING")
|
|
RETURN_NAMES = ("selected_prompt", "prompt", "lora_model", "lora_clip", "template_json")
|
|
FUNCTION = "build_prompt_lora"
|
|
CATEGORY = "template_reference"
|
|
DESCRIPTION = "Outputs selected prompt text and applies that block's LoRA to model/clip."
|
|
|
|
def build_prompt_lora(self, model, clip, prompt_template_json):
|
|
items = self._parse_items(prompt_template_json)
|
|
selected_template_id = self._selected_id(prompt_template_json)
|
|
list_hidden = self._list_hidden(prompt_template_json)
|
|
selected_prompt = ""
|
|
lora_model = model
|
|
lora_clip = clip
|
|
|
|
if items and selected_template_id:
|
|
selected = next((item for item in items if str(item.get("id") or "") == str(selected_template_id or "")), None)
|
|
if selected is not None:
|
|
selected_prompt = str(selected.get("text") or "")
|
|
lora_model, lora_clip = self._apply_lora(model, clip, selected)
|
|
|
|
normalized_json = json.dumps(
|
|
{"version": 1, "selected_id": selected_template_id, "list_hidden": list_hidden, "items": items},
|
|
ensure_ascii=False,
|
|
)
|
|
return (selected_prompt, selected_prompt, lora_model, lora_clip, normalized_json)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, model, clip, prompt_template_json):
|
|
hasher = hashlib.sha256()
|
|
hasher.update(str(prompt_template_json).encode("utf-8", errors="replace"))
|
|
return hasher.hexdigest()
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, model, clip, prompt_template_json):
|
|
return PromptTemplate.VALIDATE_INPUTS(prompt_template_json)
|
|
|
|
@staticmethod
|
|
def _apply_lora(model, clip, item):
|
|
global _LOADED_LORA_CACHE
|
|
lora = item.get("lora") if isinstance(item.get("lora"), dict) else {}
|
|
lora_name = str(lora.get("name") or "")
|
|
try:
|
|
strength = float(lora.get("strength", 1.0))
|
|
except (TypeError, ValueError):
|
|
strength = 1.0
|
|
|
|
if not lora_name or strength == 0:
|
|
return model, clip
|
|
|
|
try:
|
|
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
|
except Exception as exc:
|
|
LOGGER.warning("Prompt Template LoRA not found %r: %s", lora_name, exc)
|
|
return model, clip
|
|
loaded_lora = None
|
|
lora_metadata = None
|
|
if _LOADED_LORA_CACHE is not None:
|
|
if _LOADED_LORA_CACHE[0] == lora_path:
|
|
loaded_lora = _LOADED_LORA_CACHE[1]
|
|
lora_metadata = _LOADED_LORA_CACHE[2] if len(_LOADED_LORA_CACHE) > 2 else None
|
|
else:
|
|
_LOADED_LORA_CACHE = None
|
|
|
|
if loaded_lora is None:
|
|
loaded_lora, lora_metadata = comfy.utils.load_torch_file(lora_path, safe_load=True, return_metadata=True)
|
|
_LOADED_LORA_CACHE = (lora_path, loaded_lora, lora_metadata)
|
|
|
|
return comfy.sd.load_lora_for_models(model, clip, loaded_lora, strength, strength, lora_metadata=lora_metadata)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"TemplateReference": TemplateReference,
|
|
"PromptTemplate": PromptTemplate,
|
|
"PromptTemplateLoRA": PromptTemplateLoRA,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"TemplateReference": "Template Reference",
|
|
"PromptTemplate": "Prompt Template",
|
|
"PromptTemplateLoRA": "Prompt Template LoRA",
|
|
}
|