icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Bump version to 0.11.0 and update dependencies

Browse Source 
- Updated ONNX Runtime to 1.24.4 and replaced OpenCV with PyAV 17.0.0 in the installation environment.
- Enhanced model downloading logic to prioritize Hugging Face and fallback to OmooLab R2, with improved error handling and status reporting.
- Introduced a new environment revision file for shared metadata.
- Refactored video frame extraction to utilize the PyAV library instead of OpenCV.
- Improved installation process to handle retries across specified PyPI indexes.
- Added tests for download progress reporting and model downloading behavior.
This commit is contained in:
马南
2026-08-14 14:47:40 +08:00
parent 2ad30bf4a3
commit 5255eb4e96
12 changed files with 561 additions and 304 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/snapform/blender_manifest.toml
View File

@@ -1,7 +1,7 @@
schema_version = "1.0.0"
id = "snapform"
version = "0.10.12"
version = "0.11.0"
name = "SnapForm"
tagline = "Build spatial structures from images, fast"
maintainer = "OmooLab"

3
src/snapform/client.py
View File

@@ -18,6 +18,7 @@ from urllib.request import ProxyHandler, Request, build_opener
import bpy
from .env_revision import MANIFEST_FILENAME, MANIFEST_REVISION
from .models import (
BEN2_MODEL,
get_downloadable_model,
@@ -26,8 +27,6 @@ from .models import (
)
MANIFEST_FILENAME = "manifest.json"
MANIFEST_REVISION = 7
SERVICE_ADDRESS = "127.0.0.1"
DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"
JOB_FOLDER_NAMES = {

4
src/snapform/env_revision.py Normal file
View File

@@ -0,0 +1,4 @@
"""Shared environment metadata for the Blender UI and the install subprocess."""
MANIFEST_FILENAME = "manifest.json"
MANIFEST_REVISION = 8

95
src/snapform/install_env.py
View File

@@ -19,6 +19,7 @@ SCRIPT_DIRECTORY = Path(__file__).resolve().parent
if str(SCRIPT_DIRECTORY) not in sys.path:
sys.path.insert(0, str(SCRIPT_DIRECTORY))
from env_revision import MANIFEST_REVISION
from models import BEN2_MODEL, DOWNLOADABLE_MODELS, MOGE2_MODEL
@@ -32,10 +33,9 @@ UVICORN_VERSION = "0.51.0"
ONNXRUNTIME_VERSION = "1.24.4"
NUMPY_VERSION = "2.4.2"
PILLOW_VERSION = "12.1.1"
OPENCV_VERSION = "4.13.0.92"
AV_VERSION = "17.0.0"
TRIMESH_VERSION = "4.11.2"
HUGGINGFACE_HUB_VERSION = "1.4.1"
MANIFEST_REVISION = 7
_resolved_pypi_index = None
DEFAULT_MODEL_KEYS = (MOGE2_MODEL.key, BEN2_MODEL.key)
RUNTIME_STAGE_COUNT = 2
@@ -45,15 +45,9 @@ MODELS_SLICE_START = 0.85
MODELS_SLICE_END = 1.0
RAMP_RATE = 0.02
CREEP_CAP = 0.60
INSTALL_ESTIMATES = {
"ONNX Runtime": 3 * 60,
}
INSTALL_ESTIMATE_RANGES = {
"ONNX Runtime": AI_COMPONENTS_SLICE,
}
CHECK_CODE = (
"import json; import fastapi, uvicorn, onnxruntime; "
"import cv2, numpy, PIL, trimesh; "
"import av, numpy, PIL, trimesh; "
"print('SNAPFORM_DEVICES=' + json.dumps({"
"'onnxruntime': onnxruntime.__version__, "
"'providers': onnxruntime.get_available_providers()}))"
@@ -76,7 +70,6 @@ def write_status(
stage=None,
stages=None,
stage_label=None,
eta_seconds=None,
):
payload = {"progress": progress, "message": message}
if error:
@@ -85,8 +78,6 @@ def write_status(
payload["stage"] = stage
payload["stages"] = stages
payload["stage_label"] = stage_label
if eta_seconds is not None:
payload["eta_seconds"] = max(int(eta_seconds), 0)
Path(path).write_text(
json.dumps(payload, ensure_ascii=False),
encoding="utf-8",
@@ -103,7 +94,7 @@ def child_environment(environment=None):
if env_key not in process_environment:
process_environment[env_key] = value
if _resolved_pypi_index:
process_environment["PIP_INDEX_URL"] = _resolved_pypi_index
process_environment["UV_INDEX_URL"] = _resolved_pypi_index
if environment:
process_environment.update(environment)
return process_environment
@@ -117,13 +108,6 @@ def run(command, environment=None):
)
def format_elapsed(seconds):
seconds = int(seconds)
if seconds < 60:
return f"{seconds}s"
return f"{seconds // 60}m {seconds % 60:02d}s"
def parse_package_count(line):
try:
return int(line.split()[1])
@@ -131,26 +115,12 @@ def parse_package_count(line):
return 0
def estimated_remaining(label, progress, elapsed=0.0):
expected_seconds = INSTALL_ESTIMATES.get(label)
progress_range = INSTALL_ESTIMATE_RANGES.get(label)
if expected_seconds is None or progress_range is None:
return None
start, end = progress_range
if end <= start:
return 0.0
remaining_fraction = 1.0 - min(
max((progress - start) / (end - start), 0.0),
1.0,
)
return max(expected_seconds * remaining_fraction - elapsed, 0.0)
def resolve_indexes(args):
"""???????????????????????????????????"""
"""确定用户指定的 PyPI 源，优先命令行参数，其次 uv/pip 的环境变量。"""
global _resolved_pypi_index
_resolved_pypi_index = (
getattr(args, "pypi_index", None)
or os.environ.get("UV_INDEX_URL")
or os.environ.get("PIP_INDEX_URL")
)
@@ -206,11 +176,6 @@ def _run_install_once(
stage=stage,
stages=stages,
stage_label=label,
eta_seconds=estimated_remaining(
label,
displayed,
clock() - started,
),
)
while True:
@@ -233,7 +198,21 @@ def _run_install_once(
if empty:
if elapsed - last_heartbeat >= 5.0:
last_heartbeat = elapsed
detail("Downloading packages")
# uv 的下载阶段没有增量进度，只报告已解析的包总数，
# 让用户对规模有预期，而不承诺剩余时间或已耗时。
message = (
f"Downloading {total} packages"
if total
else "Downloading packages"
)
write_status(
status_path,
displayed,
message,
stage=stage,
stages=stages,
stage_label=label,
)
continue
line = raw_line.rstrip()
print(line)
@@ -273,13 +252,20 @@ def run_install(
end=1.0,
clock=time.monotonic,
):
"""Run a uv install command, retrying with mirrors when the default index fails."""
attempts = [environment]
if _resolved_pypi_index in (None, PYPI_DEFAULT_INDEX):
attempts.extend(
{"PIP_INDEX_URL": mirror} for mirror in PYPI_MIRRORS
)
for index, attempt_environment in enumerate(attempts):
"""Run a uv install command, retrying across PyPI indexes on failure."""
if _resolved_pypi_index is None:
# 未指定源时优先国内镜像，pypi.org 作为最终兜底。
indexes = (*PYPI_MIRRORS, PYPI_DEFAULT_INDEX)
elif _resolved_pypi_index == PYPI_DEFAULT_INDEX:
# 明确选择官方源时，失败后回退国内镜像。
indexes = (PYPI_DEFAULT_INDEX, *PYPI_MIRRORS)
else:
# 用户指定了特定源，不再回退。
indexes = (_resolved_pypi_index,)
attempts = [
dict(environment or {}, UV_INDEX_URL=index) for index in indexes
]
for position, attempt_environment in enumerate(attempts):
try:
return _run_install_once(
status_path,
@@ -293,9 +279,12 @@ def run_install(
clock,
)
except subprocess.CalledProcessError:
if index + 1 == len(attempts):
if position + 1 == len(attempts):
raise
print(f"Retrying with PyPI mirror: {PYPI_MIRRORS[0]}")
print(
"Retrying with next PyPI index: "
f"{attempts[position + 1]['UV_INDEX_URL']}"
)
def venv_python(venv_dir):
@@ -330,7 +319,7 @@ def install_stages(args, venv_dir, python, uv):
f"fastapi=={FASTAPI_VERSION}",
f"huggingface-hub=={HUGGINGFACE_HUB_VERSION}",
f"numpy=={NUMPY_VERSION}",
f"opencv-python-headless=={OPENCV_VERSION}",
f"av=={AV_VERSION}",
runtime_package,
f"pillow=={PILLOW_VERSION}",
f"trimesh=={TRIMESH_VERSION}",
@@ -465,7 +454,6 @@ def run_verify(
stage=stage,
stages=stages,
stage_label=stage_label,
eta_seconds=max(60.0 - elapsed, 0.0),
)
continue
line = raw_line.rstrip()
@@ -596,7 +584,6 @@ def run_stage(
stage=stage,
stages=stages,
stage_label=label,
eta_seconds=estimated_remaining(label, start),
)
installed = []
if kind == "install":

337
src/snapform/worker.py
View File

@@ -1,15 +1,23 @@
import argparse
import hashlib
import http.client
import json
import os
import shutil
import socket
import sys
import tempfile
import threading
import time
import sys
from pathlib import Path
from urllib.parse import quote
from urllib.request import Request, urlopen
from urllib.request import (
HTTPSHandler,
ProxyHandler,
Request,
build_opener,
urlopen,
)
SCRIPT_DIRECTORY = Path(__file__).resolve().parent
@@ -40,7 +48,29 @@ EXPORT_FORMATS = (
)
HF_ENDPOINT = "https://huggingface.co"
R2_ENDPOINT = "https://models.omoolab.xyz"
DOWNLOAD_TIMEOUT = 60.0
# 连接（TCP + TLS 握手）超时，源不可达时快速失败切换；下载阶段不设读超时。
CONNECT_TIMEOUT = 6.0
class _ShortConnectHTTPSConnection(http.client.HTTPSConnection):
"""HTTPS 连接：TCP + TLS 握手用短超时，读操作不设超时。"""
def connect(self):
try:
self.timeout = CONNECT_TIMEOUT
super().connect()
finally:
self.timeout = None
# 连接建立后读操作不设超时，下载大文件不会被读超时中断。
self.sock.settimeout(None)
class _ShortConnectHTTPSHandler(HTTPSHandler):
def https_open(self, req):
return self.do_open(_ShortConnectHTTPSConnection, req)
_download_opener = build_opener(ProxyHandler(), _ShortConnectHTTPSHandler())
class TaskCancelled(RuntimeError):
@@ -107,34 +137,32 @@ def extract_video_frames(
video_path,
output_directory,
maximum_frames,
cv2,
av,
start_frame=1,
cancel_check=None,
):
output_directory = Path(output_directory)
output_directory.mkdir(parents=True, exist_ok=True)
capture = cv2.VideoCapture(str(video_path))
if not capture.isOpened():
raise RuntimeError(f"Unable to open video: {video_path}")
try:
container = av.open(str(video_path))
except Exception as error:
raise RuntimeError(f"Unable to open video: {video_path}") from error
image_paths = []
skipped = 0
try:
while len(image_paths) < maximum_frames:
stream = container.streams.video[0]
for frame in container.decode(stream):
check_cancelled(cancel_check)
success, frame = capture.read()
if not success:
break
if skipped < start_frame - 1:
skipped += 1
continue
image_path = (
output_directory / f"frame.{len(image_paths) + 1:04d}.jpg"
)
if not cv2.imwrite(str(image_path), frame):
raise RuntimeError(f"Unable to write video frame: {image_path}")
image_path = output_directory / f"frame.{len(image_paths) + 1:04d}.jpg"
frame.to_image().save(image_path, format="JPEG")
image_paths.append(image_path)
if len(image_paths) >= maximum_frames:
break
finally:
capture.release()
container.close()
if not image_paths:
raise RuntimeError(f"Unable to read any video frame from: {video_path}")
@@ -162,7 +190,7 @@ def prepare_input(
supported = ", ".join(sorted(IMAGE_SUFFIXES | VIDEO_SUFFIXES))
raise RuntimeError(f"Unsupported input type; expected one of: {supported}")
import cv2
import av
temporary_directory = Path(
tempfile.mkdtemp(prefix="snapform-video-", dir=temporary_root)
@@ -172,7 +200,7 @@ def prepare_input(
input_path,
temporary_directory,
video_frames,
cv2,
av,
start_frame=start_frame,
cancel_check=cancel_check,
)
@@ -254,33 +282,6 @@ def unproject_camera_depth(depth, intrinsics):
return (rays * depth.reshape(-1, 1)).reshape(height, width, 3)
def smooth_depth_map(depth, smoothing):
import numpy as np
factor = min(max(float(smoothing) / 100.0, 0.0), 1.0)
if factor <= 0.0:
return np.asarray(depth, dtype=np.float32).copy()
import cv2
values = np.asarray(depth, dtype=np.float32)
valid = np.isfinite(values) & (values > 0.0)
if not valid.any():
raise ValueError("DA3 returned no valid depth values")
low, high = np.percentile(values[valid], (5.0, 95.0))
depth_range = max(float(high - low), 1e-6)
diameter = 3 + 2 * int(round(3.0 * factor))
working_values = values.copy()
working_values[~valid] = float(np.median(values[valid]))
smoothed = cv2.bilateralFilter(
working_values,
diameter,
depth_range * (0.01 + 0.09 * factor),
1.0 + 4.0 * factor,
)
smoothed[~valid] = values[~valid]
return smoothed
def compact_mesh(vertices, faces):
import numpy as np
@@ -961,8 +962,10 @@ def create_download_progress_class(
class DownloadProgress(tqdm_type):
def __init__(self, *args, **kwargs):
description = str(kwargs.get("desc", ""))
# huggingface_hub 1.x 的字节进度条描述是 "Downloading (incomplete total...)"，
# 用前缀匹配而非精确匹配，否则真实下载进度会被忽略。
self.snapform_reports_progress = (
description == "Downloading bytes"
description.startswith("Downloading")
or description.startswith("Fetching ")
)
super().__init__(*args, **kwargs)
@@ -998,6 +1001,35 @@ def progress_range(args, staged):
return 0.05, 0.90
def _discovery_path(model_dir):
"""storage_root/service.json，从 storage_root/models/<name> 反推 storage_root。"""
return Path(model_dir).parent.parent / "service.json"
def _models_from_r2(model_dir):
try:
payload = json.loads(
_discovery_path(model_dir).read_text(encoding="utf-8")
)
except (OSError, ValueError, json.JSONDecodeError):
return False
return bool(payload.get("models_from_r2"))
def _mark_models_from_r2(model_dir):
"""在 service.json 上记录 models_from_r2，保留文件里已有的其它字段。"""
path = _discovery_path(model_dir)
try:
payload = json.loads(path.read_text(encoding="utf-8"))
except (OSError, ValueError, json.JSONDecodeError):
payload = {}
payload["models_from_r2"] = True
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_name(f"{path.name}.tmp")
temporary.write_text(json.dumps(payload), encoding="utf-8")
temporary.replace(path)
def download_model(args):
label = getattr(args, "label", None) or args.repository
model_dir = Path(args.model_dir)
@@ -1009,6 +1041,77 @@ def download_model(args):
"stages": getattr(args, "stages", None),
"stage_label": getattr(args, "stage_label", None),
}
r2_directory = getattr(args, "r2_directory", "")
r2_files = getattr(args, "r2_file", [])
if r2_directory and r2_files:
_download_declared_or_fail(
args,
model_dir,
label,
start,
span,
stage_kwargs,
)
return
_download_snapshot_or_fail(args, model_dir, label, start, span, stage_kwargs)
def _download_declared_or_fail(args, model_dir, label, start, span, stage_kwargs):
"""Stream declared files from Hugging Face, then the R2 mirror, remembering fallback."""
if _models_from_r2(model_dir):
sources = (
(
R2_ENDPOINT,
download_from_r2,
f"Downloading {label} | OmooLab mirror",
),
)
else:
sources = (
(HF_ENDPOINT, download_from_hf, f"Downloading {label}"),
(
R2_ENDPOINT,
download_from_r2,
f"Downloading {label} | OmooLab mirror",
),
)
failures = []
for position, (source, downloader, message) in enumerate(sources):
write_status(args.status, start, message, **stage_kwargs)
try:
downloader(
args,
model_dir,
start=start,
span=span,
label=label,
stage_kwargs=stage_kwargs,
)
_verify_model_complete(
model_dir,
getattr(args, "required_pattern", None),
)
except Exception as error:
failures.append((source, error))
continue
# 非首个源成功意味着 hf 已失败，记住本次 service 周期后续直接走 r2。
if position > 0:
_mark_models_from_r2(model_dir)
write_status(
args.status,
start + span,
f"{label} ready",
**stage_kwargs,
)
return
detail = "; ".join(
f"{source}: {error}" for source, error in failures
)
raise RuntimeError(f"All download sources failed: {detail}")
def _download_snapshot_or_fail(args, model_dir, label, start, span, stage_kwargs):
"""Fallback download for repositories without declared files, via huggingface_hub."""
failures = []
try:
from huggingface_hub import snapshot_download
@@ -1027,7 +1130,9 @@ def download_model(args):
)
for endpoint in download_endpoints() if snapshot_download else ():
os.environ["HF_ENDPOINT"] = endpoint
message = f"Downloading {label}"
# snapshot_download 先拉取文件清单并校验 ETag，还没有字节下载，
# 用 Checking 提示当前阶段，避免进度条看起来卡住。
message = f"Checking {label} files"
if endpoint != HF_ENDPOINT:
message = f"{message} (mirror)"
write_status(args.status, start, message, **stage_kwargs)
@@ -1053,43 +1158,6 @@ def download_model(args):
**stage_kwargs,
)
return
# Hugging Face 失败后使用由项目维护的静态模型镜像。
r2_directory = getattr(args, "r2_directory", "")
r2_files = getattr(args, "r2_file", [])
if not r2_directory or not r2_files:
detail = "; ".join(
f"{source}: {error}" for source, error in failures
)
raise RuntimeError(f"All download sources failed: {detail}")
try:
write_status(
args.status,
start,
f"Downloading {label} | OmooLab mirror",
**stage_kwargs,
)
download_from_r2(
args,
model_dir,
start=start,
span=span,
label=label,
stage_kwargs=stage_kwargs,
)
_verify_model_complete(
model_dir,
getattr(args, "required_pattern", None),
)
except Exception as error:
failures.append((R2_ENDPOINT, error))
else:
write_status(
args.status,
start + span,
f"{label} ready",
**stage_kwargs,
)
return
detail = "; ".join(
f"{source}: {error}" for source, error in failures
)
@@ -1121,25 +1189,24 @@ def _verify_model_complete(model_dir, required_patterns=None):
)
def _download_r2_file(
directory,
filename,
def _download_file_from_url(
url,
destination,
expected_size,
expected_checksum,
progress=None,
retries=2,
):
"""Download and verify one OmooLab R2 model file atomically."""
url = f"{R2_ENDPOINT}/{quote(directory)}/{quote(filename)}"
"""Download and verify one model file atomically from a direct URL."""
request = Request(url, headers={"User-Agent": "SnapForm"})
opener = _download_opener.open if request.type == "https" else urlopen
last_error = None
for _attempt in range(retries):
temporary = destination.with_name(f"{destination.name}.part")
try:
checksum = hashlib.sha256()
downloaded = 0
with urlopen(request, timeout=DOWNLOAD_TIMEOUT) as response:
with opener(request) as response:
with open(temporary, "wb") as output:
while True:
chunk = response.read(1024 * 1024)
@@ -1152,11 +1219,11 @@ def _download_r2_file(
progress(len(chunk))
if downloaded != expected_size:
raise RuntimeError(
f"Size mismatch for {filename}: "
f"Size mismatch for {destination.name}: "
f"{downloaded} != {expected_size}"
)
if checksum.hexdigest().lower() != expected_checksum.lower():
raise RuntimeError(f"SHA-256 mismatch for {filename}")
raise RuntimeError(f"SHA-256 mismatch for {destination.name}")
temporary.replace(destination)
return
except Exception as error:
@@ -1168,20 +1235,37 @@ def _download_r2_file(
raise last_error
def download_from_r2(
args,
model_dir,
start=0.05,
span=0.90,
label=None,
stage_kwargs=None,
):
"""Download the declared model files from OmooLab R2."""
repository = args.repository
def r2_file_url(directory, filename):
return f"{R2_ENDPOINT}/{quote(directory)}/{quote(filename)}"
def hf_file_url(repository, filename):
"""Hugging Face resolve URL for a direct, metadata-free download."""
endpoint = (os.environ.get("HF_ENDPOINT") or HF_ENDPOINT).rstrip("/")
return f"{endpoint}/{repository}/resolve/main/{quote(filename)}"
def _parse_declared_files(args):
selected = []
for value in args.r2_file:
filename, size, checksum = value.split("|", 2)
selected.append((filename, int(size), checksum))
return selected
def _download_declared_files(
args,
model_dir,
url_builder,
start,
span,
label,
stage_kwargs,
retries=2,
):
"""Download each declared file through url_builder with aggregated progress."""
repository = args.repository
selected = _parse_declared_files(args)
reporter = DownloadProgressReporter(
args.status,
repository,
@@ -1201,16 +1285,57 @@ def download_from_r2(
for path, size, checksum in selected:
destination = Path(model_dir) / path
destination.parent.mkdir(parents=True, exist_ok=True)
_download_r2_file(
args.r2_directory,
path,
_download_file_from_url(
url_builder(path),
destination,
size,
checksum,
progress=on_chunk,
retries=retries,
)
def download_from_r2(
args,
model_dir,
start=0.05,
span=0.90,
label=None,
stage_kwargs=None,
):
"""Download the declared model files from OmooLab R2."""
_download_declared_files(
args,
model_dir,
lambda filename: r2_file_url(args.r2_directory, filename),
start,
span,
label,
stage_kwargs,
)
def download_from_hf(
args,
model_dir,
start=0.05,
span=0.90,
label=None,
stage_kwargs=None,
):
"""Download the declared model files directly from Hugging Face resolve URLs."""
_download_declared_files(
args,
model_dir,
lambda filename: hf_file_url(args.repository, filename),
start,
span,
label,
stage_kwargs,
retries=1,
)
def format_rate(rate):
size = float(rate)
if size <= 0: