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:
194
sync_models.py
Normal file
194
sync_models.py
Normal file
@@ -0,0 +1,194 @@
|
||||
"""Prepare SnapForm ONNX models and sync them to OmooLab R2."""
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
MODELS_DIR = PROJECT_ROOT / "models"
|
||||
R2_ENDPOINT = "https://models.omoolab.xyz"
|
||||
CHUNK_SIZE = 1024 * 1024
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelFile:
|
||||
folder: str
|
||||
size: int
|
||||
sha256: str
|
||||
sources: tuple[str, ...]
|
||||
|
||||
@property
|
||||
def destination(self):
|
||||
filename = "model_fp16.onnx" if self.folder.endswith("/onnx") else "model.onnx"
|
||||
return MODELS_DIR / self.folder / filename
|
||||
|
||||
|
||||
def hf_url(repository, filename):
|
||||
return f"https://huggingface.co/{repository}/resolve/main/{filename}"
|
||||
|
||||
|
||||
def r2_url(folder, filename="model.onnx"):
|
||||
return f"{R2_ENDPOINT}/{folder}/{filename}"
|
||||
|
||||
|
||||
MODEL_FILES = (
|
||||
ModelFile(
|
||||
"DA3-SMALL",
|
||||
105484029,
|
||||
"7a9d81475b3e2334ace0d6d2c2a816b3c93081437e1866cc434ce1789125f631",
|
||||
(r2_url("DA3-SMALL"),),
|
||||
),
|
||||
ModelFile(
|
||||
"DA3-BASE",
|
||||
412856697,
|
||||
"31f38fc4990539d7902ba267924b9731db1acb00d840c248a60b1865810ff9e6",
|
||||
(r2_url("DA3-BASE"),),
|
||||
),
|
||||
ModelFile(
|
||||
"BEN2-ONNX/onnx",
|
||||
219121675,
|
||||
"dfdc25f421f32a0d1268e0f2ff2153d340e8f1d52d3dd16f5dc33c1ce85cedf1",
|
||||
(
|
||||
hf_url("onnx-community/BEN2-ONNX", "onnx/model_fp16.onnx"),
|
||||
r2_url("BEN2-ONNX", "onnx/model_fp16.onnx"),
|
||||
),
|
||||
),
|
||||
ModelFile(
|
||||
"moge-2-vits-normal-onnx",
|
||||
140852051,
|
||||
"24eacb5dc7a2c54c7bc98f7de085ffbed79ad006ea5b664c2c2cdc02ff3a52f0",
|
||||
(
|
||||
hf_url("Ruicheng/moge-2-vits-normal-onnx", "model.onnx"),
|
||||
r2_url("moge-2-vits-normal-onnx"),
|
||||
),
|
||||
),
|
||||
ModelFile(
|
||||
"moge-2-vitb-normal-onnx",
|
||||
419411850,
|
||||
"bbf14e07a30f11e69d36ab861590123f5598ababcbc8946a063eb4a966f35a21",
|
||||
(
|
||||
hf_url("Ruicheng/moge-2-vitb-normal-onnx", "model.onnx"),
|
||||
r2_url("moge-2-vitb-normal-onnx"),
|
||||
),
|
||||
),
|
||||
ModelFile(
|
||||
"moge-2-vitl-normal-onnx",
|
||||
1324265014,
|
||||
"afbc4ccc3450298f3afb35b90f015f4c4f552dea21dc6470d5f7b78b77e2d751",
|
||||
(
|
||||
hf_url("Ruicheng/moge-2-vitl-normal-onnx", "model.onnx"),
|
||||
r2_url("moge-2-vitl-normal-onnx"),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def file_sha256(path):
|
||||
digest = hashlib.sha256()
|
||||
with Path(path).open("rb") as stream:
|
||||
while chunk := stream.read(CHUNK_SIZE):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def verify_existing(model):
|
||||
path = model.destination
|
||||
if not path.is_file():
|
||||
return False
|
||||
if path.stat().st_size != model.size:
|
||||
raise RuntimeError(f"Existing model has an unexpected size: {path}")
|
||||
if file_sha256(path) != model.sha256:
|
||||
raise RuntimeError(f"Existing model has an unexpected SHA-256: {path}")
|
||||
print(f"Skip existing: {path.relative_to(PROJECT_ROOT)}")
|
||||
return True
|
||||
|
||||
|
||||
def download_once(model, url):
|
||||
destination = model.destination
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = destination.with_suffix(".onnx.part")
|
||||
temporary.unlink(missing_ok=True)
|
||||
request = Request(url, headers={"User-Agent": "SnapForm model sync"})
|
||||
print(f"Download: {url}")
|
||||
downloaded = 0
|
||||
digest = hashlib.sha256()
|
||||
try:
|
||||
with urlopen(request, timeout=60) as response, temporary.open("wb") as output:
|
||||
while chunk := response.read(CHUNK_SIZE):
|
||||
output.write(chunk)
|
||||
digest.update(chunk)
|
||||
downloaded += len(chunk)
|
||||
print(
|
||||
f"\r {downloaded / 1024**2:.1f} / "
|
||||
f"{model.size / 1024**2:.1f} MiB",
|
||||
end="",
|
||||
flush=True,
|
||||
)
|
||||
print()
|
||||
if downloaded != model.size:
|
||||
raise RuntimeError(f"size mismatch: {downloaded} != {model.size}")
|
||||
if digest.hexdigest() != model.sha256:
|
||||
raise RuntimeError("SHA-256 mismatch")
|
||||
temporary.replace(destination)
|
||||
except Exception:
|
||||
temporary.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
|
||||
def prepare_model(model):
|
||||
if verify_existing(model):
|
||||
return
|
||||
failures = []
|
||||
for url in model.sources:
|
||||
try:
|
||||
download_once(model, url)
|
||||
return
|
||||
except Exception as error:
|
||||
failures.append(f"{url}: {error}")
|
||||
print(f"Source failed: {error}", file=sys.stderr)
|
||||
raise RuntimeError("All sources failed:\n" + "\n".join(failures))
|
||||
|
||||
|
||||
def sync_models():
|
||||
executable = shutil.which("rclone")
|
||||
if executable is None:
|
||||
raise RuntimeError("rclone is not available on PATH")
|
||||
subprocess.run(
|
||||
[
|
||||
executable,
|
||||
"sync",
|
||||
"models",
|
||||
"omoolab-r2:ai-models",
|
||||
"--progress",
|
||||
],
|
||||
cwd=PROJECT_ROOT,
|
||||
check=True,
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--download-only",
|
||||
action="store_true",
|
||||
help="Prepare and verify models without running rclone",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
for model in MODEL_FILES:
|
||||
prepare_model(model)
|
||||
if not args.download_only:
|
||||
sync_models()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user