| 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] |
| |
| |
| assert out.shape == (6, 6) |
| assert out[0, 0] == 0 and out[-1, -1] == 3 |
| assert out[2, 2] == 2 |
| 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]]]) |
|
|
