| import json |
| from pathlib import Path |
|
|
| import pytest |
|
|
| from h3_mlx import hub |
| from h3_mlx.config import SIGMAS |
|
|
|
|
| @pytest.mark.parametrize( |
| "name", ["/absolute", "../escape", "a/../b", "a//b", "a\\b", ""] |
| ) |
| def test_manifest_paths_cannot_escape_assembly(name): |
| with pytest.raises(ValueError): |
| hub.relative_path(name) |
|
|
|
|
| def test_prepared_catalog_does_not_attempt_downloads(tmp_path, monkeypatch): |
| import huggingface_hub |
|
|
| def forbidden(*args, **kwargs): |
| pytest.fail("Prepared releases must not attempt network access") |
|
|
| monkeypatch.setattr(huggingface_hub, "snapshot_download", forbidden) |
| catalog = tmp_path / "prepared-catalog.json" |
| catalog.write_text( |
| json.dumps( |
| { |
| "format": "h3-mlx-catalog-v1", |
| "status": "prepared", |
| "variants": {"packed12": {}}, |
| } |
| ) |
| ) |
| with pytest.raises(RuntimeError, match="publication is pending"): |
| hub.download_variant("packed12", tmp_path / "model", catalog_path=catalog) |
|
|
|
|
| def test_selective_assembly_reuses_shared_files_and_protects_existing_models( |
| tmp_path, monkeypatch |
| ): |
| import huggingface_hub |
|
|
| entries = { |
| "shared": {"repo_id": "owner/shared", "revision": "a" * 40}, |
| "q4": { |
| "repo_id": "owner/q4", |
| "revision": "b" * 40, |
| "core_format": "quantized", |
| "bits": 4, |
| }, |
| "q8": { |
| "repo_id": "owner/q8", |
| "revision": "c" * 40, |
| "core_format": "quantized", |
| "bits": 8, |
| }, |
| } |
| snapshots = {} |
| for name, entry in entries.items(): |
| snapshot = tmp_path / name |
| file = snapshot / ( |
| "model/vae/decoder.safetensors" |
| if name == "shared" |
| else "model/core/core-000.safetensors" |
| ) |
| file.parent.mkdir(parents=True) |
| file.write_bytes(name.encode()) |
| (snapshot / "config.json").write_text(json.dumps({"component": name})) |
| (snapshot / "LICENSE").write_text("retained license") |
| files = { |
| p.relative_to(snapshot).as_posix(): { |
| "bytes": p.stat().st_size, |
| "sha256": hub.sha256(p), |
| } |
| for p in (file, snapshot / "config.json", snapshot / "LICENSE") |
| } |
| (snapshot / "files.json").write_text( |
| json.dumps({"format": "h3-mlx-files-v1", "files": files}) |
| ) |
| snapshots[entry["repo_id"]] = snapshot |
| catalog = tmp_path / "catalog.json" |
| catalog.write_text( |
| json.dumps( |
| { |
| "format": "h3-mlx-catalog-v1", |
| "status": "published", |
| "source": {}, |
| "sigmas": SIGMAS, |
| "shared": entries["shared"], |
| "variants": {name: entries[name] for name in ("q4", "q8")}, |
| } |
| ) |
| ) |
| calls = [] |
|
|
| def get_file(repo, filename, **kwargs): |
| calls.append((repo, filename, kwargs)) |
| return str(snapshots[repo] / filename) |
|
|
| def snapshot(repo, **kwargs): |
| calls.append((repo, "snapshot", kwargs)) |
| return str(snapshots[repo]) |
|
|
| monkeypatch.setattr(huggingface_hub, "hf_hub_download", get_file) |
| monkeypatch.setattr(huggingface_hub, "snapshot_download", snapshot) |
| first = hub.download_variant( |
| "q4", |
| tmp_path / "model-q4", |
| catalog_path=catalog, |
| skip_encoder=True, |
| verify=True, |
| ) |
| assert {call[0] for call in calls} == {"owner/shared", "owner/q4"} |
| assert all( |
| "config.json" in call[2]["allow_patterns"] |
| for call in calls |
| if call[1] == "snapshot" |
| ) |
| assert Path(first["model_path"], "core/core-000.safetensors").read_bytes() == b"q4" |
| assert Path(first["model_path"], "licenses/shared/LICENSE").is_file() |
| assert hub.load_release(first["model_path"])["bits"] == 4 |
| calls.clear() |
| hub.download_variant( |
| "q8", tmp_path / "model-q8", catalog_path=catalog, skip_encoder=True |
| ) |
| assert (tmp_path / "model-q4/vae/decoder.safetensors").resolve() == ( |
| tmp_path / "model-q8/vae/decoder.safetensors" |
| ).resolve() |
| with pytest.raises(FileExistsError, match="separate empty"): |
| hub.download_variant( |
| "q8", tmp_path / "model-q4", catalog_path=catalog, skip_encoder=True |
| ) |
| (tmp_path / "q4/model/core/core-000.safetensors").write_bytes(b"XX") |
| with pytest.raises(ValueError, match="SHA-256 mismatch"): |
| hub.download_variant( |
| "q4", |
| tmp_path / "model-q4", |
| catalog_path=catalog, |
| skip_encoder=True, |
| verify=True, |
| ) |
|
|
|
|
| def test_unpinned_revisions_are_rejected_before_network_access(): |
| with pytest.raises(ValueError, match="immutable"): |
| hub._manifest_snapshot({"repo_id": "owner/model", "revision": "main"}) |
|
|
|
|
| def test_link_cannot_replace_an_unrelated_file(tmp_path): |
| source, destination = tmp_path / "source", tmp_path / "destination" |
| source.write_bytes(b"new") |
| destination.write_bytes(b"old") |
| with pytest.raises(FileExistsError): |
| hub._link(source, destination) |
| assert destination.read_bytes() == b"old" |
|
|
|
|
| def test_frozen_modulation_requires_its_schedule_and_precision(monkeypatch): |
| monkeypatch.setattr(hub.importlib.metadata, "version", lambda name: "0.32.4") |
| monkeypatch.setenv("MLX_ENABLE_TF32", "0") |
| release = {"sigmas": list(SIGMAS), "modulation_mlx": "0.32.4"} |
| hub.validate_modulation(release, SIGMAS) |
| with pytest.raises(ValueError, match="schedule"): |
| hub.validate_modulation(release, (1, 0)) |
| monkeypatch.setenv("MLX_ENABLE_TF32", "1") |
| with pytest.raises(ValueError, match="MLX_ENABLE_TF32"): |
| hub.validate_modulation(release, SIGMAS) |
|
|
|
|
| def test_quantized_release_moves_without_original_checkpoint(tmp_path): |
| import hashlib |
| import shutil |
|
|
| import mlx.core as mx |
|
|
| from h3_mlx.config import DiTConfig |
| from h3_mlx.dit import DiT |
|
|
| model = tmp_path / "assembled" |
| transformer = model / "transformer" |
| core = model / "core" |
| transformer.mkdir(parents=True) |
| core.mkdir() |
| config = json.dumps( |
| vars(DiTConfig(hidden_size=32, text_dim=24, token_refiner_num_layers=0)) |
| ) |
| (transformer / "config.json").write_text(config) |
| mx.save_safetensors( |
| str(transformer / "support.safetensors"), {"native": mx.ones(1)} |
| ) |
| mx.save_safetensors(str(core / "core-000.safetensors"), {"quantized": mx.ones(1)}) |
| mx.save_safetensors( |
| str(model / "pdd_fused_heads.safetensors"), {"head": mx.ones(1)} |
| ) |
| (core / "quantization.json").write_text( |
| json.dumps( |
| { |
| "format": "h3-mlx-core-v1", |
| "bits": 4, |
| "group_size": 64, |
| "source_config_sha256": hashlib.sha256(config.encode()).hexdigest(), |
| "source_fingerprint": "original-path-dependent-fingerprint-is-not-used", |
| } |
| ) |
| ) |
| (model / "release.json").write_text( |
| json.dumps( |
| { |
| "format": "h3-mlx-release-v1", |
| "core_format": "quantized", |
| "bits": 4, |
| "group_size": 64, |
| } |
| ) |
| ) |
| moved = tmp_path / "another-machine" |
| shutil.move(str(model), moved) |
| loaded = DiT(moved, cache_dir=tmp_path / "cache") |
| assert loaded.bits == 4 and loaded.core_store.path == moved / "core" |
| with pytest.raises(ValueError, match="already select"): |
| DiT(moved, bits=8) |
| metadata = hub.read_json(moved / "core/quantization.json") |
| metadata["source_config_sha256"] = "wrong" |
| (moved / "core/quantization.json").write_text(json.dumps(metadata)) |
| with pytest.raises(ValueError, match="mismatch"): |
| DiT(moved) |
|
|
