raven-image: --margin pins the border band during diffusion; measures, repaints and fails loudly (exit 65); --alpha keys the backdrop
Core nodes only (EmptyImage > VAEEncode > SetLatentNoiseMask, mask by SolidMask+MaskComposite): the sampler re-imposes the flat band at every step. The post-render step measures against the band's own decoded backdrop and is the proof, not the mechanism. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
Kiko
parent
307d9e4821
commit
d0596fdc11
4 changed files with 286 additions and 9 deletions
|
|
@ -10,7 +10,7 @@ The swarm's image service: Minds on Yggdrasil make images on Raven's GPU (ComfyU
|
|||
| `comfy-nodes/` | Custom ComfyUI nodes (`*.py`) | Raven `E:\ComfyUI\...\ComfyUI\custom_nodes\` | `\GameCult\Raven-Image-Deploy` (every 10 min) syncs them and restarts ComfyUI only when they changed and the queue is empty |
|
||||
| `deploy/` | The two deployers and their units | installed once by root | changing these is a root step; ask in #minds |
|
||||
|
||||
Client usage, flags (`--ref`, `--vision`), graphs and limits: `client/README.md`.
|
||||
Client usage, flags (`--ref`, `--vision`, `--margin`), graphs and limits: `client/README.md`.
|
||||
|
||||
Raven is Deru's gaming PC. When he closes ComfyUI the service is off, and the client exits 75 ("Raven is busy or off").
|
||||
Node changes wait for an empty queue, so a deploy never kills a running job. Concept art only until the FLUX.3 Klein
|
||||
|
|
|
|||
|
|
@ -5,8 +5,8 @@ file, standard library only, plus the ComfyUI API graph it submits.
|
|||
|
||||
**Value:** a Mind can make concept art with one command and no ComfyUI knowledge.
|
||||
**Entry points:** `raven-image` (the client), `graphs/flux3-klein.json` (the graph).
|
||||
**Live:** Klein t2i with the prompt upsampler. **Also live:** `--ref` reference-conditioned edit (below). **Not built:** other families
|
||||
(MOSS), img2img, a queue of our own, any Idunn or CultMesh surface.
|
||||
**Live:** Klein t2i with the prompt upsampler. **Also live:** `--ref` reference-conditioned edit and `--margin` clear edges (below). **Not built:** other families
|
||||
(MOSS), img2img, a custom sampler node (not needed: see `--margin`), a queue of our own, any Idunn or CultMesh surface.
|
||||
|
||||
```
|
||||
raven-image "concept art: a lighthouse on a basalt cliff at dusk" -o out.png
|
||||
|
|
@ -39,7 +39,47 @@ Test 2026-11-11, Kiko's r1a 12-cell pose sheet as the only reference, same seed:
|
|||
but also pasted three of the sheet's cells into the frame; use one cropped cell as the reference. Images and verdict:
|
||||
`/opt/commons/handoff/raven-ref-edit-20261111/`.
|
||||
|
||||
Offline test (fake ComfyUI on 127.0.0.1): `python3 scripts/test-raven-image.py` in gamecult-ops.
|
||||
## `--margin PCT`: clear edges by construction
|
||||
|
||||
```
|
||||
raven-image "a goldfish mascot on a plain pale cream background" --no-enhance --margin 6 --bg fcf7ee -o fish.png
|
||||
raven-image ... --margin 6 --bg 12fa5f --alpha -o sticker.png # RGBA, backdrop keyed out
|
||||
```
|
||||
|
||||
The border band (PCT of each axis, rounded up to the 16 px latent grid, 6% of 1024 is 64 px) is held flat during
|
||||
diffusion, not cleaned afterwards. The client replaces the empty latent with a `--bg`-coloured image run through the
|
||||
graph's own VAE (`EmptyImage` > `VAEEncode`) and puts a `SetLatentNoiseMask` on it that is 0 on the band and 1 inside
|
||||
(`SolidMask` + `MaskComposite`, nodes 30-35; node 6 goes). ComfyUI's sampler re-imposes the masked-out latent at every
|
||||
step at that step's noise level (RePaint-style), so the model composes the interior around a frame it cannot change.
|
||||
Core nodes only; no custom node, no Raven deploy. Denoise stays 1. Works on `flux3-klein` and `flux3-klein-edit`
|
||||
(with or without `--vision`). Put the same backdrop in the prompt ("plain pale cream background"); the pin sets the band,
|
||||
the prompt sets what the interior tries to match.
|
||||
|
||||
After the render the client measures the band against its own backdrop (the per-channel median of the band: the VAE
|
||||
shifts a saturated `--bg` by tens of levels, `#12fa5f` came back `#1fef6d`; cream stayed within 2), repaints the whole
|
||||
band with that colour, and prints it with the per-side count of pixels that were content:
|
||||
`margin 6.0%: backdrop #fef9ec, cleared 802 of 245760 band px (0.33%; left 646 top 0 right 156 bottom 0)`.
|
||||
If more than `--max-clip` percent of the band (default 0.5) was content the file is still written, already cleared, and
|
||||
the client exits **65**: the subject ran into the frame, so it is cut flat at the seam. Re-seed or ask for more room.
|
||||
`-n` renders all images, then exits 65 if any was dirty. `--alpha` floods from the band through backdrop-coloured pixels
|
||||
and makes them transparent (hard key; pure Python, about 1 s). The output is RGBA only with `--alpha`.
|
||||
|
||||
Measured 2026-11-11 (1024x1024, 6%, one render each, same seed and ref as the unpinned image):
|
||||
|
||||
| case | unpinned: band that was content | pinned |
|
||||
|---|---|---|
|
||||
| Aqua's facelift seed 2203 (`--ref --vision`, seaweed and bubbles in the corners) | 33% | 0.33% (corners clean; character unchanged) |
|
||||
| Kiko r1a at a desk, seed 11 (green) | 31% (the desk ran off the bottom) | 0.43% (the desk now ends inside the frame) |
|
||||
|
||||
Pinned and unpinned images and the comparison: `/opt/commons/handoff/raven-ref-edit-20261111/clear-edge/`. The pin
|
||||
did not visibly degrade either render. The residue is subject touching the seam (fin, tail, desk corner), not clutter.
|
||||
Limits: a hard key leaves a faint backdrop-coloured fringe on antialiased outlines (visible on saturated green, faint on
|
||||
cream); backdrop enclosed by the subject (the gap between a headband and the head) stays opaque; a saturated key colour
|
||||
the subject also contains (green in a green-eyed character) is untouched except where connected to the band; the band is
|
||||
a rectangle, so the pin cannot ask for a vignette or a round sticker. A subject cannot be asked to stay off the seam: only
|
||||
the exit code tells you it did not.
|
||||
|
||||
Offline test (fake ComfyUI on 127.0.0.1, 127 checks): `python3 tests/test-raven-image.py`.
|
||||
|
||||
## What it does
|
||||
|
||||
|
|
|
|||
|
|
@ -5,10 +5,12 @@ Talks to ComfyUI's own API: POST /prompt, poll GET /history/<id>, GET /view,
|
|||
POST /free when the queue is empty. No retry loop: an unreachable Raven exits 75
|
||||
with "Raven may be busy or off". See README.md next to this file.
|
||||
"""
|
||||
import argparse, json, math, os, random, stat, sys, time, urllib.error, urllib.parse, urllib.request
|
||||
import argparse, json, math, os, random, re, stat, struct, sys, time, urllib.error, urllib.parse, urllib.request, zlib
|
||||
|
||||
BASE = os.environ.get("RAVEN_URL", "http://10.77.0.4:8190")
|
||||
EX_BUSY = 75 # EX_TEMPFAIL: no answer from Raven (off, busy, a read timeout mid-render, or a lock screen after a reboot); try again later, by hand
|
||||
EX_DIRTY = 65 # EX_DATAERR: --margin was asked for and the render put subject in the band; the file is written, already cleared
|
||||
TOL = 12 # per-channel distance from the band's backdrop beyond which a pixel counts as content (VAE noise on a flat fill is a few levels)
|
||||
MAX_REF = 20 * 2 ** 20 # bytes; refs persist in Raven's input folder, so cap what we put there
|
||||
MAGIC = ((b"\x89PNG\r\n\x1a\n", None), (b"\xff\xd8\xff", None), (b"RIFF", b"WEBP")) # png, jpeg, webp (bytes 8-12)
|
||||
|
||||
|
|
@ -118,6 +120,133 @@ def build(name, text, w, h, seed, enhance, refs=(), vision=False):
|
|||
return g
|
||||
|
||||
|
||||
def margin_px(pct, w, h):
|
||||
"""Band width on each axis in pixels: pct of that axis, rounded up to the 16-px latent grid so the mask edge is exact."""
|
||||
return tuple(-(-math.ceil(pct / 100 * d) // 16) * 16 for d in (w, h))
|
||||
|
||||
|
||||
def pin_edges(g, w, h, pct, bg):
|
||||
"""Native pin: sample from a flat --bg latent with a noise mask that is 0 on the border band, so the band is
|
||||
re-imposed at every step and the model composes the interior around it. Replaces node 6 (the empty latent)."""
|
||||
mx, my = margin_px(pct, w, h)
|
||||
g["30"] = {"class_type": "EmptyImage", "inputs": {"width": w, "height": h, "batch_size": 1, "color": bg}}
|
||||
g["31"] = {"class_type": "VAEEncode", "inputs": {"pixels": ["30", 0], "vae": ["3", 0]}}
|
||||
g["32"] = {"class_type": "SolidMask", "inputs": {"value": 0.0, "width": w, "height": h}} # 0 = keep, 1 = sample
|
||||
g["33"] = {"class_type": "SolidMask", "inputs": {"value": 1.0, "width": w - 2 * mx, "height": h - 2 * my}}
|
||||
g["34"] = {"class_type": "MaskComposite", "inputs": {"destination": ["32", 0], "source": ["33", 0], "x": mx, "y": my, "operation": "add"}}
|
||||
g["35"] = {"class_type": "SetLatentNoiseMask", "inputs": {"samples": ["31", 0], "mask": ["34", 0]}}
|
||||
g["7"]["inputs"]["latent_image"] = ["35", 0]
|
||||
del g["6"]
|
||||
|
||||
|
||||
def png_decode(data):
|
||||
"""8-bit non-interlaced RGB/RGBA PNG (what ComfyUI writes) -> (w, h, bpp, rows, ancillary chunks). Stdlib only."""
|
||||
if not data.startswith(b"\x89PNG\r\n\x1a\n"):
|
||||
raise ValueError("not a png")
|
||||
pos, idat, keep, hdr = 8, [], [], None
|
||||
while pos < len(data):
|
||||
n, typ = struct.unpack(">I4s", data[pos:pos + 8])
|
||||
body = data[pos + 8:pos + 8 + n]
|
||||
pos += 12 + n
|
||||
if typ == b"IHDR":
|
||||
hdr = struct.unpack(">IIBBBBB", body)
|
||||
elif typ == b"IDAT":
|
||||
idat.append(body)
|
||||
elif typ != b"IEND":
|
||||
keep.append((typ, body))
|
||||
if hdr is None or not idat:
|
||||
raise ValueError("no IHDR or no image data")
|
||||
w, h, depth, ctype, _, _, interlace = hdr
|
||||
if depth != 8 or ctype not in (2, 6) or interlace:
|
||||
raise ValueError(f"only 8-bit non-interlaced RGB/RGBA is handled (got depth {depth}, colour type {ctype}, interlace {interlace})")
|
||||
bpp = 3 if ctype == 2 else 4
|
||||
stride, rows, prev = w * bpp, [], bytearray(w * bpp)
|
||||
raw = zlib.decompressobj().decompress(b"".join(idat), h * (stride + 1) + 1) # bounded: a bomb cannot outgrow its header
|
||||
if len(raw) != h * (stride + 1):
|
||||
raise ValueError("image data does not match the header size")
|
||||
for y in range(h):
|
||||
f, cur = raw[y * (stride + 1)], bytearray(raw[y * (stride + 1) + 1:(y + 1) * (stride + 1)])
|
||||
if f == 1:
|
||||
for i in range(bpp, stride): cur[i] = (cur[i] + cur[i - bpp]) & 255
|
||||
elif f == 2:
|
||||
for i in range(stride): cur[i] = (cur[i] + prev[i]) & 255
|
||||
elif f == 3:
|
||||
for i in range(stride): cur[i] = (cur[i] + (((cur[i - bpp] if i >= bpp else 0) + prev[i]) >> 1)) & 255
|
||||
elif f == 4:
|
||||
for i in range(stride):
|
||||
a, b, c = (cur[i - bpp] if i >= bpp else 0), prev[i], (prev[i - bpp] if i >= bpp else 0)
|
||||
pa, pb, pc = abs(b - c), abs(a - c), abs(a + b - 2 * c)
|
||||
cur[i] = (cur[i] + (a if pa <= pb and pa <= pc else b if pb <= pc else c)) & 255
|
||||
elif f != 0:
|
||||
raise ValueError(f"bad PNG filter {f}")
|
||||
rows.append(cur); prev = cur
|
||||
return w, h, bpp, rows, keep
|
||||
|
||||
|
||||
def png_encode(w, h, bpp, rows, keep):
|
||||
def chunk(t, b): return struct.pack(">I", len(b)) + t + b + struct.pack(">I", zlib.crc32(t + b))
|
||||
raw = b"".join(b"\0" + bytes(r) for r in rows)
|
||||
return (b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2 if bpp == 3 else 6, 0, 0, 0))
|
||||
+ b"".join(chunk(t, b) for t, b in keep) + chunk(b"IDAT", zlib.compress(raw, 6)) + chunk(b"IEND", b""))
|
||||
|
||||
|
||||
def key_alpha(w, h, bpp, rows, back):
|
||||
"""Make the backdrop transparent: flood fill from every border pixel through pixels within TOL of `back`, so only
|
||||
backdrop connected to the (cleared) band goes; enclosed holes and same-coloured subject parts stay opaque.
|
||||
Returns RGBA rows. Hard key, no feathering: antialiased outlines keep a faint backdrop-coloured fringe."""
|
||||
out = []
|
||||
for r in rows:
|
||||
o = bytearray(w * 4)
|
||||
for c in range(3): o[c::4] = r[c::bpp]
|
||||
o[3::4] = b"\xff" * w
|
||||
out.append(o)
|
||||
near = lambda o, x: all(abs(o[x * 4 + c] - back[c]) <= TOL for c in range(3))
|
||||
stack = [(x, y) for x in range(w) for y in (0, h - 1)] + [(x, y) for y in range(h) for x in (0, w - 1)]
|
||||
while stack:
|
||||
x, y = stack.pop()
|
||||
if 0 <= x < w and 0 <= y < h and out[y][x * 4 + 3] and near(out[y], x):
|
||||
out[y][x * 4 + 3] = 0
|
||||
stack += ((x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1))
|
||||
return out
|
||||
|
||||
|
||||
def clear_edges(png, pct, alpha=False):
|
||||
"""Measure the border band of a rendered PNG against its own backdrop (the per-channel median of the band: the VAE
|
||||
shifts a saturated --bg by tens of levels, so the asked-for colour is not what a clean band decodes to), paint the
|
||||
whole band with that backdrop, and report what was content. Returns (png bytes, backdrop (r, g, b), clipped pixels per
|
||||
side {left, top, right, bottom}, band pixel count). With alpha the PNG is RGBA with the backdrop keyed out."""
|
||||
w, h, bpp, rows, keep = png_decode(png)
|
||||
mx, my = margin_px(pct, w, h)
|
||||
spans = [(y, 0, w) if y < my or y >= h - my else None for y in range(h)]
|
||||
sample = [[], [], []]
|
||||
for y in range(0, h, 3):
|
||||
for x0, x1 in ([(0, w)] if spans[y] else [(0, mx), (w - mx, w)]):
|
||||
for x in range(x0, x1, 3):
|
||||
for c in range(3): sample[c].append(rows[y][x * bpp + c])
|
||||
back = tuple(sorted(v)[len(v) // 2] for v in sample)
|
||||
want = bytes(back + ((255,) if bpp == 4 else ()))
|
||||
clip = dict.fromkeys(("left", "top", "right", "bottom"), 0)
|
||||
|
||||
def dirty(row, x0, x1):
|
||||
n = 0
|
||||
for x in range(x0, x1):
|
||||
o = x * bpp
|
||||
if abs(row[o] - back[0]) > TOL or abs(row[o + 1] - back[1]) > TOL or abs(row[o + 2] - back[2]) > TOL: n += 1
|
||||
return n
|
||||
for y in range(h):
|
||||
row, edge = rows[y], "top" if y < my else "bottom" if y >= h - my else None
|
||||
if edge: # whole row is band; the corners belong to top/bottom
|
||||
clip[edge] += dirty(row, 0, w)
|
||||
else:
|
||||
clip["left"] += dirty(row, 0, mx); clip["right"] += dirty(row, w - mx, w)
|
||||
for x0, x1 in ((0, w),) if edge else ((0, mx), (w - mx, w)):
|
||||
row[x0 * bpp:x1 * bpp] = want * (x1 - x0)
|
||||
band = w * h - (w - 2 * mx) * (h - 2 * my)
|
||||
if alpha:
|
||||
rows, bpp = key_alpha(w, h, bpp, rows, back), 4
|
||||
return png_encode(w, h, bpp, rows, keep), back, clip, band
|
||||
|
||||
|
||||
def render(g, timeout):
|
||||
try:
|
||||
r = json.loads(call("POST", "/prompt", {"prompt": g, "client_id": "raven-image"}))
|
||||
|
|
@ -164,6 +293,12 @@ def main():
|
|||
ap.add_argument("--vision", action="store_true",
|
||||
help="with --ref: also feed the first 3 refs to the text encoder (TextEncodeQwenImageEditPlus, no vae)")
|
||||
ap.add_argument("--graph", help="graphs/<name>.json (default flux3-klein, or flux3-klein-edit with --ref)")
|
||||
ap.add_argument("--margin", type=float, metavar="PCT",
|
||||
help="keep a PCT%% border band of --bg on every side (rounded up to 16 px): the sampler is pinned to it, "
|
||||
"then the band is measured and repainted; exit 65 if more than --max-clip %% of it had content")
|
||||
ap.add_argument("--bg", default="ffffff", metavar="RRGGBB", help="backdrop colour for --margin (the prompt must ask for the same one)")
|
||||
ap.add_argument("--alpha", action="store_true", help="--margin: key the backdrop connected to the band to transparent (RGBA output)")
|
||||
ap.add_argument("--max-clip", type=float, default=0.5, metavar="PCT", help="--margin: largest %% of the band that may need clearing (default 0.5)")
|
||||
ap.add_argument("--no-enhance", action="store_true", help="send the prompt as the caption, skip the upsampler")
|
||||
ap.add_argument("--timeout", type=int, default=900, help="seconds to wait per image")
|
||||
ap.add_argument("--keep-loaded", action="store_true", help="skip POST /free at the end")
|
||||
|
|
@ -172,24 +307,48 @@ def main():
|
|||
ap.error("--vision needs at least one --ref")
|
||||
if a.graph and (a.graph == "flux3-klein-edit") != bool(a.ref):
|
||||
ap.error("--ref and graph flux3-klein-edit go together (the edit graph needs refs, other graphs cannot take them)")
|
||||
if a.margin is not None and not 0 < a.margin <= 40:
|
||||
ap.error("--margin is a percentage above 0 and at most 40")
|
||||
if a.alpha and not a.margin:
|
||||
ap.error("--alpha needs --margin")
|
||||
if not re.fullmatch(r"#?[0-9a-fA-F]{6}", a.bg):
|
||||
ap.error("--bg is six hex digits, RRGGBB")
|
||||
bg = int(a.bg.lstrip("#"), 16)
|
||||
a.graph = a.graph or ("flux3-klein-edit" if a.ref else "flux3-klein")
|
||||
ref_data = [load_ref(r) for r in a.ref] # everything refusable is refused before the first request
|
||||
w, h = (int(v) for v in a.size.lower().split("x"))
|
||||
seed0 = a.seed if a.seed is not None else random.randrange(2 ** 48)
|
||||
root, ext = os.path.splitext(a.out)
|
||||
dirty = 0
|
||||
try:
|
||||
call("GET", "/system_stats", timeout=5)
|
||||
refs = [upload(r, d) for r, d in zip(a.ref, ref_data)]
|
||||
for i in range(a.n):
|
||||
seed = seed0 + i
|
||||
t = time.time()
|
||||
png = render(build(a.graph, a.prompt, w, h, seed, not a.no_enhance, refs, a.vision), a.timeout)
|
||||
g = build(a.graph, a.prompt, w, h, seed, not a.no_enhance, refs, a.vision)
|
||||
if a.margin:
|
||||
pin_edges(g, w, h, a.margin, bg)
|
||||
png = render(g, a.timeout)
|
||||
note = ""
|
||||
if a.margin:
|
||||
try:
|
||||
png, back, clip, band = clear_edges(png, a.margin, a.alpha)
|
||||
except (ValueError, zlib.error, struct.error) as e:
|
||||
sys.exit(f"raven-image: --margin: cannot read the render as a PNG ({e}); nothing written")
|
||||
pct = 100 * sum(clip.values()) / band
|
||||
note = " margin %s%%: backdrop #%02x%02x%02x, cleared %d of %d band px (%.2f%%; %s)" % (
|
||||
a.margin, *back, sum(clip.values()), band, pct, " ".join(f"{k} {v}" for k, v in clip.items()))
|
||||
if pct > a.max_clip:
|
||||
dirty += 1; note += f" DIRTY: over --max-clip {a.max_clip}%"
|
||||
path = a.out if a.n == 1 else f"{root}-{i + 1}{ext}"
|
||||
with open(path, "wb") as f:
|
||||
f.write(png)
|
||||
print(f"{path} seed={seed} {w}x{h} {time.time() - t:.0f}s", flush=True)
|
||||
print(f"{path} seed={seed} {w}x{h} {time.time() - t:.0f}s{note}", flush=True)
|
||||
if not a.keep_loaded:
|
||||
print("freed Raven's VRAM" if free_if_idle() else "queue not empty, left models loaded")
|
||||
if dirty:
|
||||
sys.exit(EX_DIRTY)
|
||||
except Busy as e:
|
||||
print(f"raven-image: Raven may be busy or off ({BASE}: {e}). Not retrying; try again later. "
|
||||
"Most likely cause after a Raven reboot: the lock screen (the ComfyUI task needs someone logged in on Raven's desktop). "
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|||
root = sys.argv[1] if len(sys.argv) > 1 else os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
tool = os.path.join(root, "client", "raven-image")
|
||||
PNG = b"\x89PNG\r\n\x1a\n" + b"\0" * 32
|
||||
seen = {"uploads": [], "prompts": []}
|
||||
seen = {"uploads": [], "prompts": [], "view": PNG}
|
||||
|
||||
|
||||
class H(BaseHTTPRequestHandler):
|
||||
|
|
@ -19,7 +19,7 @@ class H(BaseHTTPRequestHandler):
|
|||
if self.path.startswith("/history/"):
|
||||
self.send(json.dumps({self.path[9:]: {"status": {"completed": True}, "outputs": {"9": {"images": [
|
||||
{"filename": "x.png", "subfolder": "", "type": "temp"}]}}}}).encode())
|
||||
elif self.path.startswith("/view"): self.send(PNG, "image/png")
|
||||
elif self.path.startswith("/view"): self.send(seen["view"], "image/png")
|
||||
elif self.path == "/queue": self.send(b'{"queue_running":[],"queue_pending":[]}')
|
||||
else: self.send(b"{}")
|
||||
def do_POST(self):
|
||||
|
|
@ -177,6 +177,84 @@ for nm in ('a"b\r\nX-Evil: 1\\c.png', "pass wd;$x.png", "é.png"):
|
|||
p, _ = run("cap", url="http://127.0.0.1:1")
|
||||
check(p.returncode == 75 and "may be busy or off" in p.stderr and "lock screen" in p.stderr, "unreachable: exit 75, hedged message names lock screen")
|
||||
|
||||
# --margin: the pin is in the graph, the band is measured and repainted, a dirty band is loud
|
||||
import importlib.machinery, importlib.util, struct, zlib
|
||||
mod2 = importlib.util.module_from_spec(importlib.util.spec_from_loader("ri2", importlib.machinery.SourceFileLoader("ri2", tool)))
|
||||
mod2.__spec__.loader.exec_module(mod2)
|
||||
def mkpng(w, h, px, bpp=3, filt=0):
|
||||
"""px(x, y) -> tuple of bpp bytes. filt 0..4 applies that PNG filter to every row (the decoder must undo all five)."""
|
||||
rows = [bytes(v for x in range(w) for v in px(x, y)) for y in range(h)]
|
||||
raw, prev = b"", bytes(w * bpp)
|
||||
for r in rows:
|
||||
f = bytearray()
|
||||
for i in range(w * bpp):
|
||||
a, b, c = (r[i - bpp] if i >= bpp else 0), prev[i], (prev[i - bpp] if i >= bpp else 0)
|
||||
pr = (0, a, b, (a + b) >> 1, (lambda pa, pb, pc: a if pa <= pb and pa <= pc else b if pb <= pc else c)(abs(b - c), abs(a - c), abs(a + b - 2 * c)))[filt]
|
||||
f.append((r[i] - pr) & 255)
|
||||
raw += bytes([filt]) + bytes(f); prev = r
|
||||
ck = lambda t, b: struct.pack(">I", len(b)) + t + b + struct.pack(">I", zlib.crc32(t + b))
|
||||
return (b"\x89PNG\r\n\x1a\n" + ck(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2 if bpp == 3 else 6, 0, 0, 0))
|
||||
+ ck(b"tEXt", b"prompt\0keepme") + ck(b"IDAT", zlib.compress(raw)) + ck(b"IEND", b""))
|
||||
BG = (30, 200, 90)
|
||||
# codec: every filter type round-trips, ancillary chunks survive
|
||||
for f in range(5):
|
||||
for bpp in (3, 4):
|
||||
pic = lambda x, y, bpp=bpp: ((x * 7 + y * 3) & 255, (x * y) & 255, (x ^ y) & 255) + ((200,) if bpp == 4 else ())
|
||||
w_, h_, b_, rows_, keep_ = mod2.png_decode(mkpng(37, 23, pic, bpp, f))
|
||||
check((w_, h_, b_) == (37, 23, bpp) and all(bytes(rows_[y][x * bpp:(x + 1) * bpp]) == bytes(pic(x, y)) for y in range(23) for x in range(37))
|
||||
and keep_ == [(b"tEXt", b"prompt\0keepme")], f"png codec: filter {f}, {bpp} bytes/px")
|
||||
check(mod2.png_decode(mod2.png_encode(w_, h_, b_, rows_, keep_))[3] == rows_, f"png codec: encode round-trips (filter {f}, {bpp})")
|
||||
for bad in (b"nope", mkpng(8, 8, lambda x, y: BG)[:40]):
|
||||
try: mod2.png_decode(bad); check(False, "png_decode refuses garbage")
|
||||
except (ValueError, zlib.error, struct.error): check(True, "")
|
||||
# margin_px: pct of each axis, up to the 16 px latent grid
|
||||
check(mod2.margin_px(6, 1024, 1024) == (64, 64) and mod2.margin_px(5, 1536, 1024) == (80, 64) and mod2.margin_px(0.1, 1024, 1024) == (16, 16),
|
||||
"margin_px: 16 px steps per axis")
|
||||
# the pin graph: flat-bg latent, noise mask 0 on the band and 1 inside, sampler starts from it
|
||||
for extra, base in (([], "flux3-klein.json"), (["--ref", refs[0]], "flux3-klein-edit.json")):
|
||||
p, g = run(*extra, "--margin", "5", "--bg", "#12fa5f", "--size", "1536x1024", "--no-enhance", "cap")
|
||||
check(g is not None and "6" not in g and g["7"]["inputs"]["latent_image"] == ["35", 0], f"margin {base}: sampler starts from the masked latent, empty latent gone")
|
||||
if g:
|
||||
check(g["30"]["inputs"] == {"width": 1536, "height": 1024, "batch_size": 1, "color": 0x12fa5f}, f"margin {base}: flat bg image at full size")
|
||||
check(g["31"] == {"class_type": "VAEEncode", "inputs": {"pixels": ["30", 0], "vae": ["3", 0]}}, f"margin {base}: bg encoded with the graph's vae")
|
||||
check(g["32"]["inputs"] == {"value": 0.0, "width": 1536, "height": 1024} and g["33"]["inputs"] == {"value": 1.0, "width": 1536 - 160, "height": 1024 - 128},
|
||||
f"margin {base}: mask 0 outside, 1 inside by 80x64 px")
|
||||
check(g["34"]["inputs"]["x"] == 80 and g["34"]["inputs"]["y"] == 64 and g["34"]["inputs"]["operation"] == "add"
|
||||
and g["34"]["inputs"]["destination"] == ["32", 0] and g["34"]["inputs"]["source"] == ["33", 0], f"margin {base}: interior composited at the margin offset")
|
||||
check(g["35"]["inputs"] == {"samples": ["31", 0], "mask": ["34", 0]} and g["7"]["inputs"]["denoise"] == 1, f"margin {base}: noise mask on the encoded bg, denoise stays 1")
|
||||
p, g = run("--no-enhance", "cap")
|
||||
check(g is not None and "6" in g and "35" not in g, "no --margin: graph unchanged")
|
||||
for args, needle in ((["--margin", "0"], "--margin is"), (["--margin", "41"], "--margin is"), (["--margin", "5", "--bg", "green"], "--bg is"), (["--alpha"], "--alpha needs --margin")):
|
||||
refused([*args, "cap"], needle, " ".join(args))
|
||||
# the render is measured: clean band exits 0 and is exact, content in the band is cleared and exits 65, alpha keys only connected backdrop
|
||||
W = 160 # 6% -> 16 px band
|
||||
def scene(x, y, hole=False):
|
||||
if x < 6 and y < 6 or (hole and 80 <= x < 90 and 80 <= y < 90): return (250, 20, 20) # corner speck in the band; an enclosed hole is separate below
|
||||
return BG if not (50 <= x < 110 and 50 <= y < 110) else (250, 20, 20)
|
||||
seen["view"] = mkpng(W, W, lambda x, y: (BG[0] + (x + y) % 3, BG[1], BG[2]) if not (50 <= x < 110 and 50 <= y < 110) else (250, 20, 20))
|
||||
p, g = run("--margin", "6", "--size", "160x160", "--no-enhance", "cap")
|
||||
out = open(os.path.join(tmp, "o.png"), "rb").read()
|
||||
check(p.returncode == 0 and "cleared 0 of" in p.stdout and "backdrop #1ec85a" in p.stdout, f"clean band: exit 0, nothing cleared, backdrop reported ({p.stdout.strip()[-140:]})")
|
||||
r_ = mod2.png_decode(out)[3]
|
||||
check(all(bytes(r_[y][x * 3:x * 3 + 3]) == bytes(BG) for y in range(W) for x in range(W) if x < 16 or y < 16 or x >= W - 16 or y >= W - 16), "clean band: repainted exactly to the backdrop")
|
||||
check(bytes(r_[80][80 * 3:80 * 3 + 3]) == bytes((250, 20, 20)), "clean band: interior untouched")
|
||||
seen["view"] = mkpng(W, W, lambda x, y: (250, 20, 20) if (x < 12 and 30 <= y < 130) or (50 <= x < 110 and 50 <= y < 110) else BG) # a stub entering the left band
|
||||
p, g = run("--margin", "6", "--size", "160x160", "--no-enhance", "cap")
|
||||
check(p.returncode == 65 and "DIRTY" in p.stdout and "left %d" % (12 * 100) in p.stdout and "top 0" in p.stdout, f"dirty band: exit 65 naming the side ({p.returncode} {p.stdout.strip()[-160:]})")
|
||||
r_ = mod2.png_decode(open(os.path.join(tmp, "o.png"), "rb").read())[3]
|
||||
check(all(bytes(r_[y][x * 3:x * 3 + 3]) == bytes(BG) for y in range(W) for x in range(12)), "dirty band: file is written and the stub cleared")
|
||||
p, g = run("--margin", "6", "--max-clip", "50", "--size", "160x160", "--no-enhance", "cap")
|
||||
check(p.returncode == 0, "--max-clip raises the tolerance")
|
||||
holes = lambda x, y: (250, 20, 20) if (50 <= x < 110 and 50 <= y < 110 and not (70 <= x < 90 and 70 <= y < 90)) else BG # ring around an enclosed backdrop hole
|
||||
seen["view"] = mkpng(W, W, holes)
|
||||
p, g = run("--margin", "6", "--alpha", "--size", "160x160", "--no-enhance", "cap")
|
||||
w_, h_, b_, r_, k_ = mod2.png_decode(open(os.path.join(tmp, "o.png"), "rb").read())
|
||||
check(p.returncode == 0 and b_ == 4 and r_[0][3] == 0 and r_[30][3 + 4 * 40] == 0 and r_[60][4 * 60 + 3] == 255, "--alpha: RGBA, outside keyed out, subject opaque")
|
||||
check(r_[80][4 * 80 + 3] == 255, "--alpha: enclosed backdrop hole stays opaque (connected backdrop only)")
|
||||
seen["view"] = PNG
|
||||
p, g = run("--margin", "6", "--no-enhance", "cap")
|
||||
check(p.returncode != 0 and "cannot read the render" in p.stderr and "Traceback" not in p.stderr, "undecodable render: clean error, no traceback")
|
||||
|
||||
srv.shutdown()
|
||||
print(f"{count[0] - len(fails)}/{count[0]} checks passed")
|
||||
sys.exit(1 if fails else 0)
|
||||
|
|
|
|||
Loading…
Reference in a new issue