2c5bd84c10
The int_sr module (WSR1 prior, wai-upscale:sr2/sr4 lowered to one wai-ops/1 graph), the exact_upscale module (upscale record with residual and target, presentation statement), synthesis's samples subject and SR evidence, the corpus generator, the second implementation, the training tool and its trained priors. Co-Authored-By: Claude Opus 5.5 <[email protected]>
152 lines
5.6 KiB
Python
152 lines
5.6 KiB
Python
"""The integer super-resolution derivations of extensions/exact-upscale.md §3, written from
|
||
that text in integer arithmetic of their own, standard library only: the `WSR1` prior (§3.1),
|
||
`wai-upscale:sr2` (§3.2) and `wai-upscale:sr4` with its limits (§3.3). Shared by
|
||
`upscale_verify.py` and `synthesis_verify.py`; it shares no code with wai-rs.
|
||
"""
|
||
import hashlib
|
||
import struct
|
||
|
||
SR2, SR4 = "wai-upscale:sr2", "wai-upscale:sr4"
|
||
FACTOR = {SR2: 2, SR4: 4}
|
||
CAP = 1 << 26
|
||
|
||
|
||
class Unreadable(Exception):
|
||
"""A prior that breaks §3.1, or a picture past §3.3's limits: no output."""
|
||
|
||
def __init__(self, code):
|
||
super().__init__(code)
|
||
self.code = code
|
||
|
||
|
||
def read_prior(b):
|
||
"""§3.1, in its order: `bad_magic`, `version`, `field_range`, `truncated`, `trailing`."""
|
||
def need(at, n):
|
||
if len(b) < at + n:
|
||
raise Unreadable("truncated")
|
||
need(0, 4)
|
||
if b[:4] != b"WSR1":
|
||
raise Unreadable("bad_magic")
|
||
need(4, 1)
|
||
if b[4] != 1:
|
||
raise Unreadable("version")
|
||
need(5, 3)
|
||
n, act_hi = b[5], struct.unpack_from("<H", b, 6)[0]
|
||
if not 1 <= n <= 8 or not 1 <= act_hi <= 32767:
|
||
raise Unreadable("field_range")
|
||
at, cin, layers = 8, 3, []
|
||
for i in range(n):
|
||
need(at, 4)
|
||
k, shift, cout = b[at], b[at + 1], struct.unpack_from("<H", b, at + 2)[0]
|
||
last = i + 1 == n
|
||
if k not in (1, 3, 5) or shift > 31 or (cout != 12 if last else not 1 <= cout <= 64):
|
||
raise Unreadable("field_range")
|
||
at += 4
|
||
nw = cout * cin * k * k
|
||
need(at, nw + 6 * cout)
|
||
w = list(struct.unpack_from("<%db" % nw, b, at))
|
||
at += nw
|
||
bias = list(struct.unpack_from("<%di" % cout, b, at))
|
||
at += 4 * cout
|
||
mult = list(struct.unpack_from("<%dH" % cout, b, at))
|
||
at += 2 * cout
|
||
if 0 in mult:
|
||
raise Unreadable("field_range")
|
||
layers.append({"k": k, "shift": shift, "cin": cin, "cout": cout, "w": w, "bias": bias, "mult": mult})
|
||
cin = cout
|
||
if at != len(b):
|
||
raise Unreadable("trailing")
|
||
return {"act_hi": act_hi, "layers": layers}
|
||
|
||
|
||
def up2(px, w, h):
|
||
"""§3.2 step 1: integer-operators §5E E11, the last row and column repeated."""
|
||
out = bytearray(4 * w * h * 3)
|
||
ow = 2 * w
|
||
for y in range(h):
|
||
for x in range(w):
|
||
x1, y1 = min(x + 1, w - 1), min(y + 1, h - 1)
|
||
for c in range(3):
|
||
a, b2 = px[(y * w + x) * 3 + c], px[(y * w + x1) * 3 + c]
|
||
c2, d = px[(y1 * w + x) * 3 + c], px[(y1 * w + x1) * 3 + c]
|
||
out[((2 * y) * ow + 2 * x) * 3 + c] = a
|
||
out[((2 * y) * ow + 2 * x + 1) * 3 + c] = (a + b2 + 1) >> 1
|
||
out[((2 * y + 1) * ow + 2 * x) * 3 + c] = (a + c2 + 1) >> 1
|
||
out[((2 * y + 1) * ow + 2 * x + 1) * 3 + c] = (a + b2 + c2 + d + 2) >> 2
|
||
return bytes(out)
|
||
|
||
|
||
def half_up(p, s):
|
||
"""integer-operators §2 R_half_up: ⌊(p + 2^(s−1)) / 2^s⌋, p for s = 0."""
|
||
return p if s == 0 else (p + (1 << (s - 1))) >> s
|
||
|
||
|
||
def sr2(prior, px, w, h):
|
||
"""§3.2 over a `w × h` RGB8 picture: `2w × 2h` RGB8. Refuses a picture past §3.3's limits."""
|
||
if w < 1 or h < 1 or len(px) != w * h * 3:
|
||
raise Unreadable("geometry")
|
||
widest = max(l["cout"] for l in prior["layers"])
|
||
if widest * w * h > CAP or 12 * w * h > CAP:
|
||
raise Unreadable("too_large")
|
||
n = w * h
|
||
# Step 2: centred, channel-major.
|
||
x = [[px[i * 3 + c] - 128 for i in range(n)] for c in range(3)]
|
||
for li, l in enumerate(prior["layers"]):
|
||
k, r, cin = l["k"], l["k"] // 2, l["cin"]
|
||
last = li + 1 == len(prior["layers"])
|
||
y = []
|
||
for o in range(l["cout"]):
|
||
taps = []
|
||
for ci in range(cin):
|
||
for ky in range(k):
|
||
for kx in range(k):
|
||
wt = l["w"][((o * cin + ci) * k + ky) * k + kx]
|
||
if wt:
|
||
taps.append((wt, ci, ky - r, kx - r))
|
||
plane = []
|
||
for yy in range(h):
|
||
for xx in range(w):
|
||
a = 0
|
||
for wt, ci, dy, dx in taps:
|
||
sy, sx = yy + dy, xx + dx
|
||
if 0 <= sy < h and 0 <= sx < w:
|
||
a += wt * x[ci][sy * w + sx]
|
||
v = half_up((a + l["bias"][o]) * l["mult"][o], l["shift"])
|
||
plane.append(v if last else min(max(v, 0), prior["act_hi"]))
|
||
y.append(plane)
|
||
x = y
|
||
base = up2(px, w, h)
|
||
ow, oh = 2 * w, 2 * h
|
||
out = bytearray(ow * oh * 3)
|
||
for c in range(3):
|
||
for oy in range(oh):
|
||
for ox in range(ow):
|
||
d = x[4 * c + 2 * (oy % 2) + ox % 2][(oy // 2) * w + ox // 2]
|
||
i = (oy * ow + ox) * 3 + c
|
||
out[i] = min(max(base[i] + d, 0), 255)
|
||
return bytes(out)
|
||
|
||
|
||
_CACHE = {}
|
||
|
||
|
||
def derive(derivation, prior_bytes, px, w, h):
|
||
"""`wai-upscale:sr2` or `wai-upscale:sr4` (§3.3) under a prior's bytes: the output, or
|
||
Unreadable. Results are kept by the prior's digest and the input, since a verifier meets
|
||
one derivation many times."""
|
||
key = (derivation, hashlib.sha256(prior_bytes).digest(), bytes(px), w, h)
|
||
if key not in _CACHE:
|
||
try:
|
||
p = read_prior(prior_bytes)
|
||
out = sr2(p, px, w, h)
|
||
if derivation == SR4:
|
||
out = sr2(p, out, 2 * w, 2 * h)
|
||
elif derivation != SR2:
|
||
raise Unreadable("derivation")
|
||
_CACHE[key] = out
|
||
except Unreadable as e:
|
||
_CACHE[key] = e
|
||
r = _CACHE[key]
|
||
if isinstance(r, Unreadable):
|
||
raise r
|
||
return r
|