File size: 4,548 Bytes
dd0a8f2 | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """Pure-Python release path, identity, and lineage helpers."""
from __future__ import annotations
import hashlib
import importlib.metadata
import json
import os
import platform
from pathlib import Path
from typing import Any
ROOT = Path(__file__).resolve().parents[1]
CONFIG_PATH = ROOT / "config" / "deployment.json"
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for block in iter(lambda: handle.read(16 * 1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
def load_config() -> dict[str, Any]:
return json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
def release_path(relative: str) -> Path:
candidate = (ROOT / relative).resolve()
if candidate != ROOT and ROOT not in candidate.parents:
raise ValueError(f"release path escapes package root: {relative}")
return candidate
def adapter_path(config: dict[str, Any] | None = None) -> Path:
config = config or load_config()
return release_path(config["adapter"]["path"])
def tokenizer_path(config: dict[str, Any] | None = None) -> Path:
config = config or load_config()
return release_path(config["foundation"]["tokenizer_path"])
def resolve_model_root(
explicit: str | Path | None = None,
config: dict[str, Any] | None = None,
) -> Path:
config = config or load_config()
env_name = config["foundation"]["model_root_env"]
value = explicit if explicit is not None else os.environ.get(env_name)
if not value:
raise ValueError(
f"Foundation weights are external. Set {env_name} to the model-root "
"directory described in config/foundation_models.lock.json."
)
return Path(value).expanduser().resolve()
def foundation_files(
model_root: Path,
config: dict[str, Any] | None = None,
) -> list[tuple[dict[str, Any], Path]]:
config = config or load_config()
return [
(row, (model_root / row["relative_path"]).resolve())
for row in config["foundation"]["required_files"]
]
def foundation_model_dir(
model_root: Path,
config: dict[str, Any] | None = None,
) -> Path:
config = config or load_config()
return (model_root / config["foundation"]["model_id"]).resolve()
def export_lineage(
model_root: Path,
*,
config: dict[str, Any] | None = None,
) -> dict[str, Any]:
config = config or load_config()
adapter = adapter_path(config)
base_dit = next(
path
for row, path in foundation_files(model_root, config)
if row["relative_path"].endswith("diffusion_pytorch_model.safetensors")
)
exporter = (
ROOT
/ "vendor"
/ "raven-image3f-joint"
/ "control"
/ "export_image3f_block_stack_onnx.py"
)
source_manifest = ROOT / "config" / "SOURCE_MANIFEST.sha256"
profile = {
"batch_size": config["generation_contract"]["batch_size"],
"views": config["generation_contract"]["views"],
"per_view_width": config["generation_contract"]["per_view_width"],
"per_view_height": config["generation_contract"]["per_view_height"],
"frames": config["generation_contract"]["frames"],
"reference_mode": "lag256_rgb_all_views",
"onnx_opset": 20,
}
toolchain = {
"python": platform.python_version(),
"torch": importlib.metadata.version("torch"),
"onnx": importlib.metadata.version("onnx"),
"onnxscript": importlib.metadata.version("onnxscript"),
"safetensors": importlib.metadata.version("safetensors"),
}
payload = {
"adapter_sha256": sha256(adapter),
"base_dit_sha256": sha256(base_dit),
"exporter_sha256": sha256(exporter),
"release_source_manifest_sha256": sha256(source_manifest),
"export_toolchain": toolchain,
"profile": profile,
}
canonical = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
payload["export_key"] = hashlib.sha256(canonical).hexdigest()
return payload
def generated_onnx_dir(
model_root: Path,
config: dict[str, Any] | None = None,
) -> Path:
config = config or load_config()
lineage = export_lineage(model_root, config=config)
return release_path(config["generated_layout"]["onnx_root"]) / lineage[
"export_key"
][:24]
def current_engine_dir(config: dict[str, Any] | None = None) -> Path:
config = config or load_config()
return release_path(config["generated_layout"]["engine_root"]) / "current"
|