spritegen-cli 0.1.0__py3-none-any.whl

spritegen/matting.py

@@ -0,0 +1,136 @@
+"""Cutting the background locally, on the GPU, instead of paying for it.
+
+The paid matte is a segmentation model behind an endpoint. So is this one: `rembg` runs
+BiRefNet — the same family the endpoint's `Matting` variant is — through onnxruntime, and
+on a machine with CUDA the difference is seconds and no invoice. On a machine without
+one it still runs, slower, on the CPU.
+
+**This is not the chroma cut.** `imaging.cut_chroma` removes a colour the prompt asked
+for; that is exact where the prompt was obeyed and useless where it was not. This asks a
+model what the subject is, which is what hair, a blade edge and the gap under an arm
+need — and it works on art that was never generated against a key field at all.
+
+`rembg` is an optional extra, for the reason `torch` is: it brings onnxruntime with it,
+and which build works depends on the machine's CUDA. Every other command keeps working
+without it, and the refusal names the install line.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+#: The default weights. BiRefNet general is the same family as the paid endpoint's
+#: `Matting` variant, and gives continuous alpha rather than a hard mask — which is what
+#: hair and a blade edge need.
+DEFAULT_MODEL = "birefnet-general"
+
+INSTALL = (
+    "the local matte needs rembg, which is an optional extra: "
+    'pip install "spritegen-cli[matte]" (and "rembg[gpu]" for CUDA)'
+)
+
+
+def require_rembg():
+    """`rembg`, or an error naming the one command that fixes it."""
+    try:
+        import rembg
+    except ModuleNotFoundError as exc:
+        raise NotImplementedError(INSTALL) from exc
+    return rembg
+
+
+def available() -> bool:
+    """Whether the local matte could run. Used to report, never to decide silently."""
+    try:
+        require_rembg()
+    except NotImplementedError:
+        return False
+    return True
+
+
+def providers() -> list[str]:
+    """What onnxruntime says it can run on, in its own order of preference.
+
+    Reported by the command rather than assumed: a `rembg` installed without the GPU
+    build runs on the CPU and looks like a hang on a board of six frames.
+    """
+    try:
+        import onnxruntime
+    except ModuleNotFoundError:
+        return []
+    return list(onnxruntime.get_available_providers())
+
+
+def on_gpu() -> bool:
+    return any(name.startswith(("CUDA", "ROCM")) for name in providers())
+
+
+def known_models() -> tuple[str, ...]:
+    """The model names this rembg has, or empty if it will not say."""
+    require_rembg()
+    try:
+        from rembg.sessions import sessions_names
+    except ImportError:
+        return ()
+    return tuple(sessions_names)
+
+
+def check_model(model: str) -> str:
+    """`model` if this rembg has it, and an error naming what it does have if not.
+
+    Checked here rather than left to `new_session`. The name can arrive from a `.env`
+    two directories up as easily as from the command line, it is an unconstrained string
+    either way, and what a third-party lookup does with a name it does not recognise is
+    not this tool's decision to inherit. An installed rembg that will not enumerate its
+    models gets the benefit of the doubt rather than a refusal it cannot satisfy.
+    """
+    known = known_models()
+    if known and model not in known:
+        raise ValueError(f"unknown matte model {model!r}; this rembg has: {', '.join(known)}")
+    return model
+
+
+def session(model: str | None = None):
+    """One rembg session, reused across every frame of a board.
+
+    Built once and passed in, because loading the weights is most of the cost: matting
+    six frames with six sessions pays for the model six times.
+    """
+    from . import settings
+
+    rembg = require_rembg()
+    return rembg.new_session(check_model(model or settings.load().matte_model or DEFAULT_MODEL))
+
+
+def cut(image, *, current=None, refine: bool = True, model: str | None = None):
+    """One image with its background gone. Returns RGBA.
+
+    `refine` is the edge decontamination, and it is what kills the coloured halo — the
+    same decision the paid stage makes with `--no-refine`, named the same way so the two
+    backends do not need two vocabularies. Turn it off when the figure is losing its own
+    outline rather than gaining a fringe.
+    """
+    rembg = require_rembg()
+    return rembg.remove(
+        image.convert("RGBA"),
+        session=current if current is not None else session(model),
+        post_process_mask=True,
+        decontaminate=refine,
+    )
+
+
+def cut_file(source: Path, target: Path, *, current=None, refine: bool = True) -> dict:
+    """Cut one file to another, and report what the alpha came out like.
+
+    The measurement is `imaging.measure_alpha`, the same one the paid stage reports, and
+    it is how a matte that worked is told from one that handed back the picture untouched:
+    an image with no transparent pixel in it did not get cut, whatever anything claimed.
+    """
+    from PIL import Image
+
+    from .imaging import measure_alpha
+
+    result = cut(Image.open(source), current=current, refine=refine)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    result.save(target)
+    return measure_alpha(target)

spritegen/migrate.py

@@ -0,0 +1,162 @@
+"""Move a workspace from one on-disk layout to the next — R1.5.
+
+Layout 1 was a directory per stage: `anchor/`, `board/`, `matte/`, `row/`. Layout 2 is a
+directory per artifact kind, with art directions under `sheet/`. The names changed
+because the two questions they were answering had come apart — see the spec.
+
+**Nothing migrates on its own.** A command that reorganises somebody's disk in the middle
+of another command is exactly what nobody wants to find out about afterwards, and these
+are directories that paid calls produced. So this is its own command, it prints what it
+would do before it does anything, and `--dry-run` is how you see that without the move.
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from . import workspace
+
+#: The old directory, and the artifact directory it becomes. `row` carries a variant:
+#: layout 1 had one row per sprite because it had no idea an art direction existed, and
+#: the one it has is whatever was closed — `as-is` is the honest name for that.
+MOVES: tuple[tuple[str, str], ...] = (
+    ("anchor", "anchor"),
+    ("board", "frames"),
+    ("matte", "cutout"),
+    ("row", "sheet/as-is"),
+)
+
+#: How a layout-1 stage entry named what it produced, and the artifact key it is now.
+PRODUCED: dict[str, str] = {
+    "anchor": "anchor",
+    "board": "frames",
+    "matte": "cutout",
+    "row": "sheet:as-is",
+}
+
+
+def planned(asset: str) -> list[tuple[Path, Path]]:
+    """The moves this asset needs, in order. Empty when there is nothing to do."""
+    directory = workspace.asset_dir(asset)
+    moves = []
+    for old, new in MOVES:
+        source = directory / old
+        target = directory / new
+        if source.is_dir() and source != target and not target.exists():
+            moves.append((source, target))
+    return moves
+
+
+def blocked(asset: str) -> list[tuple[Path, Path]]:
+    """The moves this asset needs and cannot make, because something is already there.
+
+    Rare, and it has to stop the migration rather than be skipped. Rewriting the state
+    while the files stay put would point `dir` at a directory holding somebody else's
+    contents: the real output is still on disk, orphaned, invisible to `show` and to
+    every stage that reads its input through the state file — and the wrong directory
+    is what the next paid stage would read.
+    """
+    directory = workspace.asset_dir(asset)
+    return [
+        (directory / old, directory / new)
+        for old, new in MOVES
+        if (directory / old).is_dir() and old != new and (directory / new).exists()
+    ]
+
+
+def rewrite_state(state: workspace.State) -> workspace.State:
+    """Point the state at the new directories and fill in `artifacts` — R1.4.
+
+    `artifacts` is built from what is on disk rather than from the old state, because the
+    old state never listed files. That makes this the one place the two records are
+    allowed to be derived from each other, and it happens once.
+    """
+    for name, entry in state.stages.items():
+        produced = entry.get("produced")
+        moved = PRODUCED.get(produced, produced)
+        if moved is None:
+            continue
+        entry["produced"] = moved
+        entry["dir"] = moved.replace(":", "/")
+
+        # A layout-1 `produced` that is none of the four falls through unchanged, and it
+        # came out of the same editable file every other recorded path does.
+        target = workspace.inside(state.name, entry["dir"])
+        state.artifacts[moved] = {
+            "dir": entry["dir"],
+            "files": sorted(path.name for path in target.glob("*") if path.is_file())
+            if target.is_dir()
+            else [],
+            "stage": name,
+            "prompt": entry.get("prompt"),
+        }
+    state.layout = workspace.LAYOUT
+    return state
+
+
+class Blocked(workspace.StageRefused):
+    """A migration that cannot run without losing track of something."""
+
+
+def run(asset: str, *, dry_run: bool) -> list[tuple[Path, Path]]:
+    """Migrate one asset. Returns the moves, whether or not they were made."""
+    state = workspace.load(asset)
+    stuck = blocked(asset)
+    if stuck:
+        where = ", ".join(target.name for _, target in stuck)
+        raise Blocked(
+            f"{asset}: {where} already exists, so the old directory cannot move there. "
+            f"Nothing was changed. Move or remove it, then migrate again."
+        )
+    moves = planned(asset)
+    if dry_run:
+        return moves
+    for source, target in moves:
+        target.parent.mkdir(parents=True, exist_ok=True)
+        source.rename(target)
+    workspace.save(rewrite_state(state))
+    return moves
+
+
+def cmd_migrate(args: argparse.Namespace) -> int:
+    """Migrate one asset, or say what every outdated one would need.
+
+    ASCII only, for the reason `workspace.cmd_status` is.
+    """
+    names = [workspace.check_name(args.name)] if args.name else workspace.known()
+    if not names:
+        print(f"no assets under {workspace.assets_root()}")
+        return 0
+
+    touched = refused = 0
+    for name in names:
+        state = workspace.load(name)
+        if not state.outdated():
+            if args.name:
+                print(f"{name}  already layout {state.layout}; nothing to migrate")
+            continue
+        touched += 1
+        try:
+            moves = run(name, dry_run=args.dry_run)
+        except Blocked:
+            # Asked about one asset, a refusal is the answer and gets the exit code a
+            # refusal gets everywhere else. Asked about all of them, one that cannot move
+            # is not a reason to fail the others: it is named, and the batch goes on.
+            if args.name:
+                raise
+            refused += 1
+            print(f"{name}  refused: it cannot move onto a directory that is already there")
+            continue
+        verb = "would move" if args.dry_run else "moved"
+        print(f"{name}  layout {state.layout} -> {workspace.LAYOUT}")
+        for source, target in moves:
+            print(f"  {verb} {source.name} -> {target.relative_to(workspace.asset_dir(name))}")
+        if not moves:
+            print("  nothing on disk to move; the state file is rewritten")
+
+    if args.dry_run and touched > refused:
+        print("nothing was moved; run again without --dry-run")
+    if refused:
+        print(f"{refused} asset(s) refused; nothing was changed for them")
+    return 0

spritegen/prompts.py

@@ -0,0 +1,96 @@
+"""The prompt that produced an artifact, kept where the artifact is — R2.1 to R2.4.
+
+An image survives its own reason. The prompt arrives through `--prompt-file`, from
+somewhere outside the workspace, or through `--prompt`, from a shell that forgets it the
+moment the command returns — and once the paid call is over, what is left on disk is a
+picture nobody can regenerate or amend. Under autonomy that is worse still: the file that
+was pointed at can be edited or deleted afterwards, so a *reference* to it is not a
+record of anything.
+
+So the text is copied in, and it is copied **before the call**: a prompt that only gets
+written when the call succeeds is missing exactly when it is most wanted.
+
+**Versioned by content, not by time.** Sending the same prompt again is not a new version
+of it — nothing changed, and a run that produced the same words should point at the words
+it produced. Editing one is, and the old text stays: `anchor.v1.md` is how you find out
+what the first attempt actually asked for.
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+#: `<kind>.v<N>.md`, and nothing else in the prompts directory is one.
+PATTERN = re.compile(r"^(?P<kind>.+)\.v(?P<version>\d+)\.md$")
+
+
+def directory(asset_dir: Path) -> Path:
+    """Where this sprite's prompts live."""
+    return asset_dir / "prompts"
+
+
+def name(kind: str, version: int) -> str:
+    return f"{kind}.v{version}.md"
+
+
+def versions(asset_dir: Path, kind: str) -> list[int]:
+    """Every version of `kind` on disk, oldest first."""
+    found = []
+    for path in directory(asset_dir).glob(f"{kind}.v*.md"):
+        match = PATTERN.match(path.name)
+        if match and match.group("kind") == kind:
+            found.append(int(match.group("version")))
+    return sorted(found)
+
+
+def read(asset_dir: Path, kind: str, version: int) -> str:
+    """One stored prompt, verbatim."""
+    path = directory(asset_dir) / name(kind, version)
+    if not path.is_file():
+        raise FileNotFoundError(f"no prompt {kind} v{version} in {directory(asset_dir)}")
+    return path.read_text(encoding="utf-8")
+
+
+def store(asset_dir: Path, kind: str, text: str) -> str:
+    """Keep `text` as a version of `kind`, and return the filename to record — R2.1, R2.2.
+
+    Unchanged text returns the version it already is. That is the whole of "versioned by
+    content": four candidates from one prompt cite one version, and the fifth run after an
+    edit cites a new one.
+
+    Compared against the newest version only. A prompt that was changed and then changed
+    back is a third version, not a re-use of the first — the history is a sequence of what
+    was asked, and collapsing a revert would hide that the middle one ever happened.
+    """
+    text = text.rstrip() + "\n"
+    existing = versions(asset_dir, kind)
+    if existing and read(asset_dir, kind, existing[-1]) == text:
+        return name(kind, existing[-1])
+
+    version = (existing[-1] + 1) if existing else 1
+    target = directory(asset_dir)
+    target.mkdir(parents=True, exist_ok=True)
+    (target / name(kind, version)).write_text(text, encoding="utf-8")
+    return name(kind, version)
+
+
+def resolve(prompt: str | None, prompt_file: str | None, what: str) -> str:
+    """The prompt text, from one source or the other and never from both.
+
+    Shared by every stage that takes a prompt, because they all had the same six lines
+    and one of them would eventually have drifted.
+    """
+    if prompt and prompt_file:
+        raise ValueError("give --prompt or --prompt-file, not both")
+    if prompt:
+        return prompt
+    if not prompt_file:
+        raise ValueError(f"{what} needs --prompt or --prompt-file")
+    path = Path(prompt_file)
+    if not path.is_file():
+        raise FileNotFoundError(f"no prompt file at {path}")
+    text = path.read_text(encoding="utf-8").strip()
+    if not text:
+        raise ValueError(f"{path} is empty")
+    return text

spritegen/rrdb.py

@@ -0,0 +1,91 @@
+"""RRDBNet, the network Real-ESRGAN's weights are for.
+
+Ported from an `upscale_atlas.py` that reimplements it rather than depending on
+`basicsr` — the upstream package pins an old torchvision and does not
+install alongside a current torch. What was two classes there is one here: `RRDBNet6B`
+differed from `RRDBNet` in the default block count and in nothing else.
+
+**This module imports torch at the top, and so nothing may import it at package level.**
+`upscale` resolves it inside the function that builds the network, which is what keeps
+`status` and `cost` answering on a machine with no torch on it.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as functional
+from torch import nn
+
+
+class ResidualDenseBlock(nn.Module):
+    """Five convolutions, each seeing everything before it, added back at 0.2."""
+
+    def __init__(self, nf: int = 64, gc: int = 32):
+        super().__init__()
+        self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1)
+        self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1)
+        self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1)
+        self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1)
+        self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1)
+        self.lrelu = nn.LeakyReLU(0.2, True)
+
+    def forward(self, x):
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        return x5 * 0.2 + x
+
+
+class RRDB(nn.Module):
+    """Three of those in a row, added back at 0.2 again."""
+
+    def __init__(self, nf: int = 64, gc: int = 32):
+        super().__init__()
+        self.rdb1 = ResidualDenseBlock(nf, gc)
+        self.rdb2 = ResidualDenseBlock(nf, gc)
+        self.rdb3 = ResidualDenseBlock(nf, gc)
+
+    def forward(self, x):
+        out = self.rdb3(self.rdb2(self.rdb1(x)))
+        return out * 0.2 + x
+
+
+class RRDBNet(nn.Module):
+    """The whole network. `nb` is 23 for the general weights and 6 for the anime ones.
+
+    The upsampling is two nearest-neighbour doublings with a convolution after each, so
+    the network is 4x and only 4x. Any other factor is a resize of what it returned.
+    """
+
+    def __init__(
+        self,
+        in_nc: int = 3,
+        out_nc: int = 3,
+        nf: int = 64,
+        nb: int = 23,
+        gc: int = 32,
+        scale: int = 4,
+    ):
+        super().__init__()
+        self.scale = scale
+        self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1)
+        self.body = nn.Sequential(*[RRDB(nf, gc) for _ in range(nb)])
+        self.conv_body = nn.Conv2d(nf, nf, 3, 1, 1)
+        self.conv_up1 = nn.Conv2d(nf, nf, 3, 1, 1)
+        self.conv_up2 = nn.Conv2d(nf, nf, 3, 1, 1)
+        self.conv_hr = nn.Conv2d(nf, nf, 3, 1, 1)
+        self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1)
+        self.lrelu = nn.LeakyReLU(0.2, True)
+
+    def forward(self, x):
+        feat = self.conv_first(x)
+        feat = feat + self.conv_body(self.body(feat))
+        feat = self.lrelu(
+            self.conv_up1(functional.interpolate(feat, scale_factor=2, mode="nearest"))
+        )
+        feat = self.lrelu(
+            self.conv_up2(functional.interpolate(feat, scale_factor=2, mode="nearest"))
+        )
+        return self.conv_last(self.lrelu(self.conv_hr(feat)))

spritegen/settings.py

@@ -0,0 +1,127 @@
+"""Every setting this tool has, in one place — the environment, a `.env`, or a default.
+
+Before this, configuration was four unrelated habits: `FAL_KEY` read out of `os.environ`
+inside `fal`, `SPRITEGEN_ASSETS` read out of it inside `workspace`, a `.env` loaded by the
+CLI into the process, and a handful of module constants nobody could override without
+editing the source. Each was defensible alone; together they meant no answer to "what can
+be configured", and two of them could disagree about whether a variable had been set.
+
+**Not cached, deliberately.** A CLI runs once and exits, so the saving would be
+microseconds; what it would cost is that a variable set after the first read stops taking
+effect, which is precisely what a test does when it points the assets somewhere else. Read
+it where you need it.
+
+The import stays out of module scope everywhere it is used, for the same reason `fal` and
+`Pillow` do: `status` and `cost` must answer without paying for anything they do not need.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Literal
+
+from pydantic import AliasChoices, Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+#: Where a downloaded model weight goes. Outside any workspace: it is the same file for
+#: every project, it is 64 MB or more, and it can always be fetched again.
+DEFAULT_CACHE = Path.home() / ".cache" / "spritegen"
+
+
+def search_from(start: Path | None = None) -> tuple[Path, ...]:
+    """The directories a search may look in: the start, and its parents up to a ceiling.
+
+    **The ceiling is the point of this.** Searching upward is right — a command is typed
+    from wherever is convenient, not from the top of the thing it concerns — but upward
+    with no stopping condition means the filesystem root, and everything between here and
+    there belongs to somebody else. A `.env` in a shared parent, on a mounted drive, or in
+    a CI runner's workspace root would be read as this tool's own configuration, and it
+    can set the credential, where the assets are, and which hosts a download may come
+    from.
+
+    So the walk stops at the first repository boundary, or at the home directory,
+    whichever comes first. Neither found — the working directory is somewhere with no
+    repository above it and outside the home directory — and the search is the working
+    directory alone, which is the conservative answer rather than the far-reaching one.
+    """
+    start = Path.cwd() if start is None else start
+    home = Path.home()
+    chain: list[Path] = []
+    for base in (start, *start.parents):
+        chain.append(base)
+        if (base / ".git").exists() or base == home:
+            return tuple(chain)
+    return (start,)
+
+
+def dotenv(start: Path | None = None) -> Path | None:
+    """The nearest `.env` at or above the working directory, within the search ceiling.
+
+    An exported variable still wins over the file — that is pydantic-settings' own order,
+    and it is what makes a one-off `FAL_KEY=... spritegen anchor ...` work while a stale
+    key sits in the file.
+    """
+    for base in search_from(start):
+        candidate = base / ".env"
+        if candidate.is_file():
+            return candidate
+    return None
+
+
+class Settings(BaseSettings):
+    """What can be configured, and what it is when nobody says.
+
+    Every field is `SPRITEGEN_<NAME>` in the environment. `fal_key` also answers to bare
+    `FAL_KEY`, which is what fal's own documentation and every existing `.env` use — an
+    alias costs one line and breaking them would cost everyone a rewrite.
+    """
+
+    model_config = SettingsConfigDict(
+        env_prefix="SPRITEGEN_",
+        env_file=None,
+        extra="ignore",
+        case_sensitive=False,
+    )
+
+    #: Where the sprites live. Unset means discover it — see `workspace.assets_root`.
+    assets: Path | None = None
+
+    #: The fal credential. Absent is not an error until something tries to spend.
+    fal_key: str | None = Field(
+        default=None, validation_alias=AliasChoices("SPRITEGEN_FAL_KEY", "FAL_KEY", "fal_key")
+    )
+
+    #: Downloaded model weights, shared across every project on this machine.
+    cache: Path = DEFAULT_CACHE
+
+    #: The sheet cell, in pixels. A project with another one can say so once here rather
+    #: than passing `--cell` to every command.
+    cell: int = 166
+
+    #: Which Real-ESRGAN weights the local upscale uses, and where it runs. `auto` means
+    #: CUDA where there is one — reported by the command, never assumed silently.
+    upscale_model: Literal["anime", "general"] = "anime"
+    upscale_device: Literal["auto", "cuda", "cpu"] = "auto"
+
+    #: Whether the background comes off locally or through a paid endpoint. Local is the
+    #: default because it costs nothing and this machine has the GPU for it; the paid one
+    #: is there for a machine that does not.
+    matte_backend: Literal["local", "fal"] = "local"
+    matte_model: str = "birefnet-general"
+
+    #: Extra hosts a downloaded URL may come from, on top of the ones built in. A
+    #: comma-separated list, because that is what an environment variable can carry.
+    allowed_hosts: str = ""
+
+    def extra_hosts(self) -> tuple[str, ...]:
+        return tuple(part.strip() for part in self.allowed_hosts.split(",") if part.strip())
+
+
+def load(**overrides) -> Settings:
+    """The settings as they are right now.
+
+    `overrides` is for a caller that already has the answer — a test, or a command-line
+    option that beats the environment. It is not a way to mutate anything: this builds a
+    new object every time and nothing is stored.
+    """
+    return Settings(_env_file=dotenv(), **overrides)