Spaces:
Running on Zero
Running on Zero
| """Check the Space's copied averaging code against the repository's original. | |
| Run from the repo root inside the project container: | |
| docker compose run --rm train bash -c "pip install -q onnxruntime gradio huggingface_hub && python space/test_space.py" | |
| """ | |
| import os | |
| import sys | |
| import time | |
| import torch | |
| from PIL import Image | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| REPO = os.path.dirname(HERE) | |
| sys.path.insert(0, HERE) | |
| os.chdir(HERE) | |
| import app # noqa: E402 | |
| import ensemble # noqa: E402 (downloads config.json + the MIX ONNX from the Hub) | |
| sys.path.insert(0, REPO) | |
| from detail_diffuser.inference.core import tta_average as ref_tta, unit_normals as ref_unit # noqa: E402 | |
| img = Image.open(os.path.join(HERE, "examples", "3185371756_ff4e9fa8a6.jpg")) | |
| from framing import letterbox # noqa: E402 | |
| sq, box = letterbox(img, app.SIZE) | |
| x = app.to_tensor(sq) | |
| assert x.shape == (1, 3, app.SIZE, app.SIZE) and x.min() >= -1 and x.max() <= 1 | |
| assert box[2] == app.SIZE or box[3] == app.SIZE, box # the long side fills the square | |
| t = time.time() | |
| fast = app.refine(x) | |
| print(f"fast: {time.time() - t:.1f}s, maps {tuple(fast.shape)}") | |
| assert fast.shape == (1, 7, app.SIZE, app.SIZE) and torch.isfinite(fast).all() | |
| # The exported graph samples its input-perturbation noise on every call, so | |
| # two runs of the real model differ slightly (mean ~1e-3). The averaging code | |
| # is therefore compared with a deterministic stand-in model, which must match | |
| # the repository's implementation exactly. | |
| def stand_in(c): | |
| return torch.cat([c * 0.7, c.mean(1, keepdim=True) * 0.5, c.flip(-1) * 0.3], 1) | |
| mine = app.unit_normals(ensemble.tta_average(stand_in, x, flip=True, rot_deg=app.ROT_DEG)) | |
| ref = ref_unit(ref_tta(lambda c, nz=None: stand_in(c), x, flip=True, rot_deg=app.ROT_DEG)) | |
| d = float((mine - ref).abs().max()) | |
| print("averaging code, max |space - repo| =", d) | |
| assert d == 0.0, d | |
| e1 = app.unit_normals(ensemble.tta_average(app.refine, x, flip=True, rot_deg=app.ROT_DEG)) | |
| assert e1.shape == (1, 7, app.SIZE, app.SIZE) and torch.isfinite(e1).all() | |
| n = e1[:, 4:7].norm(dim=1) | |
| assert float((n - 1).abs().max()) < 1e-4, "normals are not unit length" | |
| print("real model, 6 passes ok, mean |maps| =", float(e1.abs().mean())) | |
| # median_of against NumPy's nanmedian on real stand-in passes (corners are covered by fewer passes) | |
| import numpy as np # noqa: E402 | |
| passes = app.tta_passes(lambda c: stand_in(c) + 0.05 * torch.randn(1, 7, app.SIZE, app.SIZE), | |
| x, flip=True, rot_deg=app.ROT_DEG) | |
| assert len(passes) == 6 and any(float(v.min()) < 0.5 for _, v in passes), "rotated passes should have uncovered corners" | |
| P = np.stack([p.numpy() for p, _ in passes]).copy() | |
| V = np.stack([v.numpy() for _, v in passes]) > 0.5 | |
| P[~np.broadcast_to(V, P.shape)] = np.nan | |
| want = np.nanmedian(P, axis=0) | |
| got = app.median_of(passes).numpy() | |
| dm = float(np.abs(got - want).max()) | |
| print("median_of vs numpy nanmedian, max diff =", dm) | |
| assert dm < 1e-5, dm | |
| assert float((app.median_of(passes) - app.mean_of(passes)).abs().mean()) > 1e-4, "median should differ from mean on noisy passes" | |
| one = [passes[0]] | |
| assert torch.equal(app.median_of(one), app.mean_of(one)), "one pass: median == mean" | |
| out = app.predict(img, app.MODE_SIX, False) | |
| assert len(out) == 6 and out[0].size == out[1].size == img.size, (out[0].size, img.size) # whole photo, own shape | |
| assert out[4] is None and "(mean)" in out[5] | |
| print("predict ok:", out[5]) | |
| out = app.predict(img, app.MODE_SIX, True) | |
| assert out[4] is not None and out[4].size == img.size and "(median)" in out[5], out[5] | |
| print("median ok:", out[5]) | |
| out = app.predict(img, app.MODE_SINGLE, True) # median outside the six-pass mode: single pass, note says so | |
| assert out[4] is None and "1 forward pass in" in out[5] and "six-pass mode" in out[5], out[5] | |
| out = app.predict(img, app.MODE_REFINE, False) # single pass + the small refiner | |
| assert out[4] is None and "refiner" in out[5] and out[1].size == img.size, out[5] | |
| print("refiner mode ok:", out[5]) | |
| # odd shapes: a panorama (padding longer than the photo), tall, square, tiny, big | |
| for size in [(2000, 300), (300, 2000), (448, 448), (64, 48), (1500, 1100)]: | |
| im = Image.effect_noise(size, 40).convert("RGB") | |
| out = app.predict(im, app.MODE_REFINE, False) | |
| ow, oh = app.display_size(*size) | |
| assert out[0].size == out[1].size == out[2].size == out[3].size == (ow, oh), (size, out[0].size, (ow, oh)) | |
| print("odd shapes ok") | |
| print("space test ok") | |