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)