diff --git a/tensors/__init__.py b/tensors/__init__.py index 68a604f..8ca8bd6 100644 --- a/tensors/__init__.py +++ b/tensors/__init__.py @@ -1,6 +1,6 @@ """tsr: Read safetensor metadata, search and download CivitAI models.""" -__version__ = "0.1.22" +__version__ = "0.1.23" from tensors.cli import main from tensors.config import ( diff --git a/tensors/characters.py b/tensors/characters.py new file mode 100644 index 0000000..1efe03f --- /dev/null +++ b/tensors/characters.py @@ -0,0 +1,159 @@ +"""Character library: named lists of prompt elements, stored as YAML. + +Characters live in ``~/.local/share/tensors/characters/.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``. + +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. +""" + +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 + +# Storage location for character YAML files. +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 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 + + +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 + + +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 + + +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()) + + +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 + + +def resolve_character( + *, + character: str | None = None, + 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 diff --git a/tensors/cli.py b/tensors/cli.py index 9a00159..a594d41 100644 --- a/tensors/cli.py +++ b/tensors/cli.py @@ -780,6 +780,14 @@ def generate( # noqa: PLR0915 ] = None, no_quality: Annotated[bool, typer.Option("--no-quality", help="Disable auto quality tags")] = False, no_negative: Annotated[bool, typer.Option("--no-negative", help="Disable auto negative prompt")] = False, + character: Annotated[ + str | None, + typer.Option("-C", "--character", help="Saved character name (loaded from ~/.local/share/tensors/characters/)"), + ] = None, + character_prompt: Annotated[ + str | None, + typer.Option("--character-prompt", help='Inline character fragment, comma-separated (e.g. "blond hair, blue eyes")'), + ] = None, family: Annotated[ str | None, typer.Option( @@ -892,6 +900,21 @@ def generate( # noqa: PLR0915 no_quality = bool(mapped["no_quality"]) if "no_negative" in mapped and "no_negative" not in explicit: no_negative = bool(mapped["no_negative"]) + if "character" in mapped and "character" not in explicit: + # Accept either a saved-name string or an already-resolved list/tuple + # (templates may carry the resolved list inline). For lists we stage + # them into character_prompt by joining with commas so the existing + # CLI splitting/dedup path applies uniformly. + val = mapped["character"] + if isinstance(val, str): + character = val + elif isinstance(val, (list, tuple)): + 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()) + ) if "rating" in mapped and "rating" not in explicit: rating = mapped["rating"] @@ -919,6 +942,8 @@ def generate( # noqa: PLR0915 rating=rating, no_quality=no_quality, no_negative=no_negative, + character=character, + character_prompt=character_prompt, family=family, output=output, remote=remote, @@ -1007,6 +1032,8 @@ def _run_generation( # noqa: PLR0915 rating: str | None = None, no_quality: bool = False, no_negative: bool = False, + character: str | None = None, + character_prompt: str | None = None, family: str | None = None, output: Path | None = None, remote: str | None = None, @@ -1058,6 +1085,29 @@ def _run_generation( # noqa: PLR0915 if not no_quality and family_defaults.get("quality_prefix"): prompt_parts.append(family_defaults["quality_prefix"]) + # Resolve character (named lookup + inline --character-prompt, merged + deduped) + character_elements: list[str] = [] + if character or character_prompt: + from tensors.characters import resolve_character # noqa: PLC0415 + + try: + character_elements = resolve_character(character=character, character_prompt=character_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 character_elements: + prompt_parts.extend(character_elements) + if not json_output: + origin = f"'{character}'" if character else "inline" + console.print( + f"[dim]Character ({origin}, {len(character_elements)} elements): " + f"{', '.join(character_elements)}[/dim]" + ) + # Add rating tag based on model family (Pony/Illustrious) if rating: from tensors.config import get_rating_tag # noqa: PLC0415 @@ -1301,6 +1351,8 @@ _STYLE_SWEEP_TEMPLATE_KEYS = { "orientation", "no_quality", "no_negative", + "character", + "character_prompt", "rating", "family", "remote", @@ -1631,6 +1683,25 @@ 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()) + ) + # Common kwargs for every _run_generation call — extracted from the # template once, reused across sequential and parallel paths. base_gen_kwargs: dict[str, Any] = { @@ -1652,6 +1723,8 @@ def style_sweep( # noqa: PLR0915 "rating": _t("rating"), "no_quality": bool(_t("no_quality", default=False)), "no_negative": bool(_t("no_negative", default=False)), + "character": char_name, + "character_prompt": char_inline, "family": _t("family"), "remote": gen_remote, "json_output": False, @@ -1804,6 +1877,17 @@ def template( lora_strength: Annotated[float, typer.Option("--lora-strength", help="LoRA strength")] = 0.8, orientation: Annotated[str, typer.Option("-O", "--orientation", help="Resolution: square, portrait, landscape")] = "square", rating: Annotated[str | None, typer.Option("--rating", "-R", help="Content rating: safe, questionable, explicit")] = None, + character: Annotated[ + str | None, + typer.Option("-C", "--character", help="Saved character name (resolved into the `character` list field)"), + ] = None, + character_prompt: Annotated[ + str | None, + typer.Option( + "--character-prompt", + help='Inline character fragment, comma-separated (merged with --character into `character`)', + ), + ] = 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. @@ -1811,11 +1895,17 @@ 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). + 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 generate --input "$(tsr template -m ponyRealism_V22.safetensors)" "a portrait" """ from tensors.config import get_model_generation_defaults, resolve_orientation # noqa: PLC0415 @@ -1869,6 +1959,28 @@ 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). + if character or character_prompt: + from tensors.characters import resolve_character # noqa: PLC0415 + + try: + resolved = resolve_character(character=character, character_prompt=character_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: + tpl["character"] = resolved + if character: + tpl["_character_name"] = character + # Add metadata (not used by generate, but informational) tpl["_family"] = family or "unknown" if base_model_str: @@ -2352,6 +2464,129 @@ def hf_download( raise typer.Exit(1) +# ============================================================================= +# Character Commands +# ============================================================================= +# Characters are named, comma-split prompt fragments stored as YAML lists in +# ~/.local/share/tensors/characters/.yml. They are injected into the +# positive prompt by `tsr generate --character ` (or inline via +# `--character-prompt "elem1, elem2"`). + +character_app = typer.Typer( + name="character", + help="Manage saved character prompts (~/.local/share/tensors/characters/).", + no_args_is_help=True, +) +app.add_typer(character_app) + + +@character_app.command("save") +def character_save( + elements: Annotated[str, typer.Argument(help='Comma-separated prompt elements (e.g. "blond hair, blue eyes")')], + name: Annotated[str, typer.Option("-o", "--output", help="Character name (used as filename)")], + json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False, +) -> None: + """Save a character as a YAML list of prompt elements. + + Examples: + tsr character save -o cassie_cage "blond hair, broad chin, skin imperfections" + tsr character save -o elvira "long black hair, pale skin, gothic dress" + """ + from tensors.characters import parse_elements, save_character # 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_character(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}) + return + + console.print(f"[green]Saved character '{name}' ({len(parsed)} elements):[/green] {path}") + for elem in parsed: + console.print(f" • {elem}") + + +@character_app.command("list") +def character_list( + json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False, +) -> None: + """List saved characters.""" + from tensors.characters import CHARACTERS_DIR, list_characters # noqa: PLC0415 + + names = list_characters() + if json_output: + console.print_json(data={"dir": str(CHARACTERS_DIR), "characters": names}) + return + + if not names: + console.print(f"[yellow]No characters saved in {CHARACTERS_DIR}.[/yellow]") + console.print("[dim]Create one with: tsr character save -o \"elem1, elem2\"[/dim]") + return + + console.print(f"[bold]Characters[/bold] ({len(names)}) [dim]in {CHARACTERS_DIR}[/dim]") + for n in names: + console.print(f" • {n}") + + +@character_app.command("show") +def character_show( + name: Annotated[str, typer.Argument(help="Character name")], + json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False, +) -> None: + """Show a character's elements.""" + from tensors.characters import character_path, load_character # noqa: PLC0415 + + try: + elements = load_character(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(character_path(name)), "elements": elements}) + return + + console.print(f"[bold]{name}[/bold] [dim]({character_path(name)})[/dim]") + for elem in elements: + console.print(f" • {elem}") + + +@character_app.command("delete") +def character_delete( + name: Annotated[str, typer.Argument(help="Character name")], + json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False, +) -> None: + """Delete a saved character.""" + from tensors.characters import delete_character # noqa: PLC0415 + + try: + deleted = delete_character(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 character '{name}'[/green]") + else: + console.print(f"[yellow]Character '{name}' does not exist[/yellow]") + raise typer.Exit(1) + + # ============================================================================= # ComfyUI Commands # ============================================================================= diff --git a/tests/test_characters.py b/tests/test_characters.py new file mode 100644 index 0000000..28fcebf --- /dev/null +++ b/tests/test_characters.py @@ -0,0 +1,194 @@ +"""Tests for the character library (tensors.characters).""" + +from __future__ import annotations + +import pytest + + +@pytest.fixture +def char_env(tmp_path, monkeypatch): + """Redirect CHARACTERS_DIR (and DATA_DIR) at fresh tmp_path per test. + + The module captures CHARACTERS_DIR at import time, so we monkeypatch the + attribute directly rather than rely on env vars (config.DATA_DIR is also + captured at import time). + """ + from tensors import characters as char_mod + + char_dir = tmp_path / "characters" + monkeypatch.setattr(char_mod, "CHARACTERS_DIR", char_dir) + return char_mod, char_dir + + +class TestParseElements: + def test_splits_on_commas_and_trims(self): + from tensors.characters import parse_elements + + assert parse_elements("a, b , c") == ["a", "b", "c"] + + def test_drops_empty_and_duplicates(self): + from tensors.characters import parse_elements + + assert parse_elements("a, , a, b, b , c") == ["a", "b", "c"] + + def test_preserves_internal_spacing(self): + from tensors.characters import parse_elements + + assert parse_elements("blond hair, blue eyes") == ["blond hair", "blue eyes"] + + def test_empty_input(self): + from tensors.characters import parse_elements + + assert parse_elements("") == [] + assert parse_elements(" , , ") == [] + + +class TestSaveLoad: + def test_save_creates_dir_and_file(self, char_env): + char_mod, char_dir = char_env + path = char_mod.save_character("cassie", ["blond hair", "broad chin"]) + + assert path == char_dir / "cassie.yml" + assert path.is_file() + assert path.read_text() == '- "blond hair"\n- "broad chin"\n' + + def test_load_roundtrip(self, char_env): + char_mod, _ = char_env + elements = ["a", "b with spaces", "c, with, commas", 'd "quotes"'] + char_mod.save_character("mixed", elements) + + assert char_mod.load_character("mixed") == elements + + def test_load_missing_raises(self, char_env): + char_mod, _ = char_env + with pytest.raises(FileNotFoundError): + char_mod.load_character("nope") + + def test_load_tolerates_hand_edited_single_quoted_yaml(self, char_env): + char_mod, char_dir = char_env + char_dir.mkdir(parents=True) + (char_dir / "manual.yml").write_text("- 'foo'\n- 'it''s'\n- bare\n# comment\n\n") + + assert char_mod.load_character("manual") == ["foo", "it's", "bare"] + + def test_save_overwrites(self, char_env): + char_mod, _ = char_env + char_mod.save_character("c", ["a"]) + char_mod.save_character("c", ["b", "c"]) + + assert char_mod.load_character("c") == ["b", "c"] + + def test_save_empty_list_writes_empty_file(self, char_env): + char_mod, _char_dir = char_env + path = char_mod.save_character("empty", []) + + assert path.is_file() + assert path.read_text() == "" + assert char_mod.load_character("empty") == [] + + +class TestListAndDelete: + def test_list_empty_when_dir_missing(self, char_env): + char_mod, _ = char_env + assert char_mod.list_characters() == [] + + def test_list_sorted_names(self, char_env): + char_mod, _ = char_env + char_mod.save_character("zeta", ["x"]) + char_mod.save_character("alpha", ["y"]) + char_mod.save_character("mu", ["z"]) + + assert char_mod.list_characters() == ["alpha", "mu", "zeta"] + + def test_delete_existing_returns_true(self, char_env): + char_mod, _ = char_env + char_mod.save_character("doomed", ["x"]) + + assert char_mod.delete_character("doomed") is True + assert char_mod.list_characters() == [] + + def test_delete_missing_returns_false(self, char_env): + char_mod, _ = char_env + assert char_mod.delete_character("ghost") is False + + +class TestNameValidation: + @pytest.mark.parametrize("bad_name", ["", "foo/bar", "../etc", "with space", "name.yml/extra", "a$b"]) + def test_invalid_names_rejected(self, char_env, bad_name): + char_mod, _ = char_env + with pytest.raises(ValueError, match=r"Invalid character name"): + char_mod.save_character(bad_name, ["x"]) + + @pytest.mark.parametrize("good_name", ["a", "foo_bar", "foo-bar", "foo.bar", "Mixed123", "v1.2_alpha-beta"]) + def test_valid_names_accepted(self, char_env, good_name): + char_mod, _ = char_env + path = char_mod.save_character(good_name, ["x"]) + assert path.is_file() + + def test_load_validates_name(self, char_env): + char_mod, _ = char_env + with pytest.raises(ValueError, match=r"Invalid character name"): + char_mod.load_character("../escape") + + def test_delete_validates_name(self, char_env): + char_mod, _ = char_env + with pytest.raises(ValueError, match=r"Invalid character name"): + char_mod.delete_character("../escape") + + +class TestResolveCharacter: + def test_named_only(self, char_env): + char_mod, _ = char_env + char_mod.save_character("base", ["x", "y"]) + + assert char_mod.resolve_character(character="base") == ["x", "y"] + + def test_inline_only(self, char_env): + char_mod, _ = char_env + + assert char_mod.resolve_character(character_prompt="a, b, c") == ["a", "b", "c"] + + def test_extra_only(self, char_env): + char_mod, _ = char_env + + assert char_mod.resolve_character(extra=["a", "b"]) == ["a", "b"] + + def test_named_plus_inline_dedup_preserves_order(self, char_env): + char_mod, _ = char_env + char_mod.save_character("base", ["a", "b"]) + + # 'b' is shared — should appear once, in named position + assert char_mod.resolve_character(character="base", character_prompt="b, c") == ["a", "b", "c"] + + def test_all_three_sources_merged(self, char_env): + char_mod, _ = char_env + char_mod.save_character("base", ["a", "b"]) + + result = char_mod.resolve_character( + character="base", + character_prompt="c, d", + extra=["d", "e"], + ) + assert result == ["a", "b", "c", "d", "e"] + + def test_no_args_returns_empty(self, char_env): + char_mod, _ = char_env + assert char_mod.resolve_character() == [] + + def test_missing_named_raises(self, char_env): + char_mod, _ = char_env + with pytest.raises(FileNotFoundError): + char_mod.resolve_character(character="absent") + + +class TestCharacterPath: + def test_path_does_not_require_existence(self, char_env): + char_mod, char_dir = char_env + + # character_path is pure — doesn't touch disk + assert char_mod.character_path("ghost") == char_dir / "ghost.yml" + + def test_path_validates_name(self, char_env): + char_mod, _ = char_env + with pytest.raises(ValueError): + char_mod.character_path("bad name")