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:
63
pack.py
63
pack.py
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user