File size: 1,688 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 | import json
import mlx.core as mx
import numpy as np
from h3_mlx.config import DiTConfig
from h3_mlx.dit import DiT
def test_refined_conditioning_cache_is_exact_and_input_specific(tmp_path):
model_path = tmp_path / "model"
transformer = model_path / "transformer"
transformer.mkdir(parents=True)
cfg = DiTConfig(hidden_size=32, text_dim=24, token_refiner_num_layers=0)
(transformer / "config.json").write_text(json.dumps(vars(cfg)))
mx.random.seed(732)
weights = {
"condition_proj.weight": (mx.random.normal((32, 24)) * 0.1).astype(mx.bfloat16),
"condition_proj.bias": mx.zeros(32, mx.bfloat16),
"token_refiner.final_norm.weight": mx.ones(32, mx.bfloat16),
}
mx.save_safetensors(str(transformer / "model.safetensors"), weights)
mx.save_safetensors(
str(model_path / "pdd_fused_heads.safetensors"), {"unused": mx.zeros(1)}
)
model = DiT(model_path, cache_dir=tmp_path / "cache")
text = mx.random.normal((17, 24)).astype(mx.bfloat16)
first = model.refine(text)
assert not model.refiner_cache_hit
second = model.refine(text)
assert model.refiner_cache_hit
np.testing.assert_array_equal(
np.asarray(first.view(mx.uint16)), np.asarray(second.view(mx.uint16))
)
model.refine(text + mx.array(1, mx.bfloat16))
assert not model.refiner_cache_hit
assert len(list((tmp_path / "cache").glob("refined-*.safetensors"))) == 2
model.cache_refined_text = False
direct = model.refine(text)
assert not model.refiner_cache_hit
np.testing.assert_array_equal(
np.asarray(first.view(mx.uint16)), np.asarray(direct.view(mx.uint16))
)
model.clear()
|