spritegen-cli 0.1.0__py3-none-any.whl

spritegen/upscale.py

@@ -0,0 +1,444 @@
+"""Real-ESRGAN, locally and on the GPU, so a paid call never has to buy resolution.
+
+The expensive calls here are billed by pixels and by seconds. Generating at the cheap
+tier and enlarging afterwards moves the quality decision off the invoice: a `low` anchor
+and a 480p clip cost a fraction of the same thing at the top tier, and what they give up
+is resolution — which is the one thing a model like this puts back.
+
+**Ported from an `upscale_atlas.py` that reimplements RRDBNet without `basicsr`.** The
+weights are downloaded rather than vendored: the ncnn `.bin`/`.param` files that ship
+with Real-ESRGAN are for a different runtime, and its `realesrgan-ncnn-vulkan` binary is
+a Linux ELF that does not run everywhere this does.
+
+What changed in the port is the alpha, and it matters more here than it did there. A
+matted board is transparent over most of its area, and the RGB behind that transparency
+is whatever the matte left — usually black. Handing that to the network sharpens the
+black into the figure's edge, and every frame comes back with a dark rim. So the
+transparent region is filled from the nearest opaque colour before the network sees it,
+and the alpha channel is enlarged on its own.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+from pathlib import Path
+
+from . import fal
+from .fal import Untrusted
+
+#: The two the project already uses. `anime` is the 6-block model and is the one for
+#: character art; `general` is the full 23-block model, for a photographic reference.
+MODEL_URLS = {
+    "anime": (
+        "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/"
+        "RealESRGAN_x4plus_anime_6B.pth"
+    ),
+    "general": (
+        "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
+    ),
+}
+
+#: Where a weight file may come from. The URLs above are constants, so this is
+#: defence in depth rather than a live path — and the line that has to change first if
+#: they ever stop being constants.
+WEIGHT_HOSTS = ("github.com", "objects.githubusercontent.com")
+
+#: The network is 4x whatever it is given; any other factor is a resize afterwards.
+NATIVE_SCALE = 4
+
+#: Inference is tiled, so a 3732x3320 board does not have to fit in VRAM at once.
+TILE = 256
+TILE_PAD = 10
+
+#: How far the opaque colour is flooded outward before the network sees the image.
+FILL_PASSES = 8
+
+INSTALL = (
+    "the upscale needs torch, which is an optional extra: "
+    'pip install "spritegen-cli[upscale]", or uv sync --extra upscale'
+)
+
+
+def cache_dir() -> Path:
+    """Where a weight file lives — `SPRITEGEN_CACHE`, or the default beside it.
+
+    Outside any workspace on purpose. A weight file is tens of megabytes, it is the same
+    file for every project this CLI is run in, and it can always be fetched again — none
+    of which describes something that belongs in somebody's repository.
+    """
+    from . import settings
+
+    return settings.load().cache / "weights"
+
+
+def require_torch():
+    """`torch`, or an error naming the one command that fixes it.
+
+    Resolved here and never at module scope. Torch is some two gigabytes with CUDA, and
+    this is one free command of a CLI whose other commands have to keep answering on a
+    machine that never installed it — `status` and `cost` especially.
+    """
+    try:
+        import torch
+    except ModuleNotFoundError as exc:
+        raise NotImplementedError(INSTALL) from exc
+    return torch
+
+
+def available() -> bool:
+    """Whether the upscale could run at all. Used to report, never to decide silently."""
+    try:
+        require_torch()
+    except NotImplementedError:
+        return False
+    return True
+
+
+def pick_device(prefer: str = "auto"):
+    """CUDA where there is one, the CPU where there is not — or whatever was asked for.
+
+    Reported rather than assumed: the difference on a board is seconds against minutes,
+    and a run that quietly fell back to the CPU looks like a hang. Asking for `cuda`
+    where there is none is an error rather than a silent fallback, for the same reason.
+    """
+    torch = require_torch()
+    if prefer == "cpu":
+        return torch.device("cpu")
+    if prefer == "cuda":
+        if not torch.cuda.is_available():
+            raise NotImplementedError(
+                "cuda was asked for and torch reports none; leave the device on auto "
+                "for the CPU, or install a torch built against this machine's CUDA"
+            )
+        return torch.device("cuda")
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def download_weights(model: str) -> Path:
+    """The weight file, out of the cache or out of the release that holds it.
+
+    Held to the same discipline `fal.download` is, and for the same reason — it is the
+    other place in this tool that writes bytes from the network to disk:
+
+    - **HTTPS to the host the constant names**, checked rather than assumed. The URLs are
+      constants, so this is not an attacker's path today; it is what stops one becoming
+      an attacker's path the first time somebody makes them configurable.
+    - **A size ceiling**, so a redirect to something enormous fails instead of filling
+      the disk.
+    - **A digest written beside the file.** Nothing upstream publishes a hash to pin
+      against, so this cannot say the *first* download was the right one. What it can say
+      is that the file has not changed since — which is the case that matters for a cache
+      trusted forever and never looked at again.
+
+    `torch.load(..., weights_only=True)` in `load_model` is what keeps a swapped file
+    from being code execution; this keeps it from being silent.
+    """
+    if model not in MODEL_URLS:
+        known = ", ".join(sorted(MODEL_URLS))
+        raise KeyError(f"unknown upscale model {model!r}; there are: {known}")
+
+    cache = cache_dir()
+    cache.mkdir(parents=True, exist_ok=True)
+    destination = cache / f"{model}.pth"
+    digest_at = destination.with_suffix(".sha256")
+
+    if destination.exists():
+        if digest_at.is_file():
+            recorded = digest_at.read_text(encoding="utf-8").strip()
+            found = _digest(destination)
+            if found != recorded:
+                raise Untrusted(
+                    f"{destination} has changed since it was downloaded "
+                    f"({found[:12]} against {recorded[:12]}); delete it to fetch again"
+                )
+        return destination
+
+    _fetch(MODEL_URLS[model], destination)
+    digest_at.write_text(_digest(destination) + "\n", encoding="utf-8")
+    return destination
+
+
+def _digest(path: Path) -> str:
+    found = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            found.update(chunk)
+    return found.hexdigest()
+
+
+def _fetch(url: str, destination: Path) -> Path:
+    """One weight file, through the same fetch a paid result comes through.
+
+    `fal.download` with a different host set, and not a second implementation of it. The
+    first version here checked the scheme and the host of the URL it was *given* — which
+    is not the host it ends up talking to, because `urlopen` follows redirects inside the
+    library and a GitHub release URL redirects by design. `download` checks every hop and
+    where the chain lands, which is what that has to mean.
+
+    Written to a `.partial` first: a fetch interrupted halfway would otherwise leave a
+    truncated file that `exists()` accepts forever.
+    """
+    partial = destination.with_suffix(".partial")
+    try:
+        fal.download(url, partial, WEIGHT_HOSTS)
+        partial.replace(destination)
+    finally:
+        partial.unlink(missing_ok=True)
+    return destination
+
+
+def fill_transparent(image):
+    """Take the black out from behind the figure before the network sees it.
+
+    The network sees RGB and no alpha. Where the matte left black behind the
+    transparency, that black is a hard edge as far as the network is concerned and it
+    sharpens it into the figure — which is the dark rim around an enlarged sprite.
+
+    Two treatments, because the transparent area is two different things. The band just
+    outside the silhouette is where the network actually looks, and it is inpainted from
+    the art beside it, so the edge continues instead of stopping. Everything beyond that
+    band is out of reach of any kernel and only has to not be black, so it takes the
+    figure's own mean colour.
+
+    An image with no transparency comes back converted and otherwise untouched, which is
+    the anchor's case before it has been cut.
+    """
+    import cv2
+    import numpy as np
+    from PIL import Image
+
+    rgba = np.asarray(image.convert("RGBA"))
+    opaque = rgba[:, :, 3] > 0
+    if not opaque.any() or opaque.all():
+        return Image.fromarray(np.ascontiguousarray(rgba[:, :, :3]), "RGB")
+
+    rgb = np.ascontiguousarray(rgba[:, :, :3])
+    rgb[~opaque] = rgb[opaque].mean(axis=0).round().astype(np.uint8)
+
+    solid = opaque.astype(np.uint8)
+    kernel = np.ones((3, 3), np.uint8)
+    reach = cv2.dilate(solid, kernel, iterations=FILL_PASSES)
+    band = ((reach > 0) & ~opaque).astype(np.uint8)
+    if band.any():
+        rgb = cv2.inpaint(rgb, band, FILL_PASSES, cv2.INPAINT_TELEA)
+    return Image.fromarray(rgb, "RGB")
+
+
+def load_model(model: str, device):
+    """The network with its weights in it, on the device that will run it."""
+    torch = require_torch()
+    from .rrdb import RRDBNet
+
+    net = RRDBNet(in_nc=3, out_nc=3, nf=64, nb=6 if model == "anime" else 23, gc=32)
+    state = torch.load(str(download_weights(model)), map_location=device, weights_only=True)
+    # A Real-ESRGAN checkpoint wraps its weights; which key depends on the release.
+    for key in ("params_ema", "params"):
+        if key in state:
+            state = state[key]
+            break
+    net.load_state_dict(state, strict=True)
+    net.eval()
+    net.to(device)
+    return net
+
+
+def tiles(height: int, width: int, tile: int = TILE, pad: int = TILE_PAD) -> list[dict]:
+    """How a `height x width` image is cut into overlapping tiles and put back together.
+
+    Pure arithmetic, and separate from the forward pass on purpose: this is the part that
+    is easy to get wrong, the way it goes wrong is a visible seam every `tile` pixels,
+    and a seam is not something a test of the network would notice. Extracted, it can be
+    checked against what it has to be — that the kept regions tile the output exactly,
+    with no gap and no overlap.
+
+    Each entry has three boxes, in input pixels for `read` and output pixels for the
+    other two:
+
+    - `read` — what goes into the network, the tile plus `pad` of context on each side
+      that exists. The context is what stops the network seeing a hard edge where the
+      tile was cut, which is the seam.
+    - `take` — where this tile's own pixels sit inside what came back, which is the pad
+      offset scaled up.
+    - `put` — where they go in the finished image.
+
+    `take` and `put` are the same size by construction; the pad is read and thrown away.
+    """
+    found = []
+    for top in range(0, height, tile):
+        for left in range(0, width, tile):
+            y0, x0 = max(top - pad, 0), max(left - pad, 0)
+            y1, x1 = min(top + tile + pad, height), min(left + tile + pad, width)
+            keep_h = min(tile, height - top) * NATIVE_SCALE
+            keep_w = min(tile, width - left) * NATIVE_SCALE
+            found.append(
+                {
+                    "read": (y0, y1, x0, x1),
+                    "take": (
+                        (top - y0) * NATIVE_SCALE,
+                        (top - y0) * NATIVE_SCALE + keep_h,
+                        (left - x0) * NATIVE_SCALE,
+                        (left - x0) * NATIVE_SCALE + keep_w,
+                    ),
+                    "put": (
+                        top * NATIVE_SCALE,
+                        top * NATIVE_SCALE + keep_h,
+                        left * NATIVE_SCALE,
+                        left * NATIVE_SCALE + keep_w,
+                    ),
+                }
+            )
+    return found
+
+
+def _tiled(net, tensor):
+    """One forward pass, in tiles, so a large board need not fit in VRAM at once."""
+    torch = require_torch()
+    _, _, height, width = tensor.shape
+    if height <= TILE and width <= TILE:
+        return net(tensor)
+
+    output = torch.zeros((1, 3, height * NATIVE_SCALE, width * NATIVE_SCALE), device=tensor.device)
+    for one in tiles(height, width):
+        ry0, ry1, rx0, rx1 = one["read"]
+        ty0, ty1, tx0, tx1 = one["take"]
+        py0, py1, px0, px1 = one["put"]
+        done = net(tensor[:, :, ry0:ry1, rx0:rx1])
+        output[:, :, py0:py1, px0:px1] = done[:, :, ty0:ty1, tx0:tx1]
+    return output
+
+
+def enlarge(image, scale: int, model: str | None = None, device=None):
+    """One image, `scale` times larger, with the alpha kept out of the network.
+
+    The network is 4x and only 4x, so any other factor is a resize of its output. Going
+    through it is still worth it at 2x: the result carries detail the model invented,
+    which is the entire reason to run it rather than resize the input.
+    """
+    # Before torch is resolved: a scale under two is the caller's mistake either way,
+    # and being told about torch instead would send them to fix the wrong thing.
+    if scale < 2:
+        raise ValueError(f"an upscale is 2 or more; got {scale}")
+
+    torch = require_torch()
+    import numpy as np
+    from PIL import Image
+
+    from . import settings
+
+    config = settings.load()
+    model = model or config.upscale_model
+    device = pick_device(config.upscale_device) if device is None else device
+    net = load_model(model, device)
+
+    array = np.asarray(fill_transparent(image), dtype=np.float32) / 255.0
+    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0).to(device)
+    with torch.no_grad():
+        out = _tiled(net, tensor).squeeze(0).clamp(0, 1).cpu()
+
+    enlarged = Image.fromarray((out.permute(1, 2, 0).numpy() * 255).round().astype(np.uint8), "RGB")
+    target = (image.size[0] * scale, image.size[1] * scale)
+    if enlarged.size != target:
+        enlarged = enlarged.resize(target, Image.LANCZOS)
+
+    if image.mode != "RGBA":
+        return enlarged
+    result = enlarged.convert("RGBA")
+    result.putalpha(image.getchannel("A").resize(target, Image.LANCZOS))
+    return result
+
+
+def enlarge_file(path: Path, scale: int, model: str | None = None, device=None) -> dict:
+    """Enlarge one file in place, keeping the original beside it the first time.
+
+    `<name>.raw.png` is written once and never overwritten, so running the upscale a
+    second time cannot lose what came back from the paid call — and the second run
+    starts from that original rather than from an already-enlarged image, which would
+    compound the model's invention twice over.
+
+    That backup is also why this command needs no ledger line: nothing it does costs
+    anything, and nothing it does is unrecoverable.
+    """
+    from PIL import Image
+
+    raw = path.with_name(f"{path.stem}.raw{path.suffix}")
+    if not raw.exists():
+        raw.write_bytes(path.read_bytes())
+    image = Image.open(raw).convert("RGBA")
+
+    from . import settings
+
+    model = model or settings.load().upscale_model
+    result = enlarge(image, scale, model, device)
+    result.save(path)
+
+    return {
+        "file": path.name,
+        "raw": raw.name,
+        "from": f"{image.size[0]}x{image.size[1]}",
+        "to": f"{result.size[0]}x{result.size[1]}",
+        "scale": scale,
+        "model": model,
+    }
+
+
+def cmd_upscale(args: argparse.Namespace) -> int:
+    """Enlarge one artifact of one sprite, locally and for nothing.
+
+    Named by artifact rather than by stage, like everything else in the workspace: what
+    is being enlarged is the anchor, or the frames, not "whatever the video stage left".
+    Without one, it takes what the sprite produced most recently, which is almost always
+    what somebody means right after producing it.
+
+    No ledger line, because nothing was spent. The original is kept as `<name>.raw.png`
+    the first time and every later run starts from that — enlarging an already-enlarged
+    image would compound the model's invention twice over.
+
+    ASCII only, for the reason `workspace.cmd_status` is.
+    """
+    from . import workspace
+
+    state = workspace.load(workspace.check_name(args.name))
+    if state.outdated():
+        raise workspace.StageRefused(
+            f"{state.name}: this workspace is an older layout; run `spritegen migrate` first"
+        )
+
+    key = args.artifact or (state.last() and state.stages[state.last()].get("produced"))
+    if not key:
+        raise workspace.StageRefused(f"{state.name}: nothing produced yet, so nothing to enlarge")
+    entry = state.artifacts.get(key)
+    if entry is None:
+        known = ", ".join(state.artifacts) or "none"
+        raise workspace.StageRefused(f"{state.name}: no artifact {key!r}; it has: {known}")
+
+    directory = workspace.inside(state.name, entry["dir"])
+    targets = [
+        path
+        for path in sorted(directory.glob("*.png"))
+        if not path.stem.endswith(".raw") and not path.name.endswith(".raw.png")
+    ]
+    if not targets:
+        raise workspace.StageRefused(f"{state.name}: no image in {directory}")
+
+    if args.dry_run:
+        from . import settings
+
+        config = settings.load()
+        print(f"dry run: {args.scale}x {config.upscale_model} on {config.upscale_device}")
+        for path in targets:
+            print(f"  {path}")
+        return 0
+
+    device = pick_device(args.device)
+    print(f"{args.scale}x on {device}")
+    for path in targets:
+        report = enlarge_file(path, args.scale, args.model, device)
+        print(f"  {report['file']}  {report['from']} -> {report['to']}  (kept {report['raw']})")
+
+    # The files changed but nothing was produced, so `artifacts` is corrected rather than
+    # added to: `show` compares this against the disk, and leaving it stale would report
+    # drift that is this command's own doing.
+    entry["files"] = sorted(path.name for path in directory.glob("*") if path.is_file())
+    workspace.save(state)
+    return 0