MiniMax-H3-Image-mlx / tests /test_runtime_guard.py
devin-lai's picture
Release v0.4.0: selective MLX model downloads
3da88e3 verified
Raw History Blame Contribute Delete
1.45 kB
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()