diff --git a/README.md b/README.md index 5a0743d..fc49c0e 100644 --- a/README.md +++ b/README.md @@ -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)` 表示正在执行任务， diff --git a/pyproject.toml b/pyproject.toml index aa5fb44..22114fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/src/snapform/ai_host.py b/src/snapform/ai_host.py index bcd389c..be078d8 100644 --- a/src/snapform/ai_host.py +++ b/src/snapform/ai_host.py @@ -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: diff --git a/src/snapform/blender_manifest.toml b/src/snapform/blender_manifest.toml index af9ca74..d283249 100644 --- a/src/snapform/blender_manifest.toml +++ b/src/snapform/blender_manifest.toml @@ -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] diff --git a/src/snapform/bootstrap_runtime.py b/src/snapform/bootstrap_runtime.py index 1241c50..68621ff 100644 --- a/src/snapform/bootstrap_runtime.py +++ b/src/snapform/bootstrap_runtime.py @@ -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(): diff --git a/src/snapform/operators.py b/src/snapform/operators.py index 3882ed6..2c4f09e 100644 --- a/src/snapform/operators.py +++ b/src/snapform/operators.py @@ -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 diff --git a/src/snapform/preferences.py b/src/snapform/preferences.py index 346ce1c..74e5d92 100644 --- a/src/snapform/preferences.py +++ b/src/snapform/preferences.py @@ -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): diff --git a/src/snapform/properties.py b/src/snapform/properties.py index 2b66aad..9e5d478 100644 --- a/src/snapform/properties.py +++ b/src/snapform/properties.py @@ -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", diff --git a/src/snapform/runtime.py b/src/snapform/runtime.py index 0a75b82..cb0e93a 100644 --- a/src/snapform/runtime.py +++ b/src/snapform/runtime.py @@ -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]) diff --git a/src/snapform/status_bar.py b/src/snapform/status_bar.py index 7d6712b..5c35dc4 100644 --- a/src/snapform/status_bar.py +++ b/src/snapform/status_bar.py @@ -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: diff --git a/src/snapform/worker.py b/src/snapform/worker.py index 0c492cd..bc09b79 100644 --- a/src/snapform/worker.py +++ b/src/snapform/worker.py @@ -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( diff --git a/tests/test_addon.py b/tests/test_addon.py index 70b7c42..c70e2a0 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -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( diff --git a/tests/test_bootstrap_runtime.py b/tests/test_bootstrap_runtime.py index 215aa72..b50d3e0 100644 --- a/tests/test_bootstrap_runtime.py +++ b/tests/test_bootstrap_runtime.py @@ -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__": diff --git a/tests/test_download_progress.py b/tests/test_download_progress.py index 459b551..2783a24 100644 --- a/tests/test_download_progress.py +++ b/tests/test_download_progress.py @@ -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() diff --git a/tests/test_pack.py b/tests/test_pack.py index 9386ee4..2376663 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -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()) diff --git a/tests/test_worker.py b/tests/test_worker.py index 8770072..766c4ac 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -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( diff --git a/uv.lock b/uv.lock index 2e8c96e..116ff06 100644 --- a/uv.lock +++ b/uv.lock @@ -4,5 +4,5 @@ requires-python = ">=3.11" [[package]] name = "snapform" -version = "0.9.29" +version = "0.9.30" source = { virtual = "." }