diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index d293dc2..6afe024 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -24,9 +24,6 @@ jobs: - name: 安装测试依赖 run: python -m pip install numpy pytest - - name: 准备 Vendor - run: python prepare_vendor.py - - name: 运行测试 run: python -m pytest diff --git a/.gitignore b/.gitignore index fc0a0e3..7ba9ec1 100644 --- a/.gitignore +++ b/.gitignore @@ -7,5 +7,4 @@ build/ *.zip .pytest_cache/ -/vendor/ /models/ diff --git a/README.md b/README.md index 6e36d3c..297d75b 100644 --- a/README.md +++ b/README.md @@ -19,9 +19,10 @@ SnapForm 是一个 Blender 扩展,可以把一张照片、一段视频或一 最快的 DirectML/CoreML provider,否则 CPU。下载超时会自动重试国内镜像,无需手动配置 5. 打开 3D 视图右侧的 `SnapForm` 面板,点击 `Install Environment` -扩展只内置对应平台的 uv。首次使用会自动下载轻量的 ONNX Runtime 环境,随后自动下载 `MoGe-2 ViT-S Normal` 和 +扩展不内置 uv 或 AI 运行库。首次使用会先下载并校验固定版本的 uv,随后自动安装轻量的 ONNX Runtime 环境,再下载 `MoGe-2 ViT-S Normal` 和 `BEN2 Base`;不安装 PyTorch,也不下载各模型的 Python 源码包。其他 AI 模型可在 `Models` panel 中按需下载,请保持网络畅通,首次安装需要一些时间。 +安装时 uv、ONNX Environment、MoGe-2 和 BEN2 各自使用独立的 0–100% 进度;下载阶段显示速度与连续字节进度。 ## 快速上手 @@ -69,7 +70,8 @@ Bump Scale 会在材质组内乘以 `0.001`,并按透视表面的深度空间 ### 视频与序列帧 SnapForm 支持视频(mp4、mkv、mov 等)和序列帧,会逐帧处理并生成逐帧对应的立体面片。处理范围默认 -跟随 Empty 上设置的 `Image Frames` 与 `Offset`,不需要额外配置。 +跟随 Empty 上设置的 `Image Frames` 与 `Offset`,不需要额外配置。DA3 的转换任务复用同一个 ONNX +Session 逐帧推理,显存峰值不随视频帧数平方增长;Debug 多帧输出仍使用多视图联合推理。 ## 任务结果保存在哪里 diff --git a/docs/depth-plane.md b/docs/depth-plane.md index 4170cd7..a7fa719 100644 --- a/docs/depth-plane.md +++ b/docs/depth-plane.md @@ -22,6 +22,8 @@ 因此材质直接用 `.foreground` 同时提供颜色和 Mask,不再额外加载原图 Image Texture。 Image Empty 引用视频或序列帧时,DA3 与 MoGe-2 会按 Empty 的 Image Frames 数值逐帧估计 深度(未设置时为完整帧数,显式设为 1 时只处理 1 帧),起始帧跟随 Empty 的 Offset,与背景蒙版范围一致; + DA3 会复用同一个 ONNX Session 对每帧独立推理,再合并为原有序列输出,避免多帧 Attention 的平方级显存占用。 + 该行为只用于 Depth Plane 与 Surface 转换;Debug 多帧输出仍联合推理。 Offset 超出视频长度时同样把窗口收进视频内。材质节点与置换贴图的 Start 保持与原 Empty 一致;Original 模式的视频或序列帧使用 `frame.0001.jpg` 等原图序列作为颜色与深度, `foreground.0001.png` 只提供前景 Alpha。 diff --git a/docs/development.md b/docs/development.md index e22d4f4..5e620c9 100644 --- a/docs/development.md +++ b/docs/development.md @@ -3,7 +3,6 @@ ## 开发 ```bash -uv run python prepare_vendor.py uv run pytest tests -v uv run python pack.py ``` @@ -12,10 +11,7 @@ uv run python pack.py ```text . -├── vendor.lock.toml 各平台 uv 的版本、URL 与 SHA-256 -├── prepare_vendor.py 下载并校验各平台 uv ├── pack.py -├── vendor/ 不进入 Git 的 uv 下载缓存 ├── src/ │ └── snapform/ │ ├── blender_manifest.toml diff --git a/docs/features.md b/docs/features.md index 8823bed..13ea5ce 100644 --- a/docs/features.md +++ b/docs/features.md @@ -37,7 +37,8 @@ - Image Empty 右键菜单可直接运行 DA3,根据 `mini_npz` 深度生成面片 - Image Empty 右键菜单可直接运行 BEN2,保留原图并换用新建的透明背景图片 - Image Empty 支持视频与序列帧输入:`Convert to Depth Plane` 逐帧估计深度, - `Remove Background` 逐帧抠图,入口保持不变 + `Remove Background` 逐帧抠图,入口保持不变。DA3 转换时每帧独立执行 ONNX 推理并复用模型 Session; + Debug 多帧任务继续联合推理,以保留跨视图相机与深度关系 - 视频生成的序列帧以外部文件引用保留在 job 目录(Blender 无法整体 Pack 序列帧); Geometry Nodes 与材质 Image Texture 均可按场景帧播放序列帧 - 使用浮点深度贴图、`O DA3 Depth Plane` Geometry Nodes 和原图 Bump 构建立体面片 @@ -49,10 +50,9 @@ ## 使用 -1. 先运行 `uv run python prepare_vendor.py`,按 `vendor.lock.toml` 下载并校验固定版本的 - 第三方构建输入;再运行 `uv run python pack.py`,分别构建 `windows-x64`、`macos-arm64`、 - `linux-x64` 三个平台的 Extension ZIP(如 `SnapForm.v0.11.1.macos-arm64.zip`)。 - 每个包只含对应平台的 uv 二进制,不再携带 AI 项目的源码包。 +1. 运行 `uv run python pack.py`,分别构建 `windows-x64`、`macos-arm64`、 + `linux-x64` 三个平台的 Extension ZIP(如 `SnapForm.v0.11.2.macos-arm64.zip`)。 + 安装包不携带 uv 或 AI 项目源码;首次安装 Environment 时下载固定版本 uv,官方源建连超过 3 秒后回退 OmooLab 镜像,并校验固定 SHA-256。 2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。 3. 在 Extension Preferences 中确认 Storage Root 与 Device。Device 是正式功能 (Image Empty 右键菜单)使用的推理设备,默认取本机最快的 DirectML/CoreML provider,否则 CPU,且只列出 @@ -60,12 +60,15 @@ 4. 在 `AI Service` 点击 `Install Environment`。首次安装会下载 Python、ONNX Runtime、FastAPI、Uvicorn; uv 已内置在扩展中,DA3、BEN2 与 MoGe 不安装源码包, 并在 Environment 安装后自动下载 `MoGe-2 ViT-S Normal` 和 FP16 `BEN2 Base` 推荐模型。Environment 安装 - 固定显示 `Environment`、`ONNX Runtime` 两个用户阶段;FastAPI 等内部安装步骤合并在 `ONNX Runtime` 中。 - 安装完成后直接写入 Environment Manifest,不再额外导入全部组件进行验证。推荐模型属于安装 - 后的自动步骤,不计入阶段编号;`DA3 Base` 等其他模型可稍后在 `Models` 面板按需下载。 - 状态文字使用 `阶段/总数 阶段名 | 当前操作 | 预计剩余时间`。依赖安装只显示阶段与已解析的 - 包总数,不再承诺剩余时间;模型下载根据字节数和速度计算剩余时间,例如 - `Downloading DA3 Small | 14.2 MB/s | About 2 min left`。重试或切换镜像时进度不会倒退。 + 首次安装依次显示 `uv`、`ONNX Environment`、`MoGe-2 ViT-S Normal` 与 `BEN2 Base`;每个阶段 + 都独立从 0 走到 100%,进入下一阶段时进度条归零。uv 已存在时省略其阶段并重新编号。 + 安装完成后直接写入 Environment Manifest,不再额外导入全部组件进行验证;`DA3 Base` 等其他模型可稍后在 `Models` 面板按需下载。 + 状态文字使用 `阶段/总数 当前操作`,例如 + `3/4 Downloading MoGe-2 ViT-S Normal 8.4 MB/s`。依赖安装只显示阶段与已解析的 + 包总数;模型下载显示速度,例如 + `Downloading DA3 Small 14.2 MB/s`;镜像来源显示为 + `Downloading DA3 Small (mirror) 14.2 MB/s`。下载按字节连续推进,重试或切换镜像时进度不会倒退。 + Environment 安装阶段与 `Models` 面板单独下载使用相同的模型名称格式,不会改写为通用的 `Downloading packages`。 Environment 安装后即使模型下载中断也不影响 Service 使用,可稍后在 `Models` 面板单独补下。 安装不做启动网络预检:依赖安装默认依次尝试清华、阿里镜像并以 pypi.org 兜底;模型下载由 worker 自行探测回退到 OmooLab 镜像,切换结果写入 service.json 的 models_from_r2 字段, diff --git a/docs/service.md b/docs/service.md index 0ff70b7..c02cfe4 100644 --- a/docs/service.md +++ b/docs/service.md @@ -21,7 +21,7 @@ Environment、模型和任务结果默认保存在 `~/.snapform`: ## ONNX-only Environment -扩展只内置对应平台的 `uv` 可执行文件。`Install Environment` 使用它创建环境并安装: +扩展不内置 `uv` 可执行文件或 AI 运行库。`Install Environment` 首先下载固定版本的 uv,校验 SHA-256 后解压到共享 Storage Root,再用它创建环境并安装: - ONNX Runtime 1.24.4;Windows 使用 `onnxruntime-directml` - NumPy 2.4.2、Pillow 12.1.1、PyAV 17.0.0、trimesh 4.11.2 @@ -41,7 +41,7 @@ Windows 当前发行环境提供 DirectML 与 CPU;`Auto` 优先使用 `DmlExec - DA3 Small、DA3 Base - BEN2 Base -首次安装自动下载 `MoGe-2 ViT-S Normal` 与 FP16 `BEN2 Base`,其他模型在 `Models` panel 按需下载。模型按 Hugging Face 仓库末段名称保存在 `~/.snapform/models/`,并保留仓库内部结构。声明了文件清单与校验值的模型会直接走 Hugging Face `resolve` 直链流式下载,跳过文件清单与 ETag 预校验;连接源不可达时 6 秒内快速失败并回退到 `https://models.omoolab.xyz//`,并把 `models_from_r2` 写入 service.json,本次 service 周期内后续模型直接走 R2,不再重复等待 Hugging Face。R2 文件会校验大小与 SHA-256。未声明清单的仓库仍使用 `huggingface_hub` 快照下载。可用 `HF_ENDPOINT` 指定 Hugging Face 镜像。 +首次安装自动下载 `MoGe-2 ViT-S Normal` 与 FP16 `BEN2 Base`,其他模型在 `Models` panel 按需下载。模型按 Hugging Face 仓库末段名称保存在 `~/.snapform/models/`,并保留仓库内部结构。声明了文件清单与校验值的模型会直接走 Hugging Face `resolve` 直链流式下载,跳过文件清单与 ETag 预校验;连接源不可达时使用 3 秒 TCP 与 TLS 建连超时,随后回退到 `https://models.omoolab.xyz//`,并把 `models_from_r2` 写入 service.json,本次 service 周期内后续模型直接走 R2,不再重复等待 Hugging Face。建连成功后不限制模型文件读取时长。R2 文件会校验大小与 SHA-256。未声明清单的仓库仍使用 `huggingface_hub` 快照下载。可用 `HF_ENDPOINT` 指定 Hugging Face 镜像。 DA3 ONNX 使用模型的 camera decoder 恢复外参与内参;旧版 `Use Ray Pose` 参数只为任务协议兼容保留,不参与 ONNX 推理。MoGe-2 的 metric depth、normal、mask、normalized intrinsics,以及 BEN2 的 Alpha,均在 NumPy 后处理中恢复。 @@ -50,7 +50,9 @@ DA3 ONNX 使用模型的 camera decoder 恢复外参与内参;旧版 `Use Ray 安装器直接开始工作,不做阻塞式网络预检。依赖安装默认依次尝试清华、阿里云 PyPI 镜像并以 pypi.org 兜底,也可用 `UV_INDEX_URL` 显式指定。Windows 系统代理会注入子进程;同时支持 `HTTP_PROXY`、`HTTPS_PROXY` 与 `ALL_PROXY`。 -GitHub 不参与用户安装流程。`prepare_vendor.py` 只为扩展打包准备各平台的 `uv` 二进制。 +uv 默认从 Astral 的 GitHub Release 下载,TCP 与 TLS 建连超过 3 秒时回退到 `https://mirror.omoolab.xyz/uv//`;镜像建连超时为 15 秒,避免较慢的 TLS 握手再次过早失败。两个来源使用相同的下载速度与连续字节进度显示,并共用固定 SHA-256 校验,建连成功后不限制文件读取时长。下载压缩包只作临时文件,提取完成后删除。模型下载不经过 GitHub。 + +首次安装按阶段分别显示 0–100%:uv 下载与解压、ONNX Environment 创建与依赖安装、MoGe-2 下载、BEN2 下载。进入下一阶段时进度条归零;若 uv 已存在则省略该阶段并重新编号。 ## 生命周期 @@ -64,9 +66,11 @@ Service 全局同时执行一个 AI 任务。BEN2 独立缓存,深度模型只 DA3 ONNX 的注意力是朴素实现(`MatMul → Softmax → MatMul`),峰值显存与 `(帧数 × 每帧 patch 数)²` 成正比。迁移到 ONNX 前的 PyTorch 版使用 `scaled_dot_product_attention`(`XFORMERS_DISABLED=1`),按块计算注意力分数、不落地完整注意力矩阵,显存近似线性,因此视频能处理几十上百帧;ONNX 导出把 SDPA 回退成朴素注意力,多帧在中高分辨率下会显存爆炸(约 12 帧 @ 1008 即需 ~90 GB)。 -作为兜底,`da3_onnx.infer` 在推理前按 `views × patches` 预估注意力张量,超过 4 GiB 直接报错并提示减少帧数或 Resolution Scale,避免 onnxruntime 报出难以定位的分配失败或 `UnicodeDecodeError`(DirectML 在中文 Windows 下报错信息为 GBK 编码,被 UTF-8 解码掩盖)。 +Depth Plane 与 Surface 的视频、序列帧转换会复用同一个 ONNX Session 逐帧独立推理,再把结果合并到现有序列输出,因此峰值显存只取决于单帧分辨率。逐帧结果不保留 DA3 的跨视图尺度与相机位姿一致性,适用于当前逐帧动画结构。Debug 多帧输出仍使用联合推理,以保留多视图关系。 + +联合推理时,`da3_onnx.infer` 会在执行前按 `views × patches` 预估注意力张量,超过 4 GiB 直接报错并提示减少帧数或 Resolution Scale,避免 onnxruntime 报出难以定位的分配失败或 `UnicodeDecodeError`(DirectML 在中文 Windows 下报错信息为 GBK 编码,被 UTF-8 解码掩盖)。 待解决方向: - 重新导出 DA3 ONNX,使用 ORT 的 `Attention` / `MultiHeadAttention` 贡献算子,或注入内存高效的注意力算子; -- 或把视频按小窗口分批推理,再按 `ref_view_strategy` 对齐相机位姿拼接。 +- 若未来需要统一坐标系的视频重建,可按小窗口分批推理,再对齐相机位姿与深度尺度。 diff --git a/pack.py b/pack.py index 9b545df..76716dc 100644 --- a/pack.py +++ b/pack.py @@ -1,4 +1,3 @@ -import hashlib import tomllib from pathlib import Path from zipfile import ZIP_DEFLATED, ZipFile @@ -6,49 +5,8 @@ from zipfile import ZIP_DEFLATED, ZipFile PROJECT_ROOT = Path(__file__).resolve().parent SOURCE_DIR = PROJECT_ROOT / "src" / "snapform" -VENDOR_DIR = PROJECT_ROOT / "vendor" DIST_DIR = PROJECT_ROOT / "dist" -LOCK_PATH = PROJECT_ROOT / "vendor.lock.toml" - - -def load_vendor_lock(): - """Return the locked third-party artifact metadata.""" - with LOCK_PATH.open("rb") as stream: - lock = tomllib.load(stream) - if lock.get("schema_version") != 1: - raise RuntimeError("Unsupported vendor.lock.toml schema") - return lock - - -VENDOR_LOCK = load_vendor_lock() -PLATFORMS = { - platform: artifact["filename"] - for platform, artifact in VENDOR_LOCK["platforms"].items() -} - - -def file_sha256(path): - """Return the SHA-256 digest for a file.""" - digest = hashlib.sha256() - with Path(path).open("rb") as stream: - while chunk := stream.read(1024 * 1024): - digest.update(chunk) - return digest.hexdigest() - - -def validate_artifact(path, expected_sha256): - """Require an artifact to exist and match its locked digest.""" - if not path.is_file(): - raise RuntimeError( - f"Missing vendor artifact: {path}. " - "Run: uv run python prepare_vendor.py" - ) - actual = file_sha256(path) - if actual != expected_sha256: - raise RuntimeError( - f"Vendor artifact hash mismatch: {path}. " - "Run: uv run python prepare_vendor.py" - ) +PLATFORMS = ("windows-x64", "macos-arm64", "linux-x64") def manifest_for_platform(platform): @@ -73,21 +31,6 @@ def manifest_for_platform(platform): return "\n".join(lines) + "\n" -def bundled_vendor_files(platform): - return { - f"vendor/{PLATFORMS[platform]}": VENDOR_DIR / PLATFORMS[platform] - } - - -def validate_vendor_files(platform): - """Validate common archives and the selected uv archive.""" - uv_artifact = VENDOR_LOCK["platforms"][platform] - validate_artifact( - VENDOR_DIR / uv_artifact["filename"], - uv_artifact["sha256"], - ) - - def build_extension(platform): if platform not in PLATFORMS: raise ValueError(f"Unsupported platform: {platform}") @@ -96,8 +39,6 @@ def build_extension(platform): f"{manifest['name']}.v{manifest['version']}.{platform}.zip" ) DIST_DIR.mkdir(exist_ok=True) - vendor_files = bundled_vendor_files(platform) - validate_vendor_files(platform) with ZipFile(output_path, "w", ZIP_DEFLATED) as archive: for source_path in sorted(SOURCE_DIR.rglob("*")): if not source_path.is_file(): @@ -105,8 +46,6 @@ def build_extension(platform): if "__pycache__" in source_path.parts: continue relative_path = source_path.relative_to(SOURCE_DIR) - if relative_path.parts[0] == "vendor": - continue if source_path.name == "blender_manifest.toml": continue archive.write(source_path, relative_path) @@ -114,11 +53,6 @@ def build_extension(platform): "blender_manifest.toml", manifest_for_platform(platform), ) - for archive_name, source_path in vendor_files.items(): - archive.writestr( - archive_name, - source_path.read_bytes(), - ) return output_path diff --git a/prepare_vendor.py b/prepare_vendor.py deleted file mode 100644 index cacd9b4..0000000 --- a/prepare_vendor.py +++ /dev/null @@ -1,110 +0,0 @@ -import argparse -import hashlib -import sys -import tomllib -import urllib.request -from pathlib import Path - - -PROJECT_ROOT = Path(__file__).resolve().parent -LOCK_PATH = PROJECT_ROOT / "vendor.lock.toml" -VENDOR_DIR = PROJECT_ROOT / "vendor" -CHUNK_SIZE = 1024 * 1024 - - -def load_lock(path=LOCK_PATH): - """Load and validate the vendor lock file.""" - with Path(path).open("rb") as stream: - lock = tomllib.load(stream) - if lock.get("schema_version") != 1: - raise RuntimeError("Unsupported vendor.lock.toml schema") - return lock - - -def file_sha256(path): - """Return the SHA-256 digest for a file.""" - digest = hashlib.sha256() - with Path(path).open("rb") as stream: - while chunk := stream.read(CHUNK_SIZE): - digest.update(chunk) - return digest.hexdigest() - - -def download_artifact(artifact, destination_dir=VENDOR_DIR): - """Download one locked artifact and verify its content.""" - destination_dir = Path(destination_dir) - destination = destination_dir / artifact["filename"] - expected = artifact["sha256"] - if destination.is_file() and file_sha256(destination) == expected: - print(f"Ready: {destination.name}") - return destination - - destination_dir.mkdir(parents=True, exist_ok=True) - temporary = destination.with_name(f"{destination.name}.part") - temporary.unlink(missing_ok=True) - print(f"Downloading: {destination.name}") - request = urllib.request.Request( - artifact["url"], - headers={"User-Agent": "SnapForm vendor preparation"}, - ) - try: - with urllib.request.urlopen(request) as response, temporary.open("wb") as output: - while chunk := response.read(CHUNK_SIZE): - output.write(chunk) - actual = file_sha256(temporary) - if actual != expected: - raise RuntimeError( - f"SHA-256 mismatch for {destination.name}: " - f"expected {expected}, got {actual}" - ) - temporary.replace(destination) - except Exception: - temporary.unlink(missing_ok=True) - raise - return destination - - -def prepare(platforms=None, lock_path=LOCK_PATH, destination_dir=VENDOR_DIR): - """Prepare bundled uv executables for selected platforms.""" - lock = load_lock(lock_path) - available = lock["platforms"] - selected = list(available) if platforms is None else platforms - unknown = sorted(set(selected) - set(available)) - if unknown: - raise ValueError(f"Unsupported platform: {', '.join(unknown)}") - - artifacts = [available[platform] for platform in selected] - return [ - download_artifact(artifact, destination_dir) - for artifact in artifacts - ] - - -def parse_args(argv=None): - """Parse command-line options.""" - parser = argparse.ArgumentParser( - description="Download and verify third-party files for SnapForm.", - ) - parser.add_argument( - "--platform", - action="append", - dest="platforms", - help="Prepare one platform; repeat as needed. Default: all platforms.", - ) - return parser.parse_args(argv) - - -def main(argv=None): - """Run vendor preparation from the command line.""" - args = parse_args(argv) - try: - paths = prepare(args.platforms) - except (OSError, RuntimeError, ValueError) as error: - print(f"Preparation failed: {error}", file=sys.stderr) - return 1 - print(f"Prepared {len(paths)} files") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index 19a00e2..9e6f42b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "snapform" -version = "0.11.1" +version = "0.11.2" description = "Fast image-guided structure building tools for Blender" readme = "README.md" requires-python = ">=3.11" diff --git a/src/snapform/blender_manifest.toml b/src/snapform/blender_manifest.toml index c2ed409..b55382e 100644 --- a/src/snapform/blender_manifest.toml +++ b/src/snapform/blender_manifest.toml @@ -1,7 +1,7 @@ schema_version = "1.0.0" id = "snapform" -version = "0.11.1" +version = "0.11.2" name = "SnapForm" tagline = "Build spatial structures from images, fast" maintainer = "OmooLab" diff --git a/src/snapform/client.py b/src/snapform/client.py index a8724bb..fa00baa 100644 --- a/src/snapform/client.py +++ b/src/snapform/client.py @@ -269,7 +269,7 @@ def job_id_from_status(status_path): def uv_executable(): - """Return the bundled uv path inside the shared storage.""" + """Return the downloaded uv path inside the shared storage.""" return str(uv_cache_path()) @@ -833,6 +833,7 @@ def reconstruct_request( mesh_filter_edges=None, start_frame=1, device=None, + inference_mode="joint", ): snapform_selected_key = model_key or selected_model_key(settings, export_format) snapform_selected_process_res = ( @@ -867,6 +868,7 @@ def reconstruct_request( "device": (device or settings.device).lower(), "use_ray_pose": 1 if snapform_selected_ray_pose else 0, "ref_view_strategy": settings.ref_view_strategy, + "inference_mode": inference_mode, "remove_background_model_dir": "", "mesh_filter_edges": bool( getattr(settings, "mesh_filter_edges", True) diff --git a/src/snapform/install_env.py b/src/snapform/install_env.py index ba6a70c..a9599c0 100644 --- a/src/snapform/install_env.py +++ b/src/snapform/install_env.py @@ -1,4 +1,6 @@ import argparse +import hashlib +import http.client import json import os import platform @@ -36,13 +38,30 @@ PILLOW_VERSION = "12.1.1" AV_VERSION = "17.0.0" TRIMESH_VERSION = "4.11.2" HUGGINGFACE_HUB_VERSION = "1.4.1" +UV_VERSION = "0.11.3" +UV_OFFICIAL_CONNECT_TIMEOUT = 3.0 +UV_MIRROR_CONNECT_TIMEOUT = 15.0 +DOWNLOAD_CHUNK_SIZE = 256 * 1024 +UV_OFFICIAL_BASE_URL = ( + f"https://github.com/astral-sh/uv/releases/download/{UV_VERSION}" +) +UV_MIRROR_BASE_URL = f"https://mirror.omoolab.xyz/uv/{UV_VERSION}" +UV_ARTIFACTS = { + "x86_64-pc-windows-msvc": ( + "uv-x86_64-pc-windows-msvc.zip", + "ae681c0aaec7cc96af184648cb88d73f8393ed60fa5880abdd6bdb910f9b227c", + ), + "aarch64-apple-darwin": ( + "uv-aarch64-apple-darwin.tar.gz", + "2bc3d0c7bf2bd08325b1e170abac6f7e5b3346e1d4eab3370d17cefec934996f", + ), + "x86_64-unknown-linux-gnu": ( + "uv-x86_64-unknown-linux-gnu.tar.gz", + "c0f3236f146e55472663cfbcc9be3042a9f1092275bbe3fe2a56a6cbfd3da5ce", + ), +} _resolved_pypi_index = None DEFAULT_MODEL_KEYS = (MOGE2_MODEL.key, BEN2_MODEL.key) -RUNTIME_STAGE_COUNT = 2 -ENVIRONMENT_SLICE = (0.02, 0.06) -AI_COMPONENTS_SLICE = (0.10, 0.85) -MODELS_SLICE_START = 0.85 -MODELS_SLICE_END = 1.0 RAMP_RATE = 0.02 CREEP_CAP = 0.60 CHECK_CODE = ( @@ -54,6 +73,48 @@ CHECK_CODE = ( ) +class _ConnectHTTPSConnection(http.client.HTTPSConnection): + """Limit TCP/TLS setup without limiting artifact reads.""" + + connect_timeout = UV_OFFICIAL_CONNECT_TIMEOUT + + def connect(self): + try: + self.timeout = self.connect_timeout + super().connect() + finally: + self.timeout = None + self.sock.settimeout(None) + + +class _MirrorConnectHTTPSConnection(_ConnectHTTPSConnection): + connect_timeout = UV_MIRROR_CONNECT_TIMEOUT + + +class _ConnectHTTPSHandler(urllib.request.HTTPSHandler): + def __init__(self, connection_type): + super().__init__() + self.connection_type = connection_type + + def https_open(self, request): + return self.do_open(self.connection_type, request) + + +def build_uv_download_opener(connection_type): + return urllib.request.build_opener( + urllib.request.ProxyHandler(), + _ConnectHTTPSHandler(connection_type), + ) + + +_uv_download_openers = { + UV_OFFICIAL_BASE_URL: build_uv_download_opener(_ConnectHTTPSConnection), + UV_MIRROR_BASE_URL: build_uv_download_opener( + _MirrorConnectHTTPSConnection + ), +} + + def configure_standard_streams(): """Windows 图形进程可能继承 GBK stdout,安装日志统一写为 UTF-8。""" for stream in (sys.stdout, sys.stderr): @@ -484,22 +545,6 @@ def uv_download_target(): raise RuntimeError(f"Unsupported platform for uv download: {sys.platform}") -def bundled_vendor_directory(): - packaged = Path(__file__).with_name("vendor") - if packaged.is_dir(): - return packaged - project_root = Path(__file__).resolve().parents[2] - if (project_root / "pyproject.toml").is_file(): - return project_root / "vendor" - return packaged - - -def bundled_uv_archive(): - target, archive = uv_download_target() - path = bundled_vendor_directory() / f"uv-{target}.{archive}" - return path if path.is_file() else None - - def archive_is_valid(path): if zipfile.is_zipfile(path): return True @@ -526,6 +571,99 @@ def extract_uv(archive_path, uv_path): os.chmod(uv_path, 0o755) +def file_sha256(path): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + while chunk := stream.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def format_download_rate(rate): + size = float(rate) + 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 + precision = 0 if unit_index == 0 else 1 + return f"{size:.{precision}f} {units[unit_index]}" + + +def download_uv_archive( + destination, + status_path, + start, + end, + stage_fields, + clock=time.monotonic, +): + target, _archive_type = uv_download_target() + artifact = UV_ARTIFACTS.get(target) + if artifact is None: + raise RuntimeError(f"SnapForm does not provide uv for {target}") + filename, expected_checksum = artifact + failures = [] + sources = ( + (UV_OFFICIAL_BASE_URL, ""), + (UV_MIRROR_BASE_URL, " (mirror)"), + ) + for base_url, source_label in sources: + destination.unlink(missing_ok=True) + url = f"{base_url}/{filename}" + write_status( + status_path, + start, + f"Downloading uv{source_label}", + **stage_fields, + ) + request = urllib.request.Request( + url, + headers={"User-Agent": "SnapForm uv bootstrap"}, + ) + try: + with _uv_download_openers[base_url].open(request) as response: + total = int(response.headers.get("Content-Length", 0)) + downloaded = 0 + started_at = clock() + last_report_at = started_at + with destination.open("wb") as output: + while chunk := response.read(DOWNLOAD_CHUNK_SIZE): + output.write(chunk) + downloaded += len(chunk) + now = clock() + if now - last_report_at < 0.20 and downloaded < total: + continue + last_report_at = now + elapsed = max(now - started_at, 1e-6) + rate = downloaded / elapsed + message = ( + f"Downloading uv{source_label} " + f"{format_download_rate(rate)}" + ) + progress = start + if total: + progress += (end - start) * downloaded / total + write_status( + status_path, + min(progress, end), + message, + **stage_fields, + ) + actual_checksum = file_sha256(destination) + if actual_checksum != expected_checksum: + raise RuntimeError( + f"SHA-256 mismatch: {actual_checksum}" + ) + if not archive_is_valid(destination): + raise RuntimeError("Downloaded uv archive is invalid") + return destination + except Exception as error: + destination.unlink(missing_ok=True) + failures.append(f"{url}: {error}") + raise RuntimeError("Unable to download uv; " + "; ".join(failures)) + + def install_uv( uv, status_path, @@ -535,29 +673,31 @@ def install_uv( start=0.0, end=1.0, ): - """Extract the bundled uv binary; fail clearly when it is missing.""" + """Download, verify, and extract the pinned uv binary.""" stage_fields = { "stage": stage, "stages": stages, "stage_label": stage_label, } - bundled_uv = bundled_uv_archive() - if bundled_uv is None: - raise RuntimeError( - "Bundled uv is missing; reinstall the extension matching your " - f"platform (searched {bundled_vendor_directory()})" - ) - if not archive_is_valid(bundled_uv): - raise RuntimeError(f"Bundled uv archive is invalid: {bundled_uv}") - print(f"Extracting bundled uv from {bundled_uv}") - write_status( - status_path, - start, - "extracting bundled uv", - **stage_fields, - ) Path(uv).parent.mkdir(parents=True, exist_ok=True) - extract_uv(bundled_uv, uv) + archive_path = Path(uv).with_name("uv-download.tmp") + try: + download_uv_archive( + archive_path, + status_path, + start, + end, + stage_fields, + ) + write_status( + status_path, + end, + "Extracting uv", + **stage_fields, + ) + extract_uv(archive_path, uv) + finally: + archive_path.unlink(missing_ok=True) def run_stage( @@ -649,38 +789,30 @@ def install(args): stages = install_stages(args, venv_dir, python, uv) env_stages = [stage for stage in stages if stage[3] != "download"] model_stages = [stage for stage in stages if stage[3] == "download"] - model_span = ( - (MODELS_SLICE_END - MODELS_SLICE_START) / len(model_stages) - if model_stages - else 0.0 - ) + total_stages = int(uv_missing) + 1 + len(model_stages) + stage = 1 installed_packages = [] - for index, (label, command, environment, kind) in enumerate(env_stages): - if index == 0: - stage = 1 - stage_start, stage_end = ENVIRONMENT_SLICE - elif index == 1: - stage = 2 - stage_start, stage_end = AI_COMPONENTS_SLICE - else: - raise RuntimeError(f"Unexpected runtime stage: {label}") - if uv_missing and index == 0: - midpoint = stage_start + (stage_end - stage_start) * 0.5 - install_uv( - uv, - args.status, - stage=stage, - stages=RUNTIME_STAGE_COUNT, - stage_label=label, - start=stage_start, - end=midpoint, - ) - stage_start = midpoint + if uv_missing: + install_uv( + uv, + args.status, + stage=stage, + stages=total_stages, + stage_label="uv", + start=0.0, + end=1.0, + ) + stage += 1 + environment_label = "ONNX Environment" + for index, (_label, command, environment, kind) in enumerate(env_stages): + if index >= 2: + raise RuntimeError("Unexpected ONNX Environment stage") + stage_start, stage_end = ((0.0, 0.05) if index == 0 else (0.05, 1.0)) stage_installed = run_stage( args.status, stage, - RUNTIME_STAGE_COUNT, - label, + total_stages, + environment_label, command, environment, kind, @@ -689,6 +821,7 @@ def install(args): ) if isinstance(stage_installed, list): installed_packages.extend(stage_installed) + stage += 1 devices = installed_devices() venv_dir.mkdir(parents=True, exist_ok=True) manifest_path.parent.mkdir(parents=True, exist_ok=True) @@ -711,44 +844,50 @@ def install(args): ), encoding="utf-8", ) - model_offset = MODELS_SLICE_START failed_models = [] for label, command, environment, kind in model_stages: - stage_start = model_offset - stage_end = stage_start + model_span try: run_stage( args.status, - None, - None, + stage, + total_stages, label, command, environment, kind, - start=stage_start, - end=stage_end, + start=0.0, + end=1.0, ) except Exception as error: print(f"Model download failed: {label}: {error}") failed_models.append(label) write_status( args.status, - stage_end, + 1.0, f"Skipping {label}; continuing", + stage=stage, + stages=total_stages, + stage_label=label, ) - model_offset = stage_end + stage += 1 if failed_models: labels = ", ".join(failed_models) write_status( args.status, 1.0, f"Environment is installed; retry {labels} from Models", + stage=total_stages, + stages=total_stages, + stage_label=model_stages[-1][0], ) return write_status( args.status, 1.0, "Environment and recommended models ready", + stage=total_stages, + stages=total_stages, + stage_label=model_stages[-1][0] if model_stages else environment_label, ) diff --git a/src/snapform/operators.py b/src/snapform/operators.py index 6cebf66..a7d1cf4 100644 --- a/src/snapform/operators.py +++ b/src/snapform/operators.py @@ -108,6 +108,7 @@ class ActiveJob: started: bool = False start_error: str = "" start_thread: threading.Thread | None = None + progress_stage: int | None = None _active_job = None @@ -520,6 +521,11 @@ class RunJob(bpy.types.Operator): and remove_background ), device=self._task_device(settings), + inference_mode=( + "per_frame" + if self.job_type in DA3_IMAGE_JOBS + else "joint" + ), ), paths, ) @@ -612,10 +618,17 @@ class RunJob(bpy.types.Operator): reported_progress = float( status.get("progress", settings.progress) ) - settings.progress = max( - settings.progress, - min(max(reported_progress, 0.0), 1.0), - ) + reported_progress = min(max(reported_progress, 0.0), 1.0) + reported_stage = status.get("stage") + if reported_stage is not None and reported_stage != getattr( + job, + "progress_stage", + None, + ): + job.progress_stage = reported_stage + settings.progress = reported_progress + else: + settings.progress = max(settings.progress, reported_progress) settings.status = self._format_status(status, settings.status) self._redraw_ui(context, force=True) @@ -624,38 +637,19 @@ class RunJob(bpy.types.Operator): stage = status.get("stage") stages = status.get("stages") stage_label = status.get("stage_label") - parts = [] if stage and stages: - stage_text = f"{stage}/{stages}" + prefix = f"{stage}/{stages}" + normalized = str(message or "").strip() + if normalized and normalized.lower() != str( + stage_label or "" + ).lower(): + return f"{prefix} {normalized}" if stage_label: - stage_text = f"{stage_text} {stage_label}" - parts.append(stage_text) - elif stage_label: - parts.append(str(stage_label)) + return f"{prefix} {stage_label}" + return prefix if message: - normalized = str(message).strip() - if normalized.lower() != str(stage_label or "").lower(): - parts.append(normalized) - eta = self._format_eta(status.get("eta_seconds")) - if eta: - parts.append(eta) - return " | ".join(parts) if parts else (message or fallback) - - @staticmethod - def _format_eta(seconds): - if seconds is None: - return "" - seconds = max(float(seconds), 0.0) - if seconds <= 10: - return "A few seconds left" - if seconds < 60: - rounded = max(5, int((seconds + 2.5) // 5) * 5) - return f"About {rounded} sec left" - if seconds < 3600: - minutes = int((seconds + 59) // 60) - return f"About {minutes} min left" - hours = int((seconds + 3599) // 3600) - return f"About {hours} hr left" + return str(message).strip() + return str(stage_label or fallback) def _finish(self, context, job, return_code): settings = context.scene.snapform_settings diff --git a/src/snapform/worker.py b/src/snapform/worker.py index d04e6a3..1c886f3 100644 --- a/src/snapform/worker.py +++ b/src/snapform/worker.py @@ -49,7 +49,8 @@ EXPORT_FORMATS = ( HF_ENDPOINT = "https://huggingface.co" R2_ENDPOINT = "https://models.omoolab.xyz" # 连接(TCP + TLS 握手)超时,源不可达时快速失败切换;下载阶段不设读超时。 -CONNECT_TIMEOUT = 6.0 +CONNECT_TIMEOUT = 3.0 +DOWNLOAD_CHUNK_SIZE = 256 * 1024 class _ShortConnectHTTPSConnection(http.client.HTTPSConnection): @@ -898,7 +899,7 @@ class DownloadProgressReporter: with self.lock: if progress <= self.progress: return - if now - self.last_write_time < 0.20 and progress < 0.95: + if now - self.last_write_time < 0.20 and factor < 1.0: return self.progress = progress self.last_write_time = now @@ -907,20 +908,10 @@ class DownloadProgressReporter: if not effective_rate and elapsed > 0: effective_rate = current / elapsed speed = format_rate(effective_rate) if effective_rate else "" - eta_seconds = None - if effective_rate and effective_rate > 0: - eta_seconds = max(total - current, 0) / effective_rate - if self.stage is not None: - message = ( - f"Downloading packages | {speed}" - if speed - else "Downloading packages" - ) - else: - prefix = self.label or self.repository - message = f"Downloading {prefix}" - if speed: - message = f"{message} | {speed}" + prefix = self.label or self.repository + message = f"Downloading {prefix}" + if speed: + message = f"{message} {speed}" write_status( self.status_path, progress, @@ -928,7 +919,6 @@ class DownloadProgressReporter: stage=self.stage, stages=self.stages, stage_label=self.stage_label, - eta_seconds=eta_seconds, ) @@ -1063,7 +1053,7 @@ def _download_declared_or_fail(args, model_dir, label, start, span, stage_kwargs ( R2_ENDPOINT, download_from_r2, - f"Downloading {label} | OmooLab mirror", + f"Downloading {label} (mirror)", ), ) else: @@ -1072,7 +1062,7 @@ def _download_declared_or_fail(args, model_dir, label, start, span, stage_kwargs ( R2_ENDPOINT, download_from_r2, - f"Downloading {label} | OmooLab mirror", + f"Downloading {label} (mirror)", ), ) failures = [] @@ -1209,7 +1199,7 @@ def _download_file_from_url( with opener(request) as response: with open(temporary, "wb") as output: while True: - chunk = response.read(1024 * 1024) + chunk = response.read(DOWNLOAD_CHUNK_SIZE) if not chunk: break output.write(chunk) @@ -1310,7 +1300,7 @@ def download_from_r2( lambda filename: r2_file_url(args.r2_directory, filename), start, span, - label, + f"{label} (mirror)" if label else "mirror", stage_kwargs, ) @@ -1935,6 +1925,25 @@ def run_moge(args, model_cache=None, cancel_check=None): shutil.rmtree(temporary_directory, ignore_errors=True) +def merge_da3_frame_predictions(predictions): + """Combine independent single-frame DA3 results for existing exporters.""" + import numpy as np + + if not predictions: + raise ValueError("DA3 returned no frame predictions") + merged = predictions[0] + for name in ( + "depth", + "conf", + "extrinsics", + "intrinsics", + "processed_images", + ): + values = [getattr(prediction, name) for prediction in predictions] + setattr(merged, name, np.concatenate(values, axis=0)) + return merged + + def reconstruct(args, model_cache=None, cancel_check=None): write_status(args.status, 0.02, "Starting Environment") check_cancelled(cancel_check) @@ -2016,15 +2025,32 @@ def reconstruct(args, model_cache=None, cancel_check=None): 0.25, f"Estimating depth for {len(images)} frame(s)", ) + inference_options = { + "process_res": args.process_res, + "process_res_method": args.process_res_method, + "use_ray_pose": bool(args.use_ray_pose), + "ref_view_strategy": args.ref_view_strategy, + } try: - prediction = infer( - model, - [str(path) for path in images], - process_res=args.process_res, - process_res_method=args.process_res_method, - use_ray_pose=bool(args.use_ray_pose), - ref_view_strategy=args.ref_view_strategy, - ) + if getattr(args, "inference_mode", "joint") == "per_frame": + predictions = [] + for index, image in enumerate(images): + check_cancelled(cancel_check) + write_status( + args.status, + 0.25 + 0.60 * index / len(images), + f"Estimating DA3 depth ({index + 1}/{len(images)})", + ) + predictions.append( + infer(model, [str(image)], **inference_options) + ) + prediction = merge_da3_frame_predictions(predictions) + else: + prediction = infer( + model, + [str(path) for path in images], + **inference_options, + ) except UnicodeDecodeError as error: # DirectML reports failures with a locale-specific message that # ONNX Runtime cannot always decode, masking the real error. diff --git a/tests/test_addon.py b/tests/test_addon.py index cd399c5..dc8aac0 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -849,7 +849,7 @@ class AddonTest(unittest.TestCase): }, "fallback", ), - "3/8 pytorch | 2 downloaded", + "3/8 2 downloaded", ) self.assertEqual( operator._format_status( @@ -886,7 +886,7 @@ class AddonTest(unittest.TestCase): }, "fallback", ), - "2/3 PyTorch | Downloading packages | About 3 min left", + "2/3 Downloading packages", ) self.assertEqual( operator._format_status( @@ -902,6 +902,30 @@ class AddonTest(unittest.TestCase): "2/3 PyTorch", ) + def test_update_status_resets_progress_when_stage_changes(self): + operator = self.addon.RunJob() + settings = SimpleNamespace(progress=0.9, status="previous") + context = SimpleNamespace( + scene=SimpleNamespace(snapform_settings=settings) + ) + job = SimpleNamespace(status_path=Path("status.json"), progress_stage=1) + + with patch.object( + self.addon.client, + "read_status", + return_value={ + "progress": 0.1, + "message": "Downloading model", + "stage": 2, + "stages": 4, + "stage_label": "MoGe-2", + }, + ), patch.object(operator, "_redraw_ui"): + operator._update_status(context, job) + + self.assertEqual(job.progress_stage, 2) + self.assertEqual(settings.progress, 0.1) + def test_model_options_mark_only_models_that_are_not_downloaded(self): items_callback = self.addon.SnapFormSettings.__annotations__["model"][ "items" @@ -1419,6 +1443,7 @@ class AddonTest(unittest.TestCase): parameters = request.payload["parameters"] self.assertEqual(request.payload["command"], "reconstruct") + self.assertEqual(parameters["inference_mode"], "joint") self.assertEqual(parameters["export_format"], "gs_ply") for parameter in ( "process_res", @@ -1660,6 +1685,10 @@ class AddonTest(unittest.TestCase): self.fake_bpy.data.objects.get = original_get self.assertEqual(reconstruct.call_args.kwargs["start_frame"], 71) + self.assertEqual( + reconstruct.call_args.kwargs["inference_mode"], + "per_frame", + ) def test_reconstruct_request_accepts_image_depth_overrides(self): client = importlib.import_module("snapform.client") diff --git a/tests/test_download_progress.py b/tests/test_download_progress.py index 78d3f0a..6b71f01 100644 --- a/tests/test_download_progress.py +++ b/tests/test_download_progress.py @@ -110,10 +110,10 @@ class DownloadProgressTest(unittest.TestCase): status = json.loads(status_path.read_text(encoding="utf-8")) self.assertEqual( status["message"], - "Downloading da3 small | 13.4 MB/s", + "Downloading da3 small 13.4 MB/s", ) - def test_reports_estimated_time_remaining(self): + def test_does_not_report_estimated_time_remaining(self): times = iter((0.0, 10.0)) with tempfile.TemporaryDirectory() as directory: status_path = Path(directory) / "status.json" @@ -133,9 +133,9 @@ class DownloadProgressTest(unittest.TestCase): progress.update(50) status = json.loads(status_path.read_text(encoding="utf-8")) - self.assertEqual(status["eta_seconds"], 10) + self.assertNotIn("eta_seconds", status) - def test_staged_download_prefixes_speed_with_downloading(self): + def test_staged_download_keeps_model_label(self): class FakeTqdmWithRate(FakeTqdm): @property def format_dict(self): @@ -165,7 +165,7 @@ class DownloadProgressTest(unittest.TestCase): status = json.loads(status_path.read_text(encoding="utf-8")) self.assertEqual( status["message"], - "Downloading packages | 13.4 MB/s", + "Downloading da3 small 13.4 MB/s", ) self.assertEqual(status["stage"], 6) self.assertEqual(status["stages"], 8) diff --git a/tests/test_install_env.py b/tests/test_install_env.py index 730d09a..5ea2c12 100644 --- a/tests/test_install_env.py +++ b/tests/test_install_env.py @@ -1,7 +1,9 @@ import importlib.util import io +import hashlib import json import os +import shutil import subprocess import tarfile import tempfile @@ -129,8 +131,8 @@ class InstallEnvTest(unittest.TestCase): if value == "--allow-pattern" ] self.assertEqual(patterns, ["onnx/model_fp16.onnx"]) - self.assertNotIn("--stage", ben2_download) - self.assertNotIn("--stage-label", ben2_download) + self.assertIn("--stage", ben2_download) + self.assertIn("--stage-label", ben2_download) self.assertEqual( ben2_download[ben2_download.index("--label") + 1], "BEN2 Base", @@ -305,7 +307,7 @@ class InstallEnvTest(unittest.TestCase): ) self.assertEqual(called_command[-1], "da3 small") - def test_install_allocates_contiguous_progress_slices(self): + def test_install_resets_progress_for_each_stage(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) @@ -358,21 +360,8 @@ class InstallEnvTest(unittest.TestCase): ): 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.ENVIRONMENT_SLICE[0]) - runtime_ranges = stage_ranges[:2] - self.assertEqual( - [item[0] for item in runtime_ranges], - [1, 2], - ) - self.assertTrue( - all(item[0] is None for item in stage_ranges[4:]) - ) + self.assertEqual(stage_ranges[:2], [(1, 0.0, 0.05), (1, 0.05, 1.0)]) + self.assertEqual(stage_ranges[2:], [(2, 0.0, 1.0), (3, 0.0, 1.0)]) status = json.loads( (root / "status.json").read_text(encoding="utf-8") ) @@ -884,13 +873,12 @@ class InstallEnvTest(unittest.TestCase): ) ) - def test_install_fails_when_bundled_uv_missing(self): + def test_install_fails_when_all_uv_sources_are_unavailable(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() - self.create_bundled_vendor(root) uv_path = root / "uv" / "uv.exe" args = SimpleNamespace( uv=str(uv_path), @@ -916,29 +904,23 @@ class InstallEnvTest(unittest.TestCase): "rmtree", ), patch.object( bootstrap, - "bundled_uv_archive", - return_value=None, + "download_uv_archive", + side_effect=RuntimeError("all sources unavailable"), ): with self.assertRaisesRegex( RuntimeError, - "Bundled uv is missing", + "all sources unavailable", ): bootstrap.install(args) - def test_install_uses_bundled_vendor_and_uv_without_network(self): + def test_install_downloads_and_extracts_uv(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() - bundled = root / "bundled" - bundled.mkdir() - (bundled / "depth-anything-3.tar.gz").write_bytes(b"da3") - (bundled / "ben2.tar.gz").write_bytes(b"ben2") - (bundled / "utils3d.tar.gz").write_bytes(b"utils3d") - (bundled / "moge.tar.gz").write_bytes(b"moge") target, archive_name = bootstrap.uv_download_target() - archive = bundled / f"uv-{target}.{archive_name}" + archive = root / f"uv-{target}.{archive_name}" if archive_name.endswith("zip"): with zipfile.ZipFile(archive, "w") as bundle: bundle.writestr("uv/uv.exe", b"uv-binary") @@ -978,14 +960,18 @@ class InstallEnvTest(unittest.TestCase): install_commands.append(command) return [] + def fake_download(destination, *_args): + shutil.copy2(archive, destination) + return destination + with patch.object( - bootstrap, - "bundled_vendor_directory", - return_value=bundled, - ), patch.object( bootstrap, "resolve_indexes", return_value=None, + ), patch.object( + bootstrap, + "download_uv_archive", + side_effect=fake_download, ), patch.object( bootstrap, "run", @@ -1003,13 +989,6 @@ class InstallEnvTest(unittest.TestCase): self.assertEqual(uv_path.read_bytes(), b"uv-binary") all_commands = run_commands + install_commands - self.assertFalse( - any( - "github.com" in str(value) - for command in all_commands - for value in command - ) - ) runtime_command = next( command for command in install_commands @@ -1017,6 +996,9 @@ class InstallEnvTest(unittest.TestCase): ) self.assertFalse(any(" @ file:" in str(value) for value in runtime_command)) self.assertFalse(any("torch" in str(value) for value in runtime_command)) + status = json.loads(Path(args.status).read_text(encoding="utf-8")) + self.assertEqual(status["stage"], 4) + self.assertEqual(status["stages"], 4) def test_install_reuses_existing_uv(self): bootstrap = load_install_env() @@ -1062,6 +1044,78 @@ class InstallEnvTest(unittest.TestCase): self.assertEqual(uv_path.read_bytes(), b"binary") + def test_uv_download_falls_back_to_omoolab_and_verifies_checksum(self): + bootstrap = load_install_env() + content = io.BytesIO() + with zipfile.ZipFile(content, "w") as archive: + archive.writestr("uv/uv.exe", b"uv-binary") + payload = content.getvalue() + + class Response(io.BytesIO): + headers = {"Content-Length": str(len(payload))} + + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + destination = root / "uv.tmp" + status = root / "status.json" + urls = [] + clock_state = {"now": 0.0} + + def clock(): + clock_state["now"] += 0.25 + return clock_state["now"] + + def open_url(request): + urls.append(request.full_url) + if len(urls) == 1: + raise OSError("official unavailable") + return Response(payload) + + with patch.object( + bootstrap, + "uv_download_target", + return_value=("test-target", "zip"), + ), patch.dict( + bootstrap.UV_ARTIFACTS, + { + "test-target": ( + "uv-test.zip", + hashlib.sha256(payload).hexdigest(), + ) + }, + ), patch.dict( + bootstrap._uv_download_openers, + { + bootstrap.UV_OFFICIAL_BASE_URL: SimpleNamespace( + open=open_url + ), + bootstrap.UV_MIRROR_BASE_URL: SimpleNamespace( + open=open_url + ), + }, + ), patch.object( + bootstrap, + "write_status", + ) as write_status: + result = bootstrap.download_uv_archive( + destination, + status, + 0.0, + 1.0, + {}, + clock=clock, + ) + + self.assertEqual(result.read_bytes(), payload) + self.assertTrue(urls[0].startswith(bootstrap.UV_OFFICIAL_BASE_URL)) + self.assertTrue(urls[1].startswith(bootstrap.UV_MIRROR_BASE_URL)) + self.assertEqual(bootstrap.UV_OFFICIAL_CONNECT_TIMEOUT, 3.0) + self.assertEqual(bootstrap.UV_MIRROR_CONNECT_TIMEOUT, 15.0) + messages = [call.args[2] for call in write_status.call_args_list] + self.assertTrue(any("uv (mirror)" in value for value in messages)) + self.assertTrue(any("B/s" in value for value in messages)) + self.assertFalse(any("left" in value for value in messages)) + def test_install_stages_use_only_runtime_packages(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: @@ -1244,15 +1298,13 @@ class InstallEnvTest(unittest.TestCase): ) self.assertEqual(report, {}) - def test_install_uv_extracts_bundled_archive(self): + def test_install_uv_extracts_downloaded_archive(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) uv_path = root / "uv" - bundled = root / "bundled" - bundled.mkdir(parents=True) target, archive_name = bootstrap.uv_download_target() - archive = bundled / f"uv-{target}.{archive_name}" + archive = root / f"uv-{target}.{archive_name}" if archive_name.endswith("zip"): with zipfile.ZipFile(archive, "w") as bundle: bundle.writestr("uv/uv.exe", b"uv-binary") @@ -1262,30 +1314,32 @@ class InstallEnvTest(unittest.TestCase): info = tarfile.TarInfo("uv") info.size = len(content) tar.addfile(info, io.BytesIO(content)) + def fake_download(destination, *_args): + shutil.copy2(archive, destination) + return destination + with patch.object( bootstrap, - "bundled_vendor_directory", - return_value=bundled, + "download_uv_archive", + side_effect=fake_download, ): bootstrap.install_uv(uv_path, str(root / "status.json")) self.assertEqual(uv_path.read_bytes(), b"uv-binary") - def test_install_uv_raises_when_bundled_missing(self): + def test_install_uv_raises_when_download_fails(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) uv_path = root / "uv" - bundled = root / "bundled" - bundled.mkdir() with patch.object( bootstrap, - "bundled_vendor_directory", - return_value=bundled, + "download_uv_archive", + side_effect=RuntimeError("all uv sources failed"), ): with self.assertRaisesRegex( RuntimeError, - "Bundled uv is missing", + "all uv sources failed", ): bootstrap.install_uv(uv_path, str(root / "status.json")) diff --git a/tests/test_pack.py b/tests/test_pack.py index 3679248..2cd388d 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -61,21 +61,11 @@ class PackTest(unittest.TestCase): self.assertIn("service.py", names) self.assertNotIn("snapform/__init__.py", names) self.assertNotIn("importer.py", names) - self.assertIn( - f"vendor/{pack.PLATFORMS[platform]}", - names, - ) self.assertFalse( any(name.startswith("sources/") for name in names) ) - bundled_uv = { - name - for name in names - if name.startswith("vendor/uv-") - } - self.assertEqual( - bundled_uv, - {f"vendor/{pack.PLATFORMS[platform]}"}, + self.assertFalse( + any(name.startswith("vendor/") for name in names) ) self.assertIn(f'platforms = ["{platform}"]', manifest) self.assertNotIn( diff --git a/tests/test_prepare_vendor.py b/tests/test_prepare_vendor.py deleted file mode 100644 index fd8ea24..0000000 --- a/tests/test_prepare_vendor.py +++ /dev/null @@ -1,67 +0,0 @@ -import hashlib -import importlib.util -import io -import tempfile -import unittest -from pathlib import Path -from unittest.mock import patch - - -PROJECT_ROOT = Path(__file__).resolve().parents[1] - - -def load_prepare_vendor(): - specification = importlib.util.spec_from_file_location( - "snapform_prepare_vendor", - PROJECT_ROOT / "prepare_vendor.py", - ) - module = importlib.util.module_from_spec(specification) - specification.loader.exec_module(module) - return module - - -class PrepareVendorTest(unittest.TestCase): - def test_downloads_and_verifies_artifact(self): - prepare = load_prepare_vendor() - content = b"locked artifact" - artifact = { - "filename": "artifact.tar.gz", - "url": "https://example.invalid/artifact.tar.gz", - "sha256": hashlib.sha256(content).hexdigest(), - } - with tempfile.TemporaryDirectory() as directory, patch.object( - prepare.urllib.request, - "urlopen", - return_value=io.BytesIO(content), - ): - path = prepare.download_artifact(artifact, Path(directory)) - self.assertEqual(path.read_bytes(), content) - - def test_rejects_hash_mismatch_without_leaving_partial_file(self): - prepare = load_prepare_vendor() - artifact = { - "filename": "artifact.tar.gz", - "url": "https://example.invalid/artifact.tar.gz", - "sha256": hashlib.sha256(b"expected").hexdigest(), - } - with tempfile.TemporaryDirectory() as directory, patch.object( - prepare.urllib.request, - "urlopen", - return_value=io.BytesIO(b"unexpected"), - ): - destination = Path(directory) - with self.assertRaisesRegex(RuntimeError, "SHA-256 mismatch"): - prepare.download_artifact(artifact, destination) - self.assertFalse((destination / artifact["filename"]).exists()) - self.assertFalse( - (destination / f"{artifact['filename']}.part").exists() - ) - - def test_rejects_unknown_platform(self): - prepare = load_prepare_vendor() - with self.assertRaisesRegex(ValueError, "Unsupported platform"): - prepare.prepare(["windows-arm64"]) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_worker.py b/tests/test_worker.py index fd5198b..e0ef592 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -369,7 +369,7 @@ class WorkerTest(unittest.TestCase): "_verify_model_complete", return_value=None, ): - args.model_dir = str(Path(directory) / "model") + args.model_dir = str(Path(directory) / "models" / "DA3") args.status = str(Path(directory) / "status.json") self.worker.download_model(args) @@ -456,7 +456,7 @@ class WorkerTest(unittest.TestCase): "_verify_model_complete", return_value=None, ): - args.model_dir = str(Path(directory) / "model") + args.model_dir = str(Path(directory) / "models" / "DA3") args.status = str(Path(directory) / "status.json") self.worker.download_model(args) @@ -937,15 +937,21 @@ class WorkerTest(unittest.TestCase): self.assertEqual(mesh_calls[0][2:], (35.0, True, 0.2)) def test_reconstruct_copies_frames_for_multi_frame_input(self): + inference_inputs = [] fake_da3 = ModuleType("da3_onnx") fake_da3.create_session = lambda _model_dir, _device: object() - fake_da3.infer = lambda *_args, **_kwargs: SimpleNamespace( - depth=np.full((2, 2, 2), 2.0, dtype=np.float32), - conf=np.ones((2, 2, 2), dtype=np.float32), - extrinsics=np.repeat(np.eye(4, dtype=np.float32)[None, :3], 2, axis=0), - intrinsics=np.repeat(np.eye(3, dtype=np.float32)[None], 2, axis=0), - processed_images=np.zeros((2, 2, 2, 3), dtype=np.uint8), - ) + + def fake_infer(_session, images, **_options): + inference_inputs.append(images) + return SimpleNamespace( + depth=np.full((1, 2, 2), 2.0, dtype=np.float32), + conf=np.ones((1, 2, 2), dtype=np.float32), + extrinsics=np.eye(4, dtype=np.float32)[None, :3], + intrinsics=np.eye(3, dtype=np.float32)[None], + processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8), + ) + + fake_da3.infer = fake_infer fake_runtime = ModuleType("onnx_runtime") fake_runtime.provider_label = lambda _session: "CPU" @@ -979,6 +985,7 @@ class WorkerTest(unittest.TestCase): show_cameras=0, gs_views_interval=1, remove_background_model_dir="", + inference_mode="per_frame", ) self.worker.reconstruct(args) @@ -989,6 +996,10 @@ class WorkerTest(unittest.TestCase): self.assertTrue( (result_directory / "frame.0002.jpg").is_file() ) + self.assertEqual(len(inference_inputs), 2) + self.assertTrue(all(len(paths) == 1 for paths in inference_inputs)) + with np.load(result_directory / "results.npz") as result: + self.assertEqual(result["depth"].shape, (2, 2, 2)) self.assertTrue( (result_directory / "results.npz").is_file() ) diff --git a/vendor.lock.toml b/vendor.lock.toml deleted file mode 100644 index fe40004..0000000 --- a/vendor.lock.toml +++ /dev/null @@ -1,19 +0,0 @@ -schema_version = 1 - -[uv] -version = "0.11.3" - -[platforms.windows-x64] -filename = "uv-x86_64-pc-windows-msvc.zip" -url = "https://github.com/astral-sh/uv/releases/download/0.11.3/uv-x86_64-pc-windows-msvc.zip" -sha256 = "ae681c0aaec7cc96af184648cb88d73f8393ed60fa5880abdd6bdb910f9b227c" - -[platforms.macos-arm64] -filename = "uv-aarch64-apple-darwin.tar.gz" -url = "https://github.com/astral-sh/uv/releases/download/0.11.3/uv-aarch64-apple-darwin.tar.gz" -sha256 = "2bc3d0c7bf2bd08325b1e170abac6f7e5b3346e1d4eab3370d17cefec934996f" - -[platforms.linux-x64] -filename = "uv-x86_64-unknown-linux-gnu.tar.gz" -url = "https://github.com/astral-sh/uv/releases/download/0.11.3/uv-x86_64-unknown-linux-gnu.tar.gz" -sha256 = "c0f3236f146e55472663cfbcc9be3042a9f1092275bbe3fe2a56a6cbfd3da5ce"