| 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() |
|
|
