File size: 3,530 Bytes
3da88e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | import hashlib
import json
import mlx.core as mx
import numpy as np
import pytest
import torch
from PIL import Image
from h3_mlx.compare import make_comparison
from h3_mlx.reference import reflect_pad
from h3_mlx.vae import assemble_tiles, split_tiles, to_image
def test_tiling_native_overlap_and_coverage():
for length in (32, 256, 512, 1024, 2048, 4096):
starts, lens, overlaps = split_tiles(length)
assert starts[0] == 0 and starts[-1] + lens[-1] == length
assert sum(lens) - sum(overlaps) == length
assert all(v % 16 == 0 for v in overlaps)
assert split_tiles(512) == ([0, 128, 256], [256] * 3, [128, 128])
def test_native_vertical_then_horizontal_overlap_blending():
rows = [
[mx.full((1, 1, 4, 4), i * 2 + j, mx.float32) for j in range(2)]
for i in range(2)
]
out = np.asarray(assemble_tiles(rows, [2], [2]))[0, 0]
# Hand-computed native blend: row 1/col 1 at its first overlap pixel
# inherits the original tile above, then the original tile to its left.
assert out.shape == (6, 6)
assert out[0, 0] == 0 and out[-1, -1] == 3
assert out[2, 2] == 2 # horizontal blend takes the unblended left source
assert out[3, 3] == 2
def test_reflection_pad_matches_pytorch():
x = torch.arange(2 * 3 * 4 * 2, dtype=torch.float32).reshape(2, 3, 4, 2)
out = reflect_pad(mx.array(x.numpy()), 1, 1, 1, 1)
ref = torch.nn.functional.pad(
x.permute(0, 3, 1, 2), (1, 1, 1, 1), mode="reflect"
).permute(0, 2, 3, 1)
np.testing.assert_array_equal(np.asarray(out), ref.numpy())
def test_comparison_rejects_mismatched_seed_and_retains_native_crops(tmp_path):
a = tmp_path / "ref.png"
b = tmp_path / "candidate.png"
Image.new("RGB", (64, 64), (50, 70, 90)).save(a)
Image.new("RGB", (64, 64), (51, 70, 90)).save(b)
metadata = {
"prompt": "test",
"seed": 7,
"size": [64, 64],
"canvas": [64, 64],
"rng": "torch",
"reference_images": [],
}
a.with_suffix(".json").write_text(json.dumps(metadata))
b.with_suffix(".json").write_text(json.dumps({**metadata, "seed": 8}))
with pytest.raises(ValueError, match="seed"):
make_comparison(a, [b], tmp_path / "bad")
b.with_suffix(".json").write_text(json.dumps(metadata))
result = make_comparison(a, [b], tmp_path / "report", crop=(1, 2, 33, 34))
assert result["crop_xyxy"] == [1, 2, 33, 34]
assert result["rows"][1]["mae_8bit"] == pytest.approx(1 / 3)
assert (tmp_path / "report/stack.png").exists()
assert (tmp_path / "report/grid.png").exists()
a.with_suffix(".trace.safetensors").write_bytes(b"reference latent")
decoder = {
**metadata,
"validation_scope": "Same final latent; decoder only",
"source_trace_sha256": "different latent",
}
b.with_suffix(".json").write_text(json.dumps(decoder))
with pytest.raises(ValueError, match="final latent"):
make_comparison(a, [b], tmp_path / "wrong-latent")
decoder["source_trace_sha256"] = hashlib.sha256(b"reference latent").hexdigest()
b.with_suffix(".json").write_text(json.dumps(decoder))
verified = make_comparison(a, [b], tmp_path / "decoder")
assert "decoder-only" in verified["validation_scope"]
def test_output_quantization_truncates_like_native_writer():
x = mx.array([[[[0.5, 0.501]], [[0.25, 0.251]], [[0.75, 0.751]]]])
image = to_image(x, 2, 1)
np.testing.assert_array_equal(np.asarray(image), [[[127, 63, 191], [127, 64, 191]]])
|