| import json |
| from types import SimpleNamespace |
|
|
| import mlx.core as mx |
| import pytest |
|
|
| from h3_mlx.config import SIGMAS |
| from h3_mlx.runtime import Pipeline |
|
|
|
|
| def test_path_cache_arguments_can_be_saved_in_inference_json(tmp_path): |
| pipe = Pipeline(tmp_path, tmp_path, cache_dir=tmp_path / "cache") |
| record = {"dit": pipe.dit_options, "encoder": pipe.encoder_options} |
| saved = json.loads(json.dumps(record)) |
| assert saved["dit"]["cache_dir"] == str(tmp_path / "cache") |
| assert saved["encoder"]["cache_dir"] == str(tmp_path / "cache") |
|
|
|
|
| def test_nonfinite_prediction_is_rejected_before_image_write(tmp_path, monkeypatch): |
| from h3_mlx import encoder |
|
|
| class FakeEncoder: |
| def __init__(self, *args, **kwargs): |
| self.cache_hit = True |
| self.store = SimpleNamespace(clear=lambda: None) |
|
|
| def encode(self, *args, **kwargs): |
| return mx.zeros((1, 5120), mx.bfloat16), mx.ones(1, mx.int32) |
|
|
| class BadModel: |
| bits = None |
| group_size = 64 |
| block_limit = 0 |
| store = SimpleNamespace(fingerprint=lambda: "test") |
| metrics = {} |
| refiner_cache_hit = False |
|
|
| def prepare_modulation(self, *args, **kwargs): |
| pass |
|
|
| def refine(self, x): |
| return x |
|
|
| def predict(self, rows, *args, **kwargs): |
| return mx.full(rows.shape, float("nan")) |
|
|
| monkeypatch.setattr(encoder, "TextEncoder", FakeEncoder) |
| (tmp_path / "pdd_config.json").write_text(json.dumps({"sigmas": SIGMAS})) |
| pipe = Pipeline(tmp_path, tmp_path) |
| pipe.model = BadModel() |
| output = tmp_path / "invalid.png" |
| with pytest.raises(FloatingPointError, match="original head 0"): |
| pipe.generate( |
| "test", width=32, height=32, min_area=1024, rng="mlx", output=output |
| ) |
| assert not output.exists() |
|
|
