nan
2ad30bf4a3
- Added `ben2_onnx.py` for BEN2 ONNX inference with preprocessing. - Introduced `da3_onnx.py` for Depth Anything 3 inference, including camera projection calculations. - Created `moge_onnx.py` for MoGe-2 ONNX inference and post-processing. - Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`. - Added model synchronization script `sync_models.py` for downloading and verifying ONNX models. - Developed unit tests for BEN2, Depth Anything 3, and MoGe ONNX models to ensure functionality. - Updated `vendor.lock.toml` to remove unused artifacts.
111 lines
3.4 KiB
Python
111 lines
3.4 KiB
Python
import argparse
|
|
import hashlib
|
|
import sys
|
|
import tomllib
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parent
|
|
LOCK_PATH = PROJECT_ROOT / "vendor.lock.toml"
|
|
VENDOR_DIR = PROJECT_ROOT / "vendor"
|
|
CHUNK_SIZE = 1024 * 1024
|
|
|
|
|
|
def load_lock(path=LOCK_PATH):
|
|
"""Load and validate the vendor lock file."""
|
|
with Path(path).open("rb") as stream:
|
|
lock = tomllib.load(stream)
|
|
if lock.get("schema_version") != 1:
|
|
raise RuntimeError("Unsupported vendor.lock.toml schema")
|
|
return lock
|
|
|
|
|
|
def file_sha256(path):
|
|
"""Return the SHA-256 digest for a file."""
|
|
digest = hashlib.sha256()
|
|
with Path(path).open("rb") as stream:
|
|
while chunk := stream.read(CHUNK_SIZE):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def download_artifact(artifact, destination_dir=VENDOR_DIR):
|
|
"""Download one locked artifact and verify its content."""
|
|
destination_dir = Path(destination_dir)
|
|
destination = destination_dir / artifact["filename"]
|
|
expected = artifact["sha256"]
|
|
if destination.is_file() and file_sha256(destination) == expected:
|
|
print(f"Ready: {destination.name}")
|
|
return destination
|
|
|
|
destination_dir.mkdir(parents=True, exist_ok=True)
|
|
temporary = destination.with_name(f"{destination.name}.part")
|
|
temporary.unlink(missing_ok=True)
|
|
print(f"Downloading: {destination.name}")
|
|
request = urllib.request.Request(
|
|
artifact["url"],
|
|
headers={"User-Agent": "SnapForm vendor preparation"},
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(request) as response, temporary.open("wb") as output:
|
|
while chunk := response.read(CHUNK_SIZE):
|
|
output.write(chunk)
|
|
actual = file_sha256(temporary)
|
|
if actual != expected:
|
|
raise RuntimeError(
|
|
f"SHA-256 mismatch for {destination.name}: "
|
|
f"expected {expected}, got {actual}"
|
|
)
|
|
temporary.replace(destination)
|
|
except Exception:
|
|
temporary.unlink(missing_ok=True)
|
|
raise
|
|
return destination
|
|
|
|
|
|
def prepare(platforms=None, lock_path=LOCK_PATH, destination_dir=VENDOR_DIR):
|
|
"""Prepare bundled uv executables for selected platforms."""
|
|
lock = load_lock(lock_path)
|
|
available = lock["platforms"]
|
|
selected = list(available) if platforms is None else platforms
|
|
unknown = sorted(set(selected) - set(available))
|
|
if unknown:
|
|
raise ValueError(f"Unsupported platform: {', '.join(unknown)}")
|
|
|
|
artifacts = [available[platform] for platform in selected]
|
|
return [
|
|
download_artifact(artifact, destination_dir)
|
|
for artifact in artifacts
|
|
]
|
|
|
|
|
|
def parse_args(argv=None):
|
|
"""Parse command-line options."""
|
|
parser = argparse.ArgumentParser(
|
|
description="Download and verify third-party files for SnapForm.",
|
|
)
|
|
parser.add_argument(
|
|
"--platform",
|
|
action="append",
|
|
dest="platforms",
|
|
help="Prepare one platform; repeat as needed. Default: all platforms.",
|
|
)
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def main(argv=None):
|
|
"""Run vendor preparation from the command line."""
|
|
args = parse_args(argv)
|
|
try:
|
|
paths = prepare(args.platforms)
|
|
except (OSError, RuntimeError, ValueError) as error:
|
|
print(f"Preparation failed: {error}", file=sys.stderr)
|
|
return 1
|
|
print(f"Prepared {len(paths)} files")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|