From 2ad30bf4a37d1c37efde11c7c12ecd5a4840e0ce Mon Sep 17 00:00:00 2001 From: nan Date: Fri, 14 Aug 2026 10:27:09 +0800 Subject: [PATCH] feat: Implement ONNX model inference for BEN2, Depth Anything 3, and MoGe - Added `ben2_onnx.py` for BEN2 ONNX inference with preprocessing. - Introduced `da3_onnx.py` for Depth Anything 3 inference, including camera projection calculations. - Created `moge_onnx.py` for MoGe-2 ONNX inference and post-processing. - Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`. - Added model synchronization script `sync_models.py` for downloading and verifying ONNX models. - Developed unit tests for BEN2, Depth Anything 3, and MoGe ONNX models to ensure functionality. - Updated `vendor.lock.toml` to remove unused artifacts. --- .gitignore | 1 + README.md | 14 +- docs/depth-plane.md | 7 +- docs/development.md | 10 +- docs/features.md | 21 +- docs/outputs.md | 28 +- docs/service.md | 125 ++--- docs/surface.md | 10 +- pack.py | 63 +-- patches/depth-anything-3/pyproject.toml | 93 --- patches/depth-anything-3/requirements.txt | 23 - .../depth_anything_3/utils/export/__init__.py | 84 --- .../src/depth_anything_3/utils/ray_utils.py | 531 ------------------ patches/moge/pyproject.toml | 22 - patches/utils3d/pyproject.toml | 16 - prepare_vendor.py | 5 +- src/snapform/__init__.py | 4 +- src/snapform/ben2_onnx.py | 50 ++ src/snapform/client.py | 16 +- src/snapform/da3_onnx.py | 221 ++++++++ src/snapform/image_depth.py | 12 +- src/snapform/image_mesh.py | 7 +- src/snapform/install_env.py | 262 ++------- src/snapform/models.py | 131 +++-- src/snapform/moge_onnx.py | 165 ++++++ src/snapform/onnx_runtime.py | 55 ++ src/snapform/panel.py | 5 +- src/snapform/properties.py | 4 +- src/snapform/service.py | 63 +-- src/snapform/worker.py | 378 +++++++------ sync_models.py | 194 +++++++ tests/test_addon.py | 30 +- tests/test_ben2_onnx.py | 53 ++ tests/test_da3_onnx.py | 61 ++ tests/test_install_env.py | 260 +++------ tests/test_models.py | 64 ++- tests/test_moge_onnx.py | 130 +++++ tests/test_pack.py | 95 ---- tests/test_service.py | 60 +- tests/test_worker.py | 278 +++------ vendor.lock.toml | 27 - 41 files changed, 1569 insertions(+), 2109 deletions(-) delete mode 100644 patches/depth-anything-3/pyproject.toml delete mode 100644 patches/depth-anything-3/requirements.txt delete mode 100644 patches/depth-anything-3/src/depth_anything_3/utils/export/__init__.py delete mode 100644 patches/depth-anything-3/src/depth_anything_3/utils/ray_utils.py delete mode 100644 patches/moge/pyproject.toml delete mode 100644 patches/utils3d/pyproject.toml create mode 100644 src/snapform/ben2_onnx.py create mode 100644 src/snapform/da3_onnx.py create mode 100644 src/snapform/moge_onnx.py create mode 100644 src/snapform/onnx_runtime.py create mode 100644 sync_models.py create mode 100644 tests/test_ben2_onnx.py create mode 100644 tests/test_da3_onnx.py create mode 100644 tests/test_moge_onnx.py diff --git a/.gitignore b/.gitignore index efbc756..fc0a0e3 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ build/ .pytest_cache/ /vendor/ +/models/ diff --git a/README.md b/README.md index da309bc..01377f4 100644 --- a/README.md +++ b/README.md @@ -16,11 +16,11 @@ SnapForm 是一个 Blender 扩展，可以把一张照片、一段视频或一 2. 打开 Blender，进入 `Preferences > Extensions` 3. 点击 `Install from Disk`，选择下载的压缩包 4. 在扩展偏好设置中确认 Storage Root 与 Device（保持默认即可）；Device 默认取本机 - 最快的 CUDA/MPS，否则 CPU。下载超时会自动重试国内镜像，无需手动配置 + 最快的 DirectML/CoreML provider，否则 CPU。下载超时会自动重试国内镜像，无需手动配置 5. 打开 3D 视图右侧的 `SnapForm` 面板，点击 `Install Environment` -DA3、MoGe-2、utils3d、BEN2 源码包和 uv 已内置在扩展里，安装时直接从本地解包，不需要访问 GitHub。首次使用仍会自动 -下载运行环境依赖（PyTorch 等），随后自动下载 `MoGe-2 ViT-B Normal` 和 `BEN2 Base`。其他 AI 模型可在 +扩展只内置对应平台的 uv。首次使用会自动下载轻量的 ONNX Runtime 环境，随后自动下载 `MoGe-2 ViT-B Normal` 和 +`BEN2 Base`；不安装 PyTorch，也不下载各模型的 Python 源码包。其他 AI 模型可在 `Models` panel 中按需下载，请保持网络畅通，首次安装需要一些时间。 ## 快速上手 @@ -46,14 +46,18 @@ DA3、MoGe-2、utils3d、BEN2 源码包和 uv 已内置在扩展里，安装时 选中单张图片的 Image Empty，右键打开 `SnapForm` 子菜单并选择 `Convert to Surface`。可选择 MoGe-2 ViT-S、ViT-B、ViT-L 或 DA3 Small、Base，其中 MoGe-2 ViT-B 为默认模型。SnapForm 会按模型输出的 -相机内参与三维点建立透视网格，并匹配原 Empty 的可见尺寸、Offset 与 Transform。MoGe-2 直接使用官方 -`utils3d.build_mesh_from_map` 与深度断层 mask 生成干净拓扑；DA3 保留 confidence 与边拉伸剔除。两者再由独立的 Geometry Nodes +相机内参与三维点建立透视网格，并匹配原 Empty 的可见尺寸、Offset 与 Transform。MoGe-2 使用内置 NumPy +网格构建与深度断层 mask 生成干净拓扑；DA3 保留 confidence 与边拉伸剔除。两者再由独立的 Geometry Nodes 修改器完成剔面、快速重拓扑、Relax 与 UV 重投。开启 Remove Background 时，背景三角面会在转换期间 直接删除，不再交给 Geometry Nodes。匹配相机默认不创建；开启后会作为 Surface 的子级。 MoGe-2 的像素级 normal 会作为 Object Space normal 贴图接入 `O MoGe2 Surface Layer`，并与原图 Bump 合并；网格本身不写入 custom normal，因此重拓扑降低顶点数后仍能保留细腻的法线细节。 当前此功能仅支持单张图片。 +MoGe-2 ViT-S、ViT-B、ViT-L、DA3 Small、DA3 Base 与 BEN2 全部使用 ONNX。模型执行、相机与深度恢复、 +点云重投影和网格生成均不依赖 PyTorch。Windows 使用 DirectML 加速，provider 不可用时会回退到 ONNX Runtime CPU， +并在任务状态中显示实际 provider。 + 转换对话框提供模型与 `Create Matching Camera`、`Remove Background`、`Unshaded`；DA3 显示 `Resolution Scale`，MoGe-2 显示 `Resolution Level`， 以及分割线下方的 `Keep Original`。 diff --git a/docs/depth-plane.md b/docs/depth-plane.md index 636179c..640882f 100644 --- a/docs/depth-plane.md +++ b/docs/depth-plane.md @@ -4,17 +4,17 @@ 1. 在 3D View 选中 `Object Type = Empty`、`Display As = Image` 的对象。 2. 右键打开 `SnapForm` 子菜单并选择 `Convert to Depth Plane`，通过 `Depth from` 选择深度来源： - - `DA3`：显示 Model、Use Ray Pose 和 Resolution Scale。Process Res 先以原图长边乘以 + - `DA3`：显示 Model 和 Resolution Scale。ONNX 模型使用 camera decoder 恢复相机。Process Res 先以原图长边乘以 Resolution Scale，再钳制到 `252–1008`，最后对齐到最接近的 `14` 倍数。Resolution Scale 以百分比 显示，范围为 `1–100%`，默认 `100%`；界面只在标签中实时显示最终 Process Res。 转换导出 `mini_npz`，使用前需要在 `Models` panel 下载所选模型，并固定使用 `process_res_method="upper_bound_resize"`；设备和参考视图策略 继续使用场景设置。Model 下拉菜单只列出已经下载的 DA3 模型；若没有可用模型，对话框会引导到 `Models` panel 下载。 - - `MoGe-2`：使用所选模型输出 metric depth、normal、validity mask 与 normalized intrinsics。Depth Plane 不保存体积较大的 points；Surface 会在 AI 进程内用 points 和官方 `utils3d.build_mesh_from_map` 直接生成 `mesh.npz`。Debug 始终保留原始 `results.npz`，并按 Output Format 单独生成 Maps、纹理 GLB 或彩色 PLY。 + - `MoGe-2`：使用所选模型输出 metric depth、normal、validity mask 与 normalized intrinsics。Depth Plane 不保存体积较大的 points；Surface 会在 AI 进程内用内置 NumPy 网格构建生成 `mesh.npz`。Debug 始终保留原始 `results.npz`，并按 Output Format 单独生成 Maps、纹理 GLB 或彩色 PLY。 - `Original`：直接把原图作为深度贴图，不显示估算参数。仅在开启 `Remove Background` 时运行 BEN2 取得前景 Alpha；视频或序列帧按需同时抽取原始帧与前景序列，把原始帧作为 `.depth` 序列。 - `Original` 是默认深度来源。选择 `DA3` 时，其 Model、Use Ray Pose 与 Resolution Scale + `Original` 是默认深度来源。选择 `DA3` 时，其 Model 与 Resolution Scale 参数会集中显示在独立边框中。`Remove Background` 与 `Unshaded` 直接显示，不再使用 Options 折叠； `Keep Original` 默认关闭，并在分割线下方保持为最后一项。仅当开启 `Remove Background` 时才需要 `BEN2 Base`，并把 `.foreground` 的 Alpha 作为独立 @@ -63,4 +63,3 @@ 节点组根据 `Frame` 计算从 `1` 开始的 Image Texture Frame：场景帧与 Start 都为 `1` 时读取图片第 `1` 帧；非循环输入 钳制到有效范围；循环输入使用 Floored Modulo，让 `Frame Start` 前后的负、正相对帧都映射到 `1…Frames`，因此开始帧之前也会循环播放。Geometry 置换会像材质 Image Texture 一样随场景帧刷新。 - diff --git a/docs/development.md b/docs/development.md index 3938931..e22d4f4 100644 --- a/docs/development.md +++ b/docs/development.md @@ -12,11 +12,10 @@ uv run python pack.py ```text . -├── vendor.lock.toml 第三方构建输入的版本、URL 与 SHA-256 -├── prepare_vendor.py 下载并校验第三方构建输入 +├── vendor.lock.toml 各平台 uv 的版本、URL 与 SHA-256 +├── prepare_vendor.py 下载并校验各平台 uv ├── pack.py -├── patches/ 进入 Git 的 DA3 源码覆盖文件 -├── vendor/ 不进入 Git 的下载缓存；本地调试也直接读取这里 +├── vendor/ 不进入 Git 的 uv 下载缓存 ├── src/ │ └── snapform/ │ ├── blender_manifest.toml @@ -36,7 +35,6 @@ uv run python pack.py ## 上游与许可证 - Depth Anything 3 代码采用 Apache 2.0。 -- BEN2 代码与模型权重采用 MIT License。 +- BEN2 ONNX 模型采用其仓库声明的 Apache 2.0 License。 - 本 Extension 代码采用 GPL-3.0-or-later。 - 模型权重许可证由所选 Hugging Face 模型仓库决定，Extension 下载前会在界面显示。 - diff --git a/docs/features.md b/docs/features.md index ebeea6d..565ef9c 100644 --- a/docs/features.md +++ b/docs/features.md @@ -12,14 +12,14 @@ ## 当前功能 - 使用 Blender Extension manifest，可通过 `Install from Disk` 安装 -- 在独立 Environment 中通过常驻 Service 运行 Torch、DA3、MoGe-2 与 BEN2，避免污染 Blender Python 进程 +- 在独立 Environment 中通过常驻 Service 运行 ONNX Runtime；MoGe-2、DA3 与 BEN2 的模型执行及后处理完全不使用 PyTorch，避免污染 Blender Python 进程 - Service 的深度模型缓存只保留 DA3 或 MoGe-2 其中一个，切换时主动释放前一个；BEN2 使用独立缓存， 可与当前深度模型共存 - Blender 使用内置 `urllib.request` 访问仅监听 `127.0.0.1` 的 Service；请求带随机 Token -- 常驻 Service 复用 Torch、CUDA、BEN2 与最近使用的 DA3 模型，连续执行时不再重复加载 +- 常驻 Service 复用 ONNX Runtime session、BEN2 与最近使用的深度模型，连续执行时不再重复加载 - 多个 Blender 实例自动连接同一个共享 Service；`AI Service` 面板显示状态、端口、客户端与模型缓存，并提供安装、启动、清理、重启和停止控制 - Environment、模型和任务结果默认保存在 `~/.snapform`，可由多个 Blender 版本共享 -- 当前启用 DA3 Small 和 Base；Large 1.1、Giant 1.1 与 Nested Giant Large 1.1 暂停使用 +- 当前仅启用 FP32 DA3 Small 和 Base；它们不随首次安装下载，Giant 1.1 与 Nested Giant Large 1.1 暂停使用 - 支持 BEN2 Base 单图背景剔除测试，输出透明 PNG 与推理耗时 - 支持单张 JPG、JPEG、PNG、WEBP 图片或常见视频文件输入 - 视频按顺序截取前若干帧，可设置参与重建的最大帧数 @@ -52,17 +52,16 @@ 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.10.12.macos-arm64.zip`）。 - 每个包只含对应平台的 uv 二进制，并内置 DA3、BEN2 源码包。 + 每个包只含对应平台的 uv 二进制，不再携带 AI 项目的源码包。 2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。 3. 在 Extension Preferences 中确认 Storage Root 与 Device。Device 是正式功能 - （Image Empty 右键菜单）使用的推理设备，默认取本机最快的 CUDA/MPS，否则 CPU，且只列出 + （Image Empty 右键菜单）使用的推理设备，默认取本机最快的 DirectML/CoreML provider，否则 CPU，且只列出 本机支持的设备。 -4. 在 `AI Service` 点击 `Install Environment`。首次安装会下载 Python、PyTorch、FastAPI、Uvicorn； - DA3、BEN2 源码包与 uv 已内置在扩展中，安装时从 安装任务的临时目录中使用，安装结束后自动清理， - 并在 Environment 安装后自动下载 `DA3 Small` 和 `BEN2 Base` 推荐模型。Environment 安装 - 固定显示 `Environment`、`PyTorch`、`AI components` 三个用户阶段；xFormers（macOS 跳过）、 - FastAPI 等内部安装步骤合并在 `AI components` 中。耗时最长的 PyTorch 与 AI 组件占主要进度 - 区间。安装完成后直接写入 Environment Manifest，不再额外导入全部组件进行验证。推荐模型属于安装 +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` 面板按需下载。 状态文字使用 `阶段/总数 阶段名 | 当前操作 | 预计剩余时间`，依赖安装会结合预设耗时与实际 进度修正预期；模型下载根据字节数和速度计算剩余时间，例如 diff --git a/docs/outputs.md b/docs/outputs.md index 656e2ed..97cc1c3 100644 --- a/docs/outputs.md +++ b/docs/outputs.md @@ -1,31 +1,29 @@ # 输出格式 -## 输出格式说明 +## DA3 ONNX `Output Type` 提供以下结果： -- Point Cloud (GLB)：DA3 `glb` -- Gaussian Splats (PLY)：DA3 `gs_ply`，同时传入 `infer_gs=True` -- Depth Preview (Images)：DA3 `depth_vis` -- Depth Data (NPZ)：DA3 `mini_npz` -- Textured Surface (GLB)：SnapForm 根据 DA3 Prediction 构建并输出 `mesh.glb` +- Point Cloud (GLB)：由 depth、confidence、相机参数和原图生成彩色点云 +- Depth Preview (Images)：保存深度着色预览 +- Depth Data (NPZ)：保存 depth、confidence、extrinsics、intrinsics 与处理后的图像 +- Textured Surface (GLB)：SnapForm 根据 Prediction 构建纹理网格 -Point Cloud 专属设置仅在对应类型下显示，`gs_views_interval` 仅用于 Gaussian Splats。 -Textured Surface 使用 0–100% 的 Confidence Threshold 删除置信度最低的一定比例： +DA3 Small 与 Base 的 ONNX 图不包含 Gaussian head，因此不提供 Gaussian Splats。默认推理参数为： - `process_res=504` - `process_res_method="upper_bound_resize"` -- `use_ray_pose=False` - `ref_view_strategy="saddle_balanced"` - `conf_thresh_percentile=40.0` - `num_max_points=1_000_000` - `show_cameras=True` -- `export_kwargs["gs_ply"]["gs_views_interval"]=1` -DA3 官方 `glb` 是彩色点云与可选相机线框；`gs_ply` 是带 Gaussian 属性的 PLY； -`depth_vis` 是原图与着色深度图的并排预览；`mini_npz` 保存深度、置信度和可用的相机参数。 -Point Cloud GLB 任务会关闭其默认附带的深度预览，使五种输出产物保持分离。 +相机外参与内参由模型的 camera decoder 输出恢复，不运行 ray-pose 求解器。 -官方格式与限制可参阅 -[DA3 Python API 文档](https://github.com/ByteDance-Seed/Depth-Anything-3/blob/main/docs/API.md)。 +## MoGe-2 ONNX +MoGe-2 Debug 可保存原始 NPZ、depth/normal/mask maps、纹理 GLB 或彩色 PLY。Surface 输出由模型 points、mask 与 depth edge 在内部构建三角网格。 + +## BEN2 ONNX + +BEN2 输出 `foreground.png`，RGB 保留原始输入，Alpha 来自 ONNX mask 并恢复到输入分辨率。 diff --git a/docs/service.md b/docs/service.md index 6f43c28..fadc7b7 100644 --- a/docs/service.md +++ b/docs/service.md @@ -1,98 +1,61 @@ # AI Service -## AI Service +## 运行架构 -Environment 指磁盘上已安装的 Python、依赖及模型执行组件；Service 是由 Environment 启动、接收 Blender 任务的 -常驻 HTTP 进程。`manifest.json` 记录当前 Environment 的安装信息；Service 的 `/health` 返回进程状态和 -实例标识。项目发布前不维护安装清单或任务协议的兼容版本号。 +SnapForm 在独立的 Python 3.12 Environment 中运行常驻 HTTP Service，Blender 只通过标准库 `urllib.request` 提交任务，不把 AI 运行库载入 Blender Python。Service 只监听 `127.0.0.1`，请求使用随机 Token；同一 Storage Root 由 `service.lock` 保证只有一个实例。 -Extension 本身不包含 Torch。`Install Environment` 使用 `uv` 创建独立 Python 3.12 -环境，并安装。机器上即使没有安装 `uv` 或 Python 也没关系：安装程序由 Blender -自带的 Python 启动，缺 `uv` 时会解包扩展内置的对应平台二进制，再用它下载 -Python 3.12 与全部依赖。`uv` 会缓存在扩展存储目录下，后续安装直接复用。 -安装不做网络预检，直接开始；环境变量仍可切换镜像与代理：`PIP_INDEX_URL`（PyPI 镜像）、 -`SNAPFORM_TORCH_INDEX`（PyTorch 索引）、`HF_ENDPOINT`（Hugging Face 镜像）、 -`HTTP(S)_PROXY`/`ALL_PROXY`（代理）。未指定 PyPI 镜像时， -pip 安装失败会自动依次重试镜像（清华、阿里云）；模型下载不预探测，官方 Hugging Face -失败后自动回退 OmooLab R2（`models.omoolab.xyz`），并校验文件大小和 SHA-256。 -GitHub 不再参与安装：DA3、BEN2、uv 已内置，缺少内置文件会直接报错提示重新安装，不再下载兜底。 -单独下载模型时遵循同样的顺序：Hugging Face（可用 `HF_ENDPOINT` 指定其他入口）失败后自动切换到 -OmooLab R2，不再使用 ModelScope。 -Windows 上代理若只在系统设置里而未写入环境变量，安装程序会尝试自动从系统代理注入子进程。 - -- PyTorch 2.9.1 -- TorchVision 0.24.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 与 BEN2 源码包随扩展打包，安装时复制到临时目录后由 uv 本地安装，完成后立即清理 - （无需访问 GitHub）；SnapForm 直接调用 DA3 - 底层 export 接口，`gs_ply` 导出不会触发 `gs_video` 渲染，因此不需要 gsplat - 与 C++ 编译工具链 -- `vendor/` 是不进入 Git 的本地缓存；`vendor.lock.toml` 固定下载 URL、commit/version 与 - SHA-256，`patches/<包名>/` 进入 Git 并存放 SnapForm 对源码的覆盖文件（移除 xformers - 硬依赖、export 入口改懒加载、 - 射线位姿的 3×3 QR 分解在 MPS 上改走 CPU）， - `pack.py` 打包时自动套入官方包，扩展 zip 内才是修补后的版本；更新上游只需替换官方包 - 并核对覆盖文件是否仍适用 -- `addict` 2.4.0（补充 DA3 上游包元数据暂未声明的运行依赖） - -安装 Environment 时会自动选择 PyTorch 后端：检测到 NVIDIA 驱动则安装 CUDA 12.8 版本，否则安装 -CPU 版；macOS 始终安装 CPU 版 PyTorch（MPS 包含其中）。CPU 可以用于验证，但实际推理通常很慢。 -Extension Preferences 的 `Device` 决定正式功能的推理设备，默认取本机最快的 CUDA/MPS，否则 -CPU；Debug 面板的 `Debug Device` 只影响 Debug 任务。两处设备选项都只列出本机支持的设备： -Windows 不显示 MPS，机器上没有可用 NVIDIA 驱动时不显示 CUDA。Environment 安装完成时会检测真实 -可用的设备并写入 `manifest.json`，设备列表以该检测结果为准。 -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 暂不支持。DA3 射线位姿估计中的 3×3 QR 分解在 MPS -上尚未实现，SnapForm 的覆盖文件会单独把该分解放到 CPU 执行，其余推理仍走 MPS。 - -Extension 启用后只会查找已有的共享 Service，不会主动启动新进程。第一次执行 AI 操作或在侧边栏 -手动点击 `Start` 时才会启动 Service。Service 只绑定回环地址，并让操作系统分配动态端口；第一个实例把 -端口、进程和随机 Token 写入 `service.json`，其余 Blender 通过标准库 -`urllib.request` 连接同一个 Service。`service.lock` 保证同一 Storage Root 同时只有一个 Service，Token -则阻止同机其他进程直接调用接口。Service 全局同时只执行一个 AI 任务。 - -每个 Blender 会注册独立客户端并每 5 秒发送心跳。正常关闭 Blender 时先注销自身；如果当前 Service -就是由这个 Blender 启动的，会直接终止它，无需等待。异常退出的客户端会在 10 秒后过期；Service -在最后一个客户端离开后仍会保留模型缓存等待 20 秒再退出，作为 Service 由其他进程启动等场景的兜底。 -取消任务不会重启 Service，并会保留已经加载的模型； -如果取消发生在一次不可中断的 Torch 调用中，任务会在该调用返回后的最近检查点结束。 - -侧边栏 `AI Service` 面板以 `Running (Clients: n)` 表示空闲服务，以 `Running (Busy)` 表示正在执行任务， -并显示缓存模型和活动任务。Blender 完成结果导入时会立即释放当前任务对应的本地 Busy 快照，无需等待 -下一次 Service 轮询。空闲且只有当前 Blender 连接时，可以清理模型缓存、重启或停止 Service； -多客户端或任务执行期间不会允许破坏共享服务。 - -BEN2 在首次使用后保持加载；DA3 保留最近使用的一个模型。切换 DA3 模型时会释放旧模型；加载时如果 -显存不足，会先释放另一类缓存模型并重试。 - -Environment Manifest Revision 变化时，旧 Environment 需要重新执行 `Install Environment`。 -Service 的启动和崩溃日志写入 `~/.snapform/service.log`。 - -Environment、模型和任务结果默认保存在用户主目录的 `~/.snapform`： +Environment、模型和任务结果默认保存在 `~/.snapform`： ```text ~/.snapform/ ├── .venv/ ├── models/ ├── jobs/ +├── manifest.json ├── service.log ├── service.json # Service 运行时存在 └── service.lock # Service 运行时持有 ``` -`models/` 下每个模型一个文件夹，命名统一带前缀：DA3 为 `da3_*`（如 `da3_small`、`da3_base`）， -BEN2 为 `ben2_*`（如 `ben2_base`）。旧版无前缀的 DA3 文件夹会在首次使用时自动改名为新名称。 +不同 Blender 版本可共享此目录。也可以在 Extension Preferences 的 `Storage Root` 中修改位置。 -不同 Blender 版本安装相同 Extension 后会共用这个目录。也可以在 Extension -Preferences 的 `Storage Root` 中修改共享位置。不要把模型写入 Extension 安装目录， -因为该目录可能只读，并且升级时会被替换。不要同时从多个 Blender 实例执行 -`Install Environment`；Service、任务队列和模型读取可以共享。 +## ONNX-only Environment +扩展只内置对应平台的 `uv` 可执行文件。`Install Environment` 使用它创建环境并安装： + +- ONNX Runtime 1.24.4；Windows 使用 `onnxruntime-directml` +- NumPy 2.4.2、Pillow 12.1.1、OpenCV Headless 4.13.0.92、trimesh 4.11.2 +- FastAPI 0.139.2、Uvicorn 0.51.0、huggingface-hub 1.4.1 + +运行环境不安装 PyTorch、TorchVision、xFormers，也不安装 DA3、MoGe、BEN2 或 utils3d 的 Python 源码包。模型推理和 SnapForm 所需的相机、点云、网格后处理均由 ONNX Runtime、NumPy 与内部实现完成。 + +Windows 当前发行环境提供 DirectML 与 CPU；`Auto` 优先使用 `DmlExecutionProvider`。显式指定不可用设备时会安全回退到 CPU。其他平台安装标准 ONNX Runtime，并按可用 provider 选择 CoreML、OpenVINO 或 CPU。 + +`manifest.json` 记录 Environment 版本、ONNX Runtime 版本、runtime variant 与设备列表。Manifest Revision 改变后需要重新执行 `Install Environment`。 + +## 模型 + +所有正式功能都使用 ONNX： + +- MoGe-2 ViT-S、ViT-B、ViT-L Normal +- DA3 Small、DA3 Base +- BEN2 Base + +首次安装自动下载 `MoGe-2 ViT-S Normal` 与 FP16 `BEN2 Base`，其他模型在 `Models` panel 按需下载。模型按 Hugging Face 仓库末段名称保存在 `~/.snapform/models/`，并保留仓库内部结构。下载顺序为 Hugging Face，然后回退到 `https://models.omoolab.xyz//`；R2 文件会校验大小与 SHA-256。可用 `HF_ENDPOINT` 指定 Hugging Face 镜像。 + +DA3 ONNX 使用模型的 camera decoder 恢复外参与内参；旧版 `Use Ray Pose` 参数只为任务协议兼容保留，不参与 ONNX 推理。MoGe-2 的 metric depth、normal、mask、normalized intrinsics，以及 BEN2 的 Alpha，均在 NumPy 后处理中恢复。 + + +## 安装网络与镜像 + +安装器直接开始工作，不做阻塞式网络预检。依赖安装失败时会依次重试清华与阿里云 PyPI 镜像，也可用 `PIP_INDEX_URL` 或 Preferences 中的 `PyPI Index` 显式指定。Windows 系统代理会注入子进程；同时支持 `HTTP_PROXY`、`HTTPS_PROXY` 与 `ALL_PROXY`。 + +GitHub 不参与用户安装流程。`prepare_vendor.py` 只为扩展打包准备各平台的 `uv` 二进制。 + +## 生命周期 + +扩展启用后只查找已有 Service；第一次执行 AI 操作或点击 `Start` 时才启动。每个 Blender 客户端每 5 秒发送心跳，正常关闭时注销。最后一个客户端离开后，Service 保留模型缓存 20 秒再退出。 + +Service 全局同时执行一个 AI 任务。BEN2 独立缓存，深度模型只保留最近使用的一个；切换模型时释放旧会话。取消任务不会重启 Service，ONNX Runtime 当前调用返回后在最近检查点结束。 + +启动与崩溃信息写入 `~/.snapform/service.log`。 diff --git a/docs/surface.md b/docs/surface.md index 3c40eb7..b6fc842 100644 --- a/docs/surface.md +++ b/docs/surface.md @@ -25,7 +25,7 @@ - Smooth 采用 `O Smooth` 的边界检测和模糊法线权重，固定开放边界与薄面，只保留一个 `Smooth` 程度系数；节点组不再额外执行 Flip Faces； - 从修改前的原始表面以 nearest-surface 重新采样 `UVMap`，最终明确写为 Corner 域 `2D Vector`。 - MoGe-2 使用独立的 `O MoGe2 Surface`。Service 由 `utils3d.np.build_mesh_from_map(..., tri=True)` + MoGe-2 使用独立的 `O MoGe2 Surface`。Service 由内置 NumPy 网格构建函数 生成有限深度范围内的紧凑三角网格，并保留模型 mask、逐点 depth 及相邻点相对深度差。Geometry Nodes 按官方 `mask & ~depth_map_edge(depth, rtol=0.04)` 的含义，先由 `Mask` 删除 mask 外的面，再由 `Depth Continuity` 按同一份 3×3 邻域相对深度差删除断层面。DA3 与 MoGe-2 使用相同的属性、 @@ -163,7 +163,7 @@ Image Empty 使用 Packed Image 时，SnapForm 会通过 Blender 写出临时 PN 浮点深度；保存 `.blend` 后再次打开仍需对应 job 目录存在。 选择 PLY 后，模型菜单只会提供 `DA3 Giant 1.1` 和 -`DA3 Nested Giant Large 1.1`。Small、Base 和 Large 没有 Gaussian head，无法由 +`DA3 Nested Giant Large 1.1`。Small 和 Base 没有 Gaussian head，无法由 DA3 生成 `gs_ply`。 每个任务在 `~/.snapform/jobs/` 下用独立文件夹包裹。文件夹统一采用 @@ -185,7 +185,7 @@ Blender 为 Packed Image 或序列帧生成的输入副本也归档到对应 job - `remove-image-background` - `generate-original-depth-plane` -DA3 官方 exporter 产生的 `gs_ply/`、`depth_vis/`、`exports/mini_npz/` 会在完成后自动平铺， +DA3 ONNX exporter 产生的 `depth_vis/`、`exports/mini_npz/` 会在完成后自动平铺， 不会保留只含一层产物的子目录： ```text @@ -218,6 +218,4 @@ remove-background_20260807-153000-123456/ └── depth.0001.exr… AI 深度序列帧 ``` -建议从 `DA3 Base` 开始。它使用 Apache 2.0 模型权重；`DA3 Large 1.1` 权重采用 -CC BY-NC 4.0，仅限非商业用途。 - +建议从 `DA3 Base` 开始。它使用 Apache 2.0 模型权重。 diff --git a/pack.py b/pack.py index bf996d9..9b545df 100644 --- a/pack.py +++ b/pack.py @@ -1,6 +1,4 @@ import hashlib -import io -import tarfile import tomllib from pathlib import Path from zipfile import ZIP_DEFLATED, ZipFile @@ -10,7 +8,6 @@ PROJECT_ROOT = Path(__file__).resolve().parent SOURCE_DIR = PROJECT_ROOT / "src" / "snapform" VENDOR_DIR = PROJECT_ROOT / "vendor" DIST_DIR = PROJECT_ROOT / "dist" -PATCHES_DIR = PROJECT_ROOT / "patches" LOCK_PATH = PROJECT_ROOT / "vendor.lock.toml" @@ -28,15 +25,6 @@ PLATFORMS = { platform: artifact["filename"] for platform, artifact in VENDOR_LOCK["platforms"].items() } -BUNDLED_ARCHIVES = { - f"vendor/{artifact['filename']}": artifact["filename"] - for artifact in VENDOR_LOCK["artifacts"].values() -} -OVERLAYS = { - artifact["filename"]: artifact["patches"] - for artifact in VENDOR_LOCK["artifacts"].values() - if "patches" in artifact -} def file_sha256(path): @@ -63,41 +51,6 @@ def validate_artifact(path, expected_sha256): ) -def bundled_artifact_bytes(filename): - """Return a bundled artifact with its configured overlay applied.""" - archive_path = VENDOR_DIR / filename - overlay_dir = PATCHES_DIR / OVERLAYS.get(filename, filename) - if not overlay_dir.is_dir(): - return archive_path.read_bytes() - if not filename.endswith(".tar.gz"): - raise ValueError(f"Unsupported overlay archive: {filename}") - return _apply_tar_overlay(archive_path, overlay_dir) - - -def _apply_tar_overlay(archive_path, overlay_dir): - with tarfile.open(archive_path, "r:gz") as source: - members = source.getmembers() - top_level = members[0].name.split("/", 1)[0] - buffer = io.BytesIO() - with tarfile.open(fileobj=buffer, mode="w:gz") as output: - for member in members: - if member.isdir() or member.islnk() or member.issym(): - output.addfile(member) - continue - relative = member.name[len(top_level) + 1:] - overlay_file = overlay_dir / relative - if overlay_file.is_file(): - info = tarfile.TarInfo(member.name) - info.size = overlay_file.stat().st_size - info.mode = member.mode - with overlay_file.open("rb") as overlay_stream: - output.addfile(info, overlay_stream) - else: - with source.extractfile(member) as extracted: - output.addfile(member, extracted) - return buffer.getvalue() - - def manifest_for_platform(platform): """Return the manifest text declaring only the given platform.""" text = (SOURCE_DIR / "blender_manifest.toml").read_text(encoding="utf-8") @@ -121,23 +74,13 @@ def manifest_for_platform(platform): def bundled_vendor_files(platform): - entries = { - name: VENDOR_DIR / filename - for name, filename in BUNDLED_ARCHIVES.items() + return { + f"vendor/{PLATFORMS[platform]}": VENDOR_DIR / PLATFORMS[platform] } - entries[f"vendor/{PLATFORMS[platform]}"] = ( - VENDOR_DIR / PLATFORMS[platform] - ) - return entries def validate_vendor_files(platform): """Validate common archives and the selected uv archive.""" - for artifact in VENDOR_LOCK["artifacts"].values(): - validate_artifact( - VENDOR_DIR / artifact["filename"], - artifact["sha256"], - ) uv_artifact = VENDOR_LOCK["platforms"][platform] validate_artifact( VENDOR_DIR / uv_artifact["filename"], @@ -174,7 +117,7 @@ def build_extension(platform): for archive_name, source_path in vendor_files.items(): archive.writestr( archive_name, - bundled_artifact_bytes(source_path.name), + source_path.read_bytes(), ) return output_path diff --git a/patches/depth-anything-3/pyproject.toml b/patches/depth-anything-3/pyproject.toml deleted file mode 100644 index e0b6dc6..0000000 --- a/patches/depth-anything-3/pyproject.toml +++ /dev/null @@ -1,93 +0,0 @@ -[build-system] -requires = ["hatchling>=1.25", "hatch-vcs>=0.4"] -build-backend = "hatchling.build" - -[project] -name = "depth-anything-3" -version = "0.0.0" -description = "Depth Anything 3" -readme = "README.md" -requires-python = ">=3.9, <=3.13" -license = { text = "Apache-2.0" } -authors = [{ name = "Your Name" }] - -dependencies = [ - "pre-commit", - "trimesh", - "torch>=2", - "torchvision", - "einops", - "huggingface_hub", - "imageio", - "numpy<2", - "opencv-python", - "open3d", - "fastapi", - "uvicorn", - "requests", - "typer", - "pillow", - "omegaconf", - "evo", - "e3nn", - "moviepy", - "plyfile", - "pillow_heif", - "safetensors", - "uvicorn", - "moviepy==1.0.3", - "typer>=0.9.0", - "pycolmap", -] - -[project.optional-dependencies] -app = ["gradio>=5", "pillow>=9.0"] # requires that python3>=3.10 -gs = ["gsplat @ git+https://github.com/nerfstudio-project/gsplat.git@0b4dddf04cb687367602c01196913cde6a743d70"] -all = ["depth-anything-3[app,gs]"] - - -[project.scripts] -da3 = "depth_anything_3.cli:app" - -[project.urls] -Homepage = "https://github.com/ByteDance-Seed/Depth-Anything-3" - -[tool.hatch.version] -source = "vcs" - -[tool.hatch.build.targets.wheel] -packages = ["src/depth_anything_3"] - -[tool.hatch.build.targets.sdist] -include = [ - "/README.md", - "/pyproject.toml", - "/src/depth_anything_3", -] - -[tool.hatch.metadata] -allow-direct-references = true - -[tool.mypy] -plugins = ["jaxtyping.mypy_plugin"] - -[tool.black] -line-length = 99 -target-version = ['py37', 'py38', 'py39', 'py310', 'py311'] -include = '\.pyi?$' -exclude = ''' -/( - | \.git -)/ -''' - -[tool.isort] -profile = "black" -multi_line_output = 3 -include_trailing_comma = true -known_third_party = ["bson","cruise","cv2","dataloader","diffusers","omegaconf","tensorflow","torch","torchvision","transformers","gsplat"] -known_first_party = ["common", "data", "models", "projects"] -sections = ["FUTURE","STDLIB","THIRDPARTY","FIRSTPARTY","LOCALFOLDER"] -skip_gitignore = true -line_length = 99 -no_lines_before="THIRDPARTY" diff --git a/patches/depth-anything-3/requirements.txt b/patches/depth-anything-3/requirements.txt deleted file mode 100644 index a10d6b6..0000000 --- a/patches/depth-anything-3/requirements.txt +++ /dev/null @@ -1,23 +0,0 @@ -pre-commit -trimesh -torch>=2 -torchvision -einops -huggingface_hub -imageio -numpy<2 -opencv-python -open3d -fastapi -uvicorn -requests -typer -pillow -omegaconf -evo -e3nn -moviepy -plyfile -pillow_heif -safetensors -pycolmap diff --git a/patches/depth-anything-3/src/depth_anything_3/utils/export/__init__.py b/patches/depth-anything-3/src/depth_anything_3/utils/export/__init__.py deleted file mode 100644 index 9f774c1..0000000 --- a/patches/depth-anything-3/src/depth_anything_3/utils/export/__init__.py +++ /dev/null @@ -1,84 +0,0 @@ -# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from depth_anything_3.specs import Prediction - -# Canonical set of export formats this dispatcher understands. Treat this as the -# server-trusted enum for any request-facing validation (e.g. the HTTP backend) instead -# of letting a client-supplied export_format string reach the dispatcher unchecked. -SUPPORTED_EXPORT_FORMATS = frozenset( - { - "glb", - "mini_npz", - "npz", - "feat_vis", - "depth_vis", - "gs_ply", - "gs_video", - "colmap", - } -) - - -def export( - prediction: Prediction, - export_format: str, - export_dir: str, - **kwargs, -): - if "-" in export_format: - export_formats = export_format.split("-") - for export_format in export_formats: - export(prediction, export_format, export_dir, **kwargs) - return # Prevent falling through to single-format handling - - if export_format == "glb": - from .glb import export_to_glb - - export_to_glb(prediction, export_dir, **kwargs.get(export_format, {})) - elif export_format == "mini_npz": - from .npz import export_to_mini_npz - - export_to_mini_npz(prediction, export_dir) - elif export_format == "npz": - from .npz import export_to_npz - - export_to_npz(prediction, export_dir) - elif export_format == "feat_vis": - from .feat_vis import export_to_feat_vis - - export_to_feat_vis(prediction, export_dir, **kwargs.get(export_format, {})) - elif export_format == "depth_vis": - from .depth_vis import export_to_depth_vis - - export_to_depth_vis(prediction, export_dir) - elif export_format == "gs_ply": - from .gs import export_to_gs_ply - - export_to_gs_ply(prediction, export_dir, **kwargs.get(export_format, {})) - elif export_format == "gs_video": - from .gs import export_to_gs_video - - export_to_gs_video(prediction, export_dir, **kwargs.get(export_format, {})) - elif export_format == "colmap": - from .colmap import export_to_colmap - - export_to_colmap(prediction, export_dir, **kwargs.get(export_format, {})) - else: - raise ValueError(f"Unsupported export format: {export_format}") - - -__all__ = [ - export, -] diff --git a/patches/depth-anything-3/src/depth_anything_3/utils/ray_utils.py b/patches/depth-anything-3/src/depth_anything_3/utils/ray_utils.py deleted file mode 100644 index b92434f..0000000 --- a/patches/depth-anything-3/src/depth_anything_3/utils/ray_utils.py +++ /dev/null @@ -1,531 +0,0 @@ -# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import torch -from einops import repeat -from .geometry import unproject_depth - - -def compute_optimal_rotation_intrinsics_batch( - rays_origin, rays_target, z_threshold=1e-4, reproj_threshold=0.2, weights=None, - n_sample = None, - n_iter=100, - num_sample_for_ransac=8, - rand_sample_iters_idx=None, -): - """ - Args: - rays_origin (torch.Tensor): (B, N, 3) - rays_target (torch.Tensor): (B, N, 3) - z_threshold (float): Threshold for z value to be considered valid. - - Returns: - R (torch.tensor): (3, 3) - focal_length (torch.tensor): (2,) - principal_point (torch.tensor): (2,) - """ - device = rays_origin.device - B, N, _ = rays_origin.shape - z_mask = torch.logical_and( - torch.abs(rays_target[:, :, 2]) > z_threshold, torch.abs(rays_origin[:, :, 2]) > z_threshold - ) # (B, N, 1) - rays_origin = rays_origin.clone() - rays_target = rays_target.clone() - rays_origin[:, :, 0][z_mask] /= rays_origin[:, :, 2][z_mask] - rays_origin[:, :, 1][z_mask] /= rays_origin[:, :, 2][z_mask] - rays_target[:, :, 0][z_mask] /= rays_target[:, :, 2][z_mask] - rays_target[:, :, 1][z_mask] /= rays_target[:, :, 2][z_mask] - - rays_origin = rays_origin[:, :, :2] - rays_target = rays_target[:, :, :2] - assert weights is not None, "weights must be provided" - weights[~z_mask] = 0 - - A_list = [] - max_chunk_size = 2 - for i in range(0, rays_origin.shape[0], max_chunk_size): - A = ransac_find_homography_weighted_fast_batch( - rays_origin[i:i+max_chunk_size], - rays_target[i:i+max_chunk_size], - weights[i:i+max_chunk_size], - n_iter=n_iter, - n_sample = n_sample, - num_sample_for_ransac=num_sample_for_ransac, - reproj_threshold=reproj_threshold, - rand_sample_iters_idx=rand_sample_iters_idx, - max_inlier_num=8000, - ) - A = A.to(device) - A_need_inv_mask = torch.linalg.det(A) < 0 - A[A_need_inv_mask] = -A[A_need_inv_mask] - A_list.append(A) - - A = torch.cat(A_list, dim=0) - - R_list = [] - f_list = [] - pp_list = [] - for i in range(A.shape[0]): - R, L = ql_decomposition(A[i]) - L = L / L[2][2] - - f = torch.stack((L[0][0], L[1][1])) - pp = torch.stack((L[2][0], L[2][1])) - R_list.append(R) - f_list.append(f) - pp_list.append(pp) - - R = torch.stack(R_list) - f = torch.stack(f_list) - pp = torch.stack(pp_list) - - return R, f, pp - - -# https://www.reddit.com/r/learnmath/comments/v1crd7/linear_algebra_qr_to_ql_decomposition/ -def ql_decomposition(A): - device = A.device - P = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]], device=device).float() - A_tilde = torch.matmul(A, P) - # MPS 尚未实现 torch.linalg.qr，仅把 3x3 分解放到 CPU 执行后移回 - use_cpu = device.type == "mps" - if use_cpu: - A_tilde = A_tilde.cpu() - Q_tilde, R_tilde = torch.linalg.qr(A_tilde) - if use_cpu: - Q_tilde = Q_tilde.to(device=device) - R_tilde = R_tilde.to(device=device) - Q = torch.matmul(Q_tilde, P) - L = torch.matmul(torch.matmul(P, R_tilde), P) - d = torch.diag(L) - Q[:, 0] *= torch.sign(d[0]) - Q[:, 1] *= torch.sign(d[1]) - Q[:, 2] *= torch.sign(d[2]) - L[0] *= torch.sign(d[0]) - L[1] *= torch.sign(d[1]) - L[2] *= torch.sign(d[2]) - return Q, L - -def find_homography_least_squares_weighted_torch(src_pts, dst_pts, confident_weight): - """ - src_pts: (N,2) source points (torch.Tensor, float32/float64) - dst_pts: (N,2) target points (torch.Tensor, float32/float64) - confident_weight: (N,) weights (torch.Tensor) - Returns: (3,3) homography matrix H (torch.Tensor) - """ - assert src_pts.shape == dst_pts.shape - N = src_pts.shape[0] - if N < 4: - raise ValueError("At least 4 points are required to compute homography.") - assert confident_weight.shape == (N,) - - w = confident_weight.sqrt().unsqueeze(1) # (N,1) - - x = src_pts[:, 0:1] # (N,1) - y = src_pts[:, 1:2] # (N,1) - u = dst_pts[:, 0:1] - v = dst_pts[:, 1:2] - - zeros = torch.zeros_like(x) - - # Construct A matrix (2N, 9) - A1 = torch.cat([-x * w, -y * w, -w, zeros, zeros, zeros, x * u * w, y * u * w, u * w], dim=1) - A2 = torch.cat([zeros, zeros, zeros, -x * w, -y * w, -w, x * v * w, y * v * w, v * w], dim=1) - A = torch.cat([A1, A2], dim=0) # (2N, 9) - - # SVD - # Note: torch.linalg.svd returns U, S, Vh, where Vh is the transpose of V - _, _, Vh = torch.linalg.svd(A) - H = Vh[-1].reshape(3, 3) - H = H / H[-1, -1] - return H - - -def ransac_find_homography_weighted( - src_pts, - dst_pts, - confident_weight, - n_iter=100, - sample_ratio=0.2, - reproj_threshold=3.0, - num_sample_for_ransac=16, - random_seed=None, -): - """ - RANSAC version of weighted Homography estimation. - Sample 4 points from the top 50% weighted points each time. - reproj_threshold: points with reprojection error less than this value are inliers - Returns: best_H - """ - if random_seed is not None: - torch.manual_seed(random_seed) - N = src_pts.shape[0] - assert N >= 4 - # 1. Select top 50% weighted points - sorted_idx = torch.argsort(confident_weight, descending=True) - n_sample = max(num_sample_for_ransac, int(N * sample_ratio)) - candidate_idx = sorted_idx[:n_sample] - best_inlier_mask = None - best_score = 0 - for _ in range(n_iter): - # 2. Randomly sample 4 points - idx = candidate_idx[torch.randperm(n_sample)[:num_sample_for_ransac]] - # 3. Compute Homography - try: - H = find_homography_least_squares_weighted_torch( - src_pts[idx], dst_pts[idx], confident_weight[idx] - ) - except Exception: - H = torch.eye(3, dtype=src_pts.dtype, device=src_pts.device) - # 4. Compute reprojection error for all points - src_homo = torch.cat( - [src_pts, torch.ones(N, 1, dtype=src_pts.dtype, device=src_pts.device)], dim=1 - ) - proj = (H @ src_homo.T).T - proj = proj[:, :2] / proj[:, 2:3] - error = ((proj - dst_pts) ** 2).sum(dim=1).sqrt() # Euclidean distance - inlier_mask = error < reproj_threshold - total_score = (inlier_mask * confident_weight).sum().item() - n_inlier = inlier_mask.sum().item() - if n_inlier < 4: - continue # At least 4 inliers required for fitting - - if total_score > best_score: - best_score = total_score - best_inlier_mask = inlier_mask - - # 5. Refit Homography using inliers - H_inlier = find_homography_least_squares_weighted_torch( - src_pts[best_inlier_mask], dst_pts[best_inlier_mask], confident_weight[best_inlier_mask] - ) - - return H_inlier - - -def find_homography_least_squares_weighted_torch_batch( - src_pts_batch, dst_pts_batch, confident_weight_batch -): - """ - Batch version of weighted least squares Homography - src_pts_batch: (B, K, 2) - dst_pts_batch: (B, K, 2) - confident_weight_batch: (B, K) - Returns: (B, 3, 3) - """ - B, K, _ = src_pts_batch.shape - w = confident_weight_batch.sqrt().unsqueeze(2) # (B,K,1) - x = src_pts_batch[:, :, 0:1] - y = src_pts_batch[:, :, 1:2] - u = dst_pts_batch[:, :, 0:1] - v = dst_pts_batch[:, :, 1:2] - zeros = torch.zeros_like(x) - A1 = torch.cat([-x * w, -y * w, -w, zeros, zeros, zeros, x * u * w, y * u * w, u * w], dim=2) - A2 = torch.cat([zeros, zeros, zeros, -x * w, -y * w, -w, x * v * w, y * v * w, v * w], dim=2) - A = torch.cat([A1, A2], dim=1) # (B, 2K, 9) - # SVD: torch.linalg.svd supports batch - _, _, Vh = torch.linalg.svd(A) - H = Vh[:, -1].reshape(B, 3, 3) - H = H / H[:, 2:3, 2:3] - return H - - -def ransac_find_homography_weighted_fast( - src_pts, - dst_pts, - confident_weight, - n_sample, - n_iter=100, - reproj_threshold=3.0, - num_sample_for_ransac=8, - random_seed=None, - rand_sample_iters_idx=None, -): - """ - Batch version of RANSAC weighted Homography estimation. - Returns: H_inlier - """ - if random_seed is not None: - torch.manual_seed(random_seed) - N = src_pts.shape[0] - device = src_pts.device - assert N >= 4 - # 1. Select top weighted points by sample_ratio - sorted_idx = torch.argsort(confident_weight, descending=True) - candidate_idx = sorted_idx[:n_sample] # (n_sample,) - if rand_sample_iters_idx is None: - rand_sample_iters_idx = torch.stack( - [torch.randperm(n_sample, device=device)[:num_sample_for_ransac] for _ in range(n_iter)], - dim=0, - ) # (n_iter, num_sample_for_ransac) - # 2. Generate all sampling groups at once - # shape: (n_iter, num_sample_for_ransac) - rand_idx = candidate_idx[rand_sample_iters_idx] # (n_iter, num_sample_for_ransac) - # 3. Construct batch input - src_pts_batch = src_pts[rand_idx] # (n_iter, num_sample_for_ransac, 2) - dst_pts_batch = dst_pts[rand_idx] # (n_iter, num_sample_for_ransac, 2) - confident_weight_batch = confident_weight[rand_idx] # (n_iter, num_sample_for_ransac) - # 4. Batch fit Homography - H_batch = find_homography_least_squares_weighted_torch_batch( - src_pts_batch, dst_pts_batch, confident_weight_batch - ) # (n_iter, 3, 3) - # 5. Batch evaluate inliers for all H - src_homo = torch.cat( - [src_pts, torch.ones(N, 1, dtype=src_pts.dtype, device=src_pts.device)], dim=1 - ) # (N,3) - src_homo_expand = src_homo.unsqueeze(0).expand(n_iter, N, 3) # (n_iter, N, 3) - dst_pts_expand = dst_pts.unsqueeze(0).expand(n_iter, N, 2) # (n_iter, N, 2) - confident_weight_expand = confident_weight.unsqueeze(0).expand(n_iter, N) # (n_iter, N) - # H_batch: (n_iter, 3, 3) - proj = torch.bmm(src_homo_expand, H_batch.transpose(1, 2)) # (n_iter, N, 3) - proj_xy = proj[:, :, :2] / proj[:, :, 2:3] # (n_iter, N, 2) - error = ((proj_xy - dst_pts_expand) ** 2).sum(dim=2).sqrt() # (n_iter, N) - inlier_mask = error < reproj_threshold # (n_iter, N) - total_score = (inlier_mask * confident_weight_expand).sum(dim=1) # (n_iter,) - # 6. Select the sampling group with the highest score - best_idx = torch.argmax(total_score) - best_inlier_mask = inlier_mask[best_idx] # (N,) - inlier_src_pts = src_pts[best_inlier_mask] - inlier_dst_pts = dst_pts[best_inlier_mask] - inlier_confident_weight = confident_weight[best_inlier_mask] - - max_inlier_num = 10000 - sorted_idx = torch.argsort(inlier_confident_weight, descending=True) - - # method 1: sort according to confident_weight, and only keep max_inlier_num pts - # sorted_idx = sorted_idx[:max_inlier_num] - - # method 2: random choose max_inlier_num pts - sorted_idx = sorted_idx[torch.randperm(len(sorted_idx))[:max_inlier_num]] - - inlier_src_pts = inlier_src_pts[sorted_idx] - inlier_dst_pts = inlier_dst_pts[sorted_idx] - inlier_confident_weight = inlier_confident_weight[sorted_idx] - # 7. Refit Homography using inliers - H_inlier = find_homography_least_squares_weighted_torch( - inlier_src_pts, inlier_dst_pts, inlier_confident_weight - ) - return H_inlier - - -def ransac_find_homography_weighted_fast_batch( - src_pts, # (B, N, 3) - dst_pts, # (B, N, 2) - confident_weight, # (B, N) - n_sample, - n_iter=100, - reproj_threshold=3.0, - num_sample_for_ransac=8, - max_inlier_num=10000, - random_seed=None, - rand_sample_iters_idx=None, -): - """ - Batch version of RANSAC weighted Homography estimation (supports batch). - Input: - src_pts: (B, N, 2) - dst_pts: (B, N, 2) - confident_weight: (B, N) - Returns: - H_inlier: (B, 3, 3) - """ - if random_seed is not None: - torch.manual_seed(random_seed) - B, N, _ = src_pts.shape - assert N >= 4 - - device = src_pts.device - - # 1. Select top weighted points by sample_ratio - sorted_idx = torch.argsort(confident_weight, descending=True, dim=1) # (B, N) - candidate_idx = sorted_idx[:, :n_sample] # (B, n_sample) - - # 2. Generate all sampling groups at once - # rand_idx: (B, n_iter, num_sample_for_ransac) - if rand_sample_iters_idx is None: - rand_sample_iters_idx = torch.stack( - [torch.randperm(n_sample, device=device)[:num_sample_for_ransac] for _ in range(n_iter)], - dim=0, - ) # (n_iter, num_sample_for_ransac) - - rand_idx = candidate_idx[:, rand_sample_iters_idx] # (B, n_iter, num_sample_for_ransac) - - # 3. Construct batch input - # Indexing method below: (B, n_iter, num_sample_for_ransac, ...) - b_idx = torch.arange(B, device=device).view(B, 1, 1).expand(B, n_iter, num_sample_for_ransac) - src_pts_batch = src_pts[b_idx, rand_idx] # (B, n_iter, num_sample_for_ransac, 2) - dst_pts_batch = dst_pts[b_idx, rand_idx] # (B, n_iter, num_sample_for_ransac, 2) - confident_weight_batch = confident_weight[b_idx, rand_idx] # (B, n_iter, num_sample_for_ransac) - - # 4. Batch fit Homography - # Need to implement batch version that supports (B, n_iter, num_sample_for_ransac, ...) input - # Output H_batch: (B, n_iter, 3, 3) - cB, cN = src_pts_batch.shape[:2] - H_batch = find_homography_least_squares_weighted_torch_batch( - src_pts_batch.flatten(0, 1), dst_pts_batch.flatten(0, 1), confident_weight_batch.flatten(0, 1) - ) # (B, n_iter, 3, 3) - H_batch = H_batch.unflatten(0, (cB, cN)) - - # 5. Batch evaluate inliers for all H - src_homo = torch.cat( - [src_pts, torch.ones(B, N, 1, dtype=src_pts.dtype, device=src_pts.device)], dim=2 - ) # (B, N, 3) - src_homo_expand = src_homo.unsqueeze(1).expand(B, n_iter, N, 3) # (B, n_iter, N, 3) - dst_pts_expand = dst_pts.unsqueeze(1).expand(B, n_iter, N, 2) # (B, n_iter, N, 2) - confident_weight_expand = confident_weight.unsqueeze(1).expand(B, n_iter, N) # (B, n_iter, N) - - # H_batch: (B, n_iter, 3, 3) - # Need to reshape H_batch to (B*n_iter, 3, 3), src_homo_expand to (B*n_iter, N, 3) - H_batch_flat = H_batch.reshape(-1, 3, 3) - src_homo_expand_flat = src_homo_expand.reshape(-1, N, 3) - proj = torch.bmm(src_homo_expand_flat, H_batch_flat.transpose(1, 2)) # (B*n_iter, N, 3) - proj_xy = proj[:, :, :2] / proj[:, :, 2:3] # (B*n_iter, N, 2) - proj_xy = proj_xy.reshape(B, n_iter, N, 2) - error = ((proj_xy - dst_pts_expand) ** 2).sum(dim=3).sqrt() # (B, n_iter, N) - inlier_mask = error < reproj_threshold # (B, n_iter, N) - total_score = (inlier_mask * confident_weight_expand).sum(dim=2) # (B, n_iter) - - # 6. Select the sampling group with the highest score - best_idx = torch.argmax(total_score, dim=1) # (B,) - best_inlier_mask = inlier_mask[torch.arange(B, device=device), best_idx] # (B, N) - - # 7. Refit Homography using inliers - H_inlier_list = [] - for b in range(B): - mask = best_inlier_mask[b] - inlier_src_pts = src_pts[b][mask] # (?, 3) - inlier_dst_pts = dst_pts[b][mask] # (?, 2) - inlier_confident_weight = confident_weight[b][mask] # (?) - - sorted_idx = torch.argsort(inlier_confident_weight, descending=True) - # # method 1: sort according to confident_weight, and only keep max_inlier_num pts - # sorted_idx = sorted_idx[:max_inlier_num] - # method 2: random choose max_inlier_num pts - if len(sorted_idx) > max_inlier_num: - # random choose from first 95% confident pts - keep_len = max(int(len(sorted_idx) * 0.95), max_inlier_num) - sorted_idx = sorted_idx[:keep_len] - perm = torch.randperm(len(sorted_idx), device=device)[:max_inlier_num] - sorted_idx = sorted_idx[perm] - inlier_src_pts = inlier_src_pts[sorted_idx] - inlier_dst_pts = inlier_dst_pts[sorted_idx] - inlier_confident_weight = inlier_confident_weight[sorted_idx] - - H_inlier = find_homography_least_squares_weighted_torch( - inlier_src_pts, inlier_dst_pts, inlier_confident_weight - ) # (3, 3) - H_inlier_list.append(H_inlier) - H_inlier = torch.stack(H_inlier_list, dim=0) # (B, 3, 3) - return H_inlier - -def get_params_for_ransac(N, device): - n_iter=100 - sample_ratio=0.3 - num_sample_for_ransac=8 - n_sample = max(num_sample_for_ransac, int(N * sample_ratio)) - rand_sample_iters_idx = torch.stack( - [torch.randperm(n_sample, device=device)[:num_sample_for_ransac] for _ in range(n_iter)], - dim=0, - ) # (n_iter, num_sample_for_ransac) - return n_iter, num_sample_for_ransac, n_sample, rand_sample_iters_idx - - -def camray_to_caminfo(camray, confidence=None, reproj_threshold=0.2, training=False): - """ - Args: - camray: (B, S, num_patches_y, num_patches_x, 6) - confidence: (B, S, num_patches_y, num_patches_x) - Returns: - R: (B, S, 3, 3) - T: (B, S, 3) - focal_lengths: (B, S, 2) - principal_points: (B, S, 2) - """ - if confidence is None: - confidence = torch.ones_like(camray[:, :, :, :, 0]) - B, S, num_patches_y, num_patches_x, _ = camray.shape - # identity K, assume imw=imh=2.0 - I_K = torch.eye(3, dtype=camray.dtype, device=camray.device) - I_K[0, 2] = 1.0 - I_K[1, 2] = 1.0 - # repeat I_K to match camray - I_K = I_K.unsqueeze(0).unsqueeze(0).expand(B, S, -1, -1) - - cam_plane_depth = torch.ones( - B, S, num_patches_y, num_patches_x, 1, dtype=camray.dtype, device=camray.device - ) - I_cam_plane_unproj = unproject_depth( - cam_plane_depth, - I_K, - c2w=None, - ixt_normalized=True, - num_patches_x=num_patches_x, - num_patches_y=num_patches_y, - ) # (B, S, num_patches_y, num_patches_x, 3) - - camray = camray.flatten(0, 1).flatten(1, 2) # (B*S, num_patches_y*num_patches_x, 6) - I_cam_plane_unproj = I_cam_plane_unproj.flatten(0, 1).flatten( - 1, 2 - ) # (B*S, num_patches_y*num_patches_x, 3) - confidence = confidence.flatten(0, 1).flatten(1, 2) # (B*S, num_patches_y*num_patches_x) - - # Compute optimal rotation to align rays - N = camray.shape[-2] - device = camray.device - n_iter, num_sample_for_ransac, n_sample, rand_sample_iters_idx = get_params_for_ransac(N, device) - - # Use batch processing (confidence is guaranteed to be not None at this point) - if training: - camray = camray.clone().detach() - I_cam_plane_unproj = I_cam_plane_unproj.clone().detach() - confidence = confidence.clone().detach() - R, focal_lengths, principal_points = compute_optimal_rotation_intrinsics_batch( - I_cam_plane_unproj, - camray[:, :, :3], - reproj_threshold=reproj_threshold, - weights=confidence, - n_sample = n_sample, - n_iter=n_iter, - num_sample_for_ransac=num_sample_for_ransac, - rand_sample_iters_idx=rand_sample_iters_idx, - ) - - T = torch.sum(camray[:, :, 3:] * confidence.unsqueeze(-1), dim=1) / torch.sum( - confidence, dim=-1, keepdim=True - ) - - R = R.reshape(B, S, 3, 3) - T = T.reshape(B, S, 3) - focal_lengths = focal_lengths.reshape(B, S, 2) - principal_points = principal_points.reshape(B, S, 2) - - return R, T, 1.0 / focal_lengths, principal_points + 1.0 - -def get_extrinsic_from_camray(camray, conf, patch_size_y, patch_size_x, training=False): - pred_R, pred_T, pred_focal_lengths, pred_principal_points = camray_to_caminfo( - camray, confidence=conf.squeeze(-1), training=training - ) - - pred_extrinsic = torch.cat( - [ - torch.cat([pred_R, pred_T.unsqueeze(-1)], dim=-1), - repeat( - torch.tensor([0, 0, 0, 1], dtype=pred_R.dtype, device=pred_R.device), - "c -> b s 1 c", - b=pred_R.shape[0], - s=pred_R.shape[1], - ), - ], - dim=-2, - ) # B, S, 4, 4 - return pred_extrinsic, pred_focal_lengths, pred_principal_points diff --git a/patches/moge/pyproject.toml b/patches/moge/pyproject.toml deleted file mode 100644 index c6b7a3a..0000000 --- a/patches/moge/pyproject.toml +++ /dev/null @@ -1,22 +0,0 @@ -[build-system] -requires = ["setuptools>=61.0", "wheel"] -build-backend = "setuptools.build_meta" - -[project] -name = "moge" -version = "2.0.0" -description = "MoGe monocular geometry estimation" -readme = "README.md" -license = {text = "MIT"} -dependencies = [ - "huggingface_hub", - "numpy", - "pillow", - "torch>=2.0.0", - "utils3d", -] -requires-python = ">=3.9" - -[tool.setuptools.packages.find] -where = ["."] -include = ["moge*"] diff --git a/patches/utils3d/pyproject.toml b/patches/utils3d/pyproject.toml deleted file mode 100644 index d0898ed..0000000 --- a/patches/utils3d/pyproject.toml +++ /dev/null @@ -1,16 +0,0 @@ -[build-system] -requires = ["setuptools>=61.0", "wheel"] -build-backend = "setuptools.build_meta" - -[project] -name = "utils3d" -version = "1.7" -description = "Utilities for 3D computer vision" -readme = "README.md" -license = {text = "MIT"} -dependencies = ["numpy", "scipy"] -requires-python = ">=3.8" - -[tool.setuptools.packages.find] -where = ["."] -include = ["utils3d*"] diff --git a/prepare_vendor.py b/prepare_vendor.py index dd09544..cacd9b4 100644 --- a/prepare_vendor.py +++ b/prepare_vendor.py @@ -65,7 +65,7 @@ def download_artifact(artifact, destination_dir=VENDOR_DIR): def prepare(platforms=None, lock_path=LOCK_PATH, destination_dir=VENDOR_DIR): - """Prepare common third-party archives for selected platforms.""" + """Prepare bundled uv executables for selected platforms.""" lock = load_lock(lock_path) available = lock["platforms"] selected = list(available) if platforms is None else platforms @@ -73,8 +73,7 @@ def prepare(platforms=None, lock_path=LOCK_PATH, destination_dir=VENDOR_DIR): if unknown: raise ValueError(f"Unsupported platform: {', '.join(unknown)}") - artifacts = list(lock["artifacts"].values()) - artifacts.extend(available[platform] for platform in selected) + artifacts = [available[platform] for platform in selected] return [ download_artifact(artifact, destination_dir) for artifact in artifacts diff --git a/src/snapform/__init__.py b/src/snapform/__init__.py index cb4c813..8a45708 100644 --- a/src/snapform/__init__.py +++ b/src/snapform/__init__.py @@ -21,9 +21,9 @@ from .operators import ( ) from .panel import ( Ben2RemovalDebugPanel, + AiServicePanel, Da3OutputDebugPanel, DebugPanel, - AiServicePanel, ModelsPanel, Moge2ApiDebugPanel, ) @@ -50,8 +50,8 @@ CLASSES = ( AiServicePanel, ModelsPanel, DebugPanel, - Da3OutputDebugPanel, Moge2ApiDebugPanel, + Da3OutputDebugPanel, Ben2RemovalDebugPanel, ) diff --git a/src/snapform/ben2_onnx.py b/src/snapform/ben2_onnx.py new file mode 100644 index 0000000..9ab9803 --- /dev/null +++ b/src/snapform/ben2_onnx.py @@ -0,0 +1,50 @@ +"""BEN2 ONNX inference with Pillow and NumPy preprocessing.""" + +from pathlib import Path + +try: + from .onnx_runtime import create_session as create_runtime_session +except ImportError: + from onnx_runtime import create_session as create_runtime_session + + +ONNX_PATH = Path("onnx/model_fp16.onnx") +INPUT_SIZE = (1024, 1024) + + +def model_path(model_directory): + return Path(model_directory) / ONNX_PATH + + +def create_session(model_directory, requested_device="auto"): + return create_runtime_session(model_path(model_directory), requested_device) + + +def preprocess(image): + import numpy as np + from PIL import Image + + resized = image.convert("RGB").resize(INPUT_SIZE, Image.Resampling.LANCZOS) + array = np.asarray(resized, dtype=np.float32) / 255.0 + mean = np.array((0.485, 0.456, 0.406), dtype=np.float32) + std = np.array((0.229, 0.224, 0.225), dtype=np.float32) + array = (array - mean) / std + return np.ascontiguousarray(array.transpose(2, 0, 1)[None]) + + +def infer(session, image): + import numpy as np + from PIL import Image + + input_name = session.get_inputs()[0].name + output_name = session.get_outputs()[0].name + alpha = session.run([output_name], {input_name: preprocess(image)})[0] + alpha = np.asarray(alpha, dtype=np.float32).squeeze() + minimum = float(alpha.min()) + span = max(float(alpha.max()) - minimum, np.finfo(np.float32).eps) + alpha = np.clip((alpha - minimum) / span * 255.0, 0.0, 255.0) + mask = Image.fromarray(alpha.astype(np.uint8), mode="L") + mask = mask.resize(image.size, Image.Resampling.LANCZOS) + foreground = image.convert("RGBA") + foreground.putalpha(mask) + return foreground diff --git a/src/snapform/client.py b/src/snapform/client.py index a24f070..8a6f30c 100644 --- a/src/snapform/client.py +++ b/src/snapform/client.py @@ -26,10 +26,8 @@ from .models import ( ) -DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4" -BEN2_COMMIT = "2c99a5da477b5523585bfa5c893888a6e818a8f6" MANIFEST_FILENAME = "manifest.json" -MANIFEST_REVISION = 5 +MANIFEST_REVISION = 7 SERVICE_ADDRESS = "127.0.0.1" DEFAULT_STORAGE_ROOT = Path.home() / ".snapform" JOB_FOLDER_NAMES = { @@ -172,8 +170,8 @@ def env_devices(): def machine_devices(): """Devices the machine can run before environment installation.""" devices = ["cpu"] - if shutil.which("nvidia-smi"): - devices.insert(0, "cuda") + if sys.platform == "win32": + devices.insert(0, "directml") elif sys.platform == "darwin" and platform.machine().lower() in { "arm64", "aarch64", @@ -189,6 +187,7 @@ def supported_devices(): DEVICE_LABELS = { "cuda": ("CUDA", "NVIDIA CUDA"), + "directml": ("DirectML", "Windows GPU (DirectML)"), "mps": ("MPS", "Apple Metal (MPS)"), "cpu": ("CPU", "CPU"), } @@ -283,7 +282,6 @@ def uv_cache_path(): def install_env_command(status_path): executable = uv_executable() - torch_variant = "cu128" if "cuda" in machine_devices() else "cpu" script = Path(__file__).with_name("install_env.py") command = [ sys.executable, @@ -296,12 +294,6 @@ def install_env_command(status_path): str(storage_root() / MANIFEST_FILENAME), "--models-dir", str(models_directory()), - "--torch-variant", - torch_variant, - "--da3-commit", - DA3_COMMIT, - "--ben2-commit", - BEN2_COMMIT, "--status", str(status_path), ] diff --git a/src/snapform/da3_onnx.py b/src/snapform/da3_onnx.py new file mode 100644 index 0000000..20c211d --- /dev/null +++ b/src/snapform/da3_onnx.py @@ -0,0 +1,221 @@ +"""Depth Anything 3 inference with ONNX Runtime and NumPy.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np + +try: + from .onnx_runtime import create_session as create_onnx_session +except ImportError: + from onnx_runtime import create_session as create_onnx_session + + +ONNX_PATH = "model.onnx" +PATCH_SIZE = 14 +STRATEGY_CODES = { + "first": 0, + "middle": 1, + "saddle_balanced": 2, + "saddle_sim_range": 3, +} +IMAGE_MEAN = np.asarray((0.485, 0.456, 0.406), dtype=np.float32) +IMAGE_STD = np.asarray((0.229, 0.224, 0.225), dtype=np.float32) + + +@dataclass +class Prediction: + depth: np.ndarray + conf: np.ndarray + extrinsics: np.ndarray + intrinsics: np.ndarray + processed_images: np.ndarray + is_metric: int = 0 + sky: np.ndarray | None = None + gaussians: object | None = None + aux: dict | None = None + scale_factor: float | None = None + + +def create_session(model_dir, requested_device="auto"): + path = Path(model_dir) / ONNX_PATH + return create_onnx_session(path, requested_device) + + +def nearest_multiple(value, multiple=PATCH_SIZE): + lower = value // multiple * multiple + upper = lower + multiple + return upper if abs(upper - value) <= abs(value - lower) else lower + + +def target_size(size, process_res, method): + width, height = size + boundary = max(width, height) + if method.startswith("lower_bound"): + boundary = min(width, height) + scale = float(process_res) / boundary + width = max(1, int(round(width * scale))) + height = max(1, int(round(height * scale))) + if method.endswith("resize"): + return ( + max(PATCH_SIZE, nearest_multiple(width)), + max(PATCH_SIZE, nearest_multiple(height)), + ) + if method.endswith("crop"): + return ( + max(PATCH_SIZE, width // PATCH_SIZE * PATCH_SIZE), + max(PATCH_SIZE, height // PATCH_SIZE * PATCH_SIZE), + ) + raise ValueError(f"Unsupported process_res_method: {method}") + + +def resize_and_crop(image, process_res, method): + from PIL import Image + width, height = image.size + boundary = max(width, height) + if method.startswith("lower_bound"): + boundary = min(width, height) + scale = float(process_res) / boundary + resized_size = ( + max(1, int(round(width * scale))), + max(1, int(round(height * scale))), + ) + resampling = ( + Image.Resampling.BICUBIC + if scale > 1.0 + else Image.Resampling.BOX + ) + image = image.resize(resized_size, resample=resampling) + final_size = target_size((width, height), process_res, method) + if method.endswith("resize"): + if image.size != final_size: + upscale = final_size[0] > image.width or final_size[1] > image.height + resampling = ( + Image.Resampling.BICUBIC + if upscale + else Image.Resampling.BOX + ) + image = image.resize(final_size, resample=resampling) + return image + left = (image.width - final_size[0]) // 2 + top = (image.height - final_size[1]) // 2 + return image.crop((left, top, left + final_size[0], top + final_size[1])) + + +def preprocess(image_paths, process_res=504, process_res_method="upper_bound_resize"): + from PIL import Image + processed = [] + for path in image_paths: + with Image.open(path) as source: + image = resize_and_crop( + source.convert("RGB"), + process_res, + process_res_method, + ) + processed.append(np.asarray(image, dtype=np.uint8)) + min_height = min(image.shape[0] for image in processed) + min_width = min(image.shape[1] for image in processed) + cropped = [] + for image in processed: + top = (image.shape[0] - min_height) // 2 + left = (image.shape[1] - min_width) // 2 + cropped.append( + image[top : top + min_height, left : left + min_width] + ) + images = np.stack(cropped) + normalized = images.astype(np.float32) / 255.0 + normalized = (normalized - IMAGE_MEAN) / IMAGE_STD + model_input = normalized.transpose(0, 3, 1, 2)[None] + return np.ascontiguousarray(model_input), images + + +def quaternion_to_matrix(quaternion): + i, j, k, real = np.moveaxis(quaternion, -1, 0) + two_scale = 2.0 / np.sum(quaternion * quaternion, axis=-1) + values = ( + 1 - two_scale * (j * j + k * k), + two_scale * (i * j - k * real), + two_scale * (i * k + j * real), + two_scale * (i * j + k * real), + 1 - two_scale * (i * i + k * k), + two_scale * (j * k - i * real), + two_scale * (i * k - j * real), + two_scale * (j * k + i * real), + 1 - two_scale * (i * i + j * j), + ) + return np.stack(values, axis=-1).reshape(quaternion.shape[:-1] + (3, 3)) + + +def camera_from_pose(pose_encoding, height, width): + translation = pose_encoding[..., :3] + rotation = quaternion_to_matrix(pose_encoding[..., 3:7]) + rotation_inverse = np.swapaxes(rotation, -1, -2) + translation_inverse = -rotation_inverse @ translation[..., None] + extrinsics = np.concatenate( + (rotation_inverse, translation_inverse), + axis=-1, + ).astype(np.float32) + + field_of_view_height = pose_encoding[..., 7] + field_of_view_width = pose_encoding[..., 8] + focal_y = (height / 2.0) / np.maximum( + np.tan(field_of_view_height / 2.0), + 1e-6, + ) + focal_x = (width / 2.0) / np.maximum( + np.tan(field_of_view_width / 2.0), + 1e-6, + ) + intrinsics = np.zeros(pose_encoding.shape[:-1] + (3, 3), dtype=np.float32) + intrinsics[..., 0, 0] = focal_x + intrinsics[..., 1, 1] = focal_y + intrinsics[..., 0, 2] = width / 2.0 + intrinsics[..., 1, 2] = height / 2.0 + intrinsics[..., 2, 2] = 1.0 + return extrinsics, intrinsics + + +def infer( + session, + image_paths, + process_res=504, + process_res_method="upper_bound_resize", + ref_view_strategy="saddle_balanced", + use_ray_pose=False, +): + del use_ray_pose + inputs = session.get_inputs() + model_input, processed_images = preprocess( + image_paths, + process_res, + process_res_method, + ) + strategy_code = np.asarray( + [STRATEGY_CODES.get(ref_view_strategy, 2)], + dtype=np.int64, + ) + input_names = [value.name for value in inputs] + outputs = session.run( + None, + { + input_names[0]: model_input, + input_names[1]: strategy_code, + }, + ) + depth, confidence, pose_encoding = outputs[:3] + height, width = depth.shape[-2:] + extrinsics, intrinsics = camera_from_pose( + pose_encoding[0], + height, + width, + ) + return Prediction( + depth=np.asarray(depth[0], dtype=np.float32), + conf=np.asarray(confidence[0], dtype=np.float32), + extrinsics=extrinsics, + intrinsics=intrinsics, + processed_images=processed_images, + aux={}, + ) diff --git a/src/snapform/image_depth.py b/src/snapform/image_depth.py index d12ce6d..d72401d 100644 --- a/src/snapform/image_depth.py +++ b/src/snapform/image_depth.py @@ -2239,12 +2239,12 @@ class ConvertImageToDepthPlane(bpy.types.Operator): name="Depth From", description="Model or original image used to build depth", items=image_depth_model_items, - default=1, + default=0, ) snapform_use_ray_pose: bpy.props.BoolProperty( name="Use Ray Pose", - description="Use ray-based camera pose estimation", - default=True, + description="Legacy task compatibility; ONNX uses the camera decoder", + default=False, ) snapform_process_res_factor: process_resolution_scale_property() snapform_moge_resolution_level: bpy.props.IntProperty( @@ -2278,14 +2278,8 @@ class ConvertImageToDepthPlane(bpy.types.Operator): layout = self.layout layout.prop(self, "snapform_model") if self.snapform_model in {"SMALL", "BASE"}: - layout.prop(self, "snapform_use_ray_pose") draw_process_resolution(layout, self, context.object.data) elif self.snapform_model.startswith("MOGE2_"): - from .client import model_ready - - if not model_ready(self.snapform_model): - layout.label(text="MoGe-2 model is not installed", icon="ERROR") - layout.label(text="Download it in Models") layout.prop(self, "snapform_moge_resolution_level") layout.prop(self, "snapform_remove_background") layout.prop(self, "snapform_unshaded") diff --git a/src/snapform/image_mesh.py b/src/snapform/image_mesh.py index ae53e6b..3c4e3b3 100644 --- a/src/snapform/image_mesh.py +++ b/src/snapform/image_mesh.py @@ -608,12 +608,12 @@ class ConvertImageToMesh(bpy.types.Operator): name="Depth From", description="Model used to estimate image depth", items=image_surface_model_items, - default=1, + default=0, ) snapform_use_ray_pose: bpy.props.BoolProperty( name="Use Ray Pose", - description="Use ray-based camera pose estimation", - default=True, + description="Legacy task compatibility; ONNX uses the camera decoder", + default=False, ) snapform_process_res_factor: process_resolution_scale_property() snapform_moge_resolution_level: bpy.props.IntProperty( @@ -670,7 +670,6 @@ class ConvertImageToMesh(bpy.types.Operator): if self.snapform_model.startswith("MOGE2_"): layout.prop(self, "snapform_moge_resolution_level") else: - layout.prop(self, "snapform_use_ray_pose") draw_process_resolution(layout, self, context.object.data) layout.prop(self, "snapform_create_camera") layout.prop(self, "snapform_remove_background") diff --git a/src/snapform/install_env.py b/src/snapform/install_env.py index 366c236..1a6cac6 100644 --- a/src/snapform/install_env.py +++ b/src/snapform/install_env.py @@ -1,5 +1,4 @@ import argparse -import filecmp import json import os import platform @@ -28,47 +27,36 @@ PYPI_MIRRORS = ( "https://mirrors.aliyun.com/pypi/simple/", ) PYPI_DEFAULT_INDEX = "https://pypi.org/simple/" -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" -MANIFEST_REVISION = 5 -TORCH_INDEXES = { - "cu128": "https://download.pytorch.org/whl/cu128", - "cpu": "https://download.pytorch.org/whl/cpu", -} +ONNXRUNTIME_VERSION = "1.24.4" +NUMPY_VERSION = "2.4.2" +PILLOW_VERSION = "12.1.1" +OPENCV_VERSION = "4.13.0.92" +TRIMESH_VERSION = "4.11.2" +HUGGINGFACE_HUB_VERSION = "1.4.1" +MANIFEST_REVISION = 7 _resolved_pypi_index = None DEFAULT_MODEL_KEYS = (MOGE2_MODEL.key, BEN2_MODEL.key) -RUNTIME_STAGE_COUNT = 3 +RUNTIME_STAGE_COUNT = 2 ENVIRONMENT_SLICE = (0.02, 0.06) -PYTORCH_SLICE = (0.06, 0.48) -XFORMERS_SLICE = (0.48, 0.62) -AI_COMPONENTS_SLICE = (0.62, 0.85) +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 INSTALL_ESTIMATES = { - "PyTorch": 8 * 60, - "AI components": 8 * 60, + "ONNX Runtime": 3 * 60, } INSTALL_ESTIMATE_RANGES = { - "PyTorch": PYTORCH_SLICE, - "AI components": (XFORMERS_SLICE[0], AI_COMPONENTS_SLICE[1]), + "ONNX Runtime": AI_COMPONENTS_SLICE, } CHECK_CODE = ( - "import json; import torch; " - "from depth_anything_3.api import DepthAnything3; " - "from ben2 import BEN_Base; " - "from moge.model.v2 import MoGeModel; " - "import fastapi, uvicorn; " - "mps = getattr(torch.backends, 'mps', None); " + "import json; import fastapi, uvicorn, onnxruntime; " + "import cv2, numpy, PIL, trimesh; " "print('SNAPFORM_DEVICES=' + json.dumps({" - "'torch': torch.__version__, " - "'cuda': torch.cuda.is_available(), " - "'mps': bool(mps is not None and mps.is_available())}))" + "'onnxruntime': onnxruntime.__version__, " + "'providers': onnxruntime.get_available_providers()}))" ) @@ -158,10 +146,6 @@ def estimated_remaining(label, progress, elapsed=0.0): return max(expected_seconds * remaining_fraction - elapsed, 0.0) -def effective_torch_index(torch_variant): - return os.environ.get("SNAPFORM_TORCH_INDEX") or TORCH_INDEXES.get(torch_variant) - - def resolve_indexes(args): """???????????????????????????????????""" global _resolved_pypi_index @@ -320,52 +304,13 @@ def venv_python(venv_dir): return Path(venv_dir) / "bin" / "python" -def install_stages(args, venv_dir, python, uv, vendor_dir=None): - # macOS 的 PyTorch wheel 只在 PyPI 上，download.pytorch.org 没有 macosx 构建 - torch_index = ( - None - if sys.platform == "darwin" - else effective_torch_index(args.torch_variant) +def install_stages(args, venv_dir, python, uv): + """Build the lightweight ONNX-only runtime and model download stages.""" + runtime_package = ( + f"onnxruntime-directml=={ONNXRUNTIME_VERSION}" + if sys.platform == "win32" + else f"onnxruntime=={ONNXRUNTIME_VERSION}" ) - vendor_dir = Path(vendor_dir or Path(venv_dir).parent / "vendor") - - def archive_requirement(project_name, filename): - local = Path(vendor_dir) / filename - if not local.is_file(): - raise RuntimeError( - "Bundled vendor artifact is missing; reinstall the extension: " - f"{local}" - ) - return f"{project_name} @ {local.as_uri()}" - - da3_requirement = archive_requirement( - "depth-anything-3", - "depth-anything-3.tar.gz", - ) - ben2_requirement = archive_requirement( - "ben2", - "ben2.tar.gz", - ) - utils3d_requirement = archive_requirement( - "utils3d", - "utils3d.tar.gz", - ) - moge_requirement = archive_requirement( - "moge", - "moge.tar.gz", - ) - torch_command = [ - uv, - "pip", - "install", - "-v", - "--python", - python, - f"torch=={TORCH_VERSION}", - f"torchvision=={TORCHVISION_VERSION}", - ] - if torch_index is not None: - torch_command.extend(["--index-url", torch_index]) stages = [ ( "Environment", @@ -374,13 +319,7 @@ def install_stages(args, venv_dir, python, uv, vendor_dir=None): "plain", ), ( - "PyTorch", - torch_command, - None, - "install", - ), - ( - "AI components", + "ONNX Runtime", [ uv, "pip", @@ -388,38 +327,19 @@ def install_stages(args, venv_dir, python, uv, vendor_dir=None): "-v", "--python", python, - da3_requirement, - f"addict=={ADDICT_VERSION}", - ben2_requirement, - utils3d_requirement, - moge_requirement, f"fastapi=={FASTAPI_VERSION}", + f"huggingface-hub=={HUGGINGFACE_HUB_VERSION}", + f"numpy=={NUMPY_VERSION}", + f"opencv-python-headless=={OPENCV_VERSION}", + runtime_package, + f"pillow=={PILLOW_VERSION}", + f"trimesh=={TRIMESH_VERSION}", f"uvicorn=={UVICORN_VERSION}", - "setuptools", - "wheel", ], None, "install", ), ] - if sys.platform != "darwin": - xformers_command = [ - uv, - "pip", - "install", - "-v", - "--python", - python, - f"xformers=={XFORMERS_VERSION}", - ] - if args.torch_variant == "cu128": - xformers_command.extend(["--index-url", torch_index]) - else: - xformers_command.append("--no-deps") - stages.insert( - 2, - ("AI components", xformers_command, None, "install"), - ) models = DOWNLOADABLE_MODELS worker_script = Path(__file__).with_name("worker.py") models_dir = Path(args.models_dir) @@ -475,13 +395,13 @@ def env_installed(venv_dir, manifest_path=None): ) -def installed_devices(args): +def installed_devices(): devices = [] machine = platform.machine().lower() if sys.platform == "darwin" and machine in {"arm64", "aarch64"}: devices.append("mps") - elif args.torch_variant.startswith("cu"): - devices.append("cuda") + elif sys.platform == "win32": + devices.append("directml") devices.append("cpu") return devices @@ -592,82 +512,6 @@ def bundled_uv_archive(): return path if path.is_file() else None -def vendor_overlay_directory(filename): - names = { - "depth-anything-3.tar.gz": "depth-anything-3", - "moge.tar.gz": "moge", - "utils3d.tar.gz": "utils3d", - } - overlay_name = names.get(filename) - if overlay_name is None: - return None - project_root = Path(__file__).resolve().parents[2] - overlay = project_root / "patches" / overlay_name - return overlay if overlay.is_dir() else None - - -def copy_archive_with_overlay(source, target, overlay): - """开发目录使用原始 vendor 时，应用与正式打包相同的源码覆盖。""" - with tarfile.open(source, "r:gz") as input_archive: - members = input_archive.getmembers() - top_level = members[0].name.split("/", 1)[0] - with tarfile.open(target, "w:gz") as output_archive: - for member in members: - if member.isdir() or member.islnk() or member.issym(): - output_archive.addfile(member) - continue - relative = member.name[len(top_level) + 1 :] - replacement = overlay / relative - if replacement.is_file(): - info = tarfile.TarInfo(member.name) - info.size = replacement.stat().st_size - info.mode = member.mode - with replacement.open("rb") as stream: - output_archive.addfile(info, stream) - continue - with input_archive.extractfile(member) as stream: - output_archive.addfile(member, stream) - - -def copy_vendor_inputs(bundled_dir, vendor_dir): - """Copy bundled installation inputs into a temporary directory.""" - for name in ( - "depth-anything-3.tar.gz", - "ben2.tar.gz", - "utils3d.tar.gz", - "moge.tar.gz", - ): - bundled = Path(bundled_dir) / name - if not bundled.is_file(): - continue - target = Path(vendor_dir) / name - overlay = vendor_overlay_directory(name) - if overlay is None and target.is_file() and filecmp.cmp( - bundled, - target, - shallow=False, - ): - continue - Path(vendor_dir).mkdir(parents=True, exist_ok=True) - if overlay is not None and tarfile.is_tarfile(bundled): - copy_archive_with_overlay(bundled, target, overlay) - else: - shutil.copy2(bundled, target) - - -def remove_legacy_vendor(storage_root): - """Remove installation inputs persisted by older development builds.""" - root = Path(storage_root) - shutil.rmtree(root / "vendor", ignore_errors=True) - for name in ( - "depth-anything-3.tar.gz", - "ben2.tar.gz", - "utils3d.tar.gz", - "moge.tar.gz", - ): - (root / name).unlink(missing_ok=True) - - def archive_is_valid(path): if zipfile.is_zipfile(path): return True @@ -800,14 +644,6 @@ def run_stage( 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" - ) resolve_indexes(args) venv_dir = Path(args.venv_dir) manifest_path = Path( @@ -823,13 +659,7 @@ def install(args): uv = Path(resolved) if resolved else uv uv_missing = not uv.is_file() python = venv_python(venv_dir) - vendor_temp = tempfile.TemporaryDirectory(prefix="snapform_vendor_") - vendor_dir = Path(vendor_temp.name) - copy_vendor_inputs( - bundled_vendor_directory(), - vendor_dir, - ) - stages = install_stages(args, venv_dir, python, uv, vendor_dir) + 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 = ( @@ -838,23 +668,15 @@ def install(args): else 0.0 ) installed_packages = [] - component_stages = env_stages[2:] - component_ranges = ( - [XFORMERS_SLICE, AI_COMPONENTS_SLICE] - if len(component_stages) == 2 - else [(XFORMERS_SLICE[0], AI_COMPONENTS_SLICE[1])] - ) 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 = PYTORCH_SLICE + stage_start, stage_end = AI_COMPONENTS_SLICE else: - stage = 3 - component_index = index - 2 - stage_start, stage_end = component_ranges[component_index] + raise RuntimeError(f"Unexpected runtime stage: {label}") if uv_missing and index == 0: midpoint = stage_start + (stage_end - stage_start) * 0.5 install_uv( @@ -880,9 +702,7 @@ def install(args): ) if isinstance(stage_installed, list): installed_packages.extend(stage_installed) - vendor_temp.cleanup() - remove_legacy_vendor(venv_dir.parent) - devices = installed_devices(args) + devices = installed_devices() venv_dir.mkdir(parents=True, exist_ok=True) manifest_path.parent.mkdir(parents=True, exist_ok=True) ready_path = manifest_path @@ -891,14 +711,13 @@ def install(args): { "manifest_revision": MANIFEST_REVISION, "python": "3.12", - "torch": TORCH_VERSION, - "xformers": XFORMERS_VERSION, + "onnxruntime": ONNXRUNTIME_VERSION, "fastapi": FASTAPI_VERSION, "uvicorn": UVICORN_VERSION, - "torch_variant": args.torch_variant, + "runtime_variant": ( + "directml" if sys.platform == "win32" else "default" + ), "devices": devices, - "da3_commit": args.da3_commit, - "ben2_commit": args.ben2_commit, "packages": sorted(set(installed_packages)), }, indent=2, @@ -952,9 +771,6 @@ def parse_args(): parser.add_argument("--venv-dir", required=True) parser.add_argument("--manifest", required=True) parser.add_argument("--models-dir", required=True) - parser.add_argument("--torch-variant", choices=tuple(TORCH_INDEXES), required=True) - parser.add_argument("--da3-commit", required=True) - parser.add_argument("--ben2-commit", required=True) parser.add_argument("--status", required=True) parser.add_argument("--pypi-index", default=None) return parser.parse_args() diff --git a/src/snapform/models.py b/src/snapform/models.py index d9a4a45..cdabc10 100644 --- a/src/snapform/models.py +++ b/src/snapform/models.py @@ -21,42 +21,33 @@ MODELS = { "SMALL": ModelSpec( key="SMALL", label="DA3 Small", - directory_name="da3_small", + directory_name="DA3-SMALL", repository="depth-anything/DA3-SMALL", license_name="Apache 2.0", - description="Fastest any-view model with Apache 2.0 weights", + description="Fastest any-view model using ONNX Runtime", supports_gs=False, - r2_directory="DA3Small", + download_patterns=("model.onnx",), + ready_patterns=("model.onnx",), + r2_directory="DA3-SMALL", r2_files=( - ("config.json", 1202, "a486e29e82b7ab4a7d4cefc1ea4526cfe2ae438a572c8ca98917cfbcde7447d2"), - ("model.safetensors", 137248940, "364492e38a3a06d221ac75da7f6621ada3f2361cd24fde11ba79091e9f40efcf"), + ("model.onnx", 105484029, "7a9d81475b3e2334ace0d6d2c2a816b3c93081437e1866cc434ce1789125f631"), ), ), "BASE": ModelSpec( key="BASE", label="DA3 Base", - directory_name="da3_base", + directory_name="DA3-BASE", repository="depth-anything/DA3-BASE", license_name="Apache 2.0", - description="Balanced any-view model with Apache 2.0 weights", + description="Balanced any-view model using ONNX Runtime", supports_gs=False, - r2_directory="DA3Base", + download_patterns=("model.onnx",), + ready_patterns=("model.onnx",), + r2_directory="DA3-BASE", r2_files=( - ("config.json", 1205, "5e34115ebc17bd2d8d43033c5f72e9446ac8833fd61d3fa160b7e67e0bb5b7b5"), - ("configuration.json", 74, "b684d553257b2ca379fd5dfdf96a6e4a633eec3b969c97119c9875465a7c07a8"), - ("model.safetensors", 541518028, "e01067dc1659613083d9145a9a2547ccdbe6ccbbf83c4fe7b3e8a4e2bdae78b5"), + ("model.onnx", 412856697, "31f38fc4990539d7902ba267924b9731db1acb00d840c248a60b1865810ff9e6"), ), ), - "LARGE_1_1": ModelSpec( - key="LARGE_1_1", - label="DA3 Large 1.1", - directory_name="da3_large_1_1", - repository="depth-anything/DA3-LARGE-1.1", - license_name="CC BY-NC 4.0", - description="Higher quality non-commercial model with larger VRAM use", - supports_gs=False, - enabled=False, - ), "GIANT_1_1": ModelSpec( key="GIANT_1_1", label="DA3 Giant 1.1", @@ -82,52 +73,76 @@ MODELS = { BEN2_MODEL = ModelSpec( key="BEN2_BASE", label="BEN2 Base", - directory_name="ben2_base", - repository="PramaLLC/BEN2", - license_name="MIT", - description="Fast high-quality background removal model", + directory_name="BEN2-ONNX", + repository="onnx-community/BEN2-ONNX", + license_name="Apache 2.0", + description="Fast high-quality FP16 ONNX background removal model", supports_gs=False, - download_patterns=("config.json", "model.safetensors"), - r2_directory="BEN2Base", - r2_files=( - ("config.json", 124, "01ac876925577b52fa7c6af3f5b1d1420bc9b40016b6c25e2de89b38d8fff865"), - ("model.safetensors", 380577976, "ea8b7907176a09667c86343dc7d00de6a6d871076cb90bb5f753618fd6fb3ebb"), - ), + download_patterns=("onnx/model_fp16.onnx",), + ready_patterns=("onnx/model_fp16.onnx",), + r2_directory="BEN2-ONNX", + r2_files=(( + "onnx/model_fp16.onnx", + 219121675, + "dfdc25f421f32a0d1268e0f2ff2153d340e8f1d52d3dd16f5dc33c1ce85cedf1", + ),), ) -def moge2_model(key, label, directory_name, repository, description, r2_directory, size, checksum): - return ModelSpec( - key=key, - label=label, - directory_name=directory_name, - repository=repository, - license_name="MIT", - description=description, - supports_gs=False, - download_patterns=("model.pt",), - ready_patterns=("model.pt",), - r2_directory=r2_directory, - r2_files=(("model.pt", size, checksum),), - ) MOGE2_MODELS = { - "MOGE2_VITS_NORMAL": moge2_model( - "MOGE2_VITS_NORMAL", "MoGe-2 ViT-S Normal", "moge2_vits_normal", - "Ruicheng/moge-2-vits-normal", "Fastest MoGe-2 model", "MoGe2ViTSNormal", - 140550416, "79a16621928c2bf0ed04659218c55c01075e950507f40bb3332fb4c873d3e1dc", + "MOGE2_VITS_NORMAL": ModelSpec( + key="MOGE2_VITS_NORMAL", + label="MoGe-2 ViT-S Normal", + directory_name="moge-2-vits-normal-onnx", + repository="Ruicheng/moge-2-vits-normal-onnx", + license_name="MIT", + description="Fastest MoGe-2 model using ONNX Runtime", + supports_gs=False, + download_patterns=("model.onnx",), + ready_patterns=("model.onnx",), + r2_directory="moge-2-vits-normal-onnx", + r2_files=(( + "model.onnx", + 140852051, + "24eacb5dc7a2c54c7bc98f7de085ffbed79ad006ea5b664c2c2cdc02ff3a52f0", + ),), ), - "MOGE2_VITB_NORMAL": moge2_model( - "MOGE2_VITB_NORMAL", "MoGe-2 ViT-B Normal", "moge2_vitb_normal", - "Ruicheng/moge-2-vitb-normal", "Balanced MoGe-2 model", "MoGe2ViTBNormal", - 419110160, "16b8110e86d5dc5a849db120ca96ef3a223fd30b0c9146d1d81db504073da5f6", + "MOGE2_VITB_NORMAL": ModelSpec( + key="MOGE2_VITB_NORMAL", + label="MoGe-2 ViT-B Normal", + directory_name="moge-2-vitb-normal-onnx", + repository="Ruicheng/moge-2-vitb-normal-onnx", + license_name="MIT", + description="Balanced MoGe-2 model using ONNX Runtime", + supports_gs=False, + download_patterns=("model.onnx",), + ready_patterns=("model.onnx",), + r2_directory="moge-2-vitb-normal-onnx", + r2_files=(( + "model.onnx", + 419411850, + "bbf14e07a30f11e69d36ab861590123f5598ababcbc8946a063eb4a966f35a21", + ),), ), - "MOGE2_VITL_NORMAL": moge2_model( - "MOGE2_VITL_NORMAL", "MoGe-2 ViT-L Normal", "moge2_vitl_normal", - "Ruicheng/moge-2-vitl-normal", "Highest-quality MoGe-2 model", "MoGe2ViTLNormal", - 1323815904, "280741fd09bc3f403ccff9967784c2a391b52d2c0742ae3efdb21d9f90cc1a01", + "MOGE2_VITL_NORMAL": ModelSpec( + key="MOGE2_VITL_NORMAL", + label="MoGe-2 ViT-L Normal", + directory_name="moge-2-vitl-normal-onnx", + repository="Ruicheng/moge-2-vitl-normal-onnx", + license_name="MIT", + description="Highest-quality MoGe-2 model using ONNX Runtime", + supports_gs=False, + download_patterns=("model.onnx",), + ready_patterns=("model.onnx",), + r2_directory="moge-2-vitl-normal-onnx", + r2_files=(( + "model.onnx", + 1324265014, + "afbc4ccc3450298f3afb35b90f015f4c4f552dea21dc6470d5f7b78b77e2d751", + ),), ), } -MOGE2_MODEL = MOGE2_MODELS["MOGE2_VITB_NORMAL"] +MOGE2_MODEL = MOGE2_MODELS["MOGE2_VITS_NORMAL"] DOWNLOADABLE_MODELS = { **MODELS, BEN2_MODEL.key: BEN2_MODEL, diff --git a/src/snapform/moge_onnx.py b/src/snapform/moge_onnx.py new file mode 100644 index 0000000..5e5d390 --- /dev/null +++ b/src/snapform/moge_onnx.py @@ -0,0 +1,165 @@ +"""MoGe-2 ONNX inference and NumPy post-processing.""" + +from pathlib import Path +import sys + + +SCRIPT_DIRECTORY = Path(__file__).resolve().parent +if str(SCRIPT_DIRECTORY) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIRECTORY)) + +try: + from .onnx_runtime import create_session as create_runtime_session + from .onnx_runtime import provider_label, select_providers +except ImportError: + from onnx_runtime import create_session as create_runtime_session + from onnx_runtime import provider_label, select_providers + + +ONNX_FILENAME = "model.onnx" +MIN_TOKENS = 1200 +MAX_TOKENS = 3600 + + +def is_onnx_model(model_directory): + return (Path(model_directory) / ONNX_FILENAME).is_file() + + +def create_session(model_directory, requested_device="auto"): + path = Path(model_directory) / ONNX_FILENAME + return create_runtime_session(path, requested_device) + + +def normalized_view_plane_uv(width, height, dtype): + import numpy as np + + aspect_ratio = width / height + span_x = aspect_ratio / (1.0 + aspect_ratio**2) ** 0.5 + span_y = 1.0 / (1.0 + aspect_ratio**2) ** 0.5 + u = np.linspace( + -span_x * (width - 1) / width, + span_x * (width - 1) / width, + width, + dtype=dtype, + ) + v = np.linspace( + -span_y * (height - 1) / height, + span_y * (height - 1) / height, + height, + dtype=dtype, + ) + u, v = np.meshgrid(u, v, indexing="xy") + return np.stack((u, v), axis=-1) + + +def solve_focal_shift(uv, points): + import numpy as np + + uv = uv.reshape(-1, 2) + xy = points[..., :2].reshape(-1, 2) + z = points[..., 2].reshape(-1) + matrix = np.stack((xy, -uv), axis=-1).reshape(-1, 2) + target = (uv * z[:, None]).reshape(-1) + focal, shift = np.linalg.lstsq(matrix, target, rcond=None)[0] + epsilon = np.finfo(np.float32).eps + for _iteration in range(20): + denominator = z + shift + denominator = np.where( + np.abs(denominator) < epsilon, + np.copysign(epsilon, denominator), + denominator, + ) + projected = xy / denominator[:, None] + residual = (focal * projected - uv).reshape(-1) + jacobian = np.stack( + ( + projected, + -focal * xy / np.square(denominator)[:, None], + ), + axis=-1, + ).reshape(-1, 2) + delta = np.linalg.lstsq(jacobian, -residual, rcond=None)[0] + focal += delta[0] + shift += delta[1] + if np.max(np.abs(delta)) < 1e-6: + break + return np.float32(focal), shift + + +def recover_focal_shift(points, mask, downsample_size=(64, 64)): + import numpy as np + + height, width = points.shape[:2] + uv = normalized_view_plane_uv(width, height, points.dtype) + target_height, target_width = downsample_size + y = np.floor(np.arange(target_height) * height / target_height).astype(int) + x = np.floor(np.arange(target_width) * width / target_width).astype(int) + points_low = points[y[:, None], x[None, :]] + uv_low = uv[y[:, None], x[None, :]] + mask_low = mask[y[:, None], x[None, :]] + points_low = points_low[mask_low] + uv_low = uv_low[mask_low] + if len(points_low) < 2: + return np.float32(1.0), np.float32(0.0) + return solve_focal_shift(uv_low, points_low) + + +def depth_to_points(depth, intrinsics): + import numpy as np + + height, width = depth.shape + u = np.linspace(0.5 / width, 1.0 - 0.5 / width, width, dtype=depth.dtype) + v = np.linspace(0.5 / height, 1.0 - 0.5 / height, height, dtype=depth.dtype) + u, v = np.meshgrid(u, v, indexing="xy") + x = (u - intrinsics[0, 2]) / intrinsics[0, 0] * depth + y = (v - intrinsics[1, 2]) / intrinsics[1, 1] * depth + return np.stack((x, y, depth), axis=-1) + + +def postprocess(raw_output): + import numpy as np + + points = np.asarray(raw_output["points"][0], dtype=np.float32) + normal = np.asarray(raw_output["normal"][0], dtype=np.float32) + mask = np.asarray(raw_output["mask"][0] > 0.5) + metric_scale = np.float32(np.asarray(raw_output["metric_scale"]).reshape(-1)[0]) + height, width = mask.shape + aspect_ratio = width / height + focal, shift = recover_focal_shift(points, mask) + diagonal_factor = (1.0 + aspect_ratio**2) ** 0.5 + fx = focal / 2.0 * diagonal_factor / aspect_ratio + fy = focal / 2.0 * diagonal_factor + intrinsics = np.array( + ((fx, 0.0, 0.5), (0.0, fy, 0.5), (0.0, 0.0, 1.0)), + dtype=np.float32, + ) + depth = points[..., 2] + shift + mask &= depth > 0.0 + depth = depth * metric_scale + points = depth_to_points(depth, intrinsics) + return { + "points": points.astype(np.float32), + "depth": depth.astype(np.float32), + "normal": normal, + "mask": mask, + "intrinsics": intrinsics, + } + + +def infer(session, image, resolution_level): + import numpy as np + + level = min(max(int(resolution_level), 0), 9) + num_tokens = int(MIN_TOKENS + level / 9.0 * (MAX_TOKENS - MIN_TOKENS)) + image_input = np.ascontiguousarray(image.transpose(2, 0, 1)[None], dtype=np.float32) + input_feed = { + "image": image_input, + "num_tokens": np.asarray(num_tokens, dtype=np.int64), + } + output_names = {output.name for output in session.get_outputs()} + scale_name = "metric_scale" if "metric_scale" in output_names else "scale" + names = ("points", "normal", "mask", scale_name) + values = session.run(list(names), input_feed) + raw_output = dict(zip(names, values)) + raw_output["metric_scale"] = raw_output.pop(scale_name) + return postprocess(raw_output) diff --git a/src/snapform/onnx_runtime.py b/src/snapform/onnx_runtime.py new file mode 100644 index 0000000..109b8c6 --- /dev/null +++ b/src/snapform/onnx_runtime.py @@ -0,0 +1,55 @@ +"""Shared ONNX Runtime session helpers.""" + + +def select_providers(requested, available): + available = set(available) + providers_for_device = { + "cuda": ("CUDAExecutionProvider", "DmlExecutionProvider"), + "directml": ("DmlExecutionProvider",), + "mps": ("CoreMLExecutionProvider",), + "cpu": ("CPUExecutionProvider",), + } + if requested != "auto": + candidates = providers_for_device[requested] + for provider in candidates: + if provider in available: + return [provider] + if "CPUExecutionProvider" in available: + return ["CPUExecutionProvider"] + raise RuntimeError( + f"{requested.upper()} was requested but the ONNX Runtime " + f"providers {', '.join(candidates)} are not available" + ) + priorities = ( + "CUDAExecutionProvider", + "CoreMLExecutionProvider", + "DmlExecutionProvider", + "OpenVINOExecutionProvider", + "CPUExecutionProvider", + ) + providers = [provider for provider in priorities if provider in available] + if not providers: + raise RuntimeError("ONNX Runtime has no usable execution provider") + return providers + + +def create_session(path, requested_device="auto"): + import onnxruntime + + providers = select_providers( + requested_device, + onnxruntime.get_available_providers(), + ) + return onnxruntime.InferenceSession(str(path), providers=providers) + + +def provider_label(session): + provider = session.get_providers()[0] + labels = { + "CUDAExecutionProvider": "CUDA", + "CoreMLExecutionProvider": "CoreML", + "DmlExecutionProvider": "DirectML", + "OpenVINOExecutionProvider": "OpenVINO", + "CPUExecutionProvider": "CPU", + } + return labels.get(provider, provider) diff --git a/src/snapform/panel.py b/src/snapform/panel.py index fb16dc5..76d8fd4 100644 --- a/src/snapform/panel.py +++ b/src/snapform/panel.py @@ -186,7 +186,6 @@ def draw_da3_input_settings(layout, settings, model_property): layout.prop(settings, "video_frames") layout.prop(settings, "process_res") layout.prop(settings, "process_res_method") - layout.prop(settings, "use_ray_pose") layout.prop(settings, "ref_view_strategy") @@ -198,7 +197,7 @@ class Da3OutputDebugPanel(bpy.types.Panel): bl_region_type = "UI" bl_category = "SnapForm" bl_options = {"DEFAULT_CLOSED"} - bl_order = 0 + bl_order = 1 def draw(self, context): layout = self.layout @@ -248,7 +247,7 @@ class Moge2ApiDebugPanel(bpy.types.Panel): bl_region_type = "UI" bl_category = "SnapForm" bl_options = {"DEFAULT_CLOSED"} - bl_order = 1 + bl_order = 0 def draw(self, context): layout = self.layout diff --git a/src/snapform/properties.py b/src/snapform/properties.py index c093103..080928c 100644 --- a/src/snapform/properties.py +++ b/src/snapform/properties.py @@ -76,7 +76,7 @@ class SnapFormSettings(bpy.types.PropertyGroup): name="Model", description="MoGe-2 model used for depth estimation", items=moge2_model_enum_items, - default=1, + default=0, ) moge_export_format: bpy.props.EnumProperty( name="Output Format", @@ -110,7 +110,7 @@ class SnapFormSettings(bpy.types.PropertyGroup): ) use_ray_pose: bpy.props.BoolProperty( name="Use Ray Pose", - description="Use the slower and more accurate ray-based pose estimation", + description="Legacy task compatibility; ONNX uses the camera decoder", default=False, ) ref_view_strategy: bpy.props.EnumProperty( diff --git a/src/snapform/service.py b/src/snapform/service.py index dda942a..537f793 100644 --- a/src/snapform/service.py +++ b/src/snapform/service.py @@ -74,9 +74,6 @@ class ServiceLock: class ModelCache: def __init__(self): - import torch - - self.torch = torch self.ben2_model = None self.ben2_key = None self.da3_model = None @@ -85,75 +82,51 @@ class ModelCache: self.moge_key = None def get_ben2(self, model_dir, device): - key = (str(Path(model_dir).resolve()), device.type) + requested_device = getattr(device, "type", device) + key = (str(Path(model_dir).resolve()), requested_device, "onnx") if self.ben2_model is not None and self.ben2_key == key: return self.ben2_model, 0.0 self.release_ben2() - try: - return self._load_ben2(model_dir, device, key) - except self.torch.cuda.OutOfMemoryError: - self.release_da3() - self.release_moge() - return self._load_ben2(model_dir, device, key) + return self._load_ben2(model_dir, requested_device, key) - def _load_ben2(self, model_dir, device, key): - from ben2 import BEN_Base + def _load_ben2(self, model_dir, requested_device, key): + from ben2_onnx import create_session started = time.perf_counter() - model = BEN_Base.from_pretrained(model_dir) - model = model.to(device=device) - model.eval() + model = create_session(model_dir, requested_device) self.ben2_model = model self.ben2_key = key elapsed_ms = (time.perf_counter() - started) * 1000.0 return model, elapsed_ms def get_da3(self, model_dir, device): - key = (str(Path(model_dir).resolve()), device.type) + requested_device = getattr(device, "type", device) + key = (str(Path(model_dir).resolve()), requested_device, "onnx") if self.moge_model is not None: self.release_moge() if self.da3_model is not None and self.da3_key == key: return self.da3_model self.release_da3() - try: - return self._load_da3(model_dir, device, key) - except self.torch.cuda.OutOfMemoryError: - self.release_ben2() - self.release_moge() - return self._load_da3(model_dir, device, key) + return self._load_da3(model_dir, requested_device, key) - def _load_da3(self, model_dir, device, key): - from depth_anything_3.api import DepthAnything3 + def _load_da3(self, model_dir, requested_device, key): + from da3_onnx import create_session - model = DepthAnything3.from_pretrained(model_dir) - model = model.to(device=device) - model.eval() + model = create_session(model_dir, requested_device) self.da3_model = model self.da3_key = key return model - def get_moge(self, model_dir, device): - key = (str(Path(model_dir).resolve()), device.type) + def get_moge_onnx(self, model_dir, requested_device): + from moge_onnx import create_session + + key = (str(Path(model_dir).resolve()), requested_device, "onnx") if self.da3_model is not None: self.release_da3() if self.moge_model is not None and self.moge_key == key: return self.moge_model self.release_moge() - try: - return self._load_moge(model_dir, device, key) - except self.torch.cuda.OutOfMemoryError: - self.release_da3() - self.release_ben2() - return self._load_moge(model_dir, device, key) - - def _load_moge(self, model_dir, device, key): - worker.prepare_moge_runtime() - from moge.model.v2 import MoGeModel - - checkpoint = Path(model_dir) / "model.pt" - model = MoGeModel.from_pretrained(checkpoint) - model = model.to(device=device) - model.eval() + model = create_session(model_dir, requested_device) self.moge_model = model self.moge_key = key return model @@ -201,8 +174,6 @@ class ModelCache: def _release_cuda_cache(self): gc.collect() - if self.torch.cuda.is_available(): - self.torch.cuda.empty_cache() class ServiceState: diff --git a/src/snapform/worker.py b/src/snapform/worker.py index 2e24e92..895488d 100644 --- a/src/snapform/worker.py +++ b/src/snapform/worker.py @@ -6,11 +6,17 @@ import shutil import tempfile import threading import time +import sys from pathlib import Path from urllib.parse import quote from urllib.request import Request, urlopen +SCRIPT_DIRECTORY = Path(__file__).resolve().parent +if str(SCRIPT_DIRECTORY) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIRECTORY)) + + IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"} VIDEO_SUFFIXES = { ".mp4", @@ -434,6 +440,92 @@ def export_prediction_attributes(prediction, output_directory): return output_path +def export_da3_mini_npz(prediction, output_directory): + import numpy as np + + values = {"depth": np.round(prediction.depth, 8)} + for name in ("conf", "extrinsics", "intrinsics"): + value = getattr(prediction, name, None) + if value is not None: + values[name] = np.round(value, 2) if name == "conf" else value + output_path = Path(output_directory) / "results.npz" + np.savez_compressed(output_path, **values) + return output_path + + +def export_da3_depth_preview(prediction, output_directory): + import numpy as np + from PIL import Image + + target = Path(output_directory) / "depth_vis" + target.mkdir(parents=True, exist_ok=True) + for index, depth in enumerate(prediction.depth): + valid = np.isfinite(depth) & (depth > 0.0) + normalized = np.zeros(depth.shape, dtype=np.float32) + if valid.any(): + low, high = np.percentile(depth[valid], (2.0, 98.0)) + normalized[valid] = np.clip( + (depth[valid] - low) / max(float(high - low), 1e-6), + 0.0, + 1.0, + ) + color = np.stack( + ( + normalized, + 1.0 - np.abs(normalized * 2.0 - 1.0), + 1.0 - normalized, + ), + axis=-1, + ) + color = np.asarray(color * 255.0, dtype=np.uint8) + source = np.asarray(prediction.processed_images[index], dtype=np.uint8) + combined = np.concatenate((source, color), axis=1) + Image.fromarray(combined, mode="RGB").save( + target / f"{index:04d}.jpg", + quality=95, + ) + + +def export_da3_point_cloud( + prediction, + output_directory, + confidence_percentile=40.0, + num_max_points=1_000_000, +): + import numpy as np + import trimesh + + depth = np.asarray(prediction.depth, dtype=np.float32) + confidence = np.asarray(prediction.conf, dtype=np.float32) + threshold = confidence_percentile_threshold( + depth, + confidence, + confidence_percentile, + ) + point_sets = [] + color_sets = [] + for index in range(depth.shape[0]): + points = unproject_depth_frame( + depth[index], + prediction.intrinsics[index], + prediction.extrinsics[index], + ) + valid = np.isfinite(points).all(axis=-1) & (depth[index] > 0.0) + if threshold is not None: + valid &= confidence[index] >= threshold + point_sets.append(points[valid]) + color_sets.append(prediction.processed_images[index][valid]) + points = np.concatenate(point_sets, axis=0) + colors = np.concatenate(color_sets, axis=0) + maximum = max(int(num_max_points), 1) + if len(points) > maximum: + indices = np.linspace(0, len(points) - 1, maximum, dtype=np.int64) + points = points[indices] + colors = colors[indices] + scene = trimesh.Scene(trimesh.points.PointCloud(points, colors=colors)) + scene.export(Path(output_directory) / SCENE_FILENAME) + + def remove_background_faces(faces, foreground_mask): """Return only triangles whose vertices all belong to the foreground.""" foreground_points = foreground_mask.reshape(-1) @@ -1133,29 +1225,10 @@ def format_rate(rate): return f"{size:.1f} {units[unit_index]}" -def select_device(requested, torch): - if requested == "cpu": - return torch.device("cpu") - if requested == "cuda": - if not torch.cuda.is_available(): - raise RuntimeError("CUDA was requested but is not available") - return torch.device("cuda") - 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( input_path, model_dir, device, - torch, model_cache=None, cancel_check=None, ): @@ -1165,14 +1238,13 @@ def infer_foreground( check_cancelled(cancel_check) if input_path.suffix.lower() not in IMAGE_SUFFIXES: raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image") - owns_model = model_cache is None - if owns_model: - from ben2 import BEN_Base + from ben2_onnx import create_session, infer + owns_model = model_cache is None + requested_device = getattr(device, "type", device) + if owns_model: load_started = time.perf_counter() - model = BEN_Base.from_pretrained(model_dir) - model = model.to(device=device) - model.eval() + model = create_session(model_dir, requested_device) load_ms = (time.perf_counter() - load_started) * 1000.0 else: model, load_ms = model_cache.get_ben2(model_dir, device) @@ -1180,22 +1252,10 @@ def infer_foreground( check_cancelled(cancel_check) with Image.open(input_path) as opened_image: image = opened_image.convert("RGB") - if device.type == "cuda": - torch.cuda.synchronize() inference_started = time.perf_counter() - with torch.inference_mode(): - foreground = model.inference( - image, - refine_foreground=False, - ) + foreground = infer(model, image) foreground = preserve_source_rgb(image, foreground) - if device.type == "cuda": - torch.cuda.synchronize() inference_ms = (time.perf_counter() - inference_started) * 1000.0 - if owns_model: - del model - if device.type == "cuda": - torch.cuda.empty_cache() check_cancelled(cancel_check) return foreground, load_ms, inference_ms @@ -1216,7 +1276,6 @@ def infer_foreground_frames( images, model_dir, device, - torch, output_directory, status_path, start_progress, @@ -1238,7 +1297,6 @@ def infer_foreground_frames( image_path, model_dir, device, - torch, model_cache=model_cache, cancel_check=cancel_check, ) @@ -1260,14 +1318,9 @@ def remove_background(args, model_cache=None, cancel_check=None): write_status(args.status, 0.02, "Starting Environment") check_cancelled(cancel_check) - if model_cache is None: - import torch - else: - torch = model_cache.torch - output_directory = Path(args.output_dir) output_directory.mkdir(parents=True, exist_ok=True) - device = select_device(args.device, torch) + device = args.device video_frames = getattr(args, "video_frames", 16) start_frame = getattr(args, "start_frame", 1) @@ -1279,14 +1332,13 @@ def remove_background(args, model_cache=None, cancel_check=None): cancel_check=cancel_check, ) try: - write_status(args.status, 0.12, f"Loading BEN2 on {device.type}") + write_status(args.status, 0.12, f"Loading BEN2 for {device.upper()}") if len(images) == 1: write_status(args.status, 0.35, "Removing background") foreground, load_ms, inference_ms = infer_foreground( images[0], args.model_dir, device, - torch, model_cache=model_cache, cancel_check=cancel_check, ) @@ -1305,7 +1357,6 @@ def remove_background(args, model_cache=None, cancel_check=None): images, args.model_dir, device, - torch, output_directory, args.status, 0.12, @@ -1517,41 +1568,33 @@ def export_moge_debug_mesh( export_format, ): import numpy as np - import utils3d - from moge.utils.io import save_glb, save_ply + import trimesh - clean_mask = mask & ~utils3d.np.depth_map_edge(depth, rtol=0.04) - uv = utils3d.np.uv_map(*depth.shape) - result = utils3d.np.build_mesh_from_map( + clean_mask = mask & ~depth_edge_mask(depth, threshold=0.04) + faces = mesh_triangle_faces( points, - image, - uv, - normal, - mask=clean_mask, - tri=True, + depth, + np.ones(depth.shape, dtype=np.float32), + None, + False, + 0.0, + vertex_mask=clean_mask, ) - faces, vertices, colors, vertex_uvs, vertex_normals = result + vertices = points.reshape(-1, 3) + colors = np.asarray(image.reshape(-1, 3) * 255.0, dtype=np.uint8) + vertex_normals = normal.reshape(-1, 3) vertices = vertices * np.array([1.0, -1.0, -1.0], dtype=np.float32) - vertex_uvs = vertex_uvs * np.array([1.0, -1.0]) + np.array([0.0, 1.0]) vertex_normals = vertex_normals * np.array([1.0, -1.0, -1.0], dtype=np.float32) - if export_format == "glb": - save_glb( - output_directory / "mesh.glb", - vertices, - faces, - vertex_uvs, - (image * 255.0).astype(np.uint8), - vertex_normals, - ) - return - if export_format == "ply": - save_ply( - output_directory / "pointcloud.ply", - vertices, - np.zeros((0, 3), dtype=np.int32), - colors, - vertex_normals, - ) + mesh = trimesh.Trimesh( + vertices=vertices, + faces=faces, + vertex_colors=colors, + vertex_normals=vertex_normals, + process=False, + ) + mesh.remove_unreferenced_vertices() + filename = "mesh.glb" if export_format == "glb" else "pointcloud.ply" + mesh.export(Path(output_directory) / filename) def export_moge_surface_archive( @@ -1564,7 +1607,6 @@ def export_moge_surface_archive( foreground_mask=None, ): import numpy as np - import utils3d height, width = depth.shape finite_mask = np.isfinite(depth) & (depth > 0.0) @@ -1572,34 +1614,28 @@ def export_moge_surface_archive( if foreground_mask is not None: finite_mask &= foreground_mask depth_discontinuity = depth_discontinuity_ratio(depth) - result = utils3d.np.build_mesh_from_map( + faces = mesh_triangle_faces( points, - utils3d.np.uv_map(height, width), - normal, - mask.astype(np.float32), - depth_discontinuity.astype(np.float32), depth, - mask=finite_mask, - tri=True, + np.ones(depth.shape, dtype=np.float32), + None, + False, + 0.0, + vertex_mask=finite_mask, ) - ( - faces, - vertices, - vertex_uvs, - _vertex_normals, - vertex_mask, - vertex_depth_discontinuity, - point_depth, - ) = result - vertex_uvs = vertex_uvs * np.array([1.0, -1.0]) + np.array([0.0, 1.0]) + vertices = points.reshape(-1, 3) + u = np.linspace(0.0, 1.0, width, dtype=np.float32) + v = np.linspace(1.0, 0.0, height, dtype=np.float32) + uu, vv = np.meshgrid(u, v) + vertex_uvs = np.stack((uu, vv), axis=-1).reshape(-1, 2) pixel_intrinsics = intrinsics.copy() pixel_intrinsics[0] *= width pixel_intrinsics[1] *= height pixel_intrinsics[2] = (0.0, 0.0, 1.0) - point_depth = point_depth.astype(np.float32) - face_mask = vertex_mask[faces].all(axis=1) - face_depth_discontinuity = vertex_depth_discontinuity[faces].max(axis=1) - masked_depth = point_depth[vertex_mask.astype(bool)] + point_depth = depth.reshape(-1).astype(np.float32) + face_mask = mask.reshape(-1)[faces].all(axis=1) + face_depth_discontinuity = depth_discontinuity.reshape(-1)[faces].max(axis=1) + masked_depth = point_depth[mask.reshape(-1).astype(bool)] reference_depth = np.median(masked_depth if len(masked_depth) else point_depth) np.savez( output_path, @@ -1619,14 +1655,13 @@ def export_moge_surface_archive( def run_moge(args, model_cache=None, cancel_check=None): import numpy as np from PIL import Image + from moge_onnx import ( + create_session, + infer as infer_moge_onnx, + provider_label, + ) - prepare_moge_runtime() write_status(args.status, 0.02, "Starting MoGe-2") - if model_cache is None: - import torch - from moge.model.v2 import MoGeModel - else: - torch = model_cache.torch output_directory = Path(args.output_dir) output_directory.mkdir(parents=True, exist_ok=True) images, temporary_directory = prepare_input( @@ -1637,15 +1672,15 @@ def run_moge(args, model_cache=None, cancel_check=None): cancel_check=cancel_check, ) try: - device = select_device(args.device, torch) - write_status(args.status, 0.08, f"Loading MoGe-2 on {device.type}") if model_cache is None: - checkpoint = Path(args.model_dir) / "model.pt" - model = MoGeModel.from_pretrained(checkpoint) - model = model.to(device=device) - model.eval() + model = create_session(args.model_dir, args.device) else: - model = model_cache.get_moge(args.model_dir, device) + model = model_cache.get_moge_onnx(args.model_dir, args.device) + write_status( + args.status, + 0.08, + f"Loading MoGe-2 on {provider_label(model)}", + ) depths = [] normals = [] @@ -1662,21 +1697,17 @@ def run_moge(args, model_cache=None, cancel_check=None): ) with Image.open(image_path) as opened: image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0 - image_tensor = torch.from_numpy(image).permute(2, 0, 1).to(device) source_images.append(image) - with torch.inference_mode(): - prediction = model.infer( - image_tensor, - resolution_level=int(args.resolution_level), - force_projection=True, - apply_mask=False, - use_fp16=device.type == "cuda", - ) - depth = prediction["depth"].detach().float().cpu().numpy() - normal = prediction["normal"].detach().float().cpu().numpy() - mask = prediction["mask"].detach().cpu().numpy().astype(bool) - point = prediction["points"].detach().float().cpu().numpy() - intrinsic = prediction["intrinsics"].detach().float().cpu().numpy() + prediction = infer_moge_onnx( + model, + image, + int(args.resolution_level), + ) + depth = prediction["depth"] + normal = prediction["normal"] + mask = prediction["mask"].astype(bool) + point = prediction["points"] + intrinsic = prediction["intrinsics"] depths.append(depth) normals.append(normal) masks.append(mask) @@ -1727,8 +1758,7 @@ def run_moge(args, model_cache=None, cancel_check=None): foreground, _, _ = infer_foreground( images[0], remove_background_model_dir, - device, - torch, + args.device, model_cache=model_cache, cancel_check=cancel_check, ) @@ -1743,8 +1773,7 @@ def run_moge(args, model_cache=None, cancel_check=None): infer_foreground_frames( images, remove_background_model_dir, - device, - torch, + args.device, output_directory, args.status, 0.87, @@ -1781,21 +1810,11 @@ def run_moge(args, model_cache=None, cancel_check=None): shutil.rmtree(temporary_directory, ignore_errors=True) -def prepare_moge_runtime(): - """Use PyTorch SDPA because xFormers does not support every CUDA GPU.""" - os.environ["XFORMERS_DISABLED"] = "1" - - def reconstruct(args, model_cache=None, cancel_check=None): write_status(args.status, 0.02, "Starting Environment") check_cancelled(cancel_check) - - if model_cache is None: - import torch - from depth_anything_3.api import DepthAnything3 - else: - torch = model_cache.torch - from depth_anything_3.utils.export import export + from da3_onnx import create_session, infer + from onnx_runtime import provider_label output_directory = Path(args.output_dir) output_directory.mkdir(parents=True, exist_ok=True) @@ -1810,7 +1829,6 @@ def reconstruct(args, model_cache=None, cancel_check=None): try: check_cancelled(cancel_check) - device = select_device(args.device, torch) remove_background_model_dir = getattr( args, "remove_background_model_dir", @@ -1823,8 +1841,7 @@ def reconstruct(args, model_cache=None, cancel_check=None): foreground, load_ms, inference_ms = infer_foreground( images[0], remove_background_model_dir, - device, - torch, + args.device, model_cache=model_cache, cancel_check=cancel_check, ) @@ -1841,8 +1858,7 @@ def reconstruct(args, model_cache=None, cancel_check=None): load_ms, inference_ms = infer_foreground_frames( images, remove_background_model_dir, - device, - torch, + args.device, output_directory, args.status, 0.12, @@ -1858,13 +1874,16 @@ def reconstruct(args, model_cache=None, cancel_check=None): f"BEN2 inference {inference_ms:.0f} ms; " f"model load {load_ms:.0f} ms" ) - write_status(args.status, 0.15, f"Loading model on {device.type}") + write_status(args.status, 0.15, "Loading DA3 ONNX model") if model_cache is None: - model = DepthAnything3.from_pretrained(args.model_dir) - model = model.to(device=device) - model.eval() + model = create_session(args.model_dir, args.device) else: - model = model_cache.get_da3(args.model_dir, device) + model = model_cache.get_da3(args.model_dir, args.device) + write_status( + args.status, + 0.20, + f"Loaded DA3 on {provider_label(model)}", + ) check_cancelled(cancel_check) write_status( @@ -1872,27 +1891,14 @@ def reconstruct(args, model_cache=None, cancel_check=None): 0.25, f"Estimating depth for {len(images)} frame(s)", ) - export_kwargs = {} - if args.export_format == "glb": - export_kwargs["glb"] = { - "conf_thresh_percentile": args.conf_thresh_percentile, - "num_max_points": args.num_max_points, - "show_cameras": bool(args.show_cameras), - "export_depth_vis": False, - } - if args.export_format == "gs_ply": - export_kwargs["gs_ply"] = { - "gs_views_interval": args.gs_views_interval, - } - with torch.inference_mode(): - prediction = model.inference( - [str(path) for path in images], - infer_gs=args.export_format == "gs_ply", - 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, - ) + 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, + ) check_cancelled(cancel_check) write_status(args.status, 0.88, f"Exporting {args.export_format}") if args.export_format == "mesh": @@ -1921,12 +1927,20 @@ def reconstruct(args, model_cache=None, cancel_check=None): ), progress_callback=mesh_progress, ) - else: - export( + elif args.export_format == "mini_npz": + export_da3_mini_npz(prediction, output_directory) + elif args.export_format == "depth_vis": + export_da3_depth_preview(prediction, output_directory) + elif args.export_format == "glb": + export_da3_point_cloud( prediction, - args.export_format, - str(output_directory), - **export_kwargs, + output_directory, + args.conf_thresh_percentile, + args.num_max_points, + ) + else: + raise ValueError( + f"DA3 ONNX does not support export format: {args.export_format}" ) if args.export_format in {"glb", "gs_ply", "depth_vis", "mesh"}: write_status( @@ -1989,7 +2003,7 @@ def parser(): background.add_argument("--status", required=True) background.add_argument( "--device", - choices=("auto", "cuda", "mps", "cpu"), + choices=("auto", "cuda", "directml", "mps", "cpu"), required=True, ) background.add_argument("--video-frames", type=int, default=16) @@ -2040,7 +2054,7 @@ def parser(): reconstruction.add_argument("--start-frame", type=int, default=1) reconstruction.add_argument( "--device", - choices=("auto", "cuda", "mps", "cpu"), + choices=("auto", "cuda", "directml", "mps", "cpu"), required=True, ) reconstruction.add_argument( diff --git a/sync_models.py b/sync_models.py new file mode 100644 index 0000000..e39a928 --- /dev/null +++ b/sync_models.py @@ -0,0 +1,194 @@ +"""Prepare SnapForm ONNX models and sync them to OmooLab R2.""" + +import argparse +import hashlib +import shutil +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path +from urllib.request import Request, urlopen + + +PROJECT_ROOT = Path(__file__).resolve().parent +MODELS_DIR = PROJECT_ROOT / "models" +R2_ENDPOINT = "https://models.omoolab.xyz" +CHUNK_SIZE = 1024 * 1024 + + +@dataclass(frozen=True) +class ModelFile: + folder: str + size: int + sha256: str + sources: tuple[str, ...] + + @property + def destination(self): + filename = "model_fp16.onnx" if self.folder.endswith("/onnx") else "model.onnx" + return MODELS_DIR / self.folder / filename + + +def hf_url(repository, filename): + return f"https://huggingface.co/{repository}/resolve/main/{filename}" + + +def r2_url(folder, filename="model.onnx"): + return f"{R2_ENDPOINT}/{folder}/{filename}" + + +MODEL_FILES = ( + ModelFile( + "DA3-SMALL", + 105484029, + "7a9d81475b3e2334ace0d6d2c2a816b3c93081437e1866cc434ce1789125f631", + (r2_url("DA3-SMALL"),), + ), + ModelFile( + "DA3-BASE", + 412856697, + "31f38fc4990539d7902ba267924b9731db1acb00d840c248a60b1865810ff9e6", + (r2_url("DA3-BASE"),), + ), + ModelFile( + "BEN2-ONNX/onnx", + 219121675, + "dfdc25f421f32a0d1268e0f2ff2153d340e8f1d52d3dd16f5dc33c1ce85cedf1", + ( + hf_url("onnx-community/BEN2-ONNX", "onnx/model_fp16.onnx"), + r2_url("BEN2-ONNX", "onnx/model_fp16.onnx"), + ), + ), + ModelFile( + "moge-2-vits-normal-onnx", + 140852051, + "24eacb5dc7a2c54c7bc98f7de085ffbed79ad006ea5b664c2c2cdc02ff3a52f0", + ( + hf_url("Ruicheng/moge-2-vits-normal-onnx", "model.onnx"), + r2_url("moge-2-vits-normal-onnx"), + ), + ), + ModelFile( + "moge-2-vitb-normal-onnx", + 419411850, + "bbf14e07a30f11e69d36ab861590123f5598ababcbc8946a063eb4a966f35a21", + ( + hf_url("Ruicheng/moge-2-vitb-normal-onnx", "model.onnx"), + r2_url("moge-2-vitb-normal-onnx"), + ), + ), + ModelFile( + "moge-2-vitl-normal-onnx", + 1324265014, + "afbc4ccc3450298f3afb35b90f015f4c4f552dea21dc6470d5f7b78b77e2d751", + ( + hf_url("Ruicheng/moge-2-vitl-normal-onnx", "model.onnx"), + r2_url("moge-2-vitl-normal-onnx"), + ), + ), +) + + +def file_sha256(path): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + while chunk := stream.read(CHUNK_SIZE): + digest.update(chunk) + return digest.hexdigest() + + +def verify_existing(model): + path = model.destination + if not path.is_file(): + return False + if path.stat().st_size != model.size: + raise RuntimeError(f"Existing model has an unexpected size: {path}") + if file_sha256(path) != model.sha256: + raise RuntimeError(f"Existing model has an unexpected SHA-256: {path}") + print(f"Skip existing: {path.relative_to(PROJECT_ROOT)}") + return True + + +def download_once(model, url): + destination = model.destination + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_suffix(".onnx.part") + temporary.unlink(missing_ok=True) + request = Request(url, headers={"User-Agent": "SnapForm model sync"}) + print(f"Download: {url}") + downloaded = 0 + digest = hashlib.sha256() + try: + with urlopen(request, timeout=60) as response, temporary.open("wb") as output: + while chunk := response.read(CHUNK_SIZE): + output.write(chunk) + digest.update(chunk) + downloaded += len(chunk) + print( + f"\r {downloaded / 1024**2:.1f} / " + f"{model.size / 1024**2:.1f} MiB", + end="", + flush=True, + ) + print() + if downloaded != model.size: + raise RuntimeError(f"size mismatch: {downloaded} != {model.size}") + if digest.hexdigest() != model.sha256: + raise RuntimeError("SHA-256 mismatch") + temporary.replace(destination) + except Exception: + temporary.unlink(missing_ok=True) + raise + + +def prepare_model(model): + if verify_existing(model): + return + failures = [] + for url in model.sources: + try: + download_once(model, url) + return + except Exception as error: + failures.append(f"{url}: {error}") + print(f"Source failed: {error}", file=sys.stderr) + raise RuntimeError("All sources failed:\n" + "\n".join(failures)) + + +def sync_models(): + executable = shutil.which("rclone") + if executable is None: + raise RuntimeError("rclone is not available on PATH") + subprocess.run( + [ + executable, + "sync", + "models", + "omoolab-r2:ai-models", + "--progress", + ], + cwd=PROJECT_ROOT, + check=True, + ) + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--download-only", + action="store_true", + help="Prepare and verify models without running rclone", + ) + return parser.parse_args() + + +def main(): + args = parse_args() + for model in MODEL_FILES: + prepare_model(model) + if not args.download_only: + sync_models() + + +if __name__ == "__main__": + main() diff --git a/tests/test_addon.py b/tests/test_addon.py index 58da071..cd399c5 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -280,6 +280,10 @@ class AddonTest(unittest.TestCase): annotations["snapform_process_res_factor"]["subtype"], "PERCENTAGE", ) + self.assertEqual( + annotations["snapform_model"]["default"], + 0, + ) def test_image_mesh_uses_o_prefixed_attributes_and_bundled_nodes(self): image_mesh = self.addon.image_mesh @@ -644,8 +648,11 @@ class AddonTest(unittest.TestCase): def test_conversion_dialog_exposes_requested_settings(self): operator = self.addon.ConvertImageToDepthPlane - self.assertEqual(operator.__annotations__["snapform_model"]["default"], 1) - self.assertTrue(operator.__annotations__["snapform_use_ray_pose"]["default"]) + self.assertEqual( + operator.__annotations__["snapform_model"]["default"], + 0, + ) + self.assertFalse(operator.__annotations__["snapform_use_ray_pose"]["default"]) self.assertEqual( operator.__annotations__["snapform_process_res_factor"]["default"], 100.0, @@ -1094,7 +1101,6 @@ class AddonTest(unittest.TestCase): for snapform_model, snapform_ray_pose in ( ("SMALL", True), ("BASE", False), - ("LARGE_1_1", True), ("GIANT_1_1", False), ("NESTED_GIANT_LARGE_1_1", True), ): @@ -1947,8 +1953,8 @@ class AddonTest(unittest.TestCase): self.addon.DebugPanel.bl_options, ) panels = ( - self.addon.Da3OutputDebugPanel, self.addon.Moge2ApiDebugPanel, + self.addon.Da3OutputDebugPanel, self.addon.Ben2RemovalDebugPanel, ) for panel in panels: @@ -1956,7 +1962,7 @@ class AddonTest(unittest.TestCase): self.assertIn("DEFAULT_CLOSED", panel.bl_options) self.assertEqual( [panel.bl_label for panel in panels], - ["DA3", "MoGe-2", "BEN2"], + ["MoGe-2", "DA3", "BEN2"], ) def test_service_panel_exposes_shared_service_state_and_controls(self): @@ -2483,7 +2489,7 @@ class AddonTest(unittest.TestCase): for index, value in enumerate(command) if value == "--allow-pattern" ] - self.assertEqual(patterns, ["config.json", "model.safetensors"]) + self.assertEqual(patterns, ["onnx/model_fp16.onnx"]) label_index = command.index("--label") self.assertEqual(command[label_index + 1], "BEN2 Base") @@ -2512,15 +2518,12 @@ class AddonTest(unittest.TestCase): self.assertEqual(command[0], sys.executable) self.assertIn("--uv", command) self.assertIn("C:/uv/uv.exe", command) - self.assertEqual( - command[command.index("--torch-variant") + 1], - "cu128", - ) + self.assertNotIn("--torch-variant", command) self.assertTrue( any("install_env.py" in str(value) for value in command) ) - def test_install_env_command_auto_selects_torch_variant(self): + def test_install_env_command_does_not_select_legacy_runtime(self): client = importlib.import_module("snapform.client") with patch.object( client, @@ -2542,10 +2545,7 @@ class AddonTest(unittest.TestCase): command = client.install_env_command(Path("status.json")) self.assertNotIn("--pypi-index", command) - self.assertEqual( - command[command.index("--torch-variant") + 1], - "cpu", - ) + self.assertNotIn("--torch-variant", command) self.assertIn("C:/uv/uv.exe", command) def test_completed_export_opens_job_directory(self): diff --git a/tests/test_ben2_onnx.py b/tests/test_ben2_onnx.py new file mode 100644 index 0000000..3649104 --- /dev/null +++ b/tests/test_ben2_onnx.py @@ -0,0 +1,53 @@ +import sys +import unittest +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +from PIL import Image + + +SOURCE_DIR = Path(__file__).resolve().parents[1] / "src" / "snapform" +sys.path.insert(0, str(SOURCE_DIR)) + +import ben2_onnx + + +class FakeSession: + def get_inputs(self): + return [SimpleNamespace(name="input")] + + def get_outputs(self): + return [SimpleNamespace(name="alpha")] + + def run(self, output_names, feeds): + self.output_names = output_names + self.feeds = feeds + gradient = np.linspace(0.0, 1.0, 1024, dtype=np.float32) + return [np.tile(gradient, (1024, 1))[None, None]] + + +class Ben2OnnxTest(unittest.TestCase): + def test_preprocess_matches_model_layout(self): + image = Image.new("RGB", (40, 20), (64, 128, 192)) + result = ben2_onnx.preprocess(image) + + self.assertEqual(result.shape, (1, 3, 1024, 1024)) + self.assertEqual(result.dtype, np.float32) + self.assertTrue(result.flags.c_contiguous) + + def test_infer_restores_original_size_and_alpha(self): + session = FakeSession() + image = Image.new("RGB", (40, 20), (10, 20, 30)) + + result = ben2_onnx.infer(session, image) + + self.assertEqual(result.mode, "RGBA") + self.assertEqual(result.size, image.size) + self.assertEqual(session.feeds["input"].shape, (1, 3, 1024, 1024)) + alpha = np.asarray(result.getchannel("A")) + self.assertLess(int(alpha[:, 0].max()), int(alpha[:, -1].min())) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_da3_onnx.py b/tests/test_da3_onnx.py new file mode 100644 index 0000000..4d431c3 --- /dev/null +++ b/tests/test_da3_onnx.py @@ -0,0 +1,61 @@ +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +from PIL import Image + + +SOURCE_DIR = Path(__file__).resolve().parents[1] / "src" / "snapform" +sys.path.insert(0, str(SOURCE_DIR)) + +import da3_onnx + + +class FakeSession: + def get_inputs(self): + return [SimpleNamespace(name="images"), SimpleNamespace(name="strategy")] + + def run(self, output_names, feeds): + self.feeds = feeds + _, views, _, height, width = feeds["images"].shape + depth = np.ones((1, views, height, width), dtype=np.float32) + confidence = np.full_like(depth, 0.75) + pose = np.zeros((1, views, 9), dtype=np.float32) + pose[..., 6] = 1.0 + pose[..., 7:] = np.pi / 2.0 + return depth, confidence, pose + + +class Da3OnnxTest(unittest.TestCase): + def test_target_size_uses_patch_multiple(self): + self.assertEqual( + da3_onnx.target_size((800, 600), 504, "upper_bound_resize"), + (504, 378), + ) + + def test_infer_returns_camera_and_image_arrays(self): + with tempfile.TemporaryDirectory() as directory: + image_path = Path(directory) / "input.png" + Image.new("RGB", (80, 60), (10, 20, 30)).save(image_path) + session = FakeSession() + + result = da3_onnx.infer( + session, + [image_path], + process_res=56, + ref_view_strategy="middle", + ) + + self.assertEqual(result.depth.shape, (1, 42, 56)) + self.assertEqual(result.conf.shape, result.depth.shape) + self.assertEqual(result.processed_images.shape, (1, 42, 56, 3)) + self.assertEqual(result.extrinsics.shape, (1, 3, 4)) + self.assertEqual(result.intrinsics.shape, (1, 3, 3)) + self.assertEqual(int(session.feeds["strategy"][0]), 1) + np.testing.assert_allclose(result.extrinsics[0, :3, :3], np.eye(3)) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_install_env.py b/tests/test_install_env.py index a01fb48..b56b993 100644 --- a/tests/test_install_env.py +++ b/tests/test_install_env.py @@ -36,7 +36,7 @@ class InstallEnvTest(unittest.TestCase): (vendor / "moge.tar.gz").write_bytes(b"moge") return vendor - def test_installs_da3_from_archives_without_gsplat(self): + def test_installs_onnx_runtime_without_training_frameworks(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) @@ -71,77 +71,31 @@ class InstallEnvTest(unittest.TestCase): commands += [ call.args[4] for call in mocked_run_install.call_args_list ] - da3_command = next( - command - for command in commands - if any("depth-anything-3 @" in str(value) for value in command) - ) - self.assertIn( - f"addict=={bootstrap.ADDICT_VERSION}", - da3_command, - ) - self.assertTrue( - any( - "depth-anything-3 @ file:" in str(value) - and "snapform_vendor_" 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 - ) - ) - ben2_command = next( - command - for command in commands - if any("ben2 @" in str(value) for value in command) - ) - self.assertTrue( - any( - "ben2 @ file:" in str(value) - and "snapform_vendor_" in str(value) - for value in ben2_command - ) - ) - self.assertFalse( - any("git+" in str(value) for value in ben2_command) - ) service_command = next( command for command in commands if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command ) self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", service_command) - self.assertIn("setuptools", service_command) - self.assertNotIn("ninja", service_command) + self.assertIn( + f"onnxruntime-directml=={bootstrap.ONNXRUNTIME_VERSION}", + service_command, + ) + for forbidden in ("torch", "torchvision", "xformers", "ben2 @", "moge @"): + self.assertFalse( + any(forbidden in str(value) for value in service_command), + forbidden, + ) check_command = bootstrap.verify_command( bootstrap.venv_python(venv_dir) ) - self.assertTrue( - any("import torch;" in str(value) for value in check_command) - ) - self.assertTrue( - any( - "from ben2 import BEN_Base" in str(value) - for value in check_command - ) - ) + self.assertFalse(any("torch" in str(value) for value in check_command)) self.assertTrue( any("import fastapi, uvicorn" in str(value) for value in check_command) ) + self.assertTrue( + any("onnxruntime" in str(value) for value in check_command) + ) self.assertFalse( any("import gsplat" in str(value) for value in check_command) ) @@ -157,24 +111,24 @@ class InstallEnvTest(unittest.TestCase): if any( repository in command for repository in ( - "Ruicheng/moge-2-vitb-normal", - "PramaLLC/BEN2", + "Ruicheng/moge-2-vits-normal-onnx", + "onnx-community/BEN2-ONNX", ) ) ] self.assertEqual(len(default_downloads), 2) - self.assertTrue(any("Ruicheng/moge-2-vitb-normal" in map(str, command) for command in default_downloads)) + self.assertTrue(any("Ruicheng/moge-2-vits-normal-onnx" in map(str, command) for command in default_downloads)) ben2_download = next( command for command in default_downloads - if "PramaLLC/BEN2" in command + if "onnx-community/BEN2-ONNX" in command ) patterns = [ ben2_download[index + 1] for index, value in enumerate(ben2_download) if value == "--allow-pattern" ] - self.assertEqual(patterns, ["config.json", "model.safetensors"]) + self.assertEqual(patterns, ["onnx/model_fp16.onnx"]) self.assertNotIn("--stage", ben2_download) self.assertNotIn("--stage-label", ben2_download) self.assertEqual( @@ -182,7 +136,7 @@ class InstallEnvTest(unittest.TestCase): "BEN2 Base", ) - def test_darwin_install_skips_cuda_index_and_xformers(self): + def test_darwin_install_uses_default_onnx_runtime(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) @@ -224,15 +178,15 @@ class InstallEnvTest(unittest.TestCase): commands = [ call.args[4] for call in mocked_run_install.call_args_list ] - torch_command = next( + runtime_command = next( command for command in commands if any( - f"torch=={bootstrap.TORCH_VERSION}" in str(value) + f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}" in str(value) for value in command ) ) - self.assertNotIn("--index-url", torch_command) + self.assertNotIn("--index-url", runtime_command) self.assertFalse( any( "xformers" in str(value) @@ -241,30 +195,27 @@ class InstallEnvTest(unittest.TestCase): ) ) - def test_intel_mac_is_rejected_with_clear_error(self): + def test_intel_mac_uses_portable_onnx_runtime(self): bootstrap = load_install_env() - args = SimpleNamespace( - uv="uv", - venv_dir=".venv", - models_dir="models", - torch_variant="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() + ): + stages = bootstrap.install_stages( + SimpleNamespace( + models_dir="models", + status="status.json", + ), + ".venv", + "python", + "uv", + ) + runtime = stages[1][1] + self.assertIn( + f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}", + runtime, + ) def test_run_stage_writes_stage_progress(self): bootstrap = load_install_env() @@ -414,10 +365,10 @@ class InstallEnvTest(unittest.TestCase): ) self.assertEqual(ends[-1], 1.0) self.assertEqual(starts[0], bootstrap.ENVIRONMENT_SLICE[0]) - runtime_ranges = stage_ranges[:4] + runtime_ranges = stage_ranges[:2] self.assertEqual( [item[0] for item in runtime_ranges], - [1, 2, 3, 3], + [1, 2], ) self.assertTrue( all(item[0] is None for item in stage_ranges[4:]) @@ -606,19 +557,19 @@ class InstallEnvTest(unittest.TestCase): self.assertEqual(bootstrap.format_elapsed(12), "12s") self.assertEqual(bootstrap.format_elapsed(83), "1m 23s") - def test_ai_component_estimate_continues_across_internal_steps(self): + def test_onnx_runtime_estimate_decreases_with_progress(self): bootstrap = load_install_env() first = bootstrap.estimated_remaining( - "AI components", - bootstrap.XFORMERS_SLICE[0], - ) - second = bootstrap.estimated_remaining( - "AI components", + "ONNX Runtime", bootstrap.AI_COMPONENTS_SLICE[0], ) + second = bootstrap.estimated_remaining( + "ONNX Runtime", + sum(bootstrap.AI_COMPONENTS_SLICE) / 2, + ) - self.assertEqual(first, bootstrap.INSTALL_ESTIMATES["AI components"]) + self.assertEqual(first, bootstrap.INSTALL_ESTIMATES["ONNX Runtime"]) self.assertGreater(second, 0) self.assertLess(second, first) @@ -727,7 +678,7 @@ class InstallEnvTest(unittest.TestCase): ready["packages"], ["numpy==1.26.4", "torch==2.9.1+cu128"], ) - self.assertEqual(ready["devices"], ["cuda", "cpu"]) + self.assertEqual(ready["devices"], ["directml", "cpu"]) self.assertEqual( ready["manifest_revision"], bootstrap.MANIFEST_REVISION, @@ -814,7 +765,7 @@ class InstallEnvTest(unittest.TestCase): def fake_run(command, environment=None): commands.append(command) - if any("moge-2-vitb-normal" in str(value) for value in command): + if any("moge-2-vits-normal" in str(value) for value in command): raise subprocess.CalledProcessError(1, command) with tempfile.TemporaryDirectory() as directory: @@ -852,7 +803,7 @@ class InstallEnvTest(unittest.TestCase): self.assertIn("Environment is installed", status["message"]) self.assertTrue( any( - "PramaLLC/BEN2" in map(str, command) + "onnx-community/BEN2-ONNX" in map(str, command) for command in commands ) ) @@ -983,25 +934,13 @@ class InstallEnvTest(unittest.TestCase): for value in command ) ) - ai_command = next( + runtime_command = next( command for command in install_commands - if any("depth-anything-3 @" in str(value) for value in command) - ) - self.assertTrue( - any( - "depth-anything-3 @ file:" in str(value) - and "snapform_vendor_" in str(value) - for value in ai_command - ) - ) - self.assertTrue( - any( - "ben2 @ file:" in str(value) - and "snapform_vendor_" in str(value) - for value in ai_command - ) + if any("onnxruntime" in str(value) for value in command) ) + self.assertFalse(any(" @ file:" in str(value) for value in runtime_command)) + self.assertFalse(any("torch" in str(value) for value in runtime_command)) def test_install_reuses_existing_uv(self): bootstrap = load_install_env() @@ -1047,81 +986,12 @@ class InstallEnvTest(unittest.TestCase): self.assertEqual(uv_path.read_bytes(), b"binary") - def test_copy_vendor_inputs_copies_missing_archives(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - 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") - vendor = root / "vendor" - - bootstrap.copy_vendor_inputs(bundled, vendor) - - self.assertEqual( - (vendor / "depth-anything-3.tar.gz").read_bytes(), - b"da3", - ) - self.assertEqual( - (vendor / "ben2.tar.gz").read_bytes(), - b"ben2", - ) - - def test_copy_vendor_inputs_replaces_same_size_stale_archive(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - bundled = root / "bundled" - bundled.mkdir() - (bundled / "depth-anything-3.tar.gz").write_bytes(b"new") - (bundled / "ben2.tar.gz").write_bytes(b"ben2") - (bundled / "utils3d.tar.gz").write_bytes(b"utils3d") - (bundled / "moge.tar.gz").write_bytes(b"moge") - vendor = root / "vendor" - vendor.mkdir() - (vendor / "depth-anything-3.tar.gz").write_bytes(b"old") - - bootstrap.copy_vendor_inputs(bundled, vendor) - - self.assertEqual( - (vendor / "depth-anything-3.tar.gz").read_bytes(), - b"new", - ) - - def test_remove_legacy_vendor_removes_only_install_inputs(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - vendor = root / "vendor" - vendor.mkdir() - (vendor / "depth-anything-3.tar.gz").write_bytes(b"da3") - (root / "depth-anything-3.tar.gz").write_bytes(b"da3") - (root / "ben2.tar.gz").write_bytes(b"ben2") - models = root / "models" - models.mkdir() - - bootstrap.remove_legacy_vendor(root) - - self.assertFalse(vendor.exists()) - self.assertFalse((root / "depth-anything-3.tar.gz").exists()) - self.assertFalse((root / "ben2.tar.gz").exists()) - self.assertTrue(models.is_dir()) - - def test_install_stages_uses_local_source_archives(self): + def test_install_stages_use_only_runtime_packages(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() - vendor = root / "vendor" - vendor.mkdir() - (vendor / "depth-anything-3.tar.gz").write_bytes(b"da3") - (vendor / "ben2.tar.gz").write_bytes(b"ben2") - (vendor / "utils3d.tar.gz").write_bytes(b"utils3d") - (vendor / "moge.tar.gz").write_bytes(b"moge") uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( @@ -1141,16 +1011,14 @@ class InstallEnvTest(unittest.TestCase): uv_path, ) - ai_command = next( + runtime_command = next( command for _label, command, _environment, kind in stages if kind == "install" - and any("depth-anything-3" in str(value) for value in command) + and any("onnxruntime" in str(value) for value in command) ) - da3_local = (vendor / "depth-anything-3.tar.gz").as_uri() - ben2_local = (vendor / "ben2.tar.gz").as_uri() - self.assertIn(f"depth-anything-3 @ {da3_local}", ai_command) - self.assertIn(f"ben2 @ {ben2_local}", ai_command) + self.assertFalse(any(" @ file:" in str(value) for value in runtime_command)) + self.assertFalse(any("torch" in str(value) for value in runtime_command)) def test_install_stages_download_only_default_models(self): bootstrap = load_install_env() @@ -1189,13 +1057,13 @@ class InstallEnvTest(unittest.TestCase): if kind == "download" } self.assertIn( - "moge2_vitb_normal", - " ".join(map(str, downloads["MoGe-2 ViT-B Normal"])), + "moge-2-vits-normal-onnx", + " ".join(map(str, downloads["MoGe-2 ViT-S Normal"])), ) - self.assertIn("ben2_base", " ".join(map(str, downloads["BEN2 Base"]))) + self.assertIn("BEN2-ONNX", " ".join(map(str, downloads["BEN2 Base"]))) self.assertEqual( set(downloads), - {"MoGe-2 ViT-B Normal", "BEN2 Base"}, + {"MoGe-2 ViT-S Normal", "BEN2 Base"}, ) def test_install_does_not_verify_env(self): diff --git a/tests/test_models.py b/tests/test_models.py index 9770903..d865f84 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -30,12 +30,6 @@ class ModelTest(unittest.TestCase): self.assertEqual(model.repository, "depth-anything/DA3-BASE") self.assertEqual(model.license_name, "Apache 2.0") - def test_large_model_exposes_noncommercial_license(self): - model = self.models.get_model("LARGE_1_1") - - self.assertEqual(model.license_name, "CC BY-NC 4.0") - self.assertFalse(model.enabled) - def test_giant_models_expose_gaussian_head(self): for model_key in ("GIANT_1_1", "NESTED_GIANT_LARGE_1_1"): with self.subTest(model_key=model_key): @@ -54,19 +48,20 @@ class ModelTest(unittest.TestCase): def test_ben2_download_only_includes_required_hugging_face_files(self): model = self.models.get_downloadable_model("BEN2_BASE") - self.assertEqual(model.repository, "PramaLLC/BEN2") - self.assertEqual(model.license_name, "MIT") - self.assertEqual( - model.download_patterns, - ("config.json", "model.safetensors"), - ) + self.assertEqual(model.repository, "onnx-community/BEN2-ONNX") + self.assertEqual(model.license_name, "Apache 2.0") + self.assertEqual(model.download_patterns, ("onnx/model_fp16.onnx",)) + self.assertEqual(model.ready_patterns, ("onnx/model_fp16.onnx",)) + self.assertEqual(model.r2_files[0][0], "onnx/model_fp16.onnx") self.assertNotIn("BEN2_BASE", self.models.MODELS) - def test_da3_model_directories_share_the_da3_prefix(self): + def test_da3_model_directories_use_repository_names(self): for model_key, model in self.models.MODELS.items(): + if not model.enabled: + continue with self.subTest(model_key=model_key): self.assertTrue( - model.directory_name.startswith("da3_"), + model.directory_name.startswith("DA3-"), model.directory_name, ) self.assertEqual( @@ -74,20 +69,37 @@ class ModelTest(unittest.TestCase): model.directory_name, ) - def test_ben2_model_directory_uses_ben2_prefix(self): + def test_ben2_model_directory_uses_repository_name(self): self.assertEqual( self.models.model_directory_name("BEN2_BASE"), - "ben2_base", + "BEN2-ONNX", ) - def test_moge2_uses_single_checkpoint_and_mit_license(self): - model = self.models.get_downloadable_model("MOGE2_VITB_NORMAL") + def test_enabled_da3_models_use_onnx_and_r2_checksums(self): + for model in self.models.MODELS.values(): + if not model.enabled: + continue + with self.subTest(model=model.key): + self.assertTrue(model.download_patterns[0].endswith(".onnx")) + self.assertEqual(model.ready_patterns, ("model.onnx",)) + self.assertEqual(model.r2_files[0][0], "model.onnx") - self.assertEqual(model.repository, "Ruicheng/moge-2-vitb-normal") - self.assertEqual(model.license_name, "MIT") - self.assertEqual(model.download_patterns, ("model.pt",)) - self.assertEqual(model.ready_patterns, ("model.pt",)) - self.assertFalse(model.supports_gs) + def test_moge2_models_use_official_onnx_models(self): + for key, variant in (("MOGE2_VITS_NORMAL", "vits"), + ("MOGE2_VITB_NORMAL", "vitb"), + ("MOGE2_VITL_NORMAL", "vitl")): + with self.subTest(key=key): + model = self.models.get_downloadable_model(key) + self.assertEqual( + model.repository, + f"Ruicheng/moge-2-{variant}-normal-onnx", + ) + self.assertEqual(model.license_name, "MIT") + self.assertEqual(model.download_patterns, ("model.onnx",)) + self.assertEqual(model.ready_patterns, ("model.onnx",)) + self.assertEqual(model.r2_directory, model.repository.split("/")[-1]) + self.assertEqual(model.r2_files[0][0], "model.onnx") + self.assertFalse(model.supports_gs) def test_moge2_models_match_uploaded_r2_variants(self): self.assertEqual( @@ -96,5 +108,9 @@ class ModelTest(unittest.TestCase): ) self.assertEqual( [model.r2_directory for model in self.models.MOGE2_MODELS.values()], - ["MoGe2ViTSNormal", "MoGe2ViTBNormal", "MoGe2ViTLNormal"], + [ + "moge-2-vits-normal-onnx", + "moge-2-vitb-normal-onnx", + "moge-2-vitl-normal-onnx", + ], ) diff --git a/tests/test_moge_onnx.py b/tests/test_moge_onnx.py new file mode 100644 index 0000000..1c51364 --- /dev/null +++ b/tests/test_moge_onnx.py @@ -0,0 +1,130 @@ +import importlib.util +import sys +import unittest +from pathlib import Path +from unittest.mock import patch + +import numpy as np + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] + + +def load_module(): + specification = importlib.util.spec_from_file_location( + "snapform_moge_onnx", + PROJECT_ROOT / "src" / "snapform" / "moge_onnx.py", + ) + module = importlib.util.module_from_spec(specification) + specification.loader.exec_module(module) + return module + + +class FakeInput: + def __init__(self, name): + self.name = name + + +class FakeSession: + def __init__(self, outputs): + self.outputs = outputs + self.calls = [] + + def run(self, names, inputs): + self.calls.append((names, inputs)) + return [self.outputs[name] for name in names] + + def get_outputs(self): + return [FakeInput(name) for name in self.outputs] + + +class MogeOnnxTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.moge = load_module() + + def test_auto_provider_prefers_accelerated_runtime(self): + providers = self.moge.select_providers( + "auto", + ["CPUExecutionProvider", "DmlExecutionProvider"], + ) + + self.assertEqual( + providers, + ["DmlExecutionProvider", "CPUExecutionProvider"], + ) + + def test_explicit_unavailable_provider_falls_back_to_cpu(self): + providers = self.moge.select_providers( + "cuda", + ["CPUExecutionProvider"], + ) + + self.assertEqual(providers, ["CPUExecutionProvider"]) + + def test_infer_uses_dynamic_token_input_and_numpy_postprocess(self): + height, width = 8, 12 + raw = { + "points": np.zeros((1, height, width, 3), dtype=np.float32), + "normal": np.zeros((1, height, width, 3), dtype=np.float32), + "mask": np.ones((1, height, width), dtype=np.float32), + "scale": np.ones((1,), dtype=np.float32), + } + raw["points"][..., 2] = 1.0 + raw["normal"][..., 2] = 1.0 + session = FakeSession(raw) + expected = { + "points": np.zeros((height, width, 3), dtype=np.float32), + "depth": np.ones((height, width), dtype=np.float32), + "normal": raw["normal"][0], + "mask": np.ones((height, width), dtype=bool), + "intrinsics": np.eye(3, dtype=np.float32), + } + + with patch.object(self.moge, "postprocess", return_value=expected): + result = self.moge.infer( + session, + np.zeros((height, width, 3), dtype=np.float32), + resolution_level=9, + ) + + self.assertIs(result, expected) + names, inputs = session.calls[0] + self.assertEqual(names, ["points", "normal", "mask", "scale"]) + self.assertEqual(inputs["image"].shape, (1, 3, height, width)) + self.assertEqual(inputs["num_tokens"].dtype, np.int64) + self.assertEqual(inputs["num_tokens"], 3600) + + def test_postprocess_recovers_camera_projection_and_metric_scale(self): + height, width = 64, 80 + focal = np.float32(0.9) + shift = np.float32(0.3) + metric_scale = np.float32(2.0) + uv = self.moge.normalized_view_plane_uv(width, height, np.float32) + depth = np.linspace(1.0, 2.4, height * width, dtype=np.float32).reshape( + height, + width, + ) + raw_points = np.concatenate( + (uv * depth[..., None] / focal, (depth - shift)[..., None]), + axis=-1, + ) + raw = { + "points": raw_points[None], + "normal": np.broadcast_to( + np.array((0.0, 0.0, 1.0), dtype=np.float32), + (1, height, width, 3), + ), + "mask": np.ones((1, height, width), dtype=np.float32), + "metric_scale": np.array((metric_scale,), dtype=np.float32), + } + + result = self.moge.postprocess(raw) + + np.testing.assert_allclose(result["depth"], depth * metric_scale, rtol=1e-3) + np.testing.assert_allclose(result["points"][..., 2], result["depth"], rtol=1e-5) + self.assertTrue(result["mask"].all()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_pack.py b/tests/test_pack.py index 14f08ce..3679248 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -1,7 +1,4 @@ -import ast import importlib.util -import io -import tarfile import tempfile import tomllib import unittest @@ -64,8 +61,6 @@ class PackTest(unittest.TestCase): self.assertIn("service.py", names) self.assertNotIn("snapform/__init__.py", names) self.assertNotIn("importer.py", names) - self.assertIn("vendor/depth-anything-3.tar.gz", names) - self.assertIn("vendor/ben2.tar.gz", names) self.assertIn( f"vendor/{pack.PLATFORMS[platform]}", names, @@ -108,95 +103,5 @@ class PackTest(unittest.TestCase): {"windows-x64", "macos-arm64", "linux-x64"}, ) - def test_bundled_depth_anything_3_has_no_xformers_dependency(self): - pack = load_pack() - with tarfile.open( - fileobj=io.BytesIO( - pack.bundled_artifact_bytes("depth-anything-3.tar.gz") - ), - mode="r:gz", - ) as archive: - raw_pyproject = archive.extractfile( - "Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/" - "pyproject.toml" - ).read() - pyproject = raw_pyproject.decode("utf-8") - self.assertNotIn("xformers", pyproject) - self.assertFalse(raw_pyproject.startswith(b"\xef\xbb\xbf")) - - def test_bundled_depth_anything_3_exports_are_lazy(self): - pack = load_pack() - with tarfile.open( - fileobj=io.BytesIO( - pack.bundled_artifact_bytes("depth-anything-3.tar.gz") - ), - mode="r:gz", - ) as archive: - source = archive.extractfile( - "Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/" - "src/depth_anything_3/utils/export/__init__.py" - ).read().decode("utf-8") - tree = ast.parse(source) - module_imports = [ - ast.get_source_segment(source, node) - for node in tree.body - if isinstance(node, (ast.Import, ast.ImportFrom)) - ] - self.assertEqual( - module_imports, - ["from depth_anything_3.specs import Prediction"], - ) - - def test_official_depth_anything_3_baseline_is_pristine(self): - pack = load_pack() - archive_path = pack.VENDOR_DIR / "depth-anything-3.tar.gz" - with tarfile.open(archive_path, "r:gz") as archive: - raw_pyproject = archive.extractfile( - "Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/" - "pyproject.toml" - ).read() - export_source = archive.extractfile( - "Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/" - "src/depth_anything_3/utils/export/__init__.py" - ).read().decode("utf-8") - pyproject = raw_pyproject.decode("utf-8") - self.assertIn("xformers", pyproject) - self.assertFalse(raw_pyproject.startswith(b"\xef\xbb\xbf")) - self.assertIn( - "from depth_anything_3.utils.export.gs import", - export_source, - ) - - def test_bundled_depth_anything_3_qr_falls_back_on_mps(self): - pack = load_pack() - with tarfile.open( - fileobj=io.BytesIO( - pack.bundled_artifact_bytes("depth-anything-3.tar.gz") - ), - mode="r:gz", - ) as archive: - source = archive.extractfile( - "Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/" - "src/depth_anything_3/utils/ray_utils.py" - ).read().decode("utf-8") - self.assertIn('use_cpu = device.type == "mps"', source) - self.assertIn("A_tilde = A_tilde.cpu()", source) - self.assertIn("Q_tilde = Q_tilde.to(device=device)", source) - - def test_official_depth_anything_3_ray_utils_is_pristine(self): - pack = load_pack() - archive_path = pack.VENDOR_DIR / "depth-anything-3.tar.gz" - with tarfile.open(archive_path, "r:gz") as archive: - source = archive.extractfile( - "Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/" - "src/depth_anything_3/utils/ray_utils.py" - ).read().decode("utf-8") - self.assertIn( - "Q_tilde, R_tilde = torch.linalg.qr(A_tilde)", - source, - ) - self.assertNotIn("use_cpu", source) - - if __name__ == "__main__": unittest.main() diff --git a/tests/test_service.py b/tests/test_service.py index 343fbac..05cf6c5 100644 --- a/tests/test_service.py +++ b/tests/test_service.py @@ -40,7 +40,6 @@ class ServiceTest(unittest.TestCase): def create_cache(self): cache = self.service.ModelCache.__new__(self.service.ModelCache) - cache.torch = SimpleNamespace(cuda=FakeCuda()) cache.ben2_model = None cache.ben2_key = None cache.da3_model = None @@ -52,65 +51,41 @@ class ServiceTest(unittest.TestCase): def test_ben2_model_is_reused_for_matching_runtime_key(self): loads = [] - class FakeModel: - @classmethod - def from_pretrained(cls, model_dir): - loads.append(model_dir) - return cls() - - def to(self, device): - self.device = device - return self - - def eval(self): - return None - - ben2 = ModuleType("ben2") - ben2.BEN_Base = FakeModel + ben2 = ModuleType("ben2_onnx") + ben2.create_session = lambda model_dir, device: ( + loads.append((model_dir, device)) or object() + ) cache = self.create_cache() device = SimpleNamespace(type="cuda") - with patch.dict(sys.modules, {"ben2": ben2}): + with patch.dict(sys.modules, {"ben2_onnx": ben2}): first, _load_ms = cache.get_ben2("model", device) second, cached_load_ms = cache.get_ben2("model", device) self.assertIs(first, second) - self.assertEqual(loads, ["model"]) + self.assertEqual(loads, [("model", "cuda")]) self.assertEqual(cache.status()["ben2"], "model") self.assertEqual(cached_load_ms, 0.0) def test_da3_model_is_reused_for_matching_runtime_key(self): loads = [] - class FakeModel: - @classmethod - def from_pretrained(cls, model_dir): - loads.append(model_dir) - return cls() - - def to(self, device): - self.device = device - return self - - def eval(self): - return None - - package = ModuleType("depth_anything_3") - api = ModuleType("depth_anything_3.api") - api.DepthAnything3 = FakeModel + da3 = ModuleType("da3_onnx") + da3.create_session = lambda model_dir, device: ( + loads.append((model_dir, device)) or object() + ) cache = self.create_cache() device = SimpleNamespace(type="cuda") with patch.dict( sys.modules, { - "depth_anything_3": package, - "depth_anything_3.api": api, + "da3_onnx": da3, }, ): first = cache.get_da3("model", device) second = cache.get_da3("model", device) self.assertIs(first, second) - self.assertEqual(loads, ["model"]) + self.assertEqual(loads, [("model", "cuda")]) self.assertEqual(cache.status()["da3"], "model") def test_depth_model_cache_keeps_only_da3_or_moge(self): @@ -135,13 +110,10 @@ class ServiceTest(unittest.TestCase): moge = object() - def load_moge(_model_dir, _device, key): - cache.moge_model = moge - cache.moge_key = key - return moge - - with patch.object(cache, "_load_moge", side_effect=load_moge): - self.assertIs(cache.get_moge("moge", device), moge) + fake_moge = ModuleType("moge_onnx") + fake_moge.create_session = lambda _model_dir, _device: moge + with patch.dict(sys.modules, {"moge_onnx": fake_moge}): + self.assertIs(cache.get_moge_onnx("moge", "cuda"), moge) self.assertIsNone(cache.da3_model) self.assertIsNotNone(cache.ben2_model) diff --git a/tests/test_worker.py b/tests/test_worker.py index c2844c4..6895293 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -66,57 +66,6 @@ class WorkerTest(unittest.TestCase): self.worker.check_cancelled(lambda: False) - 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_reconstruct_parser_accepts_background_model_directory(self): args = self.worker.parser().parse_args( [ @@ -436,12 +385,6 @@ class WorkerTest(unittest.TestCase): with self.assertRaisesRegex(RuntimeError, "Unsupported input type"): self.worker.prepare_input(input_path, 16, directory) - def test_moge_runtime_disables_unsupported_xformers_attention(self): - with patch.dict(os.environ, {}, clear=True): - self.worker.prepare_moge_runtime() - - self.assertEqual(os.environ["XFORMERS_DISABLED"], "1") - def test_download_model_uses_actual_request_without_preflight(self): calls = [] fake_hugging_face = ModuleType("huggingface_hub") @@ -605,8 +548,11 @@ class WorkerTest(unittest.TestCase): ) fake_torch.device = lambda device_type: SimpleNamespace(type=device_type) fake_torch.inference_mode = nullcontext - fake_ben2 = ModuleType("ben2") - fake_ben2.BEN_Base = FakeModel + fake_ben2 = ModuleType("ben2_onnx") + fake_ben2.create_session = lambda _model_dir, _device: object() + fake_ben2.infer = lambda _model, image: ( + inference_options.append((image, {})) or FakeImage() + ) fake_pil = ModuleType("PIL") fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage()) @@ -614,7 +560,7 @@ class WorkerTest(unittest.TestCase): sys.modules, { "torch": fake_torch, - "ben2": fake_ben2, + "ben2_onnx": fake_ben2, "PIL": fake_pil, }, ): @@ -635,7 +581,7 @@ class WorkerTest(unittest.TestCase): self.assertEqual( inference_options[0][1], - {"refine_foreground": False}, + {}, ) self.assertEqual( save_options[0], @@ -690,8 +636,11 @@ class WorkerTest(unittest.TestCase): ) fake_torch.device = lambda device_type: SimpleNamespace(type=device_type) fake_torch.inference_mode = nullcontext - fake_ben2 = ModuleType("ben2") - fake_ben2.BEN_Base = FakeModel + fake_ben2 = ModuleType("ben2_onnx") + fake_ben2.create_session = lambda _model_dir, _device: object() + fake_ben2.infer = lambda _model, image: ( + inference_options.append((image, {})) or FakeImage() + ) fake_pil = ModuleType("PIL") fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage()) @@ -699,7 +648,7 @@ class WorkerTest(unittest.TestCase): sys.modules, { "torch": fake_torch, - "ben2": fake_ben2, + "ben2_onnx": fake_ben2, "PIL": fake_pil, }, ): @@ -734,7 +683,7 @@ class WorkerTest(unittest.TestCase): self.assertEqual(len(inference_options), 2) for _image, options in inference_options: - self.assertEqual(options, {"refine_foreground": False}) + self.assertEqual(options, {}) def test_remove_background_parser_accepts_video_frames(self): args = self.worker.parser().parse_args( @@ -813,64 +762,25 @@ class WorkerTest(unittest.TestCase): def test_uses_da3_export_parameters_for_formats_and_debug_mesh(self): inference_options = [] - export_calls = [] mesh_calls = [] - class FakeModel: - @classmethod - def from_pretrained(cls, model_dir): - if model_dir != "model": - raise AssertionError(f"Unexpected model directory: {model_dir}") - return cls() - - def to(self, device): - if device.type != "cpu": - raise AssertionError(f"Unexpected device: {device.type}") - return self - - def eval(self): - return None - - def inference(self, images, **options): - inference_options.append((images, options)) - return SimpleNamespace( - depth=np.full((1, 2, 2), 2.0, dtype=np.float32), - ) - - def fake_export(prediction, export_format, export_dir, **options): - export_calls.append((prediction, export_format, export_dir, options)) - output_directory = Path(export_dir) - if export_format == "glb": - (output_directory / "scene.glb").touch() - elif export_format == "gs_ply": - ply_directory = output_directory / "gs_ply" - ply_directory.mkdir() - (ply_directory / "0000.ply").touch() - elif export_format == "depth_vis": - depth_directory = output_directory / "depth_vis" - depth_directory.mkdir() - (depth_directory / "0000.jpg").touch() - else: - mini_directory = output_directory / "exports" / "mini_npz" - mini_directory.mkdir(parents=True) - np.savez_compressed( - mini_directory / "results.npz", - depth=np.zeros((14, 14)), - ) - - fake_torch = ModuleType("torch") - fake_torch.device = lambda device_type: SimpleNamespace( - type=device_type, + prediction = 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_torch.inference_mode = nullcontext - fake_api = ModuleType("depth_anything_3.api") - fake_api.DepthAnything3 = FakeModel - fake_export_module = ModuleType("depth_anything_3.utils.export") - fake_export_module.export = fake_export - fake_utils = ModuleType("depth_anything_3.utils") - fake_utils.__path__ = [] - fake_package = ModuleType("depth_anything_3") - fake_package.__path__ = [] + fake_da3 = ModuleType("da3_onnx") + fake_da3.create_session = lambda _model_dir, _device: object() + + def fake_infer(_session, images, **options): + inference_options.append((images, options)) + return prediction + + fake_da3.infer = fake_infer + fake_runtime = ModuleType("onnx_runtime") + fake_runtime.provider_label = lambda _session: "CPU" def fake_mesh_export( prediction, @@ -893,6 +803,20 @@ class WorkerTest(unittest.TestCase): output_path.touch() return output_path + def fake_point_cloud(_prediction, output_directory, *_args): + (Path(output_directory) / "scene.glb").touch() + + def fake_depth_preview(_prediction, output_directory): + target = Path(output_directory) / "depth_vis" + target.mkdir() + (target / "0000.jpg").touch() + + def fake_mini_npz(_prediction, output_directory): + np.savez_compressed( + Path(output_directory) / "results.npz", + depth=np.zeros((1, 2, 2)), + ) + with tempfile.TemporaryDirectory() as directory: root = Path(directory) input_path = root / "input.png" @@ -900,20 +824,28 @@ class WorkerTest(unittest.TestCase): with patch.dict( sys.modules, { - "torch": fake_torch, - "depth_anything_3": fake_package, - "depth_anything_3.api": fake_api, - "depth_anything_3.utils": fake_utils, - "depth_anything_3.utils.export": fake_export_module, + "da3_onnx": fake_da3, + "onnx_runtime": fake_runtime, }, ), patch.object( self.worker, "export_textured_mesh_glb", side_effect=fake_mesh_export, + ), patch.object( + self.worker, + "export_da3_point_cloud", + side_effect=fake_point_cloud, + ), patch.object( + self.worker, + "export_da3_depth_preview", + side_effect=fake_depth_preview, + ), patch.object( + self.worker, + "export_da3_mini_npz", + side_effect=fake_mini_npz, ): debug_formats = ( "glb", - "gs_ply", "depth_vis", "mini_npz", "mesh", @@ -945,15 +877,11 @@ class WorkerTest(unittest.TestCase): self.assertEqual(status["progress"], 1.0) self.assertIn(export_format, status["message"]) result_directory = Path(args.output_dir) - if export_format == "gs_ply": - self.assertTrue((result_directory / "0000.ply").is_file()) - self.assertFalse((result_directory / "gs_ply").exists()) if export_format == "depth_vis": self.assertTrue((result_directory / "0000.jpg").is_file()) self.assertFalse((result_directory / "depth_vis").exists()) - self.assertEqual(len(inference_options), 5) - self.assertEqual(len(export_calls), 4) + self.assertEqual(len(inference_options), 4) self.assertEqual(len(mesh_calls), 1) for images, options in inference_options: self.assertEqual(images, [str(input_path)]) @@ -964,82 +892,21 @@ class WorkerTest(unittest.TestCase): ) self.assertTrue(options["use_ray_pose"]) self.assertEqual(options["ref_view_strategy"], "middle") - for ( - _prediction, - export_format, - _export_dir, - options, - ) in export_calls: - np.testing.assert_array_equal( - _prediction.depth, - np.full((1, 2, 2), 2.0, dtype=np.float32), - ) - if export_format == "glb": - self.assertEqual( - options, - { - "glb": { - "conf_thresh_percentile": 35.0, - "num_max_points": 12345, - "show_cameras": False, - "export_depth_vis": False, - } - }, - ) - elif export_format == "gs_ply": - self.assertEqual( - options, - {"gs_ply": {"gs_views_interval": 2}}, - ) - else: - self.assertEqual(options, {}) self.assertEqual(mesh_calls[0][1], [input_path]) self.assertEqual(mesh_calls[0][2:], (35.0, True, 0.2)) - infer_gs_values = [options["infer_gs"] for _, options in inference_options] - self.assertEqual( - infer_gs_values, - [False, True, False, False, False], - ) - def test_reconstruct_copies_frames_for_multi_frame_input(self): - class FakeModel: - @classmethod - def from_pretrained(cls, model_dir): - return cls() - - def to(self, device): - return self - - def eval(self): - return None - - def inference(self, images, **options): - return SimpleNamespace( - depth=np.full((2, 2, 2), 2.0, dtype=np.float32), - ) - - def fake_export(prediction, export_format, export_dir, **options): - mini_directory = Path(export_dir) / "exports" / "mini_npz" - mini_directory.mkdir(parents=True) - np.savez_compressed( - mini_directory / "results.npz", - depth=np.zeros((2, 14, 14)), - ) - - fake_torch = ModuleType("torch") - fake_torch.device = lambda device_type: SimpleNamespace( - type=device_type, + 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), ) - fake_torch.inference_mode = nullcontext - fake_api = ModuleType("depth_anything_3.api") - fake_api.DepthAnything3 = FakeModel - fake_export_module = ModuleType("depth_anything_3.utils.export") - fake_export_module.export = fake_export - fake_utils = ModuleType("depth_anything_3.utils") - fake_utils.__path__ = [] - fake_package = ModuleType("depth_anything_3") - fake_package.__path__ = [] + fake_runtime = ModuleType("onnx_runtime") + fake_runtime.provider_label = lambda _session: "CPU" with tempfile.TemporaryDirectory() as directory: root = Path(directory) @@ -1050,11 +917,8 @@ class WorkerTest(unittest.TestCase): with patch.dict( sys.modules, { - "torch": fake_torch, - "depth_anything_3": fake_package, - "depth_anything_3.api": fake_api, - "depth_anything_3.utils": fake_utils, - "depth_anything_3.utils.export": fake_export_module, + "da3_onnx": fake_da3, + "onnx_runtime": fake_runtime, }, ): args = SimpleNamespace( diff --git a/vendor.lock.toml b/vendor.lock.toml index 0676be0..fe40004 100644 --- a/vendor.lock.toml +++ b/vendor.lock.toml @@ -1,32 +1,5 @@ schema_version = 1 -[artifacts.depth-anything-3] -revision = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4" -filename = "depth-anything-3.tar.gz" -url = "https://codeload.github.com/ByteDance-Seed/Depth-Anything-3/tar.gz/3d835ec1a5802d64a8b8b15f817a1ab54809bfe4" -sha256 = "98aa2dd53ab44b96cef5190ae4841c6ae51797d099b784e08e12a1a55bd3a69b" -patches = "depth-anything-3" - -[artifacts.ben2] -revision = "2c99a5da477b5523585bfa5c893888a6e818a8f6" -filename = "ben2.tar.gz" -url = "https://codeload.github.com/PramaLLC/BEN2/tar.gz/2c99a5da477b5523585bfa5c893888a6e818a8f6" -sha256 = "01fd26716f3ff3a3f18ad2336c768af5b759edfc7c38bb3bf733b7c86eee6282" - -[artifacts.moge] -revision = "925b8ed835a7a9cdb7578ba15c658a0afc969030" -filename = "moge.tar.gz" -url = "https://codeload.github.com/microsoft/MoGe/tar.gz/925b8ed835a7a9cdb7578ba15c658a0afc969030" -sha256 = "c6d90b2564d643f0f3ae65b5ffe163abaea3b55d194c6e2ba9976d67f552bd26" -patches = "moge" - -[artifacts.utils3d] -revision = "3fab839f0be9931dac7c8488eb0e1600c236e183" -filename = "utils3d.tar.gz" -url = "https://codeload.github.com/EasternJournalist/utils3d/tar.gz/3fab839f0be9931dac7c8488eb0e1600c236e183" -sha256 = "f1f167b1f68809bea391a66dd0d2ad0e14f7382cbddf087b655181c31a285abb" -patches = "utils3d" - [uv] version = "0.11.3"