spritegen-cli 0.1.0__py3-none-any.whl

spritegen/stages/__init__.py

@@ -0,0 +1,490 @@
+"""The stage registry — what exists, and what each one needs before it.
+
+The pipeline is not a queue. `motion` and `video` do the same job by different routes:
+one transfers movement from a driving clip, the other animates the anchor directly.
+Both hand over a board, and `matte` takes either. That is why `requires` means *any
+one of these* rather than *this one*.
+
+Order is checked by `workspace`, reading from here. Each stage is implemented by the
+module of the same name beside this file, resolved only when it is about to run: the
+registry describes the pipeline without importing Pillow, numpy or fal_client.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from importlib import import_module
+from types import ModuleType
+
+
+@dataclass(frozen=True)
+class Option:
+    """One command-line option a stage takes.
+
+    Declared here rather than by the stage module so `--help` is complete without
+    importing five modules and their image libraries. The stage reads the value off the
+    parsed namespace; nothing about how it is parsed is the stage's business.
+    """
+
+    flags: tuple[str, ...]
+    help: str
+    kind: str = "str"
+    default: object = None
+    choices: tuple = ()
+    action: str | None = None
+
+
+@dataclass(frozen=True)
+class Stage:
+    """One stage of the pipeline.
+
+    `requires` names the stages of which **one** is enough to be complete. Empty marks
+    the root.
+
+    `produces` is the **kind** of artifact, which is also the directory it goes in. Those
+    were one word for as long as a stage made one thing; they are not the same word any
+    more. `anchor` makes box art, an anchor or an icon depending on what it was asked
+    for, and `board` closes the same frames into as many art directions as anyone wants
+    — so the kind, and the directory, follow the options rather than the stage:
+
+    - `kind_option` names the option that decides the kind. Absent, the kind is
+      `produces` and nothing varies.
+    - `variant_option` names the option that decides a subdirectory under the kind. That
+      is a dimension rather than a step: `sheet/sharp/` and `sheet/pixel-art/` are the
+      same sprite closed two ways, and neither is the sheet.
+
+    Both name an option this stage declares, and `workspace.output_dir` is the only
+    place that turns them into a path.
+    """
+
+    name: str
+    requires: tuple[str, ...]
+    produces: str
+    summary: str
+    paid: bool
+    options: tuple[Option, ...] = ()
+    kind_option: str | None = None
+    variant_option: str | None = None
+
+
+STAGES: tuple[Stage, ...] = (
+    Stage(
+        name="anchor",
+        requires=(),
+        produces="anchor",
+        summary="generate a single image: an anchor, box art or an icon",
+        paid=True,
+        kind_option="kind",
+        options=(
+            Option(
+                ("--kind",),
+                "what this image is, and the directory it goes in: anchor is the neutral "
+                "sprite every later stage animates, box-art is the illustration that "
+                "settles the look, icon is a single object",
+                default="anchor",
+                choices=("anchor", "box-art", "icon"),
+            ),
+            Option(("--prompt",), "the prompt itself"),
+            Option(("--prompt-file",), "a file holding the prompt"),
+            Option(
+                ("--ref",),
+                "a reference image; repeat it, and the prompt addresses each one in order",
+                action="append",
+            ),
+            Option(
+                ("--size",),
+                "WIDTHxHEIGHT, or a fal preset (square_hd, square, portrait_4_3, "
+                "portrait_16_9, landscape_4_3, landscape_16_9, auto)",
+                default="1024x1024",
+            ),
+            Option(
+                ("--quality",),
+                "the endpoint's tier",
+                default="low",
+                choices=("low", "medium", "high"),
+            ),
+            Option(("--count",), "how many images to generate", kind="int", default=1),
+            Option(("--transparent",), "cut the chroma background to alpha", action="store_true"),
+            Option(
+                ("--chroma",),
+                "the key colour, 6 hex digits; with --transparent the prompt asks for "
+                "this field and the cut then removes it",
+                default="00b140",
+            ),
+            Option(("--tol",), "chroma tolerance", kind="float", default=60.0),
+            Option(("--feather",), "the alpha band at the edge", kind="float", default=24.0),
+            Option(
+                ("--chroma-mode",),
+                "flood takes only what the border reaches; global takes the key anywhere",
+                default="flood",
+                choices=("flood", "global"),
+            ),
+            Option(("--despill",), "pull the leaked key out of edge pixels", action="store_true"),
+            Option(
+                ("--despill-radius",),
+                "how far from transparent to despill",
+                kind="int",
+                default=2,
+            ),
+            Option(
+                ("--pixelart",),
+                "recover the native pixel grid and rewrite on it",
+                action="store_true",
+            ),
+            Option(("--colors",), "quantise to N colours with --pixelart", kind="int", default=0),
+            Option(
+                ("--scale",),
+                "integer upscale after the grid is recovered",
+                kind="int",
+                default=1,
+            ),
+        ),
+    ),
+    Stage(
+        name="motion",
+        requires=("anchor",),
+        produces="frames",
+        summary="transfer movement from a driving clip onto the anchor or a pose",
+        paid=True,
+        options=(
+            Option(
+                ("--pose",),
+                "start the clip from this animation's pose instead of the anchor, and "
+                "end it on that pose's last frame",
+            ),
+            Option(
+                ("--sheet",),
+                "the reference sheet the movement is taken from; a game asset, not the "
+                "asset's own, which is why it is named here",
+            ),
+            Option(("--row",), "which row of the sheet, one animation each", kind="int", default=0),
+            Option(
+                ("--sheet-frames",),
+                "how many cells of that row (default: the whole row)",
+                kind="int",
+                default=0,
+            ),
+            Option(("--clip",), "a driving clip to use instead of building one from a sheet"),
+            Option(
+                ("--endpoint",),
+                "which motion endpoint",
+                default="wan-motion",
+                choices=("wan-motion", "wan-animate", "one-to-all"),
+            ),
+            Option(("--prompt",), "an optional hint, where the endpoint takes one", default=""),
+            Option(("--target",), "the driving clip's square side", kind="int", default=720),
+            Option(("--drive-fps",), "the driving clip's frame rate", kind="int", default=12),
+            Option(("--seconds",), "how long the driving clip runs", kind="float", default=4.0),
+            Option(
+                ("--frames",),
+                "how many frames to pull out of the result",
+                kind="int",
+                default=6,
+            ),
+            Option(("--start",), "where extraction begins, 0..1", kind="float", default=0.0),
+            Option(("--end",), "where extraction ends, 0..1", kind="float", default=1.0),
+            Option(("--pick",), "exact frame indices, comma-separated"),
+            Option(("--cols",), "how wide to pack the board", kind="int", default=3),
+            Option(
+                ("--set",),
+                "an endpoint option, NAME=VALUE; checked against that endpoint before "
+                "anything is spent",
+                action="append",
+            ),
+        ),
+    ),
+    Stage(
+        name="pose",
+        requires=("anchor",),
+        produces="pose",
+        summary="the frame an animation starts from, and the one it ends on",
+        paid=True,
+        # Not `--name`: the asset is already the positional `name`, and argparse would
+        # write both to the same attribute — the pose would be generated for an asset
+        # called "attack".
+        variant_option="animation",
+        options=(
+            Option(
+                ("--animation",),
+                "which animation this pose is for, and the directory it goes in "
+                "under pose/ — attack, cast, death",
+                default="",
+            ),
+            Option(("--prompt",), "the starting pose"),
+            Option(("--prompt-file",), "a file holding it"),
+            Option(
+                ("--end-prompt",),
+                "a different last frame; without it the last frame is the first, so the "
+                "cycle closes where it began",
+            ),
+            Option(("--end-prompt-file",), "a file holding the last frame's prompt"),
+            Option(
+                ("--no-end",),
+                "write no last frame at all, for a movement that does not return",
+                action="store_true",
+            ),
+            Option(
+                ("--ref",),
+                "an extra reference after the anchor, which is always the first",
+                action="append",
+            ),
+            Option(("--size",), "WIDTHxHEIGHT, or a fal preset", default="1024x1024"),
+            Option(
+                ("--quality",),
+                "the endpoint's tier",
+                default="low",
+                choices=("low", "medium", "high"),
+            ),
+            Option(("--transparent",), "cut the chroma background to alpha", action="store_true"),
+            Option(("--chroma",), "the key colour, 6 hex digits", default="00b140"),
+            Option(("--tol",), "chroma tolerance", kind="float", default=60.0),
+            Option(("--feather",), "the alpha band at the edge", kind="float", default=24.0),
+            Option(
+                ("--chroma-mode",),
+                "flood takes only what the border reaches; global takes the key anywhere",
+                default="flood",
+                choices=("flood", "global"),
+            ),
+            Option(("--despill",), "pull the leaked key out of edge pixels", action="store_true"),
+            Option(
+                ("--despill-radius",),
+                "how far from transparent to despill",
+                kind="int",
+                default=2,
+            ),
+        ),
+    ),
+    Stage(
+        name="video",
+        requires=("anchor",),
+        produces="frames",
+        summary="animate the anchor or a pose directly, with no driving clip",
+        paid=True,
+        options=(
+            Option(
+                ("--pose",),
+                "start the clip from this animation's pose instead of the anchor, and "
+                "end it on that pose's last frame",
+            ),
+            Option(("--prompt",), "the movement to animate"),
+            Option(("--prompt-file",), "a file holding the prompt"),
+            Option(
+                ("--model",),
+                "which video endpoint",
+                default="grok",
+                choices=("grok", "grok-i2v", "seedance", "kling"),
+            ),
+            Option(
+                ("--ref",),
+                "an extra reference for grok, after the anchor; the prompt addresses "
+                "each by <IMAGE_0>, <IMAGE_1> in order",
+                action="append",
+            ),
+            Option(
+                ("--loop",),
+                "send the anchor as the last frame too, so the cycle closes where it began",
+                action="store_true",
+            ),
+            Option(("--duration",), "seconds of output", kind="int", default=4),
+            Option(
+                ("--resolution",),
+                "clip resolution",
+                default="480p",
+                choices=("480p", "720p", "1080p"),
+            ),
+            Option(
+                ("--aspect",),
+                "clip aspect for grok",
+                default="1:1",
+                choices=("16:9", "4:3", "3:2", "1:1", "2:3", "3:4", "9:16"),
+            ),
+            Option(
+                ("--negative",),
+                "negative prompt, where the endpoint takes one",
+                default="blur, distort, low quality, camera movement",
+            ),
+            Option(("--frames",), "how many frames to pull out", kind="int", default=6),
+            Option(
+                ("--start",),
+                "where extraction begins, 0..1 of the clip; the first frame of an "
+                "image-to-video clip is the input image standing still",
+                kind="float",
+                default=0.0,
+            ),
+            Option(("--end",), "where extraction ends, 0..1", kind="float", default=1.0),
+            Option(
+                ("--pick",),
+                "exact frame indices, comma-separated; ignores frames, start and end",
+            ),
+            Option(("--cols",), "how wide to pack the board", kind="int", default=3),
+        ),
+    ),
+    Stage(
+        name="matte",
+        requires=("motion", "video"),
+        produces="cutout",
+        summary="cut the frames' background by segmentation, locally or paid",
+        paid=True,
+        options=(
+            Option(
+                ("--backend",),
+                "local runs BiRefNet here through rembg and costs nothing; fal is the "
+                "paid endpoint. Without this, whatever settings say (local by default)",
+                choices=("local", "fal"),
+            ),
+            Option(
+                ("--model",),
+                "the rembg weights, with --backend local; without it, what settings say",
+            ),
+            Option(
+                ("--variant",),
+                "birefnet weights at the paid endpoint; Matting gives continuous alpha, "
+                "which is what hair needs",
+                default="Matting",
+                choices=(
+                    "General Use (Light)",
+                    "General Use (Light 2K)",
+                    "General Use (Heavy)",
+                    "Matting",
+                    "Portrait",
+                    "General Use (Dynamic)",
+                ),
+            ),
+            Option(
+                ("--resolution",),
+                "operating resolution",
+                default="2048x2048",
+                choices=("1024x1024", "2048x2048", "2304x2304"),
+            ),
+            Option(
+                ("--no-refine",),
+                "turn off edge colour decontamination, which is what kills the halo",
+                action="store_true",
+            ),
+            Option(("--mask",), "write the alpha mask beside the cut", action="store_true"),
+            Option(
+                ("--no-join",),
+                "matte each image on its own instead of joining them into one call",
+                action="store_true",
+            ),
+        ),
+    ),
+    Stage(
+        name="board",
+        requires=("matte",),
+        produces="sheet",
+        summary="close the frames into a sheet row, in square cells",
+        paid=False,
+        variant_option="art",
+        options=(
+            Option(("--grid",), "the board's grid, COLSxROWS", default="3x2"),
+            Option(
+                ("--frames",),
+                "how many cells to read (default: the whole grid)",
+                kind="int",
+                default=0,
+            ),
+            Option(("--cell",), "the sheet cell, in pixels", kind="int", default=166),
+            Option(
+                ("--scale",),
+                "the scale (default: the largest that fits — a whole factor while one "
+                "fits, and the exact fraction that brings a larger figure down)",
+                kind="float",
+                default=0.0,
+            ),
+            Option(
+                ("--fit",),
+                "bring a figure larger than the cell down to it; implied by "
+                "--art sharp, and off otherwise so an overflow stays visible",
+                action="store_true",
+            ),
+            Option(
+                ("--art",),
+                "the art direction to close this row in, and the directory under "
+                "sheet/ it goes in: sharp lands a crisp render on the cell, pixel-art "
+                "recovers a native grid the art already has, as-is does neither",
+                default="as-is",
+                choices=("as-is", "sharp", "pixel-art"),
+            ),
+            Option(("--colors",), "quantise the whole row to N colours", kind="int", default=0),
+            Option(
+                ("--alpha-floor",),
+                "alpha under this is background smear and goes to nothing, with --art sharp",
+                kind="int",
+                default=16,
+            ),
+            Option(
+                ("--alpha-ceil",),
+                "alpha over this is the figure and goes to solid, with --art sharp",
+                kind="int",
+                default=240,
+            ),
+            Option(("--chroma",), "cut this key off the whole board first, 6 hex digits"),
+            Option(("--tol",), "chroma tolerance", kind="float", default=60.0),
+            Option(("--feather",), "the alpha band at the edge", kind="float", default=24.0),
+            Option(
+                ("--chroma-mode",),
+                "flood takes only what the border reaches; global takes the key anywhere",
+                default="flood",
+                choices=("flood", "global"),
+            ),
+            Option(("--no-gif",), "skip the looping preview", action="store_true"),
+            Option(("--fps",), "the preview's frames per second", kind="int", default=7),
+        ),
+    ),
+)
+
+BY_NAME: dict[str, Stage] = {stage.name: stage for stage in STAGES}
+
+
+def get(name: str) -> Stage:
+    """The stage called `name`.
+
+    Fails with the full list rather than a KeyError: whoever mistyped wants to know
+    what exists, not which key was missing.
+    """
+    try:
+        return BY_NAME[name]
+    except KeyError:
+        known = ", ".join(stage.name for stage in STAGES)
+        raise KeyError(f"unknown stage {name!r}; there are: {known}") from None
+
+
+def order() -> tuple[str, ...]:
+    """Every stage, each after a prerequisite of its own.
+
+    This is presentation order, not a path to walk: `motion` and `video` are
+    alternatives and both appear.
+    """
+    done: list[str] = []
+    pending = [stage.name for stage in STAGES]
+    while pending:
+        ready = [
+            name
+            for name in pending
+            if not BY_NAME[name].requires
+            or any(dep in done for dep in BY_NAME[name].requires)
+        ]
+        if not ready:
+            raise RuntimeError(f"cycle in the stage registry, stuck on: {pending}")
+        for name in ready:
+            done.append(name)
+            pending.remove(name)
+    return tuple(done)
+
+
+def implementation(name: str) -> ModuleType:
+    """The module implementing `name`, imported now and not before.
+
+    A stage that has not been written yet fails here, naming the module that is
+    missing — not at package import, which would take the whole CLI down over one
+    stage.
+    """
+    stage = get(name)
+    try:
+        return import_module(f"spritegen.stages.{stage.name}")
+    except ModuleNotFoundError as exc:
+        raise NotImplementedError(
+            f"stage {stage.name!r} has no implementation yet "
+            f"(spritegen/stages/{stage.name}.py)"
+        ) from exc