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

Update addon tests, bootstrap runtime, download progress, and worker functionality

Browse Source 
- Modified status messages and stage labels in `test_addon.py` for clarity and consistency.
- Added tests to ensure proper host termination and command execution in `test_addon.py`.
- Enhanced `test_bootstrap_runtime.py` with additional tests for installation behavior, including handling of CUDA and MPS devices, and improved error handling for disk space and network checks.
- Updated download progress reporting in `test_download_progress.py` to include download speed and improved message formatting.
- Added new tests in `test_worker.py` for device selection, NPZ file validation, and model downloading behavior with Hugging Face endpoints.
- Bumped version in `uv.lock` to 0.9.30.
This commit is contained in:
马南
2026-08-07 13:11:04 +08:00
parent e0cb995473
commit d08d0e9c1a
17 changed files with 1799 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
README.md
View File

@@ -43,8 +43,11 @@
2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。
3. 在 Extension Preferences 中确认 `uv` 路径和 Torch Backend。
4. 在 `Host` 点击 `Install Runtime`。首次安装会下载 Python、PyTorch、DA3、BEN2、FastAPI、Uvicorn，
并默认下载 `DA3 Small`、`DA3 Base` 和 `BEN2 Base` 模型。安装按 12 个阶段依次推进，每段进度
从 0 到 100%，模型下载阶段会实时反映已下载字节。
并默认下载 `DA3 Small`、`DA3 Base` 和 `BEN2 Base` 模型。安装按 8 个阶段依次推进，每段进度
从 0 到 100%：下载 uv（如需）、创建 Python Runtime、下载 PyTorch、安装 xFormers（macOS 跳过）、
安装 AI 依赖、修补 DA3（跳过视频渲染），然后是三个模型的下载。进度文字统一为 `阶段/总数 阶段名: 状态`，
模型下载阶段会实时显示下载速度，例如 `6/8 da3 small: downloading 14.2MB/s`。
Runtime 就绪后即使模型下载中断也不影响 Host 使用，可稍后在 `Models` 面板单独补下。
5. 在 3D View 按 `N`，打开 `SnapForm` 页签。
6. 如需其他模型，展开 `Host > Models`，找到模型并点击该行的 `Download`。
7. 展开默认收起的 `Debug`，选择 `Export Format` 和本次执行使用的模型。
@@ -100,7 +103,8 @@
所有任务使用统一的状态栏进度显示。变化中的阶段说明显示在右侧的紧凑进度条内部，取消按钮
位于进度条右侧。Image Depth 的阶段包括 Runtime 启动、输入准备、模型加载、深度估算、缓存
写出和面片构建。进度更新会强制刷新状态栏与窗口，无需移动鼠标即可实时显示。
写出和面片构建。进度更新会强制刷新状态栏与窗口，无需移动鼠标即可实时显示。安装中断后残留的
半成品 Runtime 会在下次安装前自动清理，不会因为中断而无法再次安装。
导入过程精确保留 NPZ 内深度值之间的差值。例如深度 `2.0` 和 `5.5` 在
`Displacement Strength = -1` 时，沿局部 Z 轴相差 `3.5 Blender Unit`。它不会把数值压缩到
@@ -219,20 +223,38 @@ GLB 任务会关闭其默认附带的深度预览，使四种 job 产物保持
## Runtime
Extension 本身不包含 Torch。`Install Runtime` 使用 `uv` 创建独立 Python 3.12
环境，并安装：
环境，并安装。机器上即使没有安装 `uv` 或 Python 也没关系：安装程序由 Blender
自带的 Python 启动，缺 `uv` 时会自动下载官方独立二进制（约 30MB），再用它下载
Python 3.12 与全部依赖。`uv` 会缓存在扩展存储目录下，后续安装直接复用。
安装前会自动预检磁盘空间与网络连通性，未通过会直接给出英文报错，不会开始下载。
预检与安装支持环境变量切换镜像与代理：`PIP_INDEX_URL`（PyPI 镜像）、
`SNAPFORM_TORCH_INDEX`（PyTorch 索引）、`HF_ENDPOINT`（Hugging Face 镜像，
如 `https://hf-mirror.com`）、`HTTP(S)_PROXY`/`ALL_PROXY`（代理）。设置后预检会直接检查
这些地址。直连不通时预检会**自动回退到镜像源**（清华 PyPI、阿里云 PyTorch、hf-mirror），
无需手动配置，状态栏会提示 `using mirror sources`；镜像也不通才会报错。
单独下载模型时也会先检查 Hugging Face 连通性，直连不通自动切换到 `hf-mirror` 镜像。
Windows 上代理若只在系统设置里而未写入环境变量，安装程序会尝试自动从系统代理注入子进程。
- PyTorch 2.9.1
- TorchVision 0.24.1
- xFormers 0.0.33.post2（与 PyTorch 2.9.1 对应）
- xFormers 0.0.33.post2（与 PyTorch 2.9.1 对应，仅 Windows/Linux；macOS 无 wheel 自动跳过）
- 固定到已验证 commit 的 Depth Anything 3
- 固定到已验证 commit 的 BEN2
- FastAPI 0.139.2 与 Uvicorn 0.51.0
- 固定到 DA3 官方推荐 commit 的 `gsplat`；仅安装 PLY 导出所需的纯 Python 部分，
不编译本机 CUDA rasterizer
- DA3 与 BEN2 通过 GitHub 官方源码包安装（无需 Git）；SnapForm 直接调用 DA3
底层 export 接口，`gs_ply` 导出不会触发 `gs_video` 渲染，因此不需要 gsplat
与 C++ 编译工具链
- `addict` 2.4.0（补充 DA3 上游包元数据暂未声明的运行依赖）
默认 CUDA Backend 为 CUDA 12.8，也可以选择 CPU。CPU 可以用于验证，但实际推理
通常很慢。
CUDA 12.8 的 PyTorch wheel 自带 CUDA 运行库，无需安装 CUDA Toolkit；只需要
NVIDIA 驱动 ≥ 570（数据中心显卡可用更老驱动）。驱动不满足或没有 NVIDIA 显卡时，
设备会自动回退到 CPU（macOS 为 MPS），只有显式选择 CUDA 才会报错。
macOS（Apple Silicon）会自动从 PyPI 安装 CPU 版 PyTorch，推理设备默认使用 MPS
（Apple Metal）加速，也可以在设备选项里切换 CPU。PyTorch 自 2.3 起不再提供 Intel
Mac（x86_64）轮子，Intel Mac 暂不支持。
Extension 启用后只会查找已有的共享 AI Host，不会主动启动新进程。第一次执行 AI 操作或在侧边栏
手动点击 `Start` 时才会启动 Host。Host 只绑定回环地址，并让操作系统分配动态端口；第一个实例把
@@ -240,9 +262,10 @@ Extension 启用后只会查找已有的共享 AI Host，不会主动启动新
`urllib.request` 连接同一个 Host。`host.lock` 保证同一 Storage Root 同时只有一个 Host，Token
则阻止同机其他进程直接调用接口。Host 全局同时只执行一个 AI 任务。
每个 Blender 会注册独立客户端并每 5 秒发送心跳。正常关闭 Blender 时仅注销自身，不会直接终止
其他 Blender 正在使用的 Host；异常退出的客户端会在 20 秒后过期。最后一个客户端离开后，Host
保留模型缓存等待 60 秒，再在空闲时自动退出。取消任务不会重启 Host，并会保留已经加载的模型；
每个 Blender 会注册独立客户端并每 5 秒发送心跳。正常关闭 Blender 时先注销自身；如果当前 Host
就是由这个 Blender 启动的，会直接终止它，无需等待。异常退出的客户端会在 10 秒后过期；Host
在最后一个客户端离开后仍会保留模型缓存等待 20 秒再退出，作为 Host 由其他进程启动等场景的兜底。
取消任务不会重启 Host，并会保留已经加载的模型；
如果取消发生在一次不可中断的 Torch 调用中，任务会在该调用返回后的最近检查点结束。
侧边栏 `Host` 面板以 `Running (Clients: n)` 表示空闲服务，以 `Running (Busy)` 表示正在执行任务，

2
pyproject.toml
View File

@@ -1,6 +1,6 @@
[project]
name = "snapform"
version = "0.9.29"
version = "0.9.32"
description = "Fast image-guided structure building tools for Blender"
readme = "README.md"
requires-python = ">=3.11"

4
src/snapform/ai_host.py
View File

@@ -20,8 +20,8 @@ import worker
HOST_SCHEMA = 2
CLIENT_TIMEOUT_SECONDS = 20.0
IDLE_SHUTDOWN_SECONDS = 60.0
CLIENT_TIMEOUT_SECONDS = 10.0
IDLE_SHUTDOWN_SECONDS = 20.0
class HostLock:

4
src/snapform/blender_manifest.toml
View File

@@ -1,14 +1,14 @@
schema_version = "1.0.0"
id = "snapform"
version = "0.9.29"
version = "0.9.32"
name = "SnapForm"
tagline = "Build spatial structures from images, fast"
maintainer = "OmooLab"
type = "add-on"
blender_version_min = "4.5.0"
platforms = ["windows-x64"]
platforms = ["windows-x64", "macos-x64", "macos-arm64", "linux-x64"]
license = ["SPDX:GPL-3.0-or-later"]
[permissions]

677
src/snapform/bootstrap_runtime.py
View File

@@ -1,8 +1,17 @@
import argparse
import json
import os
import platform
import queue
import shutil
import subprocess
import sys
import tarfile
import tempfile
import threading
import time
import urllib.request
import zipfile
from pathlib import Path
@@ -11,21 +20,54 @@ if str(SCRIPT_DIRECTORY) not in sys.path:
sys.path.insert(0, str(SCRIPT_DIRECTORY))
from models import BEN2_MODEL, MODELS
from worker import HF_ENDPOINT, HF_MIRRORS, format_rate, reachable
UV_VERSION = "0.11.3"
PREFLIGHT_MIN_FREE_BYTES = 8 * 1024**3
PREFLIGHT_TIMEOUT = 8.0
MIRROR_TIMEOUT = 5.0
PYPI_MIRRORS = (
"https://pypi.tuna.tsinghua.edu.cn/simple",
"https://mirrors.aliyun.com/pypi/simple/",
)
TORCH_MIRRORS = {
"cu128": ("https://mirrors.aliyun.com/pytorch-wheels/cu128",),
"cpu": ("https://mirrors.aliyun.com/pytorch-wheels/cpu",),
}
TORCH_VERSION = "2.9.1"
TORCHVISION_VERSION = "0.24.1"
XFORMERS_VERSION = "0.0.33.post2"
ADDICT_VERSION = "2.4.0"
FASTAPI_VERSION = "0.139.2"
UVICORN_VERSION = "0.51.0"
GSPLAT_COMMIT = "0b4dddf04cb687367602c01196913cde6a743d70"
RUNTIME_SCHEMA = 4
TORCH_INDEXES = {
"cu128": "https://download.pytorch.org/whl/cu128",
"cpu": "https://download.pytorch.org/whl/cpu",
}
_resolved_pypi_index = None
_resolved_hf_endpoint = None
_resolved_torch_index = None
DEFAULT_MODEL_KEYS = ("SMALL", "BASE", BEN2_MODEL.key)
PREFLIGHT_END = 0.05
MIRROR_END = 0.07
RUNTIME_SLICE_START = 0.07
RUNTIME_SLICE_END = 0.88
VERIFY_SLICE_START = 0.88
VERIFY_SLICE_END = 0.90
MODELS_SLICE_START = 0.90
MODELS_SLICE_END = 1.0
PLAIN_SLICE_SPAN = 0.02
RAMP_RATE = 0.02
CREEP_CAP = 0.60
CHECK_CODE = (
"import torch; "
"from depth_anything_3.api import DepthAnything3; "
"from ben2 import BEN_Base; "
"import fastapi, uvicorn; "
"print(torch.__version__, torch.cuda.is_available())"
)
def write_status(
@@ -50,25 +92,169 @@ def write_status(
)
def child_environment(environment=None):
"""Inherit the current environment and inject system/registry proxy settings into child processes (uv only reads proxy env variables)."""
process_environment = os.environ.copy()
for key, value in urllib.request.getproxies().items():
env_key = f"{key.upper()}_PROXY"
if env_key not in process_environment:
process_environment[env_key] = value
if _resolved_pypi_index:
process_environment["PIP_INDEX_URL"] = _resolved_pypi_index
if _resolved_hf_endpoint:
process_environment["HF_ENDPOINT"] = _resolved_hf_endpoint
if environment:
process_environment.update(environment)
return process_environment
def run(command, environment=None):
print("Running:", subprocess.list2cmdline([str(value) for value in command]))
process_environment = None
if environment:
process_environment = os.environ.copy()
process_environment.update(environment)
subprocess.check_call(
[str(value) for value in command],
env=process_environment,
env=child_environment(environment),
)
def run_install(status_path, stage, stages, label, command, environment=None):
"""Run a uv install command and translate its verbose package events into progress."""
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])
except (IndexError, ValueError):
return 0
def effective_torch_index(backend):
if _resolved_torch_index:
return _resolved_torch_index
return os.environ.get("SNAPFORM_TORCH_INDEX") or TORCH_INDEXES.get(backend)
def probe_hosts(urls, timeout):
results = {}
def worker(url):
results[url] = reachable(url, timeout)
threads = [
threading.Thread(target=worker, args=(url,)) for url in urls
]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
return results
def resolve_endpoints(backend):
"""Resolve package indexes, falling back to mirrors when direct access fails; returns (resolved, failures, used_mirror)."""
preferred = {
"pypi": os.environ.get("PIP_INDEX_URL") or "https://pypi.org/simple/",
"hf": os.environ.get("HF_ENDPOINT") or HF_ENDPOINT,
"torch": os.environ.get("SNAPFORM_TORCH_INDEX")
or TORCH_INDEXES.get(backend),
"github": "https://github.com/",
}
if sys.platform == "darwin":
preferred["torch"] = None
mirror_map = {
"pypi": PYPI_MIRRORS,
"hf": HF_MIRRORS,
"torch": TORCH_MIRRORS.get(backend, ()),
"github": (),
}
results = probe_hosts(
[url for url in preferred.values() if url],
PREFLIGHT_TIMEOUT,
)
resolved = {}
failures = []
for key, url in preferred.items():
if url is None:
resolved[key] = None
continue
if results[url]:
resolved[key] = url
continue
mirror_results = probe_hosts(mirror_map[key], MIRROR_TIMEOUT)
mirror = next(
(
candidate
for candidate in mirror_map[key]
if mirror_results.get(candidate)
),
None,
)
if mirror:
resolved[key] = mirror
else:
resolved[key] = url
failures.append(url)
used_mirror = any(
key in preferred
and preferred[key] is not None
and resolved[key] != preferred[key]
for key in resolved
)
return resolved, failures, used_mirror
def preflight(args):
"""Check disk space and network reachability; switch to mirrors when direct access fails and raise a clear English error when nothing is reachable."""
global _resolved_pypi_index, _resolved_hf_endpoint, _resolved_torch_index
problems = []
free_bytes = shutil.disk_usage(Path(args.runtime_dir).parent).free
if free_bytes < PREFLIGHT_MIN_FREE_BYTES:
needed_gb = PREFLIGHT_MIN_FREE_BYTES // (1024**3)
problems.append(
f"Not enough free disk space: {needed_gb} GB required, "
f"only {free_bytes / 1024**3:.1f} GB available."
)
resolved, failures, used_mirror = resolve_endpoints(args.backend)
_resolved_pypi_index = resolved["pypi"]
_resolved_hf_endpoint = resolved["hf"]
_resolved_torch_index = resolved["torch"]
unreachable = failures
if unreachable:
proxies = urllib.request.getproxies()
proxy_hint = (
"Proxy detected: " + ", ".join(
f"{key}={value}" for key, value in proxies.items()
)
if proxies
else "No proxy detected"
)
problems.append(
"Cannot reach: " + ", ".join(unreachable)
+ f". {proxy_hint}. Check network or proxy settings and retry."
)
if problems:
raise RuntimeError("Prerequisites check failed: " + " ".join(problems))
return resolved, used_mirror
def run_install(
status_path,
stage,
stages,
label,
command,
environment=None,
start=0.0,
end=1.0,
clock=time.monotonic,
):
"""Run a uv install command and smoothly map package-count events onto the stage progress range."""
print("Running:", subprocess.list2cmdline([str(value) for value in command]))
process_environment = os.environ.copy()
process_environment = child_environment(environment)
process_environment["UV_LOG"] = "debug"
if environment:
process_environment.update(environment)
process = subprocess.Popen(
[str(value) for value in command],
env=process_environment,
@@ -78,27 +264,16 @@ def run_install(status_path, stage, stages, label, command, environment=None):
encoding="utf-8",
errors="replace",
)
progress = 0.0
downloaded = 0
for raw_line in process.stdout:
line = raw_line.rstrip()
print(line)
lowered = line.lower()
if "resolved" in lowered and "package" in lowered:
new_progress, message = 0.05, "Resolving dependencies"
elif "downloaded" in lowered and "package" not in lowered:
downloaded += 1
new_progress = min(0.85, 0.10 + downloaded * 0.08)
message = f"Downloaded {downloaded} package(s)"
elif "prepared" in lowered:
new_progress, message = 0.70, "Preparing packages"
elif "installed" in lowered or "built" in lowered:
new_progress, message = 0.95, "Installing packages"
else:
continue
if new_progress <= progress:
continue
progress = new_progress
output_queue = queue.Queue()
def read_output():
for raw_line in process.stdout:
output_queue.put(raw_line)
output_queue.put(None)
threading.Thread(target=read_output, daemon=True).start()
def write(progress, message):
write_status(
status_path,
progress,
@@ -107,9 +282,64 @@ def run_install(status_path, stage, stages, label, command, environment=None):
stages=stages,
stage_label=label,
)
def eased(fraction):
return start + min(max(fraction, 0.0), 1.0) * (end - start)
installed = []
total = 0
displayed = start
target = start
started = clock()
last_tick = started
last_heartbeat = 0.0
while True:
try:
raw_line = output_queue.get(timeout=0.2)
empty = False
except queue.Empty:
raw_line = None
empty = True
now = clock()
elapsed = now - started
dt = now - last_tick
last_tick = now
if total:
effective = max(target, eased(CREEP_CAP))
if effective > displayed:
displayed = min(effective, displayed + dt * RAMP_RATE)
if raw_line is None and not empty:
break
if empty:
if elapsed - last_heartbeat >= 5.0:
last_heartbeat = elapsed
write(displayed, f"downloading {format_elapsed(elapsed)}")
continue
line = raw_line.rstrip()
print(line)
if line.startswith("Resolved ") and "package" in line:
total = parse_package_count(line)
resolved_progress = eased(0.02)
displayed = max(displayed, resolved_progress)
target = max(target, resolved_progress)
write(displayed, f"{total} packages")
elif line.startswith("Prepared ") and "package" in line:
done = parse_package_count(line)
if total:
target = max(target, eased(0.10 + 0.60 * done / total))
write(displayed, f"{done}/{total} packages")
elif line.startswith("Installed ") and "package" in line:
done = parse_package_count(line)
if total:
target = max(target, eased(0.75 + 0.25 * done / total))
write(displayed, f"{done}/{total} packages")
elif line.startswith("+"):
entry = line[1:].strip()
installed.append(entry)
return_code = process.wait()
if return_code != 0:
raise subprocess.CalledProcessError(return_code, command)
return installed
def environment_python(runtime_dir):
@@ -118,70 +348,50 @@ def environment_python(runtime_dir):
return Path(runtime_dir) / "bin" / "python"
def install_stages(args, runtime_dir, python):
uv = args.uv
torch_index = TORCH_INDEXES[args.backend]
def install_stages(args, runtime_dir, python, uv):
# macOS 的 PyTorch wheel 只在 PyPI 上，download.pytorch.org 没有 macosx 构建
torch_index = (
None
if sys.platform == "darwin"
else effective_torch_index(args.backend)
)
def github_archive_url(repository, commit):
return f"https://github.com/{repository}/archive/{commit}.tar.gz"
da3_requirement = (
"depth-anything-3 @ "
"git+https://github.com/ByteDance-Seed/Depth-Anything-3.git"
f"@{args.da3_commit}"
+ github_archive_url("ByteDance-Seed/Depth-Anything-3", args.da3_commit)
)
ben2_requirement = (
"ben2 @ git+https://github.com/PramaLLC/BEN2.git"
f"@{args.ben2_commit}"
"ben2 @ " + github_archive_url("PramaLLC/BEN2", args.ben2_commit)
)
xformers_command = [
torch_command = [
uv,
"pip",
"install",
"-v",
"--python",
python,
f"xformers=={XFORMERS_VERSION}",
f"torch=={TORCH_VERSION}",
f"torchvision=={TORCHVISION_VERSION}",
]
if args.backend == "cu128":
xformers_command.extend(["--index-url", torch_index])
else:
xformers_command.append("--no-deps")
gsplat_requirement = (
"gsplat @ git+https://github.com/nerfstudio-project/"
f"gsplat.git@{GSPLAT_COMMIT}"
)
check_code = (
"import torch; "
"from depth_anything_3.api import DepthAnything3; "
"from ben2 import BEN_Base; "
"import fastapi, uvicorn; "
"import gsplat; "
"print(torch.__version__, torch.cuda.is_available())"
)
if torch_index is not None:
torch_command.extend(["--index-url", torch_index])
stages = [
(
"Create Python Runtime",
"uv",
[uv, "venv", "--clear", "--python", "3.12", runtime_dir],
None,
"plain",
),
(
"Install PyTorch",
[
uv,
"pip",
"install",
"-v",
"--python",
python,
f"torch=={TORCH_VERSION}",
f"torchvision=={TORCHVISION_VERSION}",
"--index-url",
torch_index,
],
"pytorch deps",
torch_command,
None,
"install",
),
("Install xFormers", xformers_command, None, "install"),
(
"Install Depth Anything 3",
"ai deps",
[
uv,
"pip",
@@ -191,73 +401,34 @@ def install_stages(args, runtime_dir, python):
python,
da3_requirement,
f"addict=={ADDICT_VERSION}",
],
None,
"install",
),
(
"Install BEN2",
[
uv,
"pip",
"install",
"-v",
"--python",
python,
ben2_requirement,
],
None,
"install",
),
(
"Install AI Host",
[
uv,
"pip",
"install",
"-v",
"--python",
python,
f"fastapi=={FASTAPI_VERSION}",
f"uvicorn=={UVICORN_VERSION}",
],
None,
"install",
),
(
"Install gsplat build tools",
[
uv,
"pip",
"install",
"-v",
"--python",
python,
"setuptools",
"wheel",
"ninja",
],
None,
"install",
),
(
"Install gsplat",
[
uv,
"pip",
"install",
"-v",
"--python",
python,
"--no-build-isolation",
gsplat_requirement,
],
# PLY 导出不调用 rasterizer，禁用本地扩展可避免 CUDA 与 MSVC 工具链耦合
{"BUILD_NO_CUDA": "1"},
"install",
),
("Check Runtime", [python, "-c", check_code], None, "plain"),
]
if sys.platform != "darwin":
xformers_command = [
uv,
"pip",
"install",
"-v",
"--python",
python,
f"xformers=={XFORMERS_VERSION}",
]
if args.backend == "cu128":
xformers_command.extend(["--index-url", torch_index])
else:
xformers_command.append("--no-deps")
stages.insert(
2,
("xformers", xformers_command, None, "install"),
)
models = {**MODELS, BEN2_MODEL.key: BEN2_MODEL}
worker_script = Path(__file__).with_name("worker.py")
models_dir = Path(args.models_dir)
@@ -278,7 +449,7 @@ def install_stages(args, runtime_dir, python):
download_command.extend(["--allow-pattern", pattern])
stages.append(
(
f"Download {model.label}",
model.label.lower(),
download_command,
None,
"download",
@@ -287,21 +458,153 @@ def install_stages(args, runtime_dir, python):
return stages
def run_stage(status_path, stage, stages, label, command, environment, kind):
def runtime_ready(runtime_dir):
ready_path = runtime_dir / "snapform_ready.json"
if not environment_python(runtime_dir).is_file() or not ready_path.is_file():
return False
try:
payload = json.loads(ready_path.read_text(encoding="utf-8"))
except (json.JSONDecodeError, OSError):
return False
return payload.get("runtime_schema") == RUNTIME_SCHEMA
def verify_command(python):
return [python, "-c", CHECK_CODE]
def uv_download_target():
machine = platform.machine().lower()
if sys.platform == "win32":
if machine in {"arm64", "aarch64"}:
return "aarch64-pc-windows-msvc", "zip"
return "x86_64-pc-windows-msvc", "zip"
if sys.platform == "darwin":
arch = "aarch64" if machine in {"arm64", "aarch64"} else "x86_64"
return f"{arch}-apple-darwin", "tar.gz"
if sys.platform.startswith("linux"):
arch = "aarch64" if machine in {"arm64", "aarch64"} else "x86_64"
return f"{arch}-unknown-linux-gnu", "tar.gz"
raise RuntimeError(f"Unsupported platform for uv download: {sys.platform}")
def uv_download_url():
target, archive = uv_download_target()
return (
"https://github.com/astral-sh/uv/releases/download/"
f"{UV_VERSION}/uv-{target}.{archive}"
)
def extract_uv(archive_path, uv_path):
executable_name = "uv.exe" if sys.platform == "win32" else "uv"
with tempfile.TemporaryDirectory() as directory:
if zipfile.is_zipfile(archive_path):
with zipfile.ZipFile(archive_path) as archive:
archive.extractall(directory)
else:
with tarfile.open(archive_path) as archive:
archive.extractall(directory)
source = next(Path(directory).rglob(executable_name))
shutil.copy2(source, uv_path)
if sys.platform != "win32":
os.chmod(uv_path, 0o755)
def download_uv(
uv_path,
status_path,
stage=None,
stages=None,
stage_label=None,
start=0.0,
end=1.0,
):
"""Download the standalone uv binary and map byte progress onto the stage progress range."""
stage_fields = {
"stage": stage,
"stages": stages,
"stage_label": stage_label,
}
def write(progress, message):
write_status(
status_path,
start + progress * (end - start),
message,
**stage_fields,
)
write(0.0, "downloading")
uv_path = Path(uv_path)
uv_path.parent.mkdir(parents=True, exist_ok=True)
archive_path = uv_path.with_name(f"{uv_path.name}.part")
request = urllib.request.Request(
uv_download_url(),
headers={"User-Agent": "SnapForm-Runtime-Installer"},
)
try:
with urllib.request.urlopen(request, timeout=60) as response:
total = int(response.headers.get("Content-Length") or 0)
downloaded = 0
started = time.monotonic()
with archive_path.open("wb") as output:
while True:
chunk = response.read(65536)
if not chunk:
break
output.write(chunk)
downloaded += len(chunk)
progress = (
min(downloaded / total, 1.0) if total else 0.0
)
speed = downloaded / max(time.monotonic() - started, 1e-6)
write(progress, f"downloading {format_rate(speed)}")
write(1.0, "preparing")
extract_uv(archive_path, uv_path)
finally:
archive_path.unlink(missing_ok=True)
def run_stage(
status_path,
stage,
stages,
label,
command,
environment,
kind,
start=0.0,
end=1.0,
):
write_status(
status_path,
0.0,
start,
"",
stage=stage,
stages=stages,
stage_label=label,
)
installed = []
if kind == "install":
run_install(status_path, stage, stages, label, command, environment)
installed = run_install(
status_path,
stage,
stages,
label,
command,
environment,
start=start,
end=end,
)
elif kind == "download":
run(
[
*command,
"--progress-start",
str(start),
"--progress-end",
str(end),
"--stage",
str(stage),
"--stages",
@@ -315,21 +618,76 @@ def run_stage(status_path, stage, stages, label, command, environment, kind):
run(command, environment)
write_status(
status_path,
1.0,
end,
"",
stage=stage,
stages=stages,
stage_label=label,
)
return installed
def install(args):
if sys.platform == "darwin" and platform.machine().lower() not in {
"arm64",
"aarch64",
}:
raise RuntimeError(
"SnapForm requires Apple Silicon macOS; PyTorch stopped "
"shipping Intel Mac (x86_64) wheels after 2.2.2"
)
write_status(args.status, 0.0, "checking prerequisites")
_resolved, used_mirror = preflight(args)
if used_mirror:
write_status(args.status, MIRROR_END, "using mirror sources")
runtime_dir = Path(args.runtime_dir)
runtime_dir.parent.mkdir(parents=True, exist_ok=True)
if not runtime_ready(runtime_dir):
shutil.rmtree(runtime_dir, ignore_errors=True)
uv = Path(args.uv)
if not uv.is_file():
resolved = shutil.which(args.uv)
uv = Path(resolved) if resolved else uv
uv_missing = not uv.is_file()
python = environment_python(runtime_dir)
stages = install_stages(args, runtime_dir, python)
for index, (label, command, environment, kind) in enumerate(stages, start=1):
run_stage(
stages = install_stages(args, runtime_dir, python, uv)
runtime_stages = [stage for stage in stages if stage[3] != "download"]
model_stages = [stage for stage in stages if stage[3] == "download"]
plain_count = sum(1 for stage in runtime_stages if stage[3] != "install")
runtime_slice = RUNTIME_SLICE_END - RUNTIME_SLICE_START
install_count = len(runtime_stages) - plain_count
install_span = (
(runtime_slice - PLAIN_SLICE_SPAN * plain_count) / install_count
if install_count
else 0.0
)
model_span = (
(MODELS_SLICE_END - MODELS_SLICE_START) / len(model_stages)
if model_stages
else 0.0
)
installed_packages = []
runtime_offset = RUNTIME_SLICE_START
for index, (label, command, environment, kind) in enumerate(
runtime_stages,
start=1,
):
stage_start = runtime_offset
stage_end = stage_start + (
install_span if kind == "install" else PLAIN_SLICE_SPAN
)
if uv_missing and index == 1:
download_uv(
uv,
args.status,
stage=index,
stages=len(stages),
stage_label=label,
start=stage_start,
end=stage_start + PLAIN_SLICE_SPAN * 0.5,
)
stage_start += PLAIN_SLICE_SPAN * 0.5
stage_installed = run_stage(
args.status,
index,
len(stages),
@@ -337,8 +695,17 @@ def install(args):
command,
environment,
kind,
start=stage_start,
end=stage_end,
)
if isinstance(stage_installed, list):
installed_packages.extend(stage_installed)
runtime_offset = stage_end
write_status(args.status, VERIFY_SLICE_START, "verifying runtime")
run(verify_command(python))
write_status(args.status, VERIFY_SLICE_END, "Runtime is ready")
runtime_dir.mkdir(parents=True, exist_ok=True)
ready_path = runtime_dir / "snapform_ready.json"
ready_path.write_text(
json.dumps(
@@ -352,14 +719,42 @@ def install(args):
"backend": args.backend,
"da3_commit": args.da3_commit,
"ben2_commit": args.ben2_commit,
"gsplat_commit": GSPLAT_COMMIT,
"gsplat_cuda": False,
"packages": sorted(set(installed_packages)),
},
indent=2,
),
encoding="utf-8",
)
write_status(args.status, 1.0, "Runtime is ready")
offset = len(runtime_stages)
model_offset = MODELS_SLICE_START
for index, (label, command, environment, kind) in enumerate(
model_stages,
start=offset + 1,
):
stage_start = model_offset
stage_end = stage_start + model_span
try:
run_stage(
args.status,
index,
len(stages),
label,
command,
environment,
kind,
start=stage_start,
end=stage_end,
)
except Exception as error:
print(f"Model download failed: {label}: {error}")
write_status(
args.status,
MODELS_SLICE_END,
f"Runtime is ready; {label} failed",
)
return
model_offset = stage_end
write_status(args.status, 1.0, "Runtime and models ready")
def parse_args():

12
src/snapform/operators.py
View File

@@ -61,8 +61,6 @@ def active_job():
def cancel_active_job():
global _active_job
if _active_job is None:
return
_active_job.cancelled = True
@@ -422,10 +420,12 @@ class RunJob(bpy.types.Operator):
stage_label = status.get("stage_label")
if not stage or not stages:
return message
stage_text = f"Step {stage}/{stages}"
if stage_label:
stage_text = f"{stage_text} {stage_label}"
return f"{stage_text}: {message}" if message else stage_text
text = (
f"{stage}/{stages} {stage_label}"
if stage_label
else f"{stage}/{stages}"
)
return f"{text}: {message}" if message else text
def _finish(self, context, job, return_code):
settings = context.scene.snapform_settings

10
src/snapform/preferences.py
View File

@@ -1,4 +1,5 @@
import shutil
import sys
from pathlib import Path
import bpy
@@ -18,9 +19,12 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
)
uv_executable: bpy.props.StringProperty(
name="uv Executable",
description="Path to uv; a command name is accepted when uv is on PATH",
description=(
"Path to uv; leave empty to auto-download a standalone uv "
"when it is missing"
),
subtype="FILE_PATH",
default=shutil.which("uv") or "uv",
default=shutil.which("uv") or "",
)
torch_backend: bpy.props.EnumProperty(
name="Torch Backend",
@@ -29,7 +33,7 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
("CU128", "CUDA 12.8", "NVIDIA CUDA 12.8 PyTorch wheels"),
("CPU", "CPU", "CPU-only PyTorch wheels"),
),
default="CU128",
default="CPU" if sys.platform == "darwin" else "CU128",
)
def draw(self, _context):

3
src/snapform/properties.py
View File

@@ -35,8 +35,9 @@ class SnapFormSettings(bpy.types.PropertyGroup):
name="Device",
description="Inference device used by the external Runtime",
items=(
("AUTO", "Auto", "Use CUDA when available, otherwise CPU"),
("AUTO", "Auto", "Use CUDA, then MPS, otherwise CPU"),
("CUDA", "CUDA", "Require an NVIDIA CUDA device"),
("MPS", "MPS", "Use Apple Metal (MPS)"),
("CPU", "CPU", "Run on CPU"),
),
default="AUTO",

20
src/snapform/runtime.py
View File

@@ -4,6 +4,7 @@ import secrets
import shutil
import socket
import subprocess
import sys
import time
import uuid
from dataclasses import dataclass
@@ -141,7 +142,13 @@ def uv_executable():
return resolved
if Path(configured).is_file():
return str(Path(configured))
raise RuntimeError("uv executable was not found")
return str(uv_cache_path())
def uv_cache_path():
if os.name == "nt":
return storage_root() / "uv" / "uv.exe"
return storage_root() / "uv" / "uv"
def install_runtime_command(status_path):
@@ -150,11 +157,7 @@ def install_runtime_command(status_path):
backend = getattr(addon_preferences, "torch_backend", "CU128")
script = Path(__file__).with_name("bootstrap_runtime.py")
return [
executable,
"run",
"--no-project",
"--python",
"3.12",
sys.executable,
str(script),
"--uv",
executable,
@@ -450,7 +453,8 @@ def detach_host_client():
pass
_host_connection = None
_registered_instance_id = None
_close_launched_host(terminate=False)
# 如果 Host 是由当前 Blender 启动的，关闭 Blender 时直接终止它，无需等待空闲超时
_close_launched_host(terminate=True)
def schedule_host_start():
@@ -598,6 +602,8 @@ def download_model_command(model_key, status_path):
str(model_directory(model_key)),
"--status",
str(status_path),
"--label",
model.label,
]
for pattern in model.download_patterns:
command.extend(["--allow-pattern", pattern])

6
src/snapform/status_bar.py
View File

@@ -1,5 +1,3 @@
import bpy
from .operators import active_job
@@ -49,10 +47,6 @@ def redraw_ui(context, force=False):
tag_ui_regions(context)
if not force:
return
try:
bpy.ops.wm.redraw_timer(type="DRAW_WIN_SWAP", iterations=1)
except (AttributeError, RuntimeError, TypeError):
pass
workspace = getattr(context, "workspace", None)
if workspace is not None:
try:

143
src/snapform/worker.py
View File

@@ -1,9 +1,12 @@
import argparse
import json
import os
import shutil
import tempfile
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path
@@ -21,6 +24,9 @@ VIDEO_SUFFIXES = {
}
SCENE_FILENAME = "scene.glb"
EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz")
HF_ENDPOINT = "https://huggingface.co/"
HF_MIRRORS = ("https://hf-mirror.com",)
HF_PROBE_TIMEOUT = 5.0
class TaskCancelled(RuntimeError):
@@ -32,6 +38,34 @@ def check_cancelled(cancel_check):
raise TaskCancelled("Task cancelled")
def reachable(url, timeout=HF_PROBE_TIMEOUT):
request = urllib.request.Request(
url,
method="HEAD",
headers={"User-Agent": "SnapForm-Runtime-Installer"},
)
try:
with urllib.request.urlopen(request, timeout=timeout):
return True
except urllib.error.HTTPError:
return True
except Exception:
return False
def resolve_hf_endpoint():
"""Probe Hugging Face reachability and fall back to a mirror when direct access fails; returns (endpoint, used_mirror)."""
configured = os.environ.get("HF_ENDPOINT")
if configured:
return configured, False
if reachable(HF_ENDPOINT):
return HF_ENDPOINT, False
for mirror in HF_MIRRORS:
if reachable(mirror):
return mirror, True
return HF_ENDPOINT, False
def write_status(
path,
progress,
@@ -159,6 +193,7 @@ class DownloadProgressReporter:
stage=None,
stages=None,
stage_label=None,
label=None,
):
self.status_path = status_path
self.repository = repository
@@ -169,10 +204,11 @@ class DownloadProgressReporter:
self.stage = stage
self.stages = stages
self.stage_label = stage_label
self.label = label
self.progress = start
self.last_write_time = 0.0
def report(self, current, total):
def report(self, current, total, rate=None):
if not total or current <= 0:
return
@@ -186,10 +222,18 @@ class DownloadProgressReporter:
return
self.progress = progress
self.last_write_time = now
speed = format_rate(rate) if rate else ""
if self.stage is not None:
message = f"downloading {speed}" if speed else "downloading"
else:
prefix = self.label or self.repository
message = f"downloading {prefix}"
if speed:
message = f"{message} {speed}"
write_status(
self.status_path,
progress,
f"Downloading {self.repository}",
message,
stage=self.stage,
stages=self.stages,
stage_label=self.stage_label,
@@ -206,6 +250,7 @@ def create_download_progress_class(
stage=None,
stages=None,
stage_label=None,
label=None,
):
if tqdm_type is None:
from tqdm.auto import tqdm as tqdm_type
@@ -219,6 +264,7 @@ def create_download_progress_class(
stage,
stages,
stage_label,
label,
)
class DownloadProgress(tqdm_type):
@@ -233,26 +279,40 @@ def create_download_progress_class(
def update(self, amount=1):
result = super().update(amount)
if self.snapform_reports_progress:
reporter.report(self.n, self.total)
reporter.report(self.n, self.total, self._download_rate())
return result
def refresh(self, *args, **kwargs):
result = super().refresh(*args, **kwargs)
if self.snapform_reports_progress:
reporter.report(self.n, self.total)
reporter.report(self.n, self.total, self._download_rate())
return result
def _download_rate(self):
try:
return getattr(self, "format_dict", {}).get("rate")
except Exception:
return None
return DownloadProgress
def download_model(args):
from huggingface_hub import snapshot_download
def progress_range(args, staged):
start = getattr(args, "progress_start", None)
end = getattr(args, "progress_end", None)
if start is not None and end is not None:
return float(start), float(end) - float(start)
if staged:
return 0.0, 1.0
return 0.05, 0.90
def download_model(args):
label = getattr(args, "label", None) or args.repository
model_dir = Path(args.model_dir)
model_dir.mkdir(parents=True, exist_ok=True)
staged = getattr(args, "stage", None) is not None
start = 0.0 if staged else 0.05
span = 1.0 if staged else 0.90
start, span = progress_range(args, staged)
stage_kwargs = {
"stage": getattr(args, "stage", None),
"stages": getattr(args, "stages", None),
@@ -261,14 +321,24 @@ def download_model(args):
write_status(
args.status,
start,
f"Downloading {args.repository}",
"checking download source",
**stage_kwargs,
)
endpoint, used_mirror = resolve_hf_endpoint()
if endpoint != os.environ.get("HF_ENDPOINT"):
os.environ["HF_ENDPOINT"] = endpoint
from huggingface_hub import snapshot_download
message = f"downloading {label}"
if used_mirror:
message = f"{message} (mirror)"
write_status(args.status, start, message, **stage_kwargs)
progress_class = create_download_progress_class(
args.status,
args.repository,
start=start,
span=span,
label=label,
**stage_kwargs,
)
snapshot_download(
@@ -279,12 +349,26 @@ def download_model(args):
)
write_status(
args.status,
1.0,
"Model is ready",
start + span,
"model ready",
**stage_kwargs,
)
def format_rate(rate):
size = float(rate)
if size <= 0:
return ""
units = ("B/s", "KB/s", "MB/s", "GB/s")
unit_index = 0
while size >= 1024 and unit_index < len(units) - 1:
size /= 1024
unit_index += 1
if unit_index == 0:
return f"{size:.0f} {units[unit_index]}"
return f"{size:.1f} {units[unit_index]}"
def select_device(requested, torch):
if requested == "cpu":
return torch.device("cpu")
@@ -292,7 +376,15 @@ def select_device(requested, torch):
if not torch.cuda.is_available():
raise RuntimeError("CUDA was requested but is not available")
return torch.device("cuda")
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
if requested == "mps":
if not torch.backends.mps.is_available():
raise RuntimeError("MPS was requested but is not available")
return torch.device("mps")
if torch.cuda.is_available():
return torch.device("cuda")
if torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")
def infer_foreground(
@@ -404,18 +496,36 @@ def exported_paths(output_directory, export_format):
return []
def npz_is_complete(path):
"""Return True only when the npz archive can be fully read."""
try:
import numpy as np
with np.load(path, allow_pickle=False) as archive:
return "depth" in archive.files
except Exception:
return False
def wait_for_exported_paths(
output_directory,
export_format,
timeout=30.0,
cancel_check=None,
):
# 上游 NPZ exporter 使用后台线程，因此需要等待文件实际落盘。
# DA3 的 NPZ exporter 在后台线程写入，文件出现时可能尚未写完整，
# 直接返回会让下游读到半截 zip 而报 EOFError，因此需等文件可完整读取
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
check_cancelled(cancel_check)
result_paths = exported_paths(output_directory, export_format)
if result_paths:
if (
export_format in ("mini_npz", "npz")
and not npz_is_complete(result_paths[0])
):
time.sleep(0.1)
continue
return result_paths
time.sleep(0.05)
return []
@@ -549,6 +659,9 @@ def parser():
download.add_argument("--stage", type=int, default=None)
download.add_argument("--stages", type=int, default=None)
download.add_argument("--stage-label", default=None)
download.add_argument("--label", default=None)
download.add_argument("--progress-start", type=float, default=None)
download.add_argument("--progress-end", type=float, default=None)
background = commands.add_parser("remove-background")
background.add_argument("--input", required=True)
@@ -557,7 +670,7 @@ def parser():
background.add_argument("--status", required=True)
background.add_argument(
"--device",
choices=("auto", "cuda", "cpu"),
choices=("auto", "cuda", "mps", "cpu"),
required=True,
)
@@ -598,7 +711,7 @@ def parser():
reconstruction.add_argument("--video-frames", type=int, required=True)
reconstruction.add_argument(
"--device",
choices=("auto", "cuda", "cpu"),
choices=("auto", "cuda", "mps", "cpu"),
required=True,
)
reconstruction.add_argument(

90
tests/test_addon.py
View File

@@ -517,27 +517,27 @@ class AddonTest(unittest.TestCase):
operator._format_status(
{
"progress": 0.5,
"message": "Downloaded 3 packages",
"message": "2 downloaded",
"stage": 3,
"stages": 12,
"stage_label": "Install PyTorch",
"stages": 8,
"stage_label": "pytorch",
},
"fallback",
),
"Step 3/12 Install PyTorch: Downloaded 3 packages",
"3/8 pytorch: 2 downloaded",
)
self.assertEqual(
operator._format_status(
{
"progress": 0.0,
"message": "",
"stage": 10,
"stages": 12,
"stage_label": "Download DA3 Small",
"stage": 6,
"stages": 8,
"stage_label": "da3 small",
},
"fallback",
),
"Step 10/12 Download DA3 Small",
"6/8 da3 small",
)
self.assertEqual(
operator._format_status(
@@ -865,6 +865,31 @@ class AddonTest(unittest.TestCase):
self.assertIn("ProxyHandler({})", inspect.getsource(runtime))
self.assertIn('return storage_root() / "host.json"', inspect.getsource(runtime.host_discovery_path))
def test_closing_blender_terminates_host_it_started(self):
runtime = importlib.import_module("snapform.runtime")
previous_connection = runtime._host_connection
previous_registered = runtime._registered_instance_id
runtime._host_connection = SimpleNamespace(port=12345)
runtime._registered_instance_id = "instance"
try:
with patch.object(
runtime,
"_host_http_request",
) as request, patch.object(
runtime,
"_close_launched_host",
) as close_host:
runtime.detach_host_client()
request.assert_called_once()
self.assertEqual(request.call_args.args[1], "DELETE")
close_host.assert_called_once_with(terminate=True)
finally:
runtime._host_connection = previous_connection
runtime._registered_instance_id = previous_registered
self.assertIsNone(runtime._host_connection)
self.assertIsNone(runtime._registered_instance_id)
def test_host_jobs_complete_from_status_state(self):
source = inspect.getsource(self.addon.RunJob.modal)
@@ -1357,6 +1382,37 @@ class AddonTest(unittest.TestCase):
if value == "--allow-pattern"
]
self.assertEqual(patterns, ["config.json", "model.safetensors"])
label_index = command.index("--label")
self.assertEqual(command[label_index + 1], "BEN2 Base")
def test_install_runtime_command_has_no_model_label_reference(self):
runtime = importlib.import_module("snapform.runtime")
with patch.object(
runtime,
"uv_executable",
return_value="C:/uv/uv.exe",
), patch.object(
runtime,
"preferences",
return_value=SimpleNamespace(torch_backend="CU128"),
), patch.object(
runtime,
"runtime_directory",
return_value=Path("runtime"),
), patch.object(
runtime,
"models_directory",
return_value=Path("models"),
):
command = runtime.install_runtime_command(Path("status.json"))
self.assertNotIn("--label", command)
self.assertEqual(command[0], sys.executable)
self.assertIn("--uv", command)
self.assertIn("C:/uv/uv.exe", command)
self.assertTrue(
any("bootstrap_runtime.py" in str(value) for value in command)
)
def test_preferences_install_button_is_disabled_while_busy(self):
draw_source = inspect.getsource(self.addon.SnapFormPreferences.draw)
@@ -1526,7 +1582,7 @@ class AddonTest(unittest.TestCase):
],
)
def test_force_flag_triggers_window_draw_and_swap(self):
def test_force_flag_refreshes_status_text(self):
redrawn = []
context = SimpleNamespace(
window_manager=SimpleNamespace(
@@ -1553,29 +1609,17 @@ class AddonTest(unittest.TestCase):
),
)
with patch.object(
self.addon.status_bar.bpy,
"ops",
SimpleNamespace(
wm=SimpleNamespace(
redraw_timer=lambda **options: redrawn.append(
("swap", options)
),
),
),
):
self.addon.status_bar.redraw_ui(context, force=True)
self.addon.status_bar.redraw_ui(context, force=True)
self.assertEqual(
redrawn,
[
"status",
("swap", {"type": "DRAW_WIN_SWAP", "iterations": 1}),
"status-text",
],
)
def test_force_flag_ignores_missing_redraw_helpers(self):
def test_force_flag_ignores_missing_workspace(self):
redrawn = []
context = SimpleNamespace(
window_manager=SimpleNamespace(

799
tests/test_bootstrap_runtime.py
View File

@@ -1,7 +1,11 @@
import importlib.util
import json
import os
import subprocess
import tempfile
import time
import unittest
import zipfile
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
@@ -21,14 +25,16 @@ def load_bootstrap_runtime():
class BootstrapRuntimeTest(unittest.TestCase):
def test_installs_da3_and_non_compiled_gsplat(self):
def test_installs_da3_from_archives_without_gsplat(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
runtime_dir = root / "runtime"
runtime_dir.mkdir()
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv="uv",
uv=str(uv_path),
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
@@ -36,7 +42,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(bootstrap, "run") as mocked_run, patch.object(
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
bootstrap,
"run",
) as mocked_run, patch.object(
bootstrap,
"run_install",
) as mocked_run_install:
@@ -55,25 +64,44 @@ class BootstrapRuntimeTest(unittest.TestCase):
f"addict=={bootstrap.ADDICT_VERSION}",
da3_command,
)
gsplat_call = next(
call
for call in mocked_run_install.call_args_list
if any("gsplat @" in str(value) for value in call.args[4])
)
gsplat_command = gsplat_call.args[4]
self.assertIn("--no-build-isolation", gsplat_command)
self.assertTrue(
any(bootstrap.GSPLAT_COMMIT in str(value) for value in gsplat_command)
any(
"https://github.com/ByteDance-Seed/Depth-Anything-3/"
"archive/test-commit.tar.gz" in str(value)
for value in da3_command
)
)
self.assertFalse(
any("git+" in str(value) for value in da3_command)
)
self.assertFalse(
any(
"gsplat" in str(value)
for command in commands
for value in command
)
)
self.assertFalse(
any(
"--no-build-isolation" in str(value)
for command in commands
for value in command
)
)
self.assertEqual(gsplat_call.args[5], {"BUILD_NO_CUDA": "1"})
ben2_command = next(
command
for command in commands
if any("ben2 @" in str(value) for value in command)
)
self.assertTrue(
any("test-ben2-commit" in str(value) for value in ben2_command)
any(
"https://github.com/PramaLLC/BEN2/archive/"
"test-ben2-commit.tar.gz" in str(value)
for value in ben2_command
)
)
self.assertFalse(
any("git+" in str(value) for value in ben2_command)
)
host_command = next(
command
@@ -81,6 +109,8 @@ class BootstrapRuntimeTest(unittest.TestCase):
if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command
)
self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", host_command)
self.assertIn("setuptools", host_command)
self.assertNotIn("ninja", host_command)
check_command = next(
command
for command in commands
@@ -95,6 +125,9 @@ class BootstrapRuntimeTest(unittest.TestCase):
self.assertTrue(
any("import fastapi, uvicorn" in str(value) for value in check_command)
)
self.assertFalse(
any("import gsplat" in str(value) for value in check_command)
)
default_downloads = [
command
for command in commands
@@ -121,7 +154,86 @@ class BootstrapRuntimeTest(unittest.TestCase):
self.assertEqual(patterns, ["config.json", "model.safetensors"])
self.assertIn("--stage", ben2_download)
self.assertIn("--stage-label", ben2_download)
self.assertEqual(ben2_download[-1], "Download BEN2 Base")
self.assertEqual(ben2_download[-1], "ben2 base")
def test_darwin_install_skips_cuda_index_and_xformers(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
runtime_dir = root / "runtime"
runtime_dir.mkdir()
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(
bootstrap.sys,
"platform",
"darwin",
), patch.object(
bootstrap.platform,
"machine",
return_value="arm64",
), patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
bootstrap,
"run",
), patch.object(
bootstrap,
"run_install",
) as mocked_run_install:
bootstrap.install(args)
commands = [
call.args[4] for call in mocked_run_install.call_args_list
]
torch_command = next(
command
for command in commands
if any(
f"torch=={bootstrap.TORCH_VERSION}" in str(value)
for value in command
)
)
self.assertNotIn("--index-url", torch_command)
self.assertFalse(
any(
"xformers" in str(value)
for command in commands
for value in command
)
)
def test_intel_mac_is_rejected_with_clear_error(self):
bootstrap = load_bootstrap_runtime()
args = SimpleNamespace(
uv="uv",
runtime_dir="runtime",
models_dir="models",
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status="status.json",
)
with patch.object(
bootstrap.sys,
"platform",
"darwin",
), patch.object(
bootstrap.platform,
"machine",
return_value="x86_64",
), patch.object(bootstrap.shutil, "rmtree") as mocked_rmtree:
with self.assertRaisesRegex(RuntimeError, "Apple Silicon"):
bootstrap.install(args)
mocked_rmtree.assert_not_called()
def test_run_stage_writes_stage_progress(self):
bootstrap = load_bootstrap_runtime()
@@ -131,8 +243,8 @@ class BootstrapRuntimeTest(unittest.TestCase):
bootstrap.run_stage(
str(status_path),
3,
12,
"Install PyTorch",
8,
"pytorch",
["uv", "pip", "install", "torch"],
None,
"plain",
@@ -145,8 +257,131 @@ class BootstrapRuntimeTest(unittest.TestCase):
status = json.loads(status_path.read_text(encoding="utf-8"))
self.assertEqual(status["progress"], 1.0)
self.assertEqual(status["stage"], 3)
self.assertEqual(status["stages"], 12)
self.assertEqual(status["stage_label"], "Install PyTorch")
self.assertEqual(status["stages"], 8)
self.assertEqual(status["stage_label"], "pytorch")
def test_run_stage_writes_progress_range(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
status_path = Path(directory) / "status.json"
with patch.object(bootstrap, "run"), patch.object(
bootstrap,
"write_status",
) as mocked_write:
bootstrap.run_stage(
str(status_path),
3,
8,
"pytorch",
["uv", "pip", "install", "torch"],
None,
"plain",
start=0.2,
end=0.4,
)
progresses = [
call.args[1] for call in mocked_write.call_args_list
]
self.assertEqual(progresses[0], 0.2)
self.assertEqual(progresses[-1], 0.4)
def test_run_stage_passes_progress_range_to_download_command(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
status_path = Path(directory) / "status.json"
command = [
"python",
"worker.py",
"download",
"--repository",
"depth-anything/DA3-SMALL",
]
with patch.object(bootstrap, "run") as mocked_run:
bootstrap.run_stage(
str(status_path),
6,
8,
"da3 small",
command,
None,
"download",
start=0.9,
end=0.9333,
)
called_command = mocked_run.call_args.args[0]
self.assertIn("--progress-start", called_command)
self.assertIn("--progress-end", called_command)
self.assertEqual(
called_command[called_command.index("--progress-start") + 1],
"0.9",
)
self.assertEqual(
called_command[called_command.index("--progress-end") + 1],
"0.9333",
)
self.assertEqual(called_command[-1], "da3 small")
def test_install_allocates_contiguous_progress_slices(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
runtime_dir = root / "runtime"
runtime_dir.mkdir()
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
stage_ranges = []
def fake_run_stage(
status_path,
stage,
stages,
label,
command,
environment,
kind,
start=0.0,
end=1.0,
):
stage_ranges.append((stage, start, end))
bootstrap.write_status(
status_path,
end,
"",
stage=stage,
stages=stages,
stage_label=label,
)
return []
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
bootstrap,
"run_stage",
side_effect=fake_run_stage,
), patch.object(bootstrap, "run"):
bootstrap.install(args)
starts = [item[1] for item in stage_ranges]
ends = [item[2] for item in stage_ranges]
self.assertTrue(
all(starts[i] < starts[i + 1] for i in range(len(starts) - 1))
)
self.assertEqual(ends[-1], 1.0)
self.assertEqual(starts[0], bootstrap.RUNTIME_SLICE_START)
status = json.loads(
(root / "status.json").read_text(encoding="utf-8")
)
self.assertEqual(status["progress"], 1.0)
def test_run_install_tracks_package_download_events(self):
bootstrap = load_bootstrap_runtime()
@@ -161,33 +396,515 @@ class BootstrapRuntimeTest(unittest.TestCase):
fake_process = FakeProcess(
[
"Resolved 3 packages in 1.0s",
"Downloading torch (10MiB)",
" Downloaded torch",
"DEBUG Adding transitive dependency for torch==2.9.1: filelock*",
"Prepared 3 packages in 2.0s",
"Installed 3 packages in 0.5s",
"+ torch==2.9.1",
"+ numpy==1.26.4",
]
)
with patch.object(
bootstrap.subprocess,
"Popen",
return_value=fake_process,
), patch.object(bootstrap, "write_status") as mocked_write:
installed = bootstrap.run_install(
"status.json",
2,
8,
"pytorch deps",
["uv", "pip", "install", "torch"],
)
messages = [call.args[2] for call in mocked_write.call_args_list]
self.assertEqual(
messages,
[
"3 packages",
"3/3 packages",
"3/3 packages",
],
)
self.assertEqual(mocked_write.call_args_list[-1].args[1], 0.02)
first = mocked_write.call_args_list[0]
self.assertEqual(first.kwargs["stage"], 2)
self.assertEqual(first.kwargs["stages"], 8)
self.assertEqual(first.kwargs["stage_label"], "pytorch deps")
self.assertEqual(installed, ["torch==2.9.1", "numpy==1.26.4"])
def test_run_install_smooths_progress_between_events(self):
bootstrap = load_bootstrap_runtime()
class SlowLines:
def __init__(self, lines, delay):
self.lines = list(lines)
self.delay = delay
def __iter__(self):
for line in self.lines:
time.sleep(self.delay)
yield line
class FakeProcess:
def __init__(self, lines):
self.stdout = lines
def wait(self):
return 0
fake_process = FakeProcess(
SlowLines(
[
"Resolved 3 packages in 1.0s",
"Prepared 3 packages in 2.0s",
"Installed 3 packages in 0.5s",
],
0.3,
)
)
with patch.object(
bootstrap.subprocess,
"Popen",
return_value=fake_process,
), patch.object(bootstrap, "write_status") as mocked_write:
bootstrap.run_install(
"status.json",
2,
8,
"pytorch deps",
["uv", "pip", "install", "torch"],
)
progresses = [call.args[1] for call in mocked_write.call_args_list]
self.assertGreater(progresses[-1], 0.02)
self.assertTrue(
all(
progresses[i] <= progresses[i + 1] + 1e-9
for i in range(len(progresses) - 1)
)
)
messages = [call.args[2] for call in mocked_write.call_args_list]
self.assertIn("3 packages", messages)
self.assertIn("3/3 packages", messages)
def test_run_install_maps_progress_into_stage_range(self):
bootstrap = load_bootstrap_runtime()
class FakeProcess:
def __init__(self, lines):
self.stdout = lines
def wait(self):
return 0
fake_process = FakeProcess(
[
"Resolved 3 packages in 1.0s",
"Prepared 3 packages in 2.0s",
"Installed 3 packages in 0.5s",
]
)
with tempfile.TemporaryDirectory() as directory:
status_path = Path(directory) / "status.json"
with patch.object(
bootstrap.subprocess,
"Popen",
return_value=fake_process,
):
bootstrap.run_install(
str(status_path),
2,
12,
"Install PyTorch",
["uv", "pip", "install", "torch"],
)
with patch.object(
bootstrap.subprocess,
"Popen",
return_value=fake_process,
), patch.object(bootstrap, "write_status") as mocked_write:
bootstrap.run_install(
"status.json",
2,
8,
"pytorch deps",
["uv", "pip", "install", "torch"],
start=0.2,
end=0.4,
)
status = json.loads(status_path.read_text(encoding="utf-8"))
self.assertEqual(status["progress"], 0.95)
self.assertEqual(status["message"], "Installing packages")
self.assertEqual(status["stage"], 2)
self.assertEqual(status["stages"], 12)
self.assertEqual(status["stage_label"], "Install PyTorch")
first = mocked_write.call_args_list[0]
self.assertAlmostEqual(first.args[1], 0.204)
self.assertEqual(first.kwargs["stage"], 2)
self.assertEqual(first.kwargs["stages"], 8)
self.assertEqual(first.kwargs["stage_label"], "pytorch deps")
def test_formats_elapsed_time(self):
bootstrap = load_bootstrap_runtime()
self.assertEqual(bootstrap.format_elapsed(12), "12s")
self.assertEqual(bootstrap.format_elapsed(83), "1m 23s")
def test_preflight_rejects_low_disk_space(self):
bootstrap = load_bootstrap_runtime()
args = SimpleNamespace(
backend="cu128",
runtime_dir=str(Path.cwd()),
)
with patch.object(
bootstrap.shutil,
"disk_usage",
return_value=SimpleNamespace(free=1024**3),
), patch.object(bootstrap, "reachable", return_value=True):
with self.assertRaisesRegex(RuntimeError, "disk space"):
bootstrap.preflight(args)
def test_preflight_rejects_unreachable_network(self):
bootstrap = load_bootstrap_runtime()
args = SimpleNamespace(
backend="cu128",
runtime_dir=str(Path.cwd()),
)
with patch.object(
bootstrap.shutil,
"disk_usage",
return_value=SimpleNamespace(free=100 * 1024**3),
), patch.object(bootstrap, "reachable", return_value=False):
with self.assertRaisesRegex(RuntimeError, "Cannot reach"):
bootstrap.preflight(args)
def test_preflight_passes_when_checks_ok(self):
bootstrap = load_bootstrap_runtime()
args = SimpleNamespace(
backend="cu128",
runtime_dir=str(Path.cwd()),
)
with patch.object(
bootstrap.shutil,
"disk_usage",
return_value=SimpleNamespace(free=100 * 1024**3),
), patch.object(bootstrap, "reachable", return_value=True):
resolved, used_mirror = bootstrap.preflight(args)
self.assertFalse(used_mirror)
self.assertEqual(
resolved["pypi"],
"https://pypi.org/simple/",
)
def test_preflight_falls_back_to_mirrors(self):
bootstrap = load_bootstrap_runtime()
args = SimpleNamespace(
backend="cu128",
runtime_dir=str(Path.cwd()),
)
def fake_reachable(url, timeout=bootstrap.PREFLIGHT_TIMEOUT):
if "github.com" in url:
return True
if "tuna" in url or "hf-mirror" in url or "aliyun" in url:
return True
return False
with patch.object(
bootstrap.shutil,
"disk_usage",
return_value=SimpleNamespace(free=100 * 1024**3),
), patch.object(
bootstrap,
"reachable",
side_effect=fake_reachable,
):
resolved, used_mirror = bootstrap.preflight(args)
self.assertTrue(used_mirror)
self.assertEqual(
resolved["pypi"],
"https://pypi.tuna.tsinghua.edu.cn/simple",
)
self.assertEqual(resolved["hf"], "https://hf-mirror.com")
self.assertEqual(
resolved["torch"],
"https://mirrors.aliyun.com/pytorch-wheels/cu128",
)
self.assertEqual(
bootstrap._resolved_torch_index,
"https://mirrors.aliyun.com/pytorch-wheels/cu128",
)
def test_preflight_uses_environment_overrides(self):
bootstrap = load_bootstrap_runtime()
args = SimpleNamespace(
backend="cu128",
runtime_dir=str(Path.cwd()),
)
with patch.dict(
os.environ,
{
"PIP_INDEX_URL": "https://pypi.example/simple",
"HF_ENDPOINT": "https://hf.example",
"SNAPFORM_TORCH_INDEX": "https://torch.example/whl/cu128",
},
), patch.object(
bootstrap.shutil,
"disk_usage",
return_value=SimpleNamespace(free=100 * 1024**3),
), patch.object(bootstrap, "reachable", return_value=True):
resolved, used_mirror = bootstrap.preflight(args)
self.assertFalse(used_mirror)
self.assertEqual(resolved["pypi"], "https://pypi.example/simple")
self.assertEqual(resolved["hf"], "https://hf.example")
self.assertEqual(
resolved["torch"],
"https://torch.example/whl/cu128",
)
def test_child_environment_injects_resolved_mirrors(self):
bootstrap = load_bootstrap_runtime()
previous_pypi = bootstrap._resolved_pypi_index
previous_hf = bootstrap._resolved_hf_endpoint
bootstrap._resolved_pypi_index = (
"https://pypi.tuna.tsinghua.edu.cn/simple"
)
bootstrap._resolved_hf_endpoint = "https://hf-mirror.com"
try:
env = bootstrap.child_environment()
finally:
bootstrap._resolved_pypi_index = previous_pypi
bootstrap._resolved_hf_endpoint = previous_hf
self.assertEqual(
env["PIP_INDEX_URL"],
"https://pypi.tuna.tsinghua.edu.cn/simple",
)
self.assertEqual(env["HF_ENDPOINT"], "https://hf-mirror.com")
def test_child_environment_injects_registry_proxy(self):
bootstrap = load_bootstrap_runtime()
with patch.object(
bootstrap.urllib.request,
"getproxies",
return_value={
"http": "http://proxy:8080",
"https": "http://proxy:8080",
},
):
env = bootstrap.child_environment()
self.assertEqual(env["HTTP_PROXY"], "http://proxy:8080")
self.assertEqual(env["HTTPS_PROXY"], "http://proxy:8080")
def test_child_environment_keeps_explicit_proxy(self):
bootstrap = load_bootstrap_runtime()
with patch.object(
bootstrap.urllib.request,
"getproxies",
return_value={"https": "http://proxy:8080"},
):
env = bootstrap.child_environment(
{"HTTPS_PROXY": "http://explicit:3128"}
)
self.assertEqual(env["HTTPS_PROXY"], "http://explicit:3128")
def test_ready_json_records_installed_packages(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
runtime_dir = root / "runtime"
runtime_dir.mkdir()
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
bootstrap,
"run",
), patch.object(
bootstrap,
"run_stage",
return_value=["torch==2.9.1+cu128", "numpy==1.26.4"],
):
bootstrap.install(args)
ready = json.loads(
(runtime_dir / "snapform_ready.json").read_text(
encoding="utf-8"
)
)
self.assertEqual(
ready["packages"],
["numpy==1.26.4", "torch==2.9.1+cu128"],
)
def test_broken_runtime_is_cleaned_before_reinstall(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
runtime_dir = root / "runtime"
runtime_dir.mkdir()
(runtime_dir / "partial-python.exe").write_text("junk")
args = SimpleNamespace(
uv="uv",
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
bootstrap,
"run",
), patch.object(
bootstrap,
"run_install",
):
bootstrap.install(args)
self.assertFalse((runtime_dir / "partial-python.exe").exists())
self.assertTrue((runtime_dir / "snapform_ready.json").is_file())
def test_ready_runtime_is_not_cleaned(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
runtime_dir = root / "runtime"
runtime_dir.mkdir(parents=True)
(runtime_dir / "Scripts").mkdir()
(runtime_dir / "Scripts" / "python.exe").write_text("python")
(runtime_dir / "snapform_ready.json").write_text(
json.dumps({"runtime_schema": bootstrap.RUNTIME_SCHEMA}),
encoding="utf-8",
)
(runtime_dir / "marker.txt").write_text("keep")
args = SimpleNamespace(
uv="uv",
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
bootstrap,
"run",
), patch.object(
bootstrap,
"run_install",
):
bootstrap.install(args)
self.assertTrue((runtime_dir / "marker.txt").is_file())
def test_model_download_failure_keeps_runtime_ready(self):
bootstrap = load_bootstrap_runtime()
def fake_run(command, environment=None):
if any("DA3-SMALL" in str(value) for value in command):
raise subprocess.CalledProcessError(1, command)
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
runtime_dir = root / "runtime"
runtime_dir.mkdir()
args = SimpleNamespace(
uv="uv",
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
bootstrap,
"run",
side_effect=fake_run,
), patch.object(bootstrap, "run_install"):
bootstrap.install(args)
self.assertTrue((runtime_dir / "snapform_ready.json").is_file())
status = json.loads(
(root / "status.json").read_text(encoding="utf-8")
)
self.assertEqual(status["progress"], 1.0)
self.assertIn("Runtime is ready", status["message"])
def test_install_downloads_missing_uv(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
runtime_dir = root / "runtime"
runtime_dir.mkdir()
uv_path = root / "uv" / "uv.exe"
args = SimpleNamespace(
uv=str(uv_path),
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
bootstrap,
"run",
), patch.object(
bootstrap,
"run_install",
), patch.object(
bootstrap.shutil,
"which",
return_value=None,
), patch.object(
bootstrap.shutil,
"rmtree",
), patch.object(bootstrap, "download_uv") as mocked_download:
bootstrap.install(args)
mocked_download.assert_called_once()
call_args = mocked_download.call_args
self.assertEqual(call_args.args[0], uv_path)
self.assertEqual(call_args.args[1], str(root / "status.json"))
self.assertEqual(call_args.kwargs["stage"], 1)
self.assertEqual(call_args.kwargs["stage_label"], "uv")
def test_install_reuses_existing_uv(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
runtime_dir = root / "runtime"
runtime_dir.mkdir()
uv_path = root / "uv.exe"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
bootstrap,
"run",
), patch.object(
bootstrap,
"run_install",
), patch.object(bootstrap, "download_uv") as mocked_download:
bootstrap.install(args)
mocked_download.assert_not_called()
def test_extract_uv_unpacks_standalone_binary(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
archive_path = root / "uv.zip"
with zipfile.ZipFile(archive_path, "w") as bundle:
bundle.writestr("uv-x86_64-pc-windows-msvc/uv.exe", "binary")
uv_path = root / "uv.exe"
bootstrap.extract_uv(archive_path, uv_path)
self.assertEqual(uv_path.read_bytes(), b"binary")
if __name__ == "__main__":

111
tests/test_download_progress.py
View File

@@ -1,8 +1,10 @@
import importlib.util
import json
import os
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
@@ -56,7 +58,70 @@ class DownloadProgressTest(unittest.TestCase):
status = json.loads(status_path.read_text(encoding="utf-8"))
self.assertAlmostEqual(status["progress"], 0.275)
self.assertEqual(status["message"], "Downloading owner/model")
self.assertEqual(status["message"], "downloading owner/model")
def test_reports_download_speed(self):
class FakeTqdmWithRate(FakeTqdm):
@property
def format_dict(self):
return {"rate": 14_000_000}
with tempfile.TemporaryDirectory() as directory:
status_path = Path(directory) / "status.json"
progress_class = self.worker.create_download_progress_class(
status_path,
"owner/model",
tqdm_type=FakeTqdmWithRate,
clock=lambda: 1.0,
label="da3 small",
)
progress = progress_class(
desc="Downloading bytes",
total=100,
initial=0,
unit="B",
)
progress.update(50)
status = json.loads(status_path.read_text(encoding="utf-8"))
self.assertEqual(
status["message"],
"downloading da3 small 13.4 MB/s",
)
def test_staged_download_prefixes_speed_with_downloading(self):
class FakeTqdmWithRate(FakeTqdm):
@property
def format_dict(self):
return {"rate": 14_000_000}
with tempfile.TemporaryDirectory() as directory:
status_path = Path(directory) / "status.json"
progress_class = self.worker.create_download_progress_class(
status_path,
"owner/model",
tqdm_type=FakeTqdmWithRate,
clock=lambda: 1.0,
label="da3 small",
stage=6,
stages=8,
stage_label="da3 small",
)
progress = progress_class(
desc="Downloading bytes",
total=100,
initial=0,
unit="B",
)
progress.update(50)
status = json.loads(status_path.read_text(encoding="utf-8"))
self.assertEqual(status["message"], "downloading 13.4 MB/s")
self.assertEqual(status["stage"], 6)
self.assertEqual(status["stages"], 8)
self.assertEqual(status["stage_label"], "da3 small")
def test_ignores_unrelated_progress_bars(self):
with tempfile.TemporaryDirectory() as directory:
@@ -78,6 +143,50 @@ class DownloadProgressTest(unittest.TestCase):
self.assertFalse(status_path.exists())
def test_resolve_hf_endpoint_keeps_direct_endpoint(self):
with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker,
"reachable",
return_value=True,
):
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
self.assertEqual(endpoint, "https://huggingface.co/")
self.assertFalse(used_mirror)
def test_resolve_hf_endpoint_falls_back_to_mirror(self):
with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker,
"reachable",
side_effect=lambda url, timeout=None: "hf-mirror" in url,
):
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
self.assertEqual(endpoint, "https://hf-mirror.com")
self.assertTrue(used_mirror)
def test_resolve_hf_endpoint_respects_environment_override(self):
with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example"}), patch.object(
self.worker,
"reachable",
) as mocked_reachable:
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
self.assertEqual(endpoint, "https://hf.example")
self.assertFalse(used_mirror)
mocked_reachable.assert_not_called()
def test_resolve_hf_endpoint_keeps_direct_when_mirrors_unreachable(self):
with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker,
"reachable",
return_value=False,
):
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
self.assertEqual(endpoint, "https://huggingface.co/")
self.assertFalse(used_mirror)
if __name__ == "__main__":
unittest.main()

2
tests/test_pack.py
View File

@@ -29,7 +29,7 @@ class PackTest(unittest.TestCase):
finally:
pack.DIST_DIR = original_dist
self.assertEqual(output_path.name, "SnapForm.v0.9.29.zip")
self.assertEqual(output_path.name, "SnapForm.v0.9.32.zip")
with ZipFile(output_path) as archive:
names = set(archive.namelist())

133
tests/test_worker.py
View File

@@ -1,6 +1,7 @@
import importlib.util
import inspect
import json
import os
import sys
import tempfile
import unittest
@@ -9,6 +10,8 @@ from pathlib import Path
from types import ModuleType, SimpleNamespace
from unittest.mock import patch
import numpy as np
PROJECT_ROOT = Path(__file__).resolve().parents[1]
@@ -61,6 +64,57 @@ class WorkerTest(unittest.TestCase):
inspect.getsource(self.worker.remove_background),
)
def test_select_device_supports_mps_and_auto_fallback(self):
mps_torch = SimpleNamespace(
cuda=SimpleNamespace(is_available=lambda: False),
backends=SimpleNamespace(
mps=SimpleNamespace(is_available=lambda: True),
),
)
mps_torch.device = lambda device_type: SimpleNamespace(
type=device_type
)
self.assertEqual(
self.worker.select_device("mps", mps_torch).type,
"mps",
)
self.assertEqual(
self.worker.select_device("auto", mps_torch).type,
"mps",
)
self.assertEqual(
self.worker.select_device("cpu", mps_torch).type,
"cpu",
)
cuda_torch = SimpleNamespace(
cuda=SimpleNamespace(is_available=lambda: True),
)
cuda_torch.device = lambda device_type: SimpleNamespace(
type=device_type
)
self.assertEqual(
self.worker.select_device("auto", cuda_torch).type,
"cuda",
)
cpu_only_torch = SimpleNamespace(
cuda=SimpleNamespace(is_available=lambda: False),
backends=SimpleNamespace(
mps=SimpleNamespace(is_available=lambda: False),
),
)
cpu_only_torch.device = lambda device_type: SimpleNamespace(
type=device_type
)
self.assertEqual(
self.worker.select_device("auto", cpu_only_torch).type,
"cpu",
)
with self.assertRaisesRegex(RuntimeError, "MPS was requested"):
self.worker.select_device("mps", cpu_only_torch)
def test_optional_background_removal_runs_before_da3_model_load(self):
source = inspect.getsource(self.worker.reconstruct)
@@ -121,6 +175,38 @@ class WorkerTest(unittest.TestCase):
self.assertEqual(images, [image_path])
self.assertIsNone(temporary_directory)
def test_npz_is_complete_requires_readable_archive(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
valid_path = root / "valid.npz"
np.savez_compressed(
valid_path,
depth=np.zeros((4, 4)),
)
self.assertTrue(self.worker.npz_is_complete(valid_path))
empty_path = root / "empty.npz"
empty_path.write_bytes(b"PK\x05\x06" + b"\x00" * 18)
self.assertFalse(self.worker.npz_is_complete(empty_path))
def test_wait_for_exported_paths_waits_for_complete_npz(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
npz_path = root / "exports" / "mini_npz" / "results.npz"
npz_path.parent.mkdir(parents=True)
npz_path.touch()
with patch.object(
self.worker,
"npz_is_complete",
side_effect=[False, True],
):
result = self.worker.wait_for_exported_paths(
root,
"mini_npz",
timeout=2.0,
)
self.assertEqual(result, [npz_path])
def test_rejects_unsupported_input_file(self):
with tempfile.TemporaryDirectory() as directory:
input_path = Path(directory) / "input.txt"
@@ -149,6 +235,10 @@ class WorkerTest(unittest.TestCase):
self.worker,
"create_download_progress_class",
return_value=object,
), patch.object(
self.worker,
"resolve_hf_endpoint",
return_value=(self.worker.HF_ENDPOINT, False),
):
args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json")
@@ -159,6 +249,44 @@ class WorkerTest(unittest.TestCase):
["config.json", "model.safetensors"],
)
def test_download_model_switches_to_mirror_before_hf_import(self):
calls = []
fake_hugging_face = ModuleType("huggingface_hub")
fake_hugging_face.snapshot_download = lambda **options: calls.append(
options
)
args = SimpleNamespace(
repository="PramaLLC/BEN2",
model_dir="model",
status="status.json",
allow_pattern=[],
)
with tempfile.TemporaryDirectory() as directory, patch.dict(
sys.modules,
{"huggingface_hub": fake_hugging_face},
), patch.object(
self.worker,
"create_download_progress_class",
return_value=object,
), patch.object(
self.worker,
"resolve_hf_endpoint",
return_value=("https://hf-mirror.com", True),
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker,
"write_status",
) as mocked_write:
args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json")
self.worker.download_model(args)
self.assertEqual(
os.environ["HF_ENDPOINT"],
"https://hf-mirror.com",
)
messages = [call.args[2] for call in mocked_write.call_args_list]
self.assertTrue(any("(mirror)" in message for message in messages))
def test_ben2_background_removal_writes_fast_transparent_png(self):
inference_options = []
save_options = []
@@ -279,7 +407,10 @@ class WorkerTest(unittest.TestCase):
else:
mini_directory = output_directory / "exports" / "mini_npz"
mini_directory.mkdir(parents=True)
(mini_directory / "results.npz").touch()
np.savez_compressed(
mini_directory / "results.npz",
depth=np.zeros((14, 14)),
)
fake_torch = ModuleType("torch")
fake_torch.device = lambda device_type: SimpleNamespace(

2
uv.lock generated
View File

@@ -4,5 +4,5 @@ requires-python = ">=3.11"
[[package]]
name = "snapform"
version = "0.9.29"
version = "0.9.30"
source = { virtual = "." }