416 lines
14 KiB
Python
416 lines
14 KiB
Python
"""
|
|
PromptStyler style pack tool.
|
|
|
|
This script helps you add more styles without editing JSON by hand.
|
|
It writes (or updates) a JSON pack under styles/packs/, e.g.:
|
|
styles/packs/99_user_custom.json
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
from collections import Counter
|
|
from typing import Any, Dict, List, Optional, Sequence
|
|
|
|
|
|
ROOT = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
|
|
if ROOT not in sys.path:
|
|
sys.path.insert(0, ROOT)
|
|
|
|
from style_library import (
|
|
CATEGORY_BASE_PREFIX,
|
|
CATEGORY_BASE_SUFFIX,
|
|
LOAD_POLICY_STRICT,
|
|
StyleLibraryError,
|
|
available_categories,
|
|
flux_join_sentences,
|
|
load_style_library,
|
|
z_join,
|
|
)
|
|
|
|
|
|
PACKS_DIR = os.path.join(ROOT, "styles", "packs")
|
|
DEFAULT_PACK_PATH = os.path.join(PACKS_DIR, "99_user_custom.json")
|
|
|
|
|
|
def _slugify(name: str) -> str:
|
|
base = "".join(ch.lower() if ch.isalnum() else "_" for ch in (name or "")).strip("_")
|
|
while "__" in base:
|
|
base = base.replace("__", "_")
|
|
return base or "style"
|
|
|
|
|
|
def _normalize_user_subcategory(value: str) -> str:
|
|
cleaned = "".join(ch if (ch.isalnum() or ch.isspace()) else " " for ch in (value or ""))
|
|
cleaned = " ".join(cleaned.split()).strip()
|
|
return cleaned or "Custom"
|
|
|
|
|
|
def _split_csv_list(value: Optional[str]) -> List[str]:
|
|
if not value:
|
|
return []
|
|
return [part.strip() for part in value.split(",") if part.strip()]
|
|
|
|
|
|
def _read_json(path: str) -> Dict[str, Any]:
|
|
with open(path, "r", encoding="utf-8") as handle:
|
|
return json.load(handle) or {}
|
|
|
|
|
|
def _write_json(path: str, data: Dict[str, Any]) -> None:
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
with open(path, "w", encoding="utf-8") as handle:
|
|
json.dump(data, handle, indent=2, ensure_ascii=True)
|
|
handle.write("\n")
|
|
|
|
|
|
def _load_library():
|
|
return load_style_library(load_policy=LOAD_POLICY_STRICT)
|
|
|
|
|
|
def _existing_identity_sets() -> tuple[set[str], set[str]]:
|
|
library = _load_library()
|
|
return ({style.id for style in library.styles}, {style.name for style in library.styles})
|
|
|
|
|
|
def _ensure_unique_id(style_id: str, existing_ids: set[str]) -> str:
|
|
if style_id not in existing_ids:
|
|
return style_id
|
|
i = 2
|
|
while True:
|
|
candidate = f"{style_id}_{i}"
|
|
if candidate not in existing_ids:
|
|
return candidate
|
|
i += 1
|
|
|
|
|
|
def _ensure_unique_name(name: str, existing_names: set[str]) -> str:
|
|
if name not in existing_names:
|
|
return name
|
|
i = 2
|
|
while True:
|
|
candidate = f"{name} ({i})"
|
|
if candidate not in existing_names:
|
|
return candidate
|
|
i += 1
|
|
|
|
|
|
def _load_or_init_pack(path: str) -> Dict[str, Any]:
|
|
if os.path.isfile(path):
|
|
data = _read_json(path)
|
|
if not isinstance(data, dict):
|
|
raise ValueError(f"Pack is not a JSON object: {path}")
|
|
if "styles" not in data:
|
|
data["styles"] = []
|
|
if "version" not in data:
|
|
data["version"] = 1
|
|
return data
|
|
return {"version": 1, "styles": []}
|
|
|
|
|
|
def _make_style_entry(
|
|
*,
|
|
name: str,
|
|
category: str,
|
|
core: Sequence[str],
|
|
details: Sequence[str],
|
|
tags: Sequence[str],
|
|
style_id: Optional[str],
|
|
id_prefix: str,
|
|
flux_suffix: Optional[str] = None,
|
|
existing_ids: Optional[set[str]] = None,
|
|
existing_names: Optional[set[str]] = None,
|
|
) -> Dict[str, Any]:
|
|
if existing_ids is None or existing_names is None:
|
|
existing_ids, existing_names = _existing_identity_sets()
|
|
|
|
base_id = style_id.strip() if style_id else f"{id_prefix}_{_slugify(name)}"
|
|
final_id = _ensure_unique_id(base_id, existing_ids)
|
|
final_name = _ensure_unique_name(name.strip(), existing_names)
|
|
existing_ids.add(final_id)
|
|
existing_names.add(final_name)
|
|
|
|
prefix = z_join(tuple(CATEGORY_BASE_PREFIX.get(category, ())) + tuple(core))
|
|
suffix = z_join(tuple(details) + tuple(CATEGORY_BASE_SUFFIX.get(category, ())))
|
|
|
|
tags2 = [tag.strip() for tag in tags if (tag or "").strip()]
|
|
if not tags2:
|
|
tags2 = ["user"]
|
|
|
|
if flux_suffix is None:
|
|
parts = [f"Style: {final_name}"]
|
|
core_hint = z_join(core[:12]) if core else ""
|
|
detail_hint = z_join(details[:12]) if details else ""
|
|
if core_hint:
|
|
parts.append(f"Core cues: {core_hint}")
|
|
if detail_hint:
|
|
parts.append(f"Details: {detail_hint}")
|
|
parts.append("Lighting: coherent and intentional")
|
|
parts.append("Mood: consistent with the user prompt")
|
|
flux_suffix2 = flux_join_sentences(parts)
|
|
else:
|
|
flux_suffix2 = str(flux_suffix or "").strip()
|
|
if flux_suffix2 and not re.search(r"[.!?]$", flux_suffix2):
|
|
flux_suffix2 += "."
|
|
|
|
return {
|
|
"id": final_id,
|
|
"name": final_name,
|
|
"category": category,
|
|
"default": {"prefix": prefix, "suffix": suffix},
|
|
"models": {"flux_2_klein": {"prefix": "", "suffix": flux_suffix2}},
|
|
"tags": tags2,
|
|
}
|
|
|
|
|
|
def _read_bulk_csv(path: str) -> List[Dict[str, str]]:
|
|
rows: List[Dict[str, str]] = []
|
|
with open(path, "r", encoding="utf-8-sig", newline="") as handle:
|
|
reader = csv.DictReader(handle)
|
|
for row in reader:
|
|
if not row:
|
|
continue
|
|
rows.append({key.strip(): (value or "").strip() for key, value in row.items() if key})
|
|
return rows
|
|
|
|
|
|
def _read_bulk_json(path: str) -> List[Dict[str, Any]]:
|
|
data = _read_json(path)
|
|
if isinstance(data, dict) and "styles" in data:
|
|
data = data["styles"]
|
|
if not isinstance(data, list):
|
|
raise ValueError("JSON bulk file must be a list (or an object with a 'styles' list).")
|
|
return [item for item in data if isinstance(item, dict)]
|
|
|
|
|
|
def cmd_categories(_args) -> int:
|
|
library = _load_library()
|
|
counts = Counter(style.category or "Uncategorized" for style in library.styles)
|
|
for category in available_categories(library.styles):
|
|
print(f"{category} ({counts.get(category, 0)})")
|
|
return 0
|
|
|
|
|
|
def cmd_stats(_args) -> int:
|
|
library = _load_library()
|
|
counts = Counter(style.category or "Uncategorized" for style in library.styles)
|
|
total = sum(counts.values())
|
|
print(f"styles {total}")
|
|
for category, count in counts.most_common():
|
|
print(f"{count:4d} {category}")
|
|
return 0
|
|
|
|
|
|
def cmd_add(args) -> int:
|
|
pack_path = os.path.abspath(args.pack)
|
|
pack = _load_or_init_pack(pack_path)
|
|
existing_ids, existing_names = _existing_identity_sets()
|
|
|
|
core = _split_csv_list(args.core) + _split_csv_list(args.prefix_extra)
|
|
details = _split_csv_list(args.details) + _split_csv_list(args.suffix_extra)
|
|
tags = _split_csv_list(args.tags)
|
|
|
|
entry = _make_style_entry(
|
|
name=args.name,
|
|
category=args.category,
|
|
core=core,
|
|
details=details,
|
|
tags=tags,
|
|
style_id=args.id,
|
|
id_prefix=args.id_prefix,
|
|
existing_ids=existing_ids,
|
|
existing_names=existing_names,
|
|
)
|
|
|
|
pack["styles"].append(entry)
|
|
pack["styles"] = sorted(
|
|
pack["styles"],
|
|
key=lambda style: (
|
|
str(style.get("category", "")).casefold(),
|
|
str(style.get("name", "")).casefold(),
|
|
str(style.get("id", "")),
|
|
),
|
|
)
|
|
_write_json(pack_path, pack)
|
|
|
|
print(f"Added: {entry['category']} | {entry['name']} | {entry['id']}")
|
|
print(f"Pack: {pack_path}")
|
|
return 0
|
|
|
|
|
|
def cmd_bulk(args) -> int:
|
|
pack_path = os.path.abspath(args.pack)
|
|
pack = _load_or_init_pack(pack_path)
|
|
existing_ids, existing_names = _existing_identity_sets()
|
|
|
|
if args.csv:
|
|
items = _read_bulk_csv(os.path.abspath(args.csv))
|
|
elif args.json:
|
|
items = _read_bulk_json(os.path.abspath(args.json))
|
|
else:
|
|
raise RuntimeError("bulk requires --csv or --json")
|
|
|
|
added = 0
|
|
for item in items:
|
|
name = (item.get("name") or "").strip()
|
|
category = (item.get("category") or "").strip()
|
|
if not name or not category:
|
|
continue
|
|
|
|
entry = _make_style_entry(
|
|
name=name,
|
|
category=category,
|
|
core=_split_csv_list(item.get("core")) + _split_csv_list(item.get("prefix_extra")),
|
|
details=_split_csv_list(item.get("details")) + _split_csv_list(item.get("suffix_extra")),
|
|
tags=_split_csv_list(item.get("tags")),
|
|
style_id=(item.get("id") or "").strip() or None,
|
|
id_prefix=args.id_prefix,
|
|
flux_suffix=(item.get("flux") or item.get("flux_suffix") or "").strip() or None,
|
|
existing_ids=existing_ids,
|
|
existing_names=existing_names,
|
|
)
|
|
pack["styles"].append(entry)
|
|
added += 1
|
|
|
|
pack["styles"] = sorted(
|
|
pack["styles"],
|
|
key=lambda style: (
|
|
str(style.get("category", "")).casefold(),
|
|
str(style.get("name", "")).casefold(),
|
|
str(style.get("id", "")),
|
|
),
|
|
)
|
|
_write_json(pack_path, pack)
|
|
print(f"Added {added} styles to {pack_path}")
|
|
return 0
|
|
|
|
|
|
def _prompt_line(label: str, *, default: Optional[str] = None) -> str:
|
|
if default is None:
|
|
return input(f"{label}: ").strip()
|
|
value = input(f"{label} [{default}]: ").strip()
|
|
return value if value else default
|
|
|
|
|
|
def cmd_wizard(args) -> int:
|
|
pack_path = os.path.abspath(args.pack)
|
|
pack = _load_or_init_pack(pack_path)
|
|
existing_ids, existing_names = _existing_identity_sets()
|
|
|
|
print("PromptStyler: User Style Wizard")
|
|
print("Tip: use comma-separated phrases (node splits default templates on ', ').")
|
|
print("")
|
|
|
|
subcat = _normalize_user_subcategory(
|
|
_prompt_line("User subcategory (creates category User/<subcategory>)", default="Custom")
|
|
)
|
|
category = f"User/{subcat}"
|
|
name = _prompt_line("Style name")
|
|
if not name.strip():
|
|
print("ERROR: style name is required")
|
|
return 2
|
|
|
|
core = _split_csv_list(_prompt_line("Core descriptors (comma-separated)", default=""))
|
|
details = _split_csv_list(_prompt_line("Detail descriptors (comma-separated)", default=""))
|
|
tags = _split_csv_list(_prompt_line("Tags (comma-separated, optional)", default=""))
|
|
tags.extend(["user", _slugify(subcat)])
|
|
flux_suffix = _prompt_line("Flux suffix (optional; blank to auto-generate)", default="").strip() or None
|
|
|
|
entry = _make_style_entry(
|
|
name=name,
|
|
category=category,
|
|
core=core,
|
|
details=details,
|
|
tags=tags,
|
|
style_id=None,
|
|
id_prefix=args.id_prefix,
|
|
flux_suffix=flux_suffix,
|
|
existing_ids=existing_ids,
|
|
existing_names=existing_names,
|
|
)
|
|
|
|
print("")
|
|
print("Preview")
|
|
print(f" Category: {entry['category']}")
|
|
print(f" Name: {entry['name']}")
|
|
print(f" ID: {entry['id']}")
|
|
print(f" Prefix: {entry['default']['prefix']}")
|
|
print(f" Suffix: {entry['default']['suffix']}")
|
|
print(f" Flux: {entry['models']['flux_2_klein']['suffix']}")
|
|
print(f" Tags: {', '.join(entry.get('tags', []))}")
|
|
print(f" Pack: {pack_path}")
|
|
|
|
confirm = _prompt_line("Write this style? (y/n)", default="y").strip().lower()
|
|
if confirm not in ("y", "yes"):
|
|
print("Canceled.")
|
|
return 0
|
|
|
|
pack["styles"].append(entry)
|
|
pack["styles"] = sorted(
|
|
pack["styles"],
|
|
key=lambda style: (
|
|
str(style.get("category", "")).casefold(),
|
|
str(style.get("name", "")).casefold(),
|
|
str(style.get("id", "")),
|
|
),
|
|
)
|
|
_write_json(pack_path, pack)
|
|
|
|
print("")
|
|
print(f"Added: {entry['category']} | {entry['name']} | {entry['id']}")
|
|
print(f"Pack: {pack_path}")
|
|
print("Tip: paste the id into the node's style_id_override to use it immediately.")
|
|
return 0
|
|
|
|
|
|
def main(argv: Optional[List[str]] = None) -> int:
|
|
parser = argparse.ArgumentParser(prog="add_styles.py", description="PromptStyler style pack helper")
|
|
sub = parser.add_subparsers(dest="cmd", required=True)
|
|
|
|
p_categories = sub.add_parser("categories", help="List available categories")
|
|
p_categories.set_defaults(func=cmd_categories)
|
|
|
|
p_stats = sub.add_parser("stats", help="Show style counts per category")
|
|
p_stats.set_defaults(func=cmd_stats)
|
|
|
|
p_add = sub.add_parser("add", help="Add a single style to a pack")
|
|
p_add.add_argument("--pack", default=DEFAULT_PACK_PATH, help="Output pack JSON path")
|
|
p_add.add_argument("--id-prefix", default="user", help="Prefix for generated ids (default: user)")
|
|
p_add.add_argument("--id", default=None, help="Optional explicit id (must be unique)")
|
|
p_add.add_argument("--name", required=True, help="Style display name")
|
|
p_add.add_argument("--category", required=True, help="Style category (as shown in the dropdown)")
|
|
p_add.add_argument("--core", default="", help="Comma-separated core descriptors")
|
|
p_add.add_argument("--details", default="", help="Comma-separated detail descriptors")
|
|
p_add.add_argument("--prefix-extra", default="", help="Alias for --core")
|
|
p_add.add_argument("--suffix-extra", default="", help="Alias for --details")
|
|
p_add.add_argument("--tags", default="", help="Comma-separated tags")
|
|
p_add.set_defaults(func=cmd_add)
|
|
|
|
p_bulk = sub.add_parser("bulk", help="Bulk add styles from CSV or JSON")
|
|
p_bulk.add_argument("--pack", default=DEFAULT_PACK_PATH, help="Output pack JSON path")
|
|
p_bulk.add_argument("--id-prefix", default="user", help="Prefix for generated ids (default: user)")
|
|
p_bulk.add_argument("--csv", default=None, help="CSV file path (columns: name,category,core,details,tags,id)")
|
|
p_bulk.add_argument("--json", default=None, help="JSON file path (list of style objects)")
|
|
p_bulk.set_defaults(func=cmd_bulk)
|
|
|
|
p_wizard = sub.add_parser("wizard", help="Interactive wizard to add a user style")
|
|
p_wizard.add_argument("--pack", default=DEFAULT_PACK_PATH, help="Output pack JSON path")
|
|
p_wizard.add_argument("--id-prefix", default="user", help="Prefix for generated ids (default: user)")
|
|
p_wizard.set_defaults(func=cmd_wizard)
|
|
|
|
args = parser.parse_args(argv)
|
|
try:
|
|
return int(args.func(args))
|
|
except StyleLibraryError as exc:
|
|
print(f"ERROR: {exc}")
|
|
return 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|