Files
AnyImage/pack.py
nan 2ad30bf4a3 feat: Implement ONNX model inference for BEN2, Depth Anything 3, and MoGe
- 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.
2026-08-14 10:27:09 +08:00

134 lines
4.3 KiB
Python

import hashlib
import tomllib
from pathlib import Path
from zipfile import ZIP_DEFLATED, ZipFile
PROJECT_ROOT = Path(__file__).resolve().parent
SOURCE_DIR = PROJECT_ROOT / "src" / "snapform"
VENDOR_DIR = PROJECT_ROOT / "vendor"
DIST_DIR = PROJECT_ROOT / "dist"
LOCK_PATH = PROJECT_ROOT / "vendor.lock.toml"
def load_vendor_lock():
"""Return the locked third-party artifact metadata."""
with LOCK_PATH.open("rb") as stream:
lock = tomllib.load(stream)
if lock.get("schema_version") != 1:
raise RuntimeError("Unsupported vendor.lock.toml schema")
return lock
VENDOR_LOCK = load_vendor_lock()
PLATFORMS = {
platform: artifact["filename"]
for platform, artifact in VENDOR_LOCK["platforms"].items()
}
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(1024 * 1024):
digest.update(chunk)
return digest.hexdigest()
def validate_artifact(path, expected_sha256):
"""Require an artifact to exist and match its locked digest."""
if not path.is_file():
raise RuntimeError(
f"Missing vendor artifact: {path}. "
"Run: uv run python prepare_vendor.py"
)
actual = file_sha256(path)
if actual != expected_sha256:
raise RuntimeError(
f"Vendor artifact hash mismatch: {path}. "
"Run: uv run python prepare_vendor.py"
)
def manifest_for_platform(platform):
"""Return the manifest text declaring only the given platform."""
text = (SOURCE_DIR / "blender_manifest.toml").read_text(encoding="utf-8")
lines = text.splitlines()
table_start = next(
(i for i, line in enumerate(lines) if line.startswith("[")),
len(lines),
)
head = lines[:table_start]
tail = lines[table_start:]
head = [line for line in head if not line.strip().startswith("platforms =")]
# 插入顶层区域,避免 platforms 被解析为后续 table 的子键。
# 放在 blender_version_min 后面,使扩展元数据保持易读。
insert_at = next(
(i + 1 for i, line in enumerate(head) if line.strip().startswith("blender_version_min")),
len(head),
)
head.insert(insert_at, f'platforms = ["{platform}"]')
lines = head + tail
return "\n".join(lines) + "\n"
def bundled_vendor_files(platform):
return {
f"vendor/{PLATFORMS[platform]}": VENDOR_DIR / PLATFORMS[platform]
}
def validate_vendor_files(platform):
"""Validate common archives and the selected uv archive."""
uv_artifact = VENDOR_LOCK["platforms"][platform]
validate_artifact(
VENDOR_DIR / uv_artifact["filename"],
uv_artifact["sha256"],
)
def build_extension(platform):
if platform not in PLATFORMS:
raise ValueError(f"Unsupported platform: {platform}")
manifest = tomllib.loads(manifest_for_platform(platform))
output_path = DIST_DIR / (
f"{manifest['name']}.v{manifest['version']}.{platform}.zip"
)
DIST_DIR.mkdir(exist_ok=True)
vendor_files = bundled_vendor_files(platform)
validate_vendor_files(platform)
with ZipFile(output_path, "w", ZIP_DEFLATED) as archive:
for source_path in sorted(SOURCE_DIR.rglob("*")):
if not source_path.is_file():
continue
if "__pycache__" in source_path.parts:
continue
relative_path = source_path.relative_to(SOURCE_DIR)
if relative_path.parts[0] == "vendor":
continue
if source_path.name == "blender_manifest.toml":
continue
archive.write(source_path, relative_path)
archive.writestr(
"blender_manifest.toml",
manifest_for_platform(platform),
)
for archive_name, source_path in vendor_files.items():
archive.writestr(
archive_name,
source_path.read_bytes(),
)
return output_path
def build_all():
return [build_extension(platform) for platform in PLATFORMS]
if __name__ == "__main__":
extension_paths = build_all()
for extension_path in extension_paths:
size_mib = extension_path.stat().st_size / (1024 * 1024)
print(f"Built {extension_path} ({size_mib:.2f} MiB)")