160 lines
5.0 KiB
Python
160 lines
5.0 KiB
Python
import io
|
|
import time
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from sglang.srt.multimodal.inkling.image_processing import (
|
|
IMAGE_MEAN,
|
|
IMAGE_STD,
|
|
PAD_NORM,
|
|
_encode_image_bytes,
|
|
_fill_patches_numba,
|
|
)
|
|
from sglang.srt.rust_extensions._multimodal import inkling as _rs_inkling
|
|
|
|
PS = 40
|
|
|
|
|
|
def ref_patchify(arr: np.ndarray) -> torch.Tensor:
|
|
h, w, _ = arr.shape
|
|
nph = (h + PS - 1) // PS
|
|
npw = w // PS + 1
|
|
patches = np.empty((nph * npw, PS, PS, 3), dtype=np.float32)
|
|
_fill_patches_numba(arr, PS, patches, IMAGE_MEAN, IMAGE_STD, PAD_NORM)
|
|
return torch.from_numpy(patches).to(torch.bfloat16)
|
|
|
|
|
|
def rs_patchify(arr: np.ndarray) -> torch.Tensor:
|
|
h, w, _ = arr.shape
|
|
nph = (h + PS - 1) // PS
|
|
npw = w // PS + 1
|
|
bits = _rs_inkling.patchify_rgb(arr, PS)
|
|
return torch.from_numpy(bits).view(torch.bfloat16).reshape(nph * npw, PS, PS, 3)
|
|
|
|
|
|
def rs_decode_patchify(data: bytes) -> torch.Tensor:
|
|
h, w, bits = _rs_inkling.decode_patchify(data, PS)
|
|
nph = (h + PS - 1) // PS
|
|
npw = w // PS + 1
|
|
return torch.from_numpy(bits).view(torch.bfloat16).reshape(nph * npw, PS, PS, 3)
|
|
|
|
|
|
def make_photo_like(h: int, w: int, seed: int = 0) -> np.ndarray:
|
|
rng = np.random.default_rng(seed)
|
|
yy, xx = np.mgrid[0:h, 0:w]
|
|
base = np.stack(
|
|
[
|
|
127 + 100 * np.sin(yy / 97.0) * np.cos(xx / 131.0),
|
|
127 + 100 * np.cos(yy / 61.0) * np.sin(xx / 89.0),
|
|
127 + 100 * np.sin((xx + yy) / 149.0),
|
|
],
|
|
axis=-1,
|
|
)
|
|
noise = rng.normal(0, 12, (h // 8 + 1, w // 8 + 1, 3))
|
|
noise = np.kron(noise, np.ones((8, 8, 1)))[:h, :w]
|
|
return np.clip(base + noise, 0, 255).astype(np.uint8)
|
|
|
|
|
|
def encode(arr: np.ndarray, fmt: str) -> bytes:
|
|
buf = io.BytesIO()
|
|
Image.fromarray(arr).save(
|
|
buf, format=fmt, **({"quality": 90} if fmt == "JPEG" else {})
|
|
)
|
|
return buf.getvalue()
|
|
|
|
|
|
def parity_a():
|
|
print("=== Parity A: patchify from decoded array (expect bit-exact) ===")
|
|
rng = np.random.default_rng(42)
|
|
for h, w in [(1080, 1920), (1920, 1080), (40, 40), (37, 53), (720, 1280), (1, 1)]:
|
|
arr = rng.integers(0, 256, (h, w, 3), dtype=np.uint8)
|
|
ref, got = ref_patchify(arr), rs_patchify(arr)
|
|
exact = torch.equal(ref.view(torch.uint16), got.view(torch.uint16))
|
|
print(f" {h}x{w}: shape {tuple(got.shape)} bit-exact={exact}")
|
|
assert exact, f"parity A failed at {h}x{w}"
|
|
|
|
|
|
def parity_b():
|
|
print("=== Parity B: full decode path ===")
|
|
arr = make_photo_like(1080, 1920)
|
|
for fmt in ["PNG", "JPEG"]:
|
|
data = encode(arr, fmt)
|
|
ref = _encode_image_bytes(
|
|
data,
|
|
patch_size=PS,
|
|
rescale_image_frac=None,
|
|
rescale_image_max_upscaled_long_edge=None,
|
|
)
|
|
got = rs_decode_patchify(data)
|
|
got2 = got.view(got.shape[0], 1, PS, PS, 3).expand(-1, 2, -1, -1, -1)
|
|
if torch.equal(
|
|
ref.contiguous().view(torch.uint16), got2.contiguous().view(torch.uint16)
|
|
):
|
|
print(f" {fmt}: bit-exact=True ({len(data) / 1e6:.2f}MB)")
|
|
else:
|
|
d = (ref.float() - got2.float()).abs()
|
|
print(
|
|
f" {fmt}: bit-exact=False max_abs={d.max():.6f} mean_abs={d.mean():.8f} "
|
|
f"(decoder difference; normalized-feature units)"
|
|
)
|
|
|
|
|
|
def bench():
|
|
print("=== Benchmark (1080p, patch_size=40) ===")
|
|
arr = make_photo_like(1080, 1920)
|
|
jpeg = encode(arr, "JPEG")
|
|
n = 30
|
|
|
|
_encode_image_bytes(
|
|
jpeg,
|
|
patch_size=PS,
|
|
rescale_image_frac=None,
|
|
rescale_image_max_upscaled_long_edge=None,
|
|
)
|
|
rs_decode_patchify(jpeg)
|
|
_rs_inkling.decode_patchify_batch([jpeg] * 5, PS)
|
|
|
|
def run(label, fn, iters=n, images_per_call=1):
|
|
t0, c0 = time.perf_counter(), time.process_time()
|
|
for _ in range(iters):
|
|
fn()
|
|
wall = (time.perf_counter() - t0) / iters / images_per_call * 1e3
|
|
cpu = (time.process_time() - c0) / iters / images_per_call * 1e3
|
|
print(f" {label:42} wall {wall:8.2f} ms/img cpu {cpu:8.2f} ms/img")
|
|
return wall, cpu
|
|
|
|
w_py, c_py = run(
|
|
"python (PIL + numba + bf16 cast)",
|
|
lambda: _encode_image_bytes(
|
|
jpeg,
|
|
patch_size=PS,
|
|
rescale_image_frac=None,
|
|
rescale_image_max_upscaled_long_edge=None,
|
|
),
|
|
)
|
|
w_rs, c_rs = run("rust decode_patchify", lambda: rs_decode_patchify(jpeg))
|
|
w_rb, c_rb = run(
|
|
"rust decode_patchify_batch (5 imgs/call)",
|
|
lambda: _rs_inkling.decode_patchify_batch([jpeg] * 5, PS),
|
|
iters=max(n // 5, 5),
|
|
images_per_call=5,
|
|
)
|
|
|
|
run("python numba patchify only", lambda: ref_patchify(arr))
|
|
run("rust patchify_rgb only", lambda: rs_patchify(arr))
|
|
|
|
print(
|
|
f"\n speedup vs python: single {w_py / w_rs:.1f}x wall / {c_py / c_rs:.1f}x cpu, "
|
|
f"batch {w_py / w_rb:.1f}x wall / {c_py / c_rb:.1f}x cpu"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
torch.set_num_threads(8)
|
|
parity_a()
|
|
parity_b()
|
|
bench()
|
|
print("\nOK")
|