spritegen-cli 0.1.0__py3-none-any.whl

spritegen/fal.py

@@ -0,0 +1,228 @@
+"""The one place a command turns into money.
+
+Every paid call goes through `call`, which does two things before anything is spent:
+it refuses without a key, and it honours a dry run by printing the payload instead of
+sending it. Both live here rather than in each stage so no stage can forget one.
+
+`_subscribe` is the seam. It is the only function that touches `fal_client`, which is
+imported inside it: the rest of the CLI — `--help`, `status`, `cost` — must not pay the
+import cost of an HTTP client, and a test must not need one.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from urllib.parse import urlparse
+
+KEY_VAR = "FAL_KEY"
+
+#: Where a result may be fetched from. The endpoint decides what URL comes back, and
+#: `urls_in` walks the whole response for anything under a `url` key — so without this,
+#: a malformed or tampered response naming `file:///…` would have a local file copied
+#: into the asset and passed off as generated art.
+ALLOWED_HOSTS = ("fal.media", "fal.ai", "fal.run", "fal-cdn.com")
+
+#: A hard ceiling on a downloaded body. A clip is megabytes; nothing here is hundreds.
+MAX_BYTES = 256 * 1024 * 1024
+
+TIMEOUT = 600
+
+
+class Untrusted(ValueError):
+    """A URL the endpoint returned that this tool will not fetch."""
+
+
+def _allowed_names() -> tuple[str, ...]:
+    from . import settings
+
+    return ALLOWED_HOSTS + settings.load().extra_hosts()
+
+
+def _bare_host(url) -> str:
+    """The hostname httpx will connect to, lowered, with no port and no brackets."""
+    return (url.host or "").lower()
+
+
+def check_url(url: str, allowed: tuple[str, ...] | None = None) -> str:
+    r"""The URL, if it is one this tool will fetch — R2.8.
+
+    `allowed` is the host set to hold it to; without one it is fal's, extended by
+    `SPRITEGEN_ALLOWED_HOSTS`. It is a parameter because this is not the only place that
+    writes bytes from the network to disk — the weight cache does too, from a different
+    set of hosts, and it needs every one of the checks here rather than its own worse
+    copy of two of them.
+
+    The host is read from `httpx.URL`, which is the parser that opens the connection,
+    and it is required to agree with `urlparse` besides. **Any authority carrying
+    userinfo is refused outright**, and that is what catches the crafted case: fal's
+    links never need userinfo, and a URL like `https://evil.com\@fal.media/x` puts the
+    part a reader's eye stops at into the userinfo instead of the host.
+    """
+    import httpx
+
+    parsed = urlparse(url)
+    if parsed.scheme != "https":
+        raise Untrusted(f"refusing to fetch {url!r}: only https")
+
+    connects = httpx.URL(url)
+    if connects.userinfo or "@" in (parsed.netloc or ""):
+        raise Untrusted(f"refusing to fetch {url!r}: the authority carries userinfo")
+
+    host = _bare_host(connects)
+    if host != (parsed.hostname or "").lower():
+        raise Untrusted(
+            f"refusing to fetch {url!r}: it parses as {parsed.hostname!r} but connects "
+            f"to {host!r}"
+        )
+
+    extensible = allowed is None
+    allowed = _allowed_names() if extensible else allowed
+    if not any(host == name or host.endswith(f".{name}") for name in allowed):
+        hint = " Set SPRITEGEN_ALLOWED_HOSTS to add one." if extensible else ""
+        raise Untrusted(f"refusing to fetch {url!r}: only {', '.join(allowed)}.{hint}")
+    return url
+
+
+#: How many redirects a fetch will follow before giving up.
+MAX_HOPS = 5
+
+
+def upload(path: Path) -> str:
+    """Put a local file in fal's storage and return the URL that reaches it.
+
+    That URL is standing access to the file for anyone holding it, which is why
+    `ledger.redact` keeps it out of the recorded payload.
+    """
+    import fal_client
+
+    return fal_client.upload_file(str(path))
+
+
+def download(url: str, dest: Path, allowed: tuple[str, ...] | None = None) -> Path:
+    """Fetch into `dest`, refusing anything this tool should not fetch.
+
+    **The redirect chain is walked here rather than inside the client.** Following it
+    automatically means a live request has already gone to every intermediate host by the
+    time the call returns, so checking where it landed checks nothing about the middle.
+    `follow_redirects=False` plus `response.next_request` makes each hop a decision this
+    function gets to refuse before it happens.
+
+    The body streams in chunks against a ceiling rather than being read whole, so an
+    oversized response fails instead of being held in memory — and the partial file is
+    removed, because a truncated image that `exists()` accepts is worse than none.
+
+    `allowed` narrows the host set for a caller that is not fetching a fal result — the
+    weight cache passes its own. Nothing else about the fetch differs, which is why it is
+    this function with a parameter and not a second one.
+    """
+    import httpx
+
+    check_url(url, allowed)
+    dest.parent.mkdir(parents=True, exist_ok=True)
+
+    with httpx.Client(follow_redirects=False, timeout=TIMEOUT) as client:
+        request = client.build_request("GET", url)
+        for _ in range(MAX_HOPS + 1):
+            with client.send(request, stream=True) as response:
+                hop = response.next_request
+                if hop is not None:
+                    check_url(str(hop.url), allowed)
+                    request = hop
+                    continue
+
+                response.raise_for_status()
+                size = 0
+                try:
+                    with dest.open("wb") as handle:
+                        for chunk in response.iter_bytes(1 << 20):
+                            size += len(chunk)
+                            if size > MAX_BYTES:
+                                raise Untrusted(
+                                    f"{url} is over the {MAX_BYTES:,} byte ceiling"
+                                )
+                            handle.write(chunk)
+                except BaseException:
+                    dest.unlink(missing_ok=True)
+                    raise
+                return dest
+
+    raise Untrusted(f"{url} redirected more than {MAX_HOPS} times")
+
+
+class MissingKey(ValueError):
+    """No usable `FAL_KEY`.
+
+    A ValueError so the CLI turns it into an exit code and a line on stderr: an absent
+    key is a thing to go fix, not a bug to report.
+    """
+
+
+def require_key() -> str:
+    """The key, or an error naming what to set — R2.2.
+
+    Blank counts as absent. An empty environment variable is how a key goes missing in
+    practice, and letting it through only moves the failure to the endpoint, after the
+    request was built.
+    """
+    from . import settings
+
+    key = settings.load().fal_key or ""
+    if not key.strip():
+        raise MissingKey(f"{KEY_VAR} is not set; a paid call needs it")
+    return key
+
+
+def show(endpoint: str, payload: dict) -> None:
+    """Print the call that would be made — R2.3.
+
+    The payload is printed as it would be sent, not summarised: the point of a dry run
+    is to check the thing itself, and a summary is where the difference hides.
+    """
+    print(f"dry run: {endpoint}")
+    print(json.dumps(payload, indent=2, ensure_ascii=False, sort_keys=True))
+
+
+def _subscribe(endpoint: str, payload: dict) -> dict:
+    """The call itself. Patched out in tests; nothing else imports `fal_client`."""
+    import fal_client
+
+    return fal_client.subscribe(endpoint, arguments=payload, with_logs=False)
+
+
+def call(endpoint: str, payload: dict, *, dry_run: bool = False) -> dict:
+    """Run one paid call, or show what it would be — R2.2, R2.3.
+
+    A dry run returns before the key is checked, on purpose: it issues no paid call, so
+    a payload can be inspected on a machine that has no credentials at all.
+    """
+    if dry_run:
+        show(endpoint, payload)
+        return {"dry_run": True, "endpoint": endpoint, "payload": payload}
+
+    require_key()
+    return _subscribe(endpoint, payload)
+
+
+def urls_in(result: dict) -> list[str]:
+    """Every URL an endpoint returned, in the order found.
+
+    The endpoints disagree about where they put it — `video.url`, `images[].url`,
+    `url` — so this walks the result instead of encoding one shape per endpoint and
+    getting it wrong for the next one.
+    """
+    found: list[str] = []
+
+    def walk(node) -> None:
+        if isinstance(node, dict):
+            for key, value in node.items():
+                if key == "url" and isinstance(value, str):
+                    found.append(value)
+                else:
+                    walk(value)
+        elif isinstance(node, list):
+            for item in node:
+                walk(item)
+
+    walk(result)
+    return found

spritegen/imaging.py

@@ -0,0 +1,219 @@
+"""The local, unpaid transforms — chroma cutting and pixel-grid recovery.
+
+Ported from the `gen_image.py` this grew out of, and pinned by `tests/parity`. Nothing here
+calls an endpoint: given the same bytes in, these produce the same bytes out, which is
+why they can be re-run for free after a paid call and why the fixtures can hold them to
+the behaviour the scripts had.
+
+Two stages use them, which is the reason they are not inside either: `anchor` cuts the
+chroma off what it generated, and `board` recovers the grid across a whole board.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+
+def parse_chroma(text: str) -> tuple[int, int, int]:
+    """`00b140` to `(0, 177, 64)`."""
+    raw = text.strip().lstrip("#")
+    if len(raw) != 6:
+        raise ValueError("a chroma key is 6 hex digits, e.g. 00b140")
+    try:
+        value = int(raw, 16)
+    except ValueError:
+        raise ValueError(f"chroma key {text!r} is not hex") from None
+    return (value >> 16) & 0xFF, (value >> 8) & 0xFF, value & 0xFF
+
+
+def _near_transparent(alpha, radius: int):
+    """Opaque pixels within `radius` of a transparent one."""
+    import numpy as np
+
+    near = alpha < 1.0
+    grown = near.copy()
+    for _ in range(radius):
+        padded = np.pad(grown, 1, constant_values=False)
+        grown = padded[:-2, 1:-1] | padded[2:, 1:-1] | padded[1:-1, :-2] | padded[1:-1, 2:] | grown
+    return grown & (alpha > 0.0)
+
+
+def cut_chroma(
+    path: Path,
+    key: tuple[int, int, int],
+    tol: float,
+    feather: float,
+    despill: bool,
+    despill_radius: int,
+    mode: str = "flood",
+) -> dict:
+    """Replace a chroma background with alpha, in place, and report what it measured.
+
+    `flood` counts as background only what the border can reach, so a green the artist
+    put inside the figure survives. `global` takes every pixel near the key wherever it
+    is, which is the only thing that clears islanded chroma — the gap between an arm
+    and the body, the space through hair — at the cost of eating art that is
+    legitimately that colour.
+    """
+    import numpy as np
+    from PIL import Image, ImageDraw
+
+    rgb = np.asarray(Image.open(path).convert("RGB"), dtype=np.float32)
+    height, width = rgb.shape[:2]
+
+    distance = np.sqrt(((rgb - np.array(key, dtype=np.float32)) ** 2).sum(axis=2))
+
+    # Background candidate: near the key, with the feather band still inside.
+    candidate = distance <= (tol + feather)
+
+    if mode == "global":
+        reached = candidate
+    else:
+        # A one-pixel frame that is guaranteed to be a candidate gives the flood a
+        # single starting point, whatever the art does at the edge.
+        flood = Image.new("L", (width + 2, height + 2), 255)
+        flood.paste(Image.fromarray(np.where(candidate, 255, 0).astype(np.uint8)), (1, 1))
+        ImageDraw.floodfill(flood, (0, 0), 128, thresh=0)
+        reached = np.asarray(flood, dtype=np.uint8)[1 : height + 1, 1 : width + 1] == 128
+
+    # Inside the reached background, alpha ramps 0 to 1 across the feather band.
+    if feather > 0:
+        ramp = np.clip((distance - tol) / feather, 0.0, 1.0)
+    else:
+        ramp = (distance > tol).astype(np.float32)
+    alpha = np.where(reached, ramp, 1.0)
+
+    if despill and despill_radius > 0:
+        # Leaked key lives in two places: along the contour, where the chroma painted
+        # over the edge, and islanded inside the silhouette, where the border flood
+        # never arrives and the pixel stays opaque and green. Both go, because without
+        # the second a downscale spreads each island over a whole block.
+        dominant = int(np.argmax(key))
+        others = [channel for channel in range(3) if channel != dominant]
+        ceiling = np.maximum(rgb[:, :, others[0]], rgb[:, :, others[1]])
+        target = _near_transparent(alpha, despill_radius) | (~reached & candidate)
+        spilled = target & (alpha > 0.0) & (rgb[:, :, dominant] > ceiling)
+        rgb[:, :, dominant] = np.where(spilled, ceiling, rgb[:, :, dominant])
+        spill_px = int(spilled.sum())
+    else:
+        spill_px = 0
+
+    out = np.dstack([rgb, alpha * 255.0]).round().clip(0, 255).astype(np.uint8)
+    Image.fromarray(out, mode="RGBA").save(path)
+
+    return {
+        "transparent_px": int((alpha <= 0.0).sum()),
+        "partial_px": int(((alpha > 0.0) & (alpha < 1.0)).sum()),
+        "opaque_px": int((alpha >= 1.0).sum()),
+        "despilled_px": spill_px,
+    }
+
+
+#: Seeds OpenCV's RNG before pixelfixer reconstructs. See `to_pixelart`.
+PIXELART_SEED = 0
+
+INSTALL_PIXELFIXER = (
+    "recovering a native pixel grid needs pixelfixer, which is on no package registry: "
+    'pip install "pixelfixer @ '
+    'git+https://github.com/Retro-Diffusion/pixel-art-fixer.git#subdirectory=python"'
+)
+
+
+def require_pixelfixer():
+    """`pixelfixer`, or an error naming the one command that fixes it.
+
+    Not a declared dependency, and it cannot be one: it is on no registry, and a direct
+    URL in `dependencies` or an extra makes a distribution PyPI refuses — so declaring it
+    would trade an uninstallable package for an unpublishable one.
+
+    It reaches exactly two commands, `anchor --pixelart` and `board --art pixel-art`, so
+    the tool without it is the whole tool minus one art direction. That is the same shape
+    `torch` and `rembg` have, and the third time this project has met it.
+    """
+    try:
+        import pixelfixer
+    except ModuleNotFoundError as exc:
+        raise NotImplementedError(INSTALL_PIXELFIXER) from exc
+    return pixelfixer
+
+
+def pixelfixer_available() -> bool:
+    """Whether the grid recovery could run. Used to report, never to decide silently."""
+    try:
+        require_pixelfixer()
+    except NotImplementedError:
+        return False
+    return True
+
+
+def to_pixelart(path: Path, colors: int, scale: int, seed: int = PIXELART_SEED) -> dict:
+    """Recover the native pixel grid the art already has and rewrite the image on it.
+
+    The detector votes on a cell from the art itself. It does not invent one: on a
+    smooth painting with no blocks in it the vote returns nonsense, which is why the
+    prompt has to ask for flat blocks and why `cols` and `rows` come back in the report
+    — they are how you check a grid was there at all.
+
+    **Seeded, unlike the script this came from.** pixelfixer reconstructs each block by
+    k-means on OpenCV's global RNG, so the original ran three times on one input and
+    produced three different images — same grid, same colour count, different pixels.
+    Seeding costs nothing and is what makes a re-run comparable to the run before it,
+    and a paid call's output worth re-deriving instead of re-buying.
+    """
+    require_pixelfixer()
+
+    import cv2
+    import numpy as np
+    from PIL import Image
+    from pixelfixer import detect
+    from pixelfixer.reconstruct import reconstruct
+
+    cv2.setRNGSeed(seed)
+    rgba = np.array(Image.open(path).convert("RGBA"))
+    found = detect(rgba)
+    out = reconstruct(rgba, found["step_x"], found["step_y"], found["cols"], found["rows"])
+    image = Image.fromarray(out)
+
+    if colors:
+        # No dither: it scatters half-tone noise across blocks the grid just flattened.
+        quantized = image.convert("RGB").quantize(colors=colors, dither=Image.Dither.NONE)
+        merged = quantized.convert("RGBA")
+        merged.putalpha(image.getchannel("A"))
+        image = merged
+
+    native = image.size
+    if scale > 1:
+        image = image.resize((native[0] * scale, native[1] * scale), Image.NEAREST)
+
+    image.save(path)
+    used = len({tuple(pixel) for pixel in np.array(image.convert("RGBA")).reshape(-1, 4)})
+    return {
+        "cols": int(found["cols"]),
+        "rows": int(found["rows"]),
+        "step_x": round(float(found["step_x"]), 3),
+        "step_y": round(float(found["step_y"]), 3),
+        "native": f"{native[0]}x{native[1]}",
+        "scale": scale,
+        "seed": seed,
+        "colours": used,
+    }
+
+
+def measure_alpha(path: Path) -> dict:
+    """What the alpha channel of a cut image looks like.
+
+    Ported from `bg_remove.py`'s `measure`. It is how you tell a matte that worked from
+    one that returned the picture untouched: an image with no transparent pixels came
+    back opaque, whatever the endpoint reported.
+    """
+    import numpy as np
+    from PIL import Image
+
+    image = Image.open(path).convert("RGBA")
+    alpha = np.asarray(image)[:, :, 3]
+    return {
+        "size": "x".join(str(value) for value in image.size),
+        "transparent_px": int((alpha == 0).sum()),
+        "partial_px": int(((alpha > 0) & (alpha < 255)).sum()),
+        "opaque_px": int((alpha == 255).sum()),
+    }

spritegen/ledger.py

@@ -0,0 +1,143 @@
+"""One line per paid call, appended, never rewritten.
+
+This is the only record that a call happened and what it cost in artifacts. It is
+append-only JSONL rather than a document for one reason: a run that dies halfway has
+still spent the money, and a format that has to be re-serialised whole can lose the
+line that says so.
+
+What goes in a line is fixed by R2.1 — the endpoint, the exact payload sent, the URLs
+that came back, and the sha256 of every file written. That is enough to answer "was
+this file paid for, and by which call" months later, and enough for `cost` to report
+an asset whose outputs were long since deleted.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from datetime import UTC, datetime
+from pathlib import Path
+
+CHUNK = 1 << 20
+
+#: What an uploaded file's URL is replaced by. `fal_client.upload_file` returns standing
+#: access to that file for anyone holding the link, so recording it verbatim would turn
+#: this ledger into a key to every reference image, anchor and clip ever sent — art that
+#: has not shipped included. What the record needs is that a file went up and which one,
+#: not a way back to it.
+REDACTED = "<uploaded>"
+
+#: Payload keys that carry one, or a list of them.
+UPLOAD_KEYS = frozenset(
+    {
+        "image_url",
+        "image_urls",
+        "video_url",
+        "reference_image_urls",
+        "end_image_url",
+        "tail_image_url",
+    }
+)
+
+
+def redact(payload: dict) -> dict:
+    """The payload as it is recorded — everything, minus the upload URLs.
+
+    R2.1 asks for the exact payload sent, and this is the one deliberate exception:
+    every other field is what was asked for, and these are handles to files that are
+    already on this disk anyway.
+    """
+    out = {}
+    for key, value in payload.items():
+        if key not in UPLOAD_KEYS:
+            out[key] = value
+        elif isinstance(value, list):
+            out[key] = [REDACTED] * len(value)
+        else:
+            out[key] = REDACTED
+    return out
+
+
+def sha256_of(path: Path) -> str:
+    """The file's digest, read in chunks — a board is tens of megabytes."""
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while chunk := handle.read(CHUNK):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def describe(path: Path) -> dict:
+    """One written file, as it appears in a ledger line."""
+    return {"path": path.name, "bytes": path.stat().st_size, "sha256": sha256_of(path)}
+
+
+def entry(
+    *,
+    stage: str,
+    endpoint: str,
+    payload: dict,
+    urls: list[str] | None = None,
+    files: list[Path] | None = None,
+    at: str | None = None,
+) -> dict:
+    """A ledger line, built but not written.
+
+    Separate from `append` so a dry run can show exactly what would be recorded, and so
+    a test does not need a filesystem to check the shape.
+    """
+    return {
+        "at": at or datetime.now(UTC).isoformat(timespec="seconds"),
+        "stage": stage,
+        "endpoint": endpoint,
+        "payload": redact(payload),
+        "urls": list(urls or []),
+        "files": [describe(path) for path in files or []],
+    }
+
+
+def append(path: Path, line: dict) -> dict:
+    """Append one line, creating the ledger if this is the asset's first paid call."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("a", encoding="utf-8") as handle:
+        handle.write(json.dumps(line, ensure_ascii=False) + "\n")
+    return line
+
+
+def read(path: Path) -> list[dict]:
+    """Every line, in order.
+
+    A blank line is skipped and a malformed one is fatal: the ledger is the record of
+    spending, and quietly dropping a line it could not parse would understate it.
+    """
+    if not path.exists():
+        return []
+    lines = []
+    for number, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
+        if not raw.strip():
+            continue
+        try:
+            lines.append(json.loads(raw))
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"{path}:{number} is not a ledger line: {exc}") from exc
+    return lines
+
+
+def summarise(lines: list[dict]) -> list[dict]:
+    """The ledger folded to one row per stage and endpoint, in first-call order — R2.4.
+
+    Grouped by both because a stage can be re-run against a different endpoint — that
+    is the whole point of having three motion endpoints — and a row that merged them
+    would hide which one the money went to.
+    """
+    rows: dict[tuple[str, str], dict] = {}
+    for line in lines:
+        key = (line.get("stage", "?"), line.get("endpoint", "?"))
+        row = rows.setdefault(
+            key, {"stage": key[0], "endpoint": key[1], "calls": 0, "files": 0, "bytes": 0}
+        )
+        row["calls"] += 1
+        for written in line.get("files", []):
+            row["files"] += 1
+            row["bytes"] += int(written.get("bytes", 0))
+    return list(rows.values())