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.
This commit is contained in:
2026-08-14 10:27:09 +08:00
parent 13fab9ce88
commit 2ad30bf4a3
41 changed files with 1569 additions and 2109 deletions

63
pack.py
View File

@@ -1,6 +1,4 @@
import hashlib
import io
import tarfile
import tomllib
from pathlib import Path
from zipfile import ZIP_DEFLATED, ZipFile
@@ -10,7 +8,6 @@ PROJECT_ROOT = Path(__file__).resolve().parent
SOURCE_DIR = PROJECT_ROOT / "src" / "snapform"
VENDOR_DIR = PROJECT_ROOT / "vendor"
DIST_DIR = PROJECT_ROOT / "dist"
PATCHES_DIR = PROJECT_ROOT / "patches"
LOCK_PATH = PROJECT_ROOT / "vendor.lock.toml"
@@ -28,15 +25,6 @@ PLATFORMS = {
platform: artifact["filename"]
for platform, artifact in VENDOR_LOCK["platforms"].items()
}
BUNDLED_ARCHIVES = {
f"vendor/{artifact['filename']}": artifact["filename"]
for artifact in VENDOR_LOCK["artifacts"].values()
}
OVERLAYS = {
artifact["filename"]: artifact["patches"]
for artifact in VENDOR_LOCK["artifacts"].values()
if "patches" in artifact
}
def file_sha256(path):
@@ -63,41 +51,6 @@ def validate_artifact(path, expected_sha256):
)
def bundled_artifact_bytes(filename):
"""Return a bundled artifact with its configured overlay applied."""
archive_path = VENDOR_DIR / filename
overlay_dir = PATCHES_DIR / OVERLAYS.get(filename, filename)
if not overlay_dir.is_dir():
return archive_path.read_bytes()
if not filename.endswith(".tar.gz"):
raise ValueError(f"Unsupported overlay archive: {filename}")
return _apply_tar_overlay(archive_path, overlay_dir)
def _apply_tar_overlay(archive_path, overlay_dir):
with tarfile.open(archive_path, "r:gz") as source:
members = source.getmembers()
top_level = members[0].name.split("/", 1)[0]
buffer = io.BytesIO()
with tarfile.open(fileobj=buffer, mode="w:gz") as output:
for member in members:
if member.isdir() or member.islnk() or member.issym():
output.addfile(member)
continue
relative = member.name[len(top_level) + 1:]
overlay_file = overlay_dir / relative
if overlay_file.is_file():
info = tarfile.TarInfo(member.name)
info.size = overlay_file.stat().st_size
info.mode = member.mode
with overlay_file.open("rb") as overlay_stream:
output.addfile(info, overlay_stream)
else:
with source.extractfile(member) as extracted:
output.addfile(member, extracted)
return buffer.getvalue()
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")
@@ -121,23 +74,13 @@ def manifest_for_platform(platform):
def bundled_vendor_files(platform):
entries = {
name: VENDOR_DIR / filename
for name, filename in BUNDLED_ARCHIVES.items()
return {
f"vendor/{PLATFORMS[platform]}": VENDOR_DIR / PLATFORMS[platform]
}
entries[f"vendor/{PLATFORMS[platform]}"] = (
VENDOR_DIR / PLATFORMS[platform]
)
return entries
def validate_vendor_files(platform):
"""Validate common archives and the selected uv archive."""
for artifact in VENDOR_LOCK["artifacts"].values():
validate_artifact(
VENDOR_DIR / artifact["filename"],
artifact["sha256"],
)
uv_artifact = VENDOR_LOCK["platforms"][platform]
validate_artifact(
VENDOR_DIR / uv_artifact["filename"],
@@ -174,7 +117,7 @@ def build_extension(platform):
for archive_name, source_path in vendor_files.items():
archive.writestr(
archive_name,
bundled_artifact_bytes(source_path.name),
source_path.read_bytes(),
)
return output_path