feat(scenes): add scene fragment library mirroring characters
Extracts the per-name YAML-list storage logic out of tensors/characters.py into a new generic tensors.fragments.FragmentLibrary class, then layers two kind-specific function-style modules on top: - tensors.characters (unchanged public API, now thin wrapper) - tensors.scenes (new): save_scene / load_scene / list_scenes / delete_scene / scene_path / resolve_scene, stored in ~/.local/share/tensors/scenes/ CLI additions: - tsr scene save -o <name> 'elem1, elem2, ...' - tsr scene list / show / delete Wired into generate + template + style-sweep: - tsr generate -S <name> / --scene-prompt 'elems' injects scene elements into the positive prompt right after the character block (order: quality_prefix -> character -> scene -> rating -> user) - tsr template -S <name> / --scene-prompt 'elems' embeds a 'scene' list field in the dumped JSON template (named + inline merged, deduped) - style-sweep templates accept 'scene' (str|list) and 'scene_prompt' keys - generate --input JSON honors both 'scene' (str=name or list=inline) and 'scene_prompt' keys FragmentLibrary uses kind-aware error messages (e.g. "Invalid scene name" vs "Invalid character name") so users see the right context. 28 new tests cover the generic library + scenes. Bumps version to 0.1.24 (supersedes v0.1.23 which failed to publish due to a transient Sigstore network error).
This commit is contained in:
+1
-1
@@ -1,6 +1,6 @@
|
||||
"""tsr: Read safetensor metadata, search and download CivitAI models."""
|
||||
|
||||
__version__ = "0.1.23"
|
||||
__version__ = "0.1.24"
|
||||
|
||||
from tensors.cli import main
|
||||
from tensors.config import (
|
||||
|
||||
+37
-115
@@ -1,126 +1,68 @@
|
||||
"""Character library: named lists of prompt elements, stored as YAML.
|
||||
"""Character library: named lists of prompt elements for the *who*.
|
||||
|
||||
Characters live in ``~/.local/share/tensors/characters/<name>.yml`` and contain a
|
||||
flat YAML list of strings describing a subject (appearance, identity, outfit
|
||||
fragments). They are injected into a generation prompt via ``--character`` or
|
||||
``--character-prompt`` on ``tsr generate``.
|
||||
``--character-prompt`` on ``tsr generate`` and embedded in the ``character``
|
||||
field of ``tsr template`` JSON output.
|
||||
|
||||
Format is JSON-compatible YAML — each line is ``- "value"`` so the files round-trip
|
||||
through both ``json`` and any YAML parser. Manual hand-editing with plain or
|
||||
single-quoted scalars is also accepted on read.
|
||||
This module exposes a function-style API on top of
|
||||
:class:`tensors.fragments.FragmentLibrary`. Each call instantiates a library
|
||||
rooted at the current module-level ``CHARACTERS_DIR`` so tests can monkeypatch
|
||||
the directory without re-importing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path # noqa: TC003 # used in runtime return annotations exposed to typer
|
||||
from pathlib import Path # noqa: TC003 # used in runtime return annotations
|
||||
|
||||
from tensors.config import DATA_DIR
|
||||
from tensors.fragments import FragmentLibrary, parse_elements
|
||||
|
||||
# Storage location for character YAML files.
|
||||
__all__ = [
|
||||
"CHARACTERS_DIR",
|
||||
"character_path",
|
||||
"delete_character",
|
||||
"list_characters",
|
||||
"load_character",
|
||||
"parse_elements",
|
||||
"resolve_character",
|
||||
"save_character",
|
||||
]
|
||||
|
||||
# Default storage location. Tests may monkeypatch this attribute; every helper
|
||||
# below dereferences it via globals() so the override is picked up live.
|
||||
CHARACTERS_DIR = DATA_DIR / "characters"
|
||||
|
||||
# Restrict character names to a safe subset so they can't escape CHARACTERS_DIR
|
||||
# via path traversal and so file listings stay tidy.
|
||||
_NAME_RE = re.compile(r"^[A-Za-z0-9_.-]+$")
|
||||
|
||||
# Minimum length for a quoted YAML scalar: opening + closing quote.
|
||||
_MIN_QUOTED_SCALAR_LEN = 2
|
||||
|
||||
|
||||
def _validate_name(name: str) -> None:
|
||||
if not name or not _NAME_RE.match(name):
|
||||
raise ValueError(f"Invalid character name {name!r}: only letters, digits, '.', '_', '-' allowed")
|
||||
def _lib() -> FragmentLibrary:
|
||||
"""Build a fresh library bound to the current ``CHARACTERS_DIR`` value."""
|
||||
return FragmentLibrary("characters", base_dir=globals()["CHARACTERS_DIR"])
|
||||
|
||||
|
||||
def character_path(name: str) -> Path:
|
||||
"""Return the on-disk path for a character name (without ensuring it exists)."""
|
||||
_validate_name(name)
|
||||
return CHARACTERS_DIR / f"{name}.yml"
|
||||
|
||||
|
||||
def parse_elements(value: str) -> list[str]:
|
||||
"""Split a comma-separated prompt fragment into clean, order-preserving elements.
|
||||
|
||||
Empty pieces and duplicates are dropped. Used by ``tsr character save`` and
|
||||
by the ``--character-prompt`` CLI flag so both share identical splitting
|
||||
semantics.
|
||||
"""
|
||||
parts = [p.strip() for p in value.split(",")]
|
||||
seen: set[str] = set()
|
||||
out: list[str] = []
|
||||
for p in parts:
|
||||
if p and p not in seen:
|
||||
seen.add(p)
|
||||
out.append(p)
|
||||
return out
|
||||
"""Return the on-disk path for a character name (without ensuring existence)."""
|
||||
return _lib().path(name)
|
||||
|
||||
|
||||
def save_character(name: str, elements: list[str]) -> Path:
|
||||
"""Persist a character's elements to disk and return the file path.
|
||||
|
||||
Overwrites any existing file. Each element is written on its own line as a
|
||||
JSON-encoded YAML scalar (``- "value"``), which keeps embedded commas,
|
||||
quotes, and unicode safe.
|
||||
"""
|
||||
CHARACTERS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
path = character_path(name)
|
||||
body = "\n".join(f"- {json.dumps(e, ensure_ascii=False)}" for e in elements)
|
||||
path.write_text(body + "\n" if body else "")
|
||||
return path
|
||||
"""Persist a character's elements to disk and return the file path."""
|
||||
return _lib().save(name, elements)
|
||||
|
||||
|
||||
def load_character(name: str) -> list[str]:
|
||||
"""Load a character's elements. Raises ``FileNotFoundError`` if missing.
|
||||
|
||||
Accepts both JSON-quoted scalars (``- "value"``), single-quoted YAML scalars
|
||||
(``- 'value'``) and bare scalars (``- value``). Blank lines and ``#`` comments
|
||||
are ignored.
|
||||
"""
|
||||
path = character_path(name)
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(f"Character {name!r} not found at {path}")
|
||||
|
||||
elements: list[str] = []
|
||||
for raw in path.read_text().splitlines():
|
||||
line = raw.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
if not line.startswith("-"):
|
||||
# Skip any non-list-item lines (e.g. a YAML document header users
|
||||
# might add manually); we only consume flat list entries.
|
||||
continue
|
||||
item = line[1:].lstrip()
|
||||
if not item:
|
||||
continue
|
||||
# Prefer strict JSON decode (covers our own writer output and any
|
||||
# double-quoted YAML scalar). Fall back to single-quoted YAML, then
|
||||
# bare scalar.
|
||||
try:
|
||||
value = json.loads(item)
|
||||
if not isinstance(value, str):
|
||||
value = str(value)
|
||||
except json.JSONDecodeError:
|
||||
value = item[1:-1].replace("''", "'") if len(item) >= _MIN_QUOTED_SCALAR_LEN and item[0] == item[-1] == "'" else item
|
||||
elements.append(value)
|
||||
return elements
|
||||
"""Load a character's elements. Raises ``FileNotFoundError`` if missing."""
|
||||
return _lib().load(name)
|
||||
|
||||
|
||||
def list_characters() -> list[str]:
|
||||
"""Return sorted character names. Empty list if no characters dir exists yet."""
|
||||
if not CHARACTERS_DIR.is_dir():
|
||||
return []
|
||||
return sorted(p.stem for p in CHARACTERS_DIR.glob("*.yml") if p.is_file())
|
||||
"""Return sorted character names."""
|
||||
return _lib().list()
|
||||
|
||||
|
||||
def delete_character(name: str) -> bool:
|
||||
"""Delete a character file. Returns True on success, False if it was missing."""
|
||||
path = character_path(name)
|
||||
if not path.is_file():
|
||||
return False
|
||||
path.unlink()
|
||||
return True
|
||||
"""Delete a character file. Returns True on success, False if missing."""
|
||||
return _lib().delete(name)
|
||||
|
||||
|
||||
def resolve_character(
|
||||
@@ -129,25 +71,5 @@ def resolve_character(
|
||||
character_prompt: str | None = None,
|
||||
extra: list[str] | None = None,
|
||||
) -> list[str]:
|
||||
"""Merge a named character with an inline ``--character-prompt`` and optional extras.
|
||||
|
||||
Resolution order (first match wins per duplicate): named character →
|
||||
``--character-prompt`` elements → ``extra`` list. The result preserves order
|
||||
and drops duplicates, mirroring ``parse_elements``.
|
||||
"""
|
||||
merged: list[str] = []
|
||||
seen: set[str] = set()
|
||||
|
||||
def _push(items: list[str]) -> None:
|
||||
for item in items:
|
||||
if item and item not in seen:
|
||||
seen.add(item)
|
||||
merged.append(item)
|
||||
|
||||
if character:
|
||||
_push(load_character(character))
|
||||
if character_prompt:
|
||||
_push(parse_elements(character_prompt))
|
||||
if extra:
|
||||
_push(extra)
|
||||
return merged
|
||||
"""Merge a named character with an inline ``--character-prompt`` and extras."""
|
||||
return _lib().resolve(name=character, inline=character_prompt, extra=extra)
|
||||
|
||||
+247
-33
@@ -788,6 +788,14 @@ def generate( # noqa: PLR0915
|
||||
str | None,
|
||||
typer.Option("--character-prompt", help='Inline character fragment, comma-separated (e.g. "blond hair, blue eyes")'),
|
||||
] = None,
|
||||
scene: Annotated[
|
||||
str | None,
|
||||
typer.Option("-S", "--scene", help="Saved scene name (loaded from ~/.local/share/tensors/scenes/)"),
|
||||
] = None,
|
||||
scene_prompt: Annotated[
|
||||
str | None,
|
||||
typer.Option("--scene-prompt", help='Inline scene fragment, comma-separated (e.g. "luxury penthouse, volumetric lighting")'),
|
||||
] = None,
|
||||
family: Annotated[
|
||||
str | None,
|
||||
typer.Option(
|
||||
@@ -912,7 +920,20 @@ def generate( # noqa: PLR0915
|
||||
character_prompt = ", ".join(str(x) for x in val if str(x).strip())
|
||||
if "character_prompt" in mapped and "character_prompt" not in explicit:
|
||||
cp_val = mapped["character_prompt"]
|
||||
character_prompt = cp_val if isinstance(cp_val, str) else ", ".join(str(x) for x in cp_val if str(x).strip())
|
||||
character_prompt = (
|
||||
cp_val if isinstance(cp_val, str) else ", ".join(str(x) for x in cp_val if str(x).strip())
|
||||
)
|
||||
if "scene" in mapped and "scene" not in explicit:
|
||||
sv = mapped["scene"]
|
||||
if isinstance(sv, str):
|
||||
scene = sv
|
||||
elif isinstance(sv, (list, tuple)):
|
||||
scene_prompt = ", ".join(str(x) for x in sv if str(x).strip())
|
||||
if "scene_prompt" in mapped and "scene_prompt" not in explicit:
|
||||
sp_val = mapped["scene_prompt"]
|
||||
scene_prompt = (
|
||||
sp_val if isinstance(sp_val, str) else ", ".join(str(x) for x in sp_val if str(x).strip())
|
||||
)
|
||||
if "rating" in mapped and "rating" not in explicit:
|
||||
rating = mapped["rating"]
|
||||
|
||||
@@ -942,6 +963,8 @@ def generate( # noqa: PLR0915
|
||||
no_negative=no_negative,
|
||||
character=character,
|
||||
character_prompt=character_prompt,
|
||||
scene=scene,
|
||||
scene_prompt=scene_prompt,
|
||||
family=family,
|
||||
output=output,
|
||||
remote=remote,
|
||||
@@ -1032,6 +1055,8 @@ def _run_generation( # noqa: PLR0915
|
||||
no_negative: bool = False,
|
||||
character: str | None = None,
|
||||
character_prompt: str | None = None,
|
||||
scene: str | None = None,
|
||||
scene_prompt: str | None = None,
|
||||
family: str | None = None,
|
||||
output: Path | None = None,
|
||||
remote: str | None = None,
|
||||
@@ -1105,6 +1130,31 @@ def _run_generation( # noqa: PLR0915
|
||||
f"[dim]Character ({origin}, {len(character_elements)} elements): {', '.join(character_elements)}[/dim]"
|
||||
)
|
||||
|
||||
# Resolve scene (named lookup + inline --scene-prompt, merged + deduped).
|
||||
# Scene sits between character (who) and rating/user prompt (what's happening)
|
||||
# so the natural reading order is: quality → character → scene → rating → user.
|
||||
scene_elements: list[str] = []
|
||||
if scene or scene_prompt:
|
||||
from tensors.scenes import resolve_scene # noqa: PLC0415
|
||||
|
||||
try:
|
||||
scene_elements = resolve_scene(scene=scene, scene_prompt=scene_prompt)
|
||||
except FileNotFoundError as e:
|
||||
console.print(f"[red]{e}[/red]")
|
||||
raise typer.Exit(1) from e
|
||||
except ValueError as e:
|
||||
console.print(f"[red]{e}[/red]")
|
||||
raise typer.Exit(1) from e
|
||||
|
||||
if scene_elements:
|
||||
prompt_parts.extend(scene_elements)
|
||||
if not json_output:
|
||||
origin = f"'{scene}'" if scene else "inline"
|
||||
console.print(
|
||||
f"[dim]Scene ({origin}, {len(scene_elements)} elements): "
|
||||
f"{', '.join(scene_elements)}[/dim]"
|
||||
)
|
||||
|
||||
# Add rating tag based on model family (Pony/Illustrious)
|
||||
if rating:
|
||||
from tensors.config import get_rating_tag # noqa: PLC0415
|
||||
@@ -1350,6 +1400,8 @@ _STYLE_SWEEP_TEMPLATE_KEYS = {
|
||||
"no_negative",
|
||||
"character",
|
||||
"character_prompt",
|
||||
"scene",
|
||||
"scene_prompt",
|
||||
"rating",
|
||||
"family",
|
||||
"remote",
|
||||
@@ -1680,22 +1732,27 @@ def style_sweep( # noqa: PLR0915
|
||||
|
||||
pending_tasks.append((i, entry, result, out_path))
|
||||
|
||||
# Character resolution: templates may carry either a name string (look up
|
||||
# at run-time) or an inline list of resolved elements (e.g. produced by
|
||||
# `tsr template -C ...`). Lists are joined into `character_prompt` so
|
||||
# _run_generation sees a uniform CSV string and skips the disk lookup.
|
||||
char_val = tpl_data.get("character")
|
||||
char_prompt_val = tpl_data.get("character_prompt")
|
||||
char_name: str | None = None
|
||||
char_inline: str | None = None
|
||||
if isinstance(char_val, str):
|
||||
char_name = char_val
|
||||
elif isinstance(char_val, (list, tuple)):
|
||||
char_inline = ", ".join(str(x) for x in char_val if str(x).strip())
|
||||
if char_prompt_val is not None:
|
||||
char_inline = (
|
||||
char_prompt_val if isinstance(char_prompt_val, str) else ", ".join(str(x) for x in char_prompt_val if str(x).strip())
|
||||
)
|
||||
# Character / scene resolution: templates may carry either a name string
|
||||
# (look up at run-time) or an inline list of resolved elements (e.g. produced
|
||||
# by `tsr template -C ... -S ...`). Lists are joined into the *_prompt arg
|
||||
# so _run_generation sees a uniform CSV string and skips the disk lookup.
|
||||
def _split_fragment(name_val: Any, prompt_val: Any) -> tuple[str | None, str | None]:
|
||||
name_out: str | None = None
|
||||
inline_out: str | None = None
|
||||
if isinstance(name_val, str):
|
||||
name_out = name_val
|
||||
elif isinstance(name_val, (list, tuple)):
|
||||
inline_out = ", ".join(str(x) for x in name_val if str(x).strip())
|
||||
if prompt_val is not None:
|
||||
inline_out = (
|
||||
prompt_val
|
||||
if isinstance(prompt_val, str)
|
||||
else ", ".join(str(x) for x in prompt_val if str(x).strip())
|
||||
)
|
||||
return name_out, inline_out
|
||||
|
||||
char_name, char_inline = _split_fragment(tpl_data.get("character"), tpl_data.get("character_prompt"))
|
||||
scene_name, scene_inline = _split_fragment(tpl_data.get("scene"), tpl_data.get("scene_prompt"))
|
||||
|
||||
# Common kwargs for every _run_generation call — extracted from the
|
||||
# template once, reused across sequential and parallel paths.
|
||||
@@ -1720,6 +1777,8 @@ def style_sweep( # noqa: PLR0915
|
||||
"no_negative": bool(_t("no_negative", default=False)),
|
||||
"character": char_name,
|
||||
"character_prompt": char_inline,
|
||||
"scene": scene_name,
|
||||
"scene_prompt": scene_inline,
|
||||
"family": _t("family"),
|
||||
"remote": gen_remote,
|
||||
"json_output": False,
|
||||
@@ -1883,6 +1942,17 @@ def template(
|
||||
help="Inline character fragment, comma-separated (merged with --character into `character`)",
|
||||
),
|
||||
] = None,
|
||||
scene: Annotated[
|
||||
str | None,
|
||||
typer.Option("-S", "--scene", help="Saved scene name (resolved into the `scene` list field)"),
|
||||
] = None,
|
||||
scene_prompt: Annotated[
|
||||
str | None,
|
||||
typer.Option(
|
||||
"--scene-prompt",
|
||||
help='Inline scene fragment, comma-separated (merged with --scene into `scene`)',
|
||||
),
|
||||
] = None,
|
||||
output: Annotated[Path | None, typer.Option("-o", "--output", help="Save template to file")] = None,
|
||||
) -> None:
|
||||
"""Dump a JSON generation template with resolved defaults for a model.
|
||||
@@ -1890,17 +1960,18 @@ def template(
|
||||
Outputs a ready-to-use JSON object with all parameters auto-resolved from the
|
||||
checkpoint family. Pipe to 'tsr generate --input' or save to a file for reuse.
|
||||
|
||||
``--character`` and ``--character-prompt`` append a ``character`` list to the
|
||||
template (saved-name elements first, inline elements appended, deduped).
|
||||
``--character`` / ``--character-prompt`` append a ``character`` list to the
|
||||
template; ``--scene`` / ``--scene-prompt`` append a ``scene`` list (named
|
||||
elements first, inline elements appended, deduped within each list).
|
||||
|
||||
Examples:
|
||||
tsr template -m ponyDiffusionV6XL_v6StartWithThisOne.safetensors
|
||||
tsr template -m beautifulRealistic_v7.safetensors -O portrait
|
||||
tsr template -m waiIllustriousSDXL_v160.safetensors -l "Elvira iIlluLoRA.safetensors"
|
||||
tsr template -m ponyRealism_V22.safetensors -o pony_preset.json
|
||||
tsr template -m flux1-dev.safetensors -C cassie_cage # embeds saved character
|
||||
tsr template -m flux1-dev.safetensors --character-prompt "blond hair, blue eyes"
|
||||
tsr template -m flux1-dev.safetensors -C cassie --character-prompt "wet skin"
|
||||
tsr template -m flux1-dev.safetensors -C cassie_cage -S penthouse
|
||||
tsr template -m flux1-dev.safetensors --character-prompt "blond hair, blue eyes" \\
|
||||
--scene-prompt "luxury penthouse, volumetric lighting, Canon R5"
|
||||
tsr generate --input "$(tsr template -m ponyRealism_V22.safetensors)" "a portrait"
|
||||
"""
|
||||
from tensors.config import get_model_generation_defaults, resolve_orientation # noqa: PLC0415
|
||||
@@ -1954,11 +2025,10 @@ def template(
|
||||
tpl["lora"] = lora
|
||||
tpl["lora_strength"] = lora_strength
|
||||
|
||||
# Resolve character into a flat list embedded in the template. When the
|
||||
# template is later fed to `tsr generate --input`, _run_generation will
|
||||
# treat a list under `character` as inline elements (no re-lookup needed).
|
||||
# --character and --character-prompt merge in that order (named first,
|
||||
# inline appended, duplicates dropped).
|
||||
# Resolve character / scene into flat lists embedded in the template. When
|
||||
# the template is later fed to `tsr generate --input`, _run_generation will
|
||||
# treat the lists under `character` / `scene` as inline elements (no
|
||||
# re-lookup needed). The `_*_name` fields are informational only.
|
||||
if character or character_prompt:
|
||||
from tensors.characters import resolve_character # noqa: PLC0415
|
||||
|
||||
@@ -1976,6 +2046,23 @@ def template(
|
||||
if character:
|
||||
tpl["_character_name"] = character
|
||||
|
||||
if scene or scene_prompt:
|
||||
from tensors.scenes import resolve_scene # noqa: PLC0415
|
||||
|
||||
try:
|
||||
resolved_scene = resolve_scene(scene=scene, scene_prompt=scene_prompt)
|
||||
except FileNotFoundError as e:
|
||||
console.print(f"[red]{e}[/red]")
|
||||
raise typer.Exit(1) from e
|
||||
except ValueError as e:
|
||||
console.print(f"[red]{e}[/red]")
|
||||
raise typer.Exit(1) from e
|
||||
|
||||
if resolved_scene:
|
||||
tpl["scene"] = resolved_scene
|
||||
if scene:
|
||||
tpl["_scene_name"] = scene
|
||||
|
||||
# Add metadata (not used by generate, but informational)
|
||||
tpl["_family"] = family or "unknown"
|
||||
if base_model_str:
|
||||
@@ -2460,12 +2547,19 @@ def hf_download(
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Character Commands
|
||||
# Fragment Commands (character + scene)
|
||||
# =============================================================================
|
||||
# Characters are named, comma-split prompt fragments stored as YAML lists in
|
||||
# ~/.local/share/tensors/characters/<name>.yml. They are injected into the
|
||||
# positive prompt by `tsr generate --character <name>` (or inline via
|
||||
# `--character-prompt "elem1, elem2"`).
|
||||
# Characters and scenes are named, comma-split prompt fragments stored as YAML
|
||||
# lists in ~/.local/share/tensors/<kind>/<name>.yml. They are injected into the
|
||||
# positive prompt by `tsr generate --character <name>` / `--scene <name>` (or
|
||||
# inline via `--character-prompt "..."` / `--scene-prompt`).
|
||||
#
|
||||
# Both subcommand groups share the underlying tensors.fragments.FragmentLibrary,
|
||||
# but the CLI commands are spelled out per-kind to keep Typer's signature
|
||||
# introspection happy under `from __future__ import annotations` (closures
|
||||
# referencing per-kind labels break typer's eval_str=True resolution).
|
||||
|
||||
# ---- character ----
|
||||
|
||||
character_app = typer.Typer(
|
||||
name="character",
|
||||
@@ -2501,7 +2595,7 @@ def character_save(
|
||||
raise typer.Exit(1) from e
|
||||
|
||||
if json_output:
|
||||
console.print_json(data={"name": name, "path": str(path), "elements": parsed})
|
||||
console.print_json(data={"name": name, "path": str(path), "elements": parsed, "kind": "characters"})
|
||||
return
|
||||
|
||||
console.print(f"[green]Saved character '{name}' ({len(parsed)} elements):[/green] {path}")
|
||||
@@ -2582,6 +2676,126 @@ def character_delete(
|
||||
raise typer.Exit(1)
|
||||
|
||||
|
||||
# ---- scene ----
|
||||
|
||||
scene_app = typer.Typer(
|
||||
name="scene",
|
||||
help="Manage saved scene prompts (~/.local/share/tensors/scenes/).",
|
||||
no_args_is_help=True,
|
||||
)
|
||||
app.add_typer(scene_app)
|
||||
|
||||
|
||||
@scene_app.command("save")
|
||||
def scene_save(
|
||||
elements: Annotated[
|
||||
str, typer.Argument(help='Comma-separated prompt elements (e.g. "luxury penthouse, volumetric lighting")')
|
||||
],
|
||||
name: Annotated[str, typer.Option("-o", "--output", help="Scene name (used as filename)")],
|
||||
json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False,
|
||||
) -> None:
|
||||
"""Save a scene as a YAML list of prompt elements.
|
||||
|
||||
Examples:
|
||||
tsr scene save -o penthouse "luxury penthouse, volumetric lighting, Canon R5, 85mm"
|
||||
tsr scene save -o forest "deep forest, dappled sunlight, moss-covered rocks"
|
||||
"""
|
||||
from tensors.fragments import parse_elements # noqa: PLC0415
|
||||
from tensors.scenes import save_scene # noqa: PLC0415
|
||||
|
||||
parsed = parse_elements(elements)
|
||||
if not parsed:
|
||||
console.print("[red]No usable elements after splitting on commas[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
try:
|
||||
path = save_scene(name, parsed)
|
||||
except ValueError as e:
|
||||
console.print(f"[red]{e}[/red]")
|
||||
raise typer.Exit(1) from e
|
||||
|
||||
if json_output:
|
||||
console.print_json(data={"name": name, "path": str(path), "elements": parsed, "kind": "scenes"})
|
||||
return
|
||||
|
||||
console.print(f"[green]Saved scene '{name}' ({len(parsed)} elements):[/green] {path}")
|
||||
for elem in parsed:
|
||||
console.print(f" • {elem}")
|
||||
|
||||
|
||||
@scene_app.command("list")
|
||||
def scene_list(
|
||||
json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False,
|
||||
) -> None:
|
||||
"""List saved scenes."""
|
||||
from tensors.scenes import SCENES_DIR, list_scenes # noqa: PLC0415
|
||||
|
||||
names = list_scenes()
|
||||
if json_output:
|
||||
console.print_json(data={"dir": str(SCENES_DIR), "scenes": names})
|
||||
return
|
||||
|
||||
if not names:
|
||||
console.print(f"[yellow]No scenes saved in {SCENES_DIR}.[/yellow]")
|
||||
console.print('[dim]Create one with: tsr scene save -o <name> "elem1, elem2"[/dim]')
|
||||
return
|
||||
|
||||
console.print(f"[bold]Scenes[/bold] ({len(names)}) [dim]in {SCENES_DIR}[/dim]")
|
||||
for n in names:
|
||||
console.print(f" • {n}")
|
||||
|
||||
|
||||
@scene_app.command("show")
|
||||
def scene_show(
|
||||
name: Annotated[str, typer.Argument(help="Scene name")],
|
||||
json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False,
|
||||
) -> None:
|
||||
"""Show a scene's elements."""
|
||||
from tensors.scenes import load_scene, scene_path # noqa: PLC0415
|
||||
|
||||
try:
|
||||
elements = load_scene(name)
|
||||
except FileNotFoundError as e:
|
||||
console.print(f"[red]{e}[/red]")
|
||||
raise typer.Exit(1) from e
|
||||
except ValueError as e:
|
||||
console.print(f"[red]{e}[/red]")
|
||||
raise typer.Exit(1) from e
|
||||
|
||||
if json_output:
|
||||
console.print_json(data={"name": name, "path": str(scene_path(name)), "elements": elements})
|
||||
return
|
||||
|
||||
console.print(f"[bold]{name}[/bold] [dim]({scene_path(name)})[/dim]")
|
||||
for elem in elements:
|
||||
console.print(f" • {elem}")
|
||||
|
||||
|
||||
@scene_app.command("delete")
|
||||
def scene_delete(
|
||||
name: Annotated[str, typer.Argument(help="Scene name")],
|
||||
json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False,
|
||||
) -> None:
|
||||
"""Delete a saved scene."""
|
||||
from tensors.scenes import delete_scene # noqa: PLC0415
|
||||
|
||||
try:
|
||||
deleted = delete_scene(name)
|
||||
except ValueError as e:
|
||||
console.print(f"[red]{e}[/red]")
|
||||
raise typer.Exit(1) from e
|
||||
|
||||
if json_output:
|
||||
console.print_json(data={"name": name, "deleted": deleted})
|
||||
return
|
||||
|
||||
if deleted:
|
||||
console.print(f"[green]Deleted scene '{name}'[/green]")
|
||||
else:
|
||||
console.print(f"[yellow]Scene '{name}' does not exist[/yellow]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# ComfyUI Commands
|
||||
# =============================================================================
|
||||
|
||||
@@ -0,0 +1,178 @@
|
||||
"""Generic prompt-fragment library.
|
||||
|
||||
A *fragment* is a named, ordered list of comma-style prompt elements
|
||||
(e.g. ``["blond hair", "broad chin"]``) stored as a flat YAML list on disk.
|
||||
Different *kinds* of fragments (characters, scenes, …) each get their own
|
||||
subdirectory under ``~/.local/share/tensors/<kind>/`` and their own
|
||||
``FragmentLibrary`` instance.
|
||||
|
||||
Files are written as JSON-encoded YAML scalars (``- "value"``) so they round-trip
|
||||
through both ``json`` and any YAML parser. Hand-edited single-quoted or bare
|
||||
scalars are accepted on read.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path # noqa: TC003 # used in runtime return annotations exposed to typer
|
||||
|
||||
from tensors.config import DATA_DIR
|
||||
|
||||
# Restrict fragment names to a safe subset so they can't escape the storage dir
|
||||
# via path traversal and so file listings stay tidy.
|
||||
_NAME_RE = re.compile(r"^[A-Za-z0-9_.-]+$")
|
||||
|
||||
# Minimum length for a quoted YAML scalar: opening + closing quote.
|
||||
_MIN_QUOTED_SCALAR_LEN = 2
|
||||
|
||||
|
||||
class FragmentLibrary:
|
||||
"""A named collection of prompt-fragment YAML files of a single *kind*.
|
||||
|
||||
Each ``FragmentLibrary`` is rooted at ``<DATA_DIR>/<kind>/`` (overridable for
|
||||
tests). Instance methods are stateless wrappers around that directory.
|
||||
"""
|
||||
|
||||
def __init__(self, kind: str, base_dir: Path | None = None) -> None:
|
||||
"""Create a library for ``kind`` (e.g. ``"characters"`` or ``"scenes"``).
|
||||
|
||||
``base_dir`` defaults to ``DATA_DIR / kind`` and is recomputed lazily so
|
||||
tests can monkeypatch ``DATA_DIR`` *or* the ``base_dir`` attribute
|
||||
directly without re-importing.
|
||||
"""
|
||||
if not kind or not _NAME_RE.match(kind):
|
||||
raise ValueError(f"Invalid library kind {kind!r}")
|
||||
self.kind = kind
|
||||
self.base_dir = base_dir if base_dir is not None else DATA_DIR / kind
|
||||
|
||||
# ---------- internals ----------
|
||||
|
||||
@property
|
||||
def _singular(self) -> str:
|
||||
"""Human-readable singular form of ``kind`` used in error messages."""
|
||||
return self.kind[:-1] if self.kind.endswith("s") else self.kind
|
||||
|
||||
def _validate_name(self, name: str) -> None:
|
||||
if not name or not _NAME_RE.match(name):
|
||||
raise ValueError(
|
||||
f"Invalid {self._singular} name {name!r}: only letters, digits, '.', '_', '-' allowed"
|
||||
)
|
||||
|
||||
def path(self, name: str) -> Path:
|
||||
"""Return the on-disk path for ``name`` (without ensuring it exists)."""
|
||||
self._validate_name(name)
|
||||
return self.base_dir / f"{name}.yml"
|
||||
|
||||
# ---------- CRUD ----------
|
||||
|
||||
def save(self, name: str, elements: list[str]) -> Path:
|
||||
"""Persist ``elements`` to disk and return the file path.
|
||||
|
||||
Overwrites any existing file. Each element is written on its own line as
|
||||
a JSON-encoded YAML scalar so embedded commas, quotes and unicode are
|
||||
safe.
|
||||
"""
|
||||
self.base_dir.mkdir(parents=True, exist_ok=True)
|
||||
path = self.path(name)
|
||||
body = "\n".join(f"- {json.dumps(e, ensure_ascii=False)}" for e in elements)
|
||||
path.write_text(body + "\n" if body else "")
|
||||
return path
|
||||
|
||||
def load(self, name: str) -> list[str]:
|
||||
"""Load a fragment. Raises ``FileNotFoundError`` if missing.
|
||||
|
||||
Accepts JSON-quoted scalars (``- "value"``), single-quoted YAML scalars
|
||||
(``- 'value'``) and bare scalars (``- value``). Blank lines and ``#``
|
||||
comments are ignored.
|
||||
"""
|
||||
path = self.path(name)
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(f"{self._singular.capitalize()} {name!r} not found at {path}")
|
||||
|
||||
elements: list[str] = []
|
||||
for raw in path.read_text().splitlines():
|
||||
line = raw.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
if not line.startswith("-"):
|
||||
# Skip non-list lines (e.g. a YAML document header users might
|
||||
# add manually); we only consume flat list entries.
|
||||
continue
|
||||
item = line[1:].lstrip()
|
||||
if not item:
|
||||
continue
|
||||
try:
|
||||
value = json.loads(item)
|
||||
if not isinstance(value, str):
|
||||
value = str(value)
|
||||
except json.JSONDecodeError:
|
||||
value = (
|
||||
item[1:-1].replace("''", "'")
|
||||
if len(item) >= _MIN_QUOTED_SCALAR_LEN and item[0] == item[-1] == "'"
|
||||
else item
|
||||
)
|
||||
elements.append(value)
|
||||
return elements
|
||||
|
||||
def list(self) -> list[str]:
|
||||
"""Return sorted fragment names. Empty list if the dir doesn't exist yet."""
|
||||
if not self.base_dir.is_dir():
|
||||
return []
|
||||
return sorted(p.stem for p in self.base_dir.glob("*.yml") if p.is_file())
|
||||
|
||||
def delete(self, name: str) -> bool:
|
||||
"""Delete a fragment. Returns True on success, False if it was missing."""
|
||||
path = self.path(name)
|
||||
if not path.is_file():
|
||||
return False
|
||||
path.unlink()
|
||||
return True
|
||||
|
||||
# ---------- helpers ----------
|
||||
|
||||
def resolve(
|
||||
self,
|
||||
*,
|
||||
name: str | None = None,
|
||||
inline: str | None = None,
|
||||
extra: list[str] | None = None,
|
||||
) -> list[str]:
|
||||
"""Merge a named fragment with an inline CSV string and optional extras.
|
||||
|
||||
Resolution order (first match wins per duplicate): named → inline → extra.
|
||||
Result preserves order and drops duplicates and empty pieces.
|
||||
"""
|
||||
merged: list[str] = []
|
||||
seen: set[str] = set()
|
||||
|
||||
def _push(items: list[str]) -> None:
|
||||
for item in items:
|
||||
if item and item not in seen:
|
||||
seen.add(item)
|
||||
merged.append(item)
|
||||
|
||||
if name:
|
||||
_push(self.load(name))
|
||||
if inline:
|
||||
_push(parse_elements(inline))
|
||||
if extra:
|
||||
_push(extra)
|
||||
return merged
|
||||
|
||||
|
||||
def parse_elements(value: str) -> list[str]:
|
||||
"""Split a comma-separated prompt fragment into clean, order-preserving elements.
|
||||
|
||||
Empty pieces and duplicates are dropped. Shared by ``tsr character|scene save``
|
||||
and the ``--character-prompt`` / ``--scene-prompt`` CLI flags so the splitting
|
||||
semantics stay identical across surfaces.
|
||||
"""
|
||||
parts = [p.strip() for p in value.split(",")]
|
||||
seen: set[str] = set()
|
||||
out: list[str] = []
|
||||
for p in parts:
|
||||
if p and p not in seen:
|
||||
seen.add(p)
|
||||
out.append(p)
|
||||
return out
|
||||
@@ -0,0 +1,75 @@
|
||||
"""Scene library: named lists of prompt elements for the *where*.
|
||||
|
||||
Scenes live in ``~/.local/share/tensors/scenes/<name>.yml`` and contain a flat
|
||||
YAML list of strings describing a setting (location, lighting, camera, etc.).
|
||||
They are injected into a generation prompt via ``--scene`` or ``--scene-prompt``
|
||||
on ``tsr generate`` and embedded in the ``scene`` field of ``tsr template``
|
||||
JSON output.
|
||||
|
||||
This module exposes a function-style API on top of
|
||||
:class:`tensors.fragments.FragmentLibrary`. Each call instantiates a library
|
||||
rooted at the current module-level ``SCENES_DIR`` so tests can monkeypatch the
|
||||
directory without re-importing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path # noqa: TC003 # used in runtime return annotations
|
||||
|
||||
from tensors.config import DATA_DIR
|
||||
from tensors.fragments import FragmentLibrary, parse_elements
|
||||
|
||||
__all__ = [
|
||||
"SCENES_DIR",
|
||||
"delete_scene",
|
||||
"list_scenes",
|
||||
"load_scene",
|
||||
"parse_elements",
|
||||
"resolve_scene",
|
||||
"save_scene",
|
||||
"scene_path",
|
||||
]
|
||||
|
||||
# Default storage location. Tests may monkeypatch this attribute; every helper
|
||||
# below dereferences it via globals() so the override is picked up live.
|
||||
SCENES_DIR = DATA_DIR / "scenes"
|
||||
|
||||
|
||||
def _lib() -> FragmentLibrary:
|
||||
"""Build a fresh library bound to the current ``SCENES_DIR`` value."""
|
||||
return FragmentLibrary("scenes", base_dir=globals()["SCENES_DIR"])
|
||||
|
||||
|
||||
def scene_path(name: str) -> Path:
|
||||
"""Return the on-disk path for a scene name (without ensuring existence)."""
|
||||
return _lib().path(name)
|
||||
|
||||
|
||||
def save_scene(name: str, elements: list[str]) -> Path:
|
||||
"""Persist a scene's elements to disk and return the file path."""
|
||||
return _lib().save(name, elements)
|
||||
|
||||
|
||||
def load_scene(name: str) -> list[str]:
|
||||
"""Load a scene's elements. Raises ``FileNotFoundError`` if missing."""
|
||||
return _lib().load(name)
|
||||
|
||||
|
||||
def list_scenes() -> list[str]:
|
||||
"""Return sorted scene names."""
|
||||
return _lib().list()
|
||||
|
||||
|
||||
def delete_scene(name: str) -> bool:
|
||||
"""Delete a scene file. Returns True on success, False if missing."""
|
||||
return _lib().delete(name)
|
||||
|
||||
|
||||
def resolve_scene(
|
||||
*,
|
||||
scene: str | None = None,
|
||||
scene_prompt: str | None = None,
|
||||
extra: list[str] | None = None,
|
||||
) -> list[str]:
|
||||
"""Merge a named scene with an inline ``--scene-prompt`` and extras."""
|
||||
return _lib().resolve(name=scene, inline=scene_prompt, extra=extra)
|
||||
@@ -0,0 +1,91 @@
|
||||
"""Tests for the generic FragmentLibrary in tensors.fragments.
|
||||
|
||||
The character/scene-specific modules are thin wrappers over this class; their
|
||||
behavioral coverage lives in test_characters.py and test_scenes.py. Here we
|
||||
focus on cross-kind invariants (kind-aware error messages, name validation,
|
||||
parse helper).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def lib_factory(tmp_path):
|
||||
"""Return a callable that builds an isolated FragmentLibrary per kind."""
|
||||
from tensors.fragments import FragmentLibrary
|
||||
|
||||
def _make(kind: str = "characters"):
|
||||
return FragmentLibrary(kind, base_dir=tmp_path / kind)
|
||||
|
||||
return _make
|
||||
|
||||
|
||||
class TestParseElements:
|
||||
def test_splits_trims_and_dedupes(self):
|
||||
from tensors.fragments import parse_elements
|
||||
|
||||
assert parse_elements("a, b , a, c, b") == ["a", "b", "c"]
|
||||
|
||||
def test_empty_returns_empty(self):
|
||||
from tensors.fragments import parse_elements
|
||||
|
||||
assert parse_elements("") == []
|
||||
assert parse_elements(" , , ") == []
|
||||
|
||||
|
||||
class TestKindAwareMessages:
|
||||
def test_singular_in_error_for_plural_kind(self, lib_factory):
|
||||
lib = lib_factory("scenes")
|
||||
with pytest.raises(ValueError, match=r"Invalid scene name"):
|
||||
lib.save("bad name", ["x"])
|
||||
|
||||
def test_singular_in_error_for_already_singular_kind(self, lib_factory):
|
||||
lib = lib_factory("style") # already singular — should NOT strip the 'e'
|
||||
with pytest.raises(ValueError, match=r"Invalid style name"):
|
||||
lib.save("bad name", ["x"])
|
||||
|
||||
def test_load_missing_uses_kind_singular(self, lib_factory):
|
||||
lib = lib_factory("scenes")
|
||||
with pytest.raises(FileNotFoundError, match=r"Scene 'ghost' not found"):
|
||||
lib.load("ghost")
|
||||
|
||||
|
||||
class TestLibraryConstructor:
|
||||
def test_rejects_invalid_kind(self):
|
||||
from tensors.fragments import FragmentLibrary
|
||||
|
||||
with pytest.raises(ValueError, match=r"Invalid library kind"):
|
||||
FragmentLibrary("bad kind")
|
||||
|
||||
def test_accepts_dotted_alphanumeric_kind(self):
|
||||
from tensors.fragments import FragmentLibrary
|
||||
|
||||
lib = FragmentLibrary("custom_v2.beta")
|
||||
assert lib.kind == "custom_v2.beta"
|
||||
|
||||
|
||||
class TestResolve:
|
||||
def test_only_named(self, lib_factory):
|
||||
lib = lib_factory()
|
||||
lib.save("base", ["x", "y"])
|
||||
assert lib.resolve(name="base") == ["x", "y"]
|
||||
|
||||
def test_only_inline(self, lib_factory):
|
||||
lib = lib_factory()
|
||||
assert lib.resolve(inline="a, b, c") == ["a", "b", "c"]
|
||||
|
||||
def test_named_inline_extras_merge_deduped_in_order(self, lib_factory):
|
||||
lib = lib_factory()
|
||||
lib.save("base", ["a", "b"])
|
||||
assert lib.resolve(name="base", inline="b, c", extra=["c", "d"]) == ["a", "b", "c", "d"]
|
||||
|
||||
def test_no_args_empty(self, lib_factory):
|
||||
lib = lib_factory()
|
||||
assert lib.resolve() == []
|
||||
|
||||
def test_named_missing_raises(self, lib_factory):
|
||||
lib = lib_factory()
|
||||
with pytest.raises(FileNotFoundError):
|
||||
lib.resolve(name="absent")
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Tests for the scene library (tensors.scenes).
|
||||
|
||||
Mirrors test_characters.py — the underlying implementation is shared so the
|
||||
focus here is verifying that scene-specific exports route through correctly
|
||||
and that monkeypatching SCENES_DIR works the same way.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def scene_env(tmp_path, monkeypatch):
|
||||
"""Redirect SCENES_DIR at fresh tmp_path per test."""
|
||||
from tensors import scenes as scene_mod
|
||||
|
||||
scene_dir = tmp_path / "scenes"
|
||||
monkeypatch.setattr(scene_mod, "SCENES_DIR", scene_dir)
|
||||
return scene_mod, scene_dir
|
||||
|
||||
|
||||
class TestSaveLoad:
|
||||
def test_save_creates_dir_and_file(self, scene_env):
|
||||
scene_mod, scene_dir = scene_env
|
||||
path = scene_mod.save_scene("penthouse", ["luxury penthouse", "volumetric lighting"])
|
||||
|
||||
assert path == scene_dir / "penthouse.yml"
|
||||
assert path.is_file()
|
||||
assert path.read_text() == '- "luxury penthouse"\n- "volumetric lighting"\n'
|
||||
|
||||
def test_load_roundtrip(self, scene_env):
|
||||
scene_mod, _ = scene_env
|
||||
elements = ["a", "b with spaces", "c, with, commas", 'd "quotes"']
|
||||
scene_mod.save_scene("mixed", elements)
|
||||
|
||||
assert scene_mod.load_scene("mixed") == elements
|
||||
|
||||
def test_load_missing_raises(self, scene_env):
|
||||
scene_mod, _ = scene_env
|
||||
with pytest.raises(FileNotFoundError, match=r"Scene 'nope' not found"):
|
||||
scene_mod.load_scene("nope")
|
||||
|
||||
|
||||
class TestListAndDelete:
|
||||
def test_list_empty_when_dir_missing(self, scene_env):
|
||||
scene_mod, _ = scene_env
|
||||
assert scene_mod.list_scenes() == []
|
||||
|
||||
def test_list_sorted_names(self, scene_env):
|
||||
scene_mod, _ = scene_env
|
||||
scene_mod.save_scene("zeta", ["x"])
|
||||
scene_mod.save_scene("alpha", ["y"])
|
||||
|
||||
assert scene_mod.list_scenes() == ["alpha", "zeta"]
|
||||
|
||||
def test_delete_existing(self, scene_env):
|
||||
scene_mod, _ = scene_env
|
||||
scene_mod.save_scene("doomed", ["x"])
|
||||
|
||||
assert scene_mod.delete_scene("doomed") is True
|
||||
assert scene_mod.list_scenes() == []
|
||||
|
||||
def test_delete_missing(self, scene_env):
|
||||
scene_mod, _ = scene_env
|
||||
assert scene_mod.delete_scene("ghost") is False
|
||||
|
||||
|
||||
class TestResolveScene:
|
||||
def test_named_plus_inline_dedup(self, scene_env):
|
||||
scene_mod, _ = scene_env
|
||||
scene_mod.save_scene("base", ["a", "b"])
|
||||
|
||||
assert scene_mod.resolve_scene(scene="base", scene_prompt="b, c") == ["a", "b", "c"]
|
||||
|
||||
def test_extra_appends_last(self, scene_env):
|
||||
scene_mod, _ = scene_env
|
||||
|
||||
assert scene_mod.resolve_scene(scene_prompt="a, b", extra=["b", "c"]) == ["a", "b", "c"]
|
||||
|
||||
def test_no_args_empty(self, scene_env):
|
||||
scene_mod, _ = scene_env
|
||||
assert scene_mod.resolve_scene() == []
|
||||
|
||||
|
||||
class TestNameValidation:
|
||||
@pytest.mark.parametrize("bad_name", ["", "foo/bar", "../etc", "with space", "a$b"])
|
||||
def test_invalid_names_rejected(self, scene_env, bad_name):
|
||||
scene_mod, _ = scene_env
|
||||
with pytest.raises(ValueError, match=r"Invalid scene name"):
|
||||
scene_mod.save_scene(bad_name, ["x"])
|
||||
|
||||
def test_scene_path_validates(self, scene_env):
|
||||
scene_mod, _ = scene_env
|
||||
with pytest.raises(ValueError, match=r"Invalid scene name"):
|
||||
scene_mod.scene_path("bad name")
|
||||
Reference in New Issue
Block a user