File size: 1,445 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 | 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_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()
|