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

feat: Implement ONNX model inference for BEN2, Depth Anything 3, and MoGe

Browse Source 
- 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.
This commit is contained in:
马南
2026-08-14 10:27:09 +08:00
parent 13fab9ce88
commit 2ad30bf4a3
41 changed files with 1569 additions and 2109 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@@ -8,3 +8,4 @@ build/
.pytest_cache/ .pytest_cache/
/vendor/ /vendor/
/models/

14
README.md
View File

@@ -16,11 +16,11 @@ SnapForm 是一个 Blender 扩展，可以把一张照片、一段视频或一
2. 打开 Blender，进入 `Preferences > Extensions` 2. 打开 Blender，进入 `Preferences > Extensions`
3. 点击 `Install from Disk`，选择下载的压缩包 3. 点击 `Install from Disk`，选择下载的压缩包
4. 在扩展偏好设置中确认 Storage Root 与 Device（保持默认即可）；Device 默认取本机 4. 在扩展偏好设置中确认 Storage Root 与 Device（保持默认即可）；Device 默认取本机
最快的 CUDA/MPS，否则 CPU。下载超时会自动重试国内镜像，无需手动配置 最快的 DirectML/CoreML provider，否则 CPU。下载超时会自动重试国内镜像，无需手动配置
5. 打开 3D 视图右侧的 `SnapForm` 面板，点击 `Install Environment` 5. 打开 3D 视图右侧的 `SnapForm` 面板，点击 `Install Environment`
DA3、MoGe-2、utils3d、BEN2 源码包和 uv 已内置在扩展里，安装时直接从本地解包，不需要访问 GitHub。首次使用仍会自动 扩展只内置对应平台的 uv。首次使用会自动下载轻量的 ONNX Runtime 环境，随后自动下载 `MoGe-2 ViT-B Normal` 和
下载运行环境依赖（PyTorch 等），随后自动下载 `MoGe-2 ViT-B Normal` 和 `BEN2 Base`。其他 AI 模型可在 `BEN2 Base`；不安装 PyTorch，也不下载各模型的 Python 源码包。其他 AI 模型可在
`Models` panel 中按需下载，请保持网络畅通，首次安装需要一些时间。 `Models` panel 中按需下载，请保持网络畅通，首次安装需要一些时间。
## 快速上手 ## 快速上手
@@ -46,14 +46,18 @@ DA3、MoGe-2、utils3d、BEN2 源码包和 uv 已内置在扩展里，安装时
选中单张图片的 Image Empty，右键打开 `SnapForm` 子菜单并选择 `Convert to Surface`。可选择 MoGe-2 选中单张图片的 Image Empty，右键打开 `SnapForm` 子菜单并选择 `Convert to Surface`。可选择 MoGe-2
ViT-S、ViT-B、ViT-L 或 DA3 Small、Base，其中 MoGe-2 ViT-B 为默认模型。SnapForm 会按模型输出的 ViT-S、ViT-B、ViT-L 或 DA3 Small、Base，其中 MoGe-2 ViT-B 为默认模型。SnapForm 会按模型输出的
相机内参与三维点建立透视网格，并匹配原 Empty 的可见尺寸、Offset 与 Transform。MoGe-2 直接使用官方 相机内参与三维点建立透视网格，并匹配原 Empty 的可见尺寸、Offset 与 Transform。MoGe-2 使用内置 NumPy
`utils3d.build_mesh_from_map` 与深度断层 mask 生成干净拓扑；DA3 保留 confidence 与边拉伸剔除。两者再由独立的 Geometry Nodes 网格构建与深度断层 mask 生成干净拓扑；DA3 保留 confidence 与边拉伸剔除。两者再由独立的 Geometry Nodes
修改器完成剔面、快速重拓扑、Relax 与 UV 重投。开启 Remove Background 时，背景三角面会在转换期间 修改器完成剔面、快速重拓扑、Relax 与 UV 重投。开启 Remove Background 时，背景三角面会在转换期间
直接删除，不再交给 Geometry Nodes。匹配相机默认不创建；开启后会作为 Surface 的子级。 直接删除，不再交给 Geometry Nodes。匹配相机默认不创建；开启后会作为 Surface 的子级。
MoGe-2 的像素级 normal 会作为 Object Space normal 贴图接入 `O MoGe2 Surface Layer`，并与原图 MoGe-2 的像素级 normal 会作为 Object Space normal 贴图接入 `O MoGe2 Surface Layer`，并与原图
Bump 合并；网格本身不写入 custom normal，因此重拓扑降低顶点数后仍能保留细腻的法线细节。 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 显示 转换对话框提供模型与 `Create Matching Camera`、`Remove Background`、`Unshaded`；DA3 显示
`Resolution Scale`，MoGe-2 显示 `Resolution Level`， `Resolution Scale`，MoGe-2 显示 `Resolution Level`，
以及分割线下方的 `Keep Original`。 以及分割线下方的 `Keep Original`。

7
docs/depth-plane.md
View File

@@ -4,17 +4,17 @@
1. 在 3D View 选中 `Object Type = Empty`、`Display As = Image` 的对象。 1. 在 3D View 选中 `Object Type = Empty`、`Display As = Image` 的对象。
2. 右键打开 `SnapForm` 子菜单并选择 `Convert to Depth Plane`，通过 `Depth from` 选择深度来源： 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 以百分比 Resolution Scale，再钳制到 `252–1008`，最后对齐到最接近的 `14` 倍数。Resolution Scale 以百分比
显示，范围为 `1–100%`，默认 `100%`；界面只在标签中实时显示最终 Process Res。 显示，范围为 `1–100%`，默认 `100%`；界面只在标签中实时显示最终 Process Res。
转换导出 `mini_npz`，使用前需要在 转换导出 `mini_npz`，使用前需要在
`Models` panel 下载所选模型，并固定使用 `process_res_method="upper_bound_resize"`；设备和参考视图策略 `Models` panel 下载所选模型，并固定使用 `process_res_method="upper_bound_resize"`；设备和参考视图策略
继续使用场景设置。Model 下拉菜单只列出已经下载的 DA3 模型；若没有可用模型，对话框会引导到 继续使用场景设置。Model 下拉菜单只列出已经下载的 DA3 模型；若没有可用模型，对话框会引导到
`Models` panel 下载。 `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 - `Original`：直接把原图作为深度贴图，不显示估算参数。仅在开启 `Remove Background` 时运行 BEN2
取得前景 Alpha；视频或序列帧按需同时抽取原始帧与前景序列，把原始帧作为 `.depth` 序列。 取得前景 Alpha；视频或序列帧按需同时抽取原始帧与前景序列，把原始帧作为 `.depth` 序列。
`Original` 是默认深度来源。选择 `DA3` 时，其 Model、Use Ray Pose 与 Resolution Scale `Original` 是默认深度来源。选择 `DA3` 时，其 Model 与 Resolution Scale
参数会集中显示在独立边框中。`Remove Background` 与 `Unshaded` 直接显示，不再使用 Options 折叠； 参数会集中显示在独立边框中。`Remove Background` 与 `Unshaded` 直接显示，不再使用 Options 折叠；
`Keep Original` 默认关闭，并在分割线下方保持为最后一项。仅当开启 `Remove Background` 时才需要 `Keep Original` 默认关闭，并在分割线下方保持为最后一项。仅当开启 `Remove Background` 时才需要
`BEN2 Base`，并把 `.foreground` 的 Alpha 作为独立 `BEN2 Base`，并把 `.foreground` 的 Alpha 作为独立
@@ -63,4 +63,3 @@
节点组根据 `Frame` 计算从 `1` 开始的 Image Texture Frame：场景帧与 Start 都为 `1` 时读取图片第 `1` 帧；非循环输入 节点组根据 `Frame` 计算从 `1` 开始的 Image Texture Frame：场景帧与 Start 都为 `1` 时读取图片第 `1` 帧；非循环输入
钳制到有效范围；循环输入使用 Floored Modulo，让 `Frame Start` 前后的负、正相对帧都映射到 钳制到有效范围；循环输入使用 Floored Modulo，让 `Frame Start` 前后的负、正相对帧都映射到
`1…Frames`，因此开始帧之前也会循环播放。Geometry 置换会像材质 Image Texture 一样随场景帧刷新。 `1…Frames`，因此开始帧之前也会循环播放。Geometry 置换会像材质 Image Texture 一样随场景帧刷新。

10
docs/development.md
View File

@@ -12,11 +12,10 @@ uv run python pack.py
```text ```text
. .
├── vendor.lock.toml 第三方构建输入的版本、URL 与 SHA-256 ├── vendor.lock.toml 各平台 uv 的版本、URL 与 SHA-256
├── prepare_vendor.py 下载并校验第三方构建输入 ├── prepare_vendor.py 下载并校验各平台 uv
├── pack.py ├── pack.py
├── patches/ 进入 Git 的 DA3 源码覆盖文件 ├── vendor/ 不进入 Git 的 uv 下载缓存
├── vendor/ 不进入 Git 的下载缓存；本地调试也直接读取这里
├── src/ ├── src/
│ └── snapform/ │ └── snapform/
│ ├── blender_manifest.toml │ ├── blender_manifest.toml
@@ -36,7 +35,6 @@ uv run python pack.py
## 上游与许可证 ## 上游与许可证
- Depth Anything 3 代码采用 Apache 2.0。 - Depth Anything 3 代码采用 Apache 2.0。
- BEN2 代码与模型权重采用 MIT License。 - BEN2 ONNX 模型采用其仓库声明的 Apache 2.0 License。
- 本 Extension 代码采用 GPL-3.0-or-later。 - 本 Extension 代码采用 GPL-3.0-or-later。
- 模型权重许可证由所选 Hugging Face 模型仓库决定，Extension 下载前会在界面显示。 - 模型权重许可证由所选 Hugging Face 模型仓库决定，Extension 下载前会在界面显示。

21
docs/features.md
View File

@@ -12,14 +12,14 @@
## 当前功能 ## 当前功能
- 使用 Blender Extension manifest，可通过 `Install from Disk` 安装 - 使用 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 使用独立缓存， - Service 的深度模型缓存只保留 DA3 或 MoGe-2 其中一个，切换时主动释放前一个；BEN2 使用独立缓存，
可与当前深度模型共存 可与当前深度模型共存
- Blender 使用内置 `urllib.request` 访问仅监听 `127.0.0.1` 的 Service；请求带随机 Token - Blender 使用内置 `urllib.request` 访问仅监听 `127.0.0.1` 的 Service；请求带随机 Token
- 常驻 Service 复用 Torch、CUDA、BEN2 与最近使用的 DA3 模型，连续执行时不再重复加载 - 常驻 Service 复用 ONNX Runtime session、BEN2 与最近使用的深度模型，连续执行时不再重复加载
- 多个 Blender 实例自动连接同一个共享 Service；`AI Service` 面板显示状态、端口、客户端与模型缓存，并提供安装、启动、清理、重启和停止控制 - 多个 Blender 实例自动连接同一个共享 Service；`AI Service` 面板显示状态、端口、客户端与模型缓存，并提供安装、启动、清理、重启和停止控制
- Environment、模型和任务结果默认保存在 `~/.snapform`，可由多个 Blender 版本共享 - 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 与推理耗时 - 支持 BEN2 Base 单图背景剔除测试，输出透明 PNG 与推理耗时
- 支持单张 JPG、JPEG、PNG、WEBP 图片或常见视频文件输入 - 支持单张 JPG、JPEG、PNG、WEBP 图片或常见视频文件输入
- 视频按顺序截取前若干帧，可设置参与重建的最大帧数 - 视频按顺序截取前若干帧，可设置参与重建的最大帧数
@@ -52,17 +52,16 @@
1. 先运行 `uv run python prepare_vendor.py`，按 `vendor.lock.toml` 下载并校验固定版本的 1. 先运行 `uv run python prepare_vendor.py`，按 `vendor.lock.toml` 下载并校验固定版本的
第三方构建输入；再运行 `uv run python pack.py`，分别构建 `windows-x64`、`macos-arm64`、 第三方构建输入；再运行 `uv run python pack.py`，分别构建 `windows-x64`、`macos-arm64`、
`linux-x64` 三个平台的 Extension ZIP（如 `SnapForm.v0.10.12.macos-arm64.zip`）。 `linux-x64` 三个平台的 Extension ZIP（如 `SnapForm.v0.10.12.macos-arm64.zip`）。
每个包只含对应平台的 uv 二进制，并内置 DA3、BEN2 源码包。 每个包只含对应平台的 uv 二进制，不再携带 AI 项目的源码包。
2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。 2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。
3. 在 Extension Preferences 中确认 Storage Root 与 Device。Device 是正式功能 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； 4. 在 `AI Service` 点击 `Install Environment`。首次安装会下载 Python、ONNX Runtime、FastAPI、Uvicorn；
DA3、BEN2 源码包与 uv 已内置在扩展中，安装时从 安装任务的临时目录中使用，安装结束后自动清理， uv 已内置在扩展中，DA3、BEN2 与 MoGe 不安装源码包，
并在 Environment 安装后自动下载 `DA3 Small` 和 `BEN2 Base` 推荐模型。Environment 安装 并在 Environment 安装后自动下载 `MoGe-2 ViT-S Normal` 和 FP16 `BEN2 Base` 推荐模型。Environment 安装
固定显示 `Environment`、`PyTorch`、`AI components` 三个用户阶段；xFormers（macOS 跳过）、 固定显示 `Environment`、`ONNX Runtime` 两个用户阶段；FastAPI 等内部安装步骤合并在 `ONNX Runtime` 中。
FastAPI 等内部安装步骤合并在 `AI components` 中。耗时最长的 PyTorch 与 AI 组件占主要进度 安装完成后直接写入 Environment Manifest，不再额外导入全部组件进行验证。推荐模型属于安装
区间。安装完成后直接写入 Environment Manifest，不再额外导入全部组件进行验证。推荐模型属于安装
后的自动步骤，不计入阶段编号；`DA3 Base` 等其他模型可稍后在 `Models` 面板按需下载。 后的自动步骤，不计入阶段编号；`DA3 Base` 等其他模型可稍后在 `Models` 面板按需下载。
状态文字使用 `阶段/总数 阶段名 | 当前操作 | 预计剩余时间`，依赖安装会结合预设耗时与实际 状态文字使用 `阶段/总数 阶段名 | 当前操作 | 预计剩余时间`，依赖安装会结合预设耗时与实际
进度修正预期；模型下载根据字节数和速度计算剩余时间，例如 进度修正预期；模型下载根据字节数和速度计算剩余时间，例如

28
docs/outputs.md
View File

@@ -1,31 +1,29 @@
# 输出格式 # 输出格式
## 输出格式说明 ## DA3 ONNX
`Output Type` 提供以下结果： `Output Type` 提供以下结果：
- Point Cloud (GLB)：DA3 `glb` - Point Cloud (GLB)：由 depth、confidence、相机参数和原图生成彩色点云
- Gaussian Splats (PLY)：DA3 `gs_ply`，同时传入 `infer_gs=True` - Depth Preview (Images)：保存深度着色预览
- Depth Preview (Images)：DA3 `depth_vis` - Depth Data (NPZ)：保存 depth、confidence、extrinsics、intrinsics 与处理后的图像
- Depth Data (NPZ)：DA3 `mini_npz` - Textured Surface (GLB)：SnapForm 根据 Prediction 构建纹理网格
- Textured Surface (GLB)：SnapForm 根据 DA3 Prediction 构建并输出 `mesh.glb`
Point Cloud 专属设置仅在对应类型下显示，`gs_views_interval` 仅用于 Gaussian Splats。 DA3 Small 与 Base 的 ONNX 图不包含 Gaussian head，因此不提供 Gaussian Splats。默认推理参数为：
Textured Surface 使用 0–100% 的 Confidence Threshold 删除置信度最低的一定比例：
- `process_res=504` - `process_res=504`
- `process_res_method="upper_bound_resize"` - `process_res_method="upper_bound_resize"`
- `use_ray_pose=False`
- `ref_view_strategy="saddle_balanced"` - `ref_view_strategy="saddle_balanced"`
- `conf_thresh_percentile=40.0` - `conf_thresh_percentile=40.0`
- `num_max_points=1_000_000` - `num_max_points=1_000_000`
- `show_cameras=True` - `show_cameras=True`
- `export_kwargs["gs_ply"]["gs_views_interval"]=1`
DA3 官方 `glb` 是彩色点云与可选相机线框；`gs_ply` 是带 Gaussian 属性的 PLY； 相机外参与内参由模型的 camera decoder 输出恢复，不运行 ray-pose 求解器。
`depth_vis` 是原图与着色深度图的并排预览；`mini_npz` 保存深度、置信度和可用的相机参数。
Point Cloud GLB 任务会关闭其默认附带的深度预览，使五种输出产物保持分离。
官方格式与限制可参阅 ## MoGe-2 ONNX
[DA3 Python API 文档](https://github.com/ByteDance-Seed/Depth-Anything-3/blob/main/docs/API.md)。
MoGe-2 Debug 可保存原始 NPZ、depth/normal/mask maps、纹理 GLB 或彩色 PLY。Surface 输出由模型 points、mask 与 depth edge 在内部构建三角网格。
## BEN2 ONNX
BEN2 输出 `foreground.png`，RGB 保留原始输入，Alpha 来自 ONNX mask 并恢复到输入分辨率。

125
docs/service.md
View File

@@ -1,98 +1,61 @@
# AI Service # AI Service
## AI Service ## 运行架构
Environment 指磁盘上已安装的 Python、依赖及模型执行组件；Service 是由 Environment 启动、接收 Blender 任务的 SnapForm 在独立的 Python 3.12 Environment 中运行常驻 HTTP Service，Blender 只通过标准库 `urllib.request` 提交任务，不把 AI 运行库载入 Blender Python。Service 只监听 `127.0.0.1`，请求使用随机 Token；同一 Storage Root 由 `service.lock` 保证只有一个实例。
常驻 HTTP 进程。`manifest.json` 记录当前 Environment 的安装信息；Service 的 `/health` 返回进程状态和
实例标识。项目发布前不维护安装清单或任务协议的兼容版本号。
Extension 本身不包含 Torch。`Install Environment` 使用 `uv` 创建独立 Python 3.12 Environment、模型和任务结果默认保存在 `~/.snapform`：
环境，并安装。机器上即使没有安装 `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`：
```text ```text
~/.snapform/ ~/.snapform/
├── .venv/ ├── .venv/
├── models/ ├── models/
├── jobs/ ├── jobs/
├── manifest.json
├── service.log ├── service.log
├── service.json # Service 运行时存在 ├── service.json # Service 运行时存在
└── service.lock # Service 运行时持有 └── service.lock # Service 运行时持有
``` ```
`models/` 下每个模型一个文件夹，命名统一带前缀：DA3 为 `da3_*`（如 `da3_small`、`da3_base`）， 不同 Blender 版本可共享此目录。也可以在 Extension Preferences 的 `Storage Root` 中修改位置。
BEN2 为 `ben2_*`（如 `ben2_base`）。旧版无前缀的 DA3 文件夹会在首次使用时自动改名为新名称。
不同 Blender 版本安装相同 Extension 后会共用这个目录。也可以在 Extension ## ONNX-only Environment
Preferences 的 `Storage Root` 中修改共享位置。不要把模型写入 Extension 安装目录，
因为该目录可能只读，并且升级时会被替换。不要同时从多个 Blender 实例执行
`Install Environment`；Service、任务队列和模型读取可以共享。
扩展只内置对应平台的 `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/<repository-name>/<relative-path>`；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`。

10
docs/surface.md
View File

@@ -25,7 +25,7 @@
- Smooth 采用 `O Smooth` 的边界检测和模糊法线权重，固定开放边界与薄面，只保留一个 - Smooth 采用 `O Smooth` 的边界检测和模糊法线权重，固定开放边界与薄面，只保留一个
`Smooth` 程度系数；节点组不再额外执行 Flip Faces； `Smooth` 程度系数；节点组不再额外执行 Flip Faces；
- 从修改前的原始表面以 nearest-surface 重新采样 `UVMap`，最终明确写为 Corner 域 `2D Vector`。 - 从修改前的原始表面以 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 及相邻点相对深度差。Geometry Nodes
按官方 `mask & ~depth_map_edge(depth, rtol=0.04)` 的含义，先由 `Mask` 删除 mask 外的面，再由 按官方 `mask & ~depth_map_edge(depth, rtol=0.04)` 的含义，先由 `Mask` 删除 mask 外的面，再由
`Depth Continuity` 按同一份 3×3 邻域相对深度差删除断层面。DA3 与 MoGe-2 使用相同的属性、 `Depth Continuity` 按同一份 3×3 邻域相对深度差删除断层面。DA3 与 MoGe-2 使用相同的属性、
@@ -163,7 +163,7 @@ Image Empty 使用 Packed Image 时，SnapForm 会通过 Blender 写出临时 PN
浮点深度；保存 `.blend` 后再次打开仍需对应 job 目录存在。 浮点深度；保存 `.blend` 后再次打开仍需对应 job 目录存在。
选择 PLY 后，模型菜单只会提供 `DA3 Giant 1.1` 和 选择 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`。 DA3 生成 `gs_ply`。
每个任务在 `~/.snapform/jobs/` 下用独立文件夹包裹。文件夹统一采用 每个任务在 `~/.snapform/jobs/` 下用独立文件夹包裹。文件夹统一采用
@@ -185,7 +185,7 @@ Blender 为 Packed Image 或序列帧生成的输入副本也归档到对应 job
- `remove-image-background` - `remove-image-background`
- `generate-original-depth-plane` - `generate-original-depth-plane`
DA3 官方 exporter 产生的 `gs_ply/`、`depth_vis/`、`exports/mini_npz/` 会在完成后自动平铺， DA3 ONNX exporter 产生的 `depth_vis/`、`exports/mini_npz/` 会在完成后自动平铺，
不会保留只含一层产物的子目录： 不会保留只含一层产物的子目录：
```text ```text
@@ -218,6 +218,4 @@ remove-background_20260807-153000-123456/
└── depth.0001.exr… AI 深度序列帧 └── depth.0001.exr… AI 深度序列帧
``` ```
建议从 `DA3 Base` 开始。它使用 Apache 2.0 模型权重；`DA3 Large 1.1` 权重采用 建议从 `DA3 Base` 开始。它使用 Apache 2.0 模型权重。
CC BY-NC 4.0，仅限非商业用途。

63
pack.py
View File

@@ -1,6 +1,4 @@
import hashlib import hashlib
import io
import tarfile
import tomllib import tomllib
from pathlib import Path from pathlib import Path
from zipfile import ZIP_DEFLATED, ZipFile from zipfile import ZIP_DEFLATED, ZipFile
@@ -10,7 +8,6 @@ PROJECT_ROOT = Path(__file__).resolve().parent
SOURCE_DIR = PROJECT_ROOT / "src" / "snapform" SOURCE_DIR = PROJECT_ROOT / "src" / "snapform"
VENDOR_DIR = PROJECT_ROOT / "vendor" VENDOR_DIR = PROJECT_ROOT / "vendor"
DIST_DIR = PROJECT_ROOT / "dist" DIST_DIR = PROJECT_ROOT / "dist"
PATCHES_DIR = PROJECT_ROOT / "patches"
LOCK_PATH = PROJECT_ROOT / "vendor.lock.toml" LOCK_PATH = PROJECT_ROOT / "vendor.lock.toml"
@@ -28,15 +25,6 @@ PLATFORMS = {
platform: artifact["filename"] platform: artifact["filename"]
for platform, artifact in VENDOR_LOCK["platforms"].items() 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): 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): def manifest_for_platform(platform):
"""Return the manifest text declaring only the given platform.""" """Return the manifest text declaring only the given platform."""
text = (SOURCE_DIR / "blender_manifest.toml").read_text(encoding="utf-8") 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): def bundled_vendor_files(platform):
entries = { return {
name: VENDOR_DIR / filename f"vendor/{PLATFORMS[platform]}": VENDOR_DIR / PLATFORMS[platform]
for name, filename in BUNDLED_ARCHIVES.items()
} }
entries[f"vendor/{PLATFORMS[platform]}"] = (
VENDOR_DIR / PLATFORMS[platform]
)
return entries
def validate_vendor_files(platform): def validate_vendor_files(platform):
"""Validate common archives and the selected uv archive.""" """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] uv_artifact = VENDOR_LOCK["platforms"][platform]
validate_artifact( validate_artifact(
VENDOR_DIR / uv_artifact["filename"], VENDOR_DIR / uv_artifact["filename"],
@@ -174,7 +117,7 @@ def build_extension(platform):
for archive_name, source_path in vendor_files.items(): for archive_name, source_path in vendor_files.items():
archive.writestr( archive.writestr(
archive_name, archive_name,
bundled_artifact_bytes(source_path.name), source_path.read_bytes(),
) )
return output_path return output_path

93
patches/depth-anything-3/pyproject.toml
View File

@@ -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"

23
patches/depth-anything-3/requirements.txt
View File

@@ -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

84
patches/depth-anything-3/src/depth_anything_3/utils/export/__init__.py
View File

@@ -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,
]

531
patches/depth-anything-3/src/depth_anything_3/utils/ray_utils.py
View File

@@ -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

22
patches/moge/pyproject.toml
View File

@@ -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*"]

16
patches/utils3d/pyproject.toml
View File

@@ -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*"]

5
prepare_vendor.py
View File

@@ -65,7 +65,7 @@ def download_artifact(artifact, destination_dir=VENDOR_DIR):
def prepare(platforms=None, lock_path=LOCK_PATH, 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) lock = load_lock(lock_path)
available = lock["platforms"] available = lock["platforms"]
selected = list(available) if platforms is None else 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: if unknown:
raise ValueError(f"Unsupported platform: {', '.join(unknown)}") raise ValueError(f"Unsupported platform: {', '.join(unknown)}")
artifacts = list(lock["artifacts"].values()) artifacts = [available[platform] for platform in selected]
artifacts.extend(available[platform] for platform in selected)
return [ return [
download_artifact(artifact, destination_dir) download_artifact(artifact, destination_dir)
for artifact in artifacts for artifact in artifacts

4
src/snapform/__init__.py
View File

@@ -21,9 +21,9 @@ from .operators import (
) )
from .panel import ( from .panel import (
Ben2RemovalDebugPanel, Ben2RemovalDebugPanel,
AiServicePanel,
Da3OutputDebugPanel, Da3OutputDebugPanel,
DebugPanel, DebugPanel,
AiServicePanel,
ModelsPanel, ModelsPanel,
Moge2ApiDebugPanel, Moge2ApiDebugPanel,
) )
@@ -50,8 +50,8 @@ CLASSES = (
AiServicePanel, AiServicePanel,
ModelsPanel, ModelsPanel,
DebugPanel, DebugPanel,
Da3OutputDebugPanel,
Moge2ApiDebugPanel, Moge2ApiDebugPanel,
Da3OutputDebugPanel,
Ben2RemovalDebugPanel, Ben2RemovalDebugPanel,
) )

50
src/snapform/ben2_onnx.py Normal file
View File

@@ -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

16
src/snapform/client.py
View File

@@ -26,10 +26,8 @@ from .models import (
) )
DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4"
BEN2_COMMIT = "2c99a5da477b5523585bfa5c893888a6e818a8f6"
MANIFEST_FILENAME = "manifest.json" MANIFEST_FILENAME = "manifest.json"
MANIFEST_REVISION = 5 MANIFEST_REVISION = 7
SERVICE_ADDRESS = "127.0.0.1" SERVICE_ADDRESS = "127.0.0.1"
DEFAULT_STORAGE_ROOT = Path.home() / ".snapform" DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"
JOB_FOLDER_NAMES = { JOB_FOLDER_NAMES = {
@@ -172,8 +170,8 @@ def env_devices():
def machine_devices(): def machine_devices():
"""Devices the machine can run before environment installation.""" """Devices the machine can run before environment installation."""
devices = ["cpu"] devices = ["cpu"]
if shutil.which("nvidia-smi"): if sys.platform == "win32":
devices.insert(0, "cuda") devices.insert(0, "directml")
elif sys.platform == "darwin" and platform.machine().lower() in { elif sys.platform == "darwin" and platform.machine().lower() in {
"arm64", "arm64",
"aarch64", "aarch64",
@@ -189,6 +187,7 @@ def supported_devices():
DEVICE_LABELS = { DEVICE_LABELS = {
"cuda": ("CUDA", "NVIDIA CUDA"), "cuda": ("CUDA", "NVIDIA CUDA"),
"directml": ("DirectML", "Windows GPU (DirectML)"),
"mps": ("MPS", "Apple Metal (MPS)"), "mps": ("MPS", "Apple Metal (MPS)"),
"cpu": ("CPU", "CPU"), "cpu": ("CPU", "CPU"),
} }
@@ -283,7 +282,6 @@ def uv_cache_path():
def install_env_command(status_path): def install_env_command(status_path):
executable = uv_executable() executable = uv_executable()
torch_variant = "cu128" if "cuda" in machine_devices() else "cpu"
script = Path(__file__).with_name("install_env.py") script = Path(__file__).with_name("install_env.py")
command = [ command = [
sys.executable, sys.executable,
@@ -296,12 +294,6 @@ def install_env_command(status_path):
str(storage_root() / MANIFEST_FILENAME), str(storage_root() / MANIFEST_FILENAME),
"--models-dir", "--models-dir",
str(models_directory()), str(models_directory()),
"--torch-variant",
torch_variant,
"--da3-commit",
DA3_COMMIT,
"--ben2-commit",
BEN2_COMMIT,
"--status", "--status",
str(status_path), str(status_path),
] ]

221
src/snapform/da3_onnx.py Normal file
View File

@@ -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={},
)

12
src/snapform/image_depth.py
View File

@@ -2239,12 +2239,12 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
name="Depth From", name="Depth From",
description="Model or original image used to build depth", description="Model or original image used to build depth",
items=image_depth_model_items, items=image_depth_model_items,
default=1, default=0,
) )
snapform_use_ray_pose: bpy.props.BoolProperty( snapform_use_ray_pose: bpy.props.BoolProperty(
name="Use Ray Pose", name="Use Ray Pose",
description="Use ray-based camera pose estimation", description="Legacy task compatibility; ONNX uses the camera decoder",
default=True, default=False,
) )
snapform_process_res_factor: process_resolution_scale_property() snapform_process_res_factor: process_resolution_scale_property()
snapform_moge_resolution_level: bpy.props.IntProperty( snapform_moge_resolution_level: bpy.props.IntProperty(
@@ -2278,14 +2278,8 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
layout = self.layout layout = self.layout
layout.prop(self, "snapform_model") layout.prop(self, "snapform_model")
if self.snapform_model in {"SMALL", "BASE"}: if self.snapform_model in {"SMALL", "BASE"}:
layout.prop(self, "snapform_use_ray_pose")
draw_process_resolution(layout, self, context.object.data) draw_process_resolution(layout, self, context.object.data)
elif self.snapform_model.startswith("MOGE2_"): 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_moge_resolution_level")
layout.prop(self, "snapform_remove_background") layout.prop(self, "snapform_remove_background")
layout.prop(self, "snapform_unshaded") layout.prop(self, "snapform_unshaded")

7
src/snapform/image_mesh.py
View File

@@ -608,12 +608,12 @@ class ConvertImageToMesh(bpy.types.Operator):
name="Depth From", name="Depth From",
description="Model used to estimate image depth", description="Model used to estimate image depth",
items=image_surface_model_items, items=image_surface_model_items,
default=1, default=0,
) )
snapform_use_ray_pose: bpy.props.BoolProperty( snapform_use_ray_pose: bpy.props.BoolProperty(
name="Use Ray Pose", name="Use Ray Pose",
description="Use ray-based camera pose estimation", description="Legacy task compatibility; ONNX uses the camera decoder",
default=True, default=False,
) )
snapform_process_res_factor: process_resolution_scale_property() snapform_process_res_factor: process_resolution_scale_property()
snapform_moge_resolution_level: bpy.props.IntProperty( snapform_moge_resolution_level: bpy.props.IntProperty(
@@ -670,7 +670,6 @@ class ConvertImageToMesh(bpy.types.Operator):
if self.snapform_model.startswith("MOGE2_"): if self.snapform_model.startswith("MOGE2_"):
layout.prop(self, "snapform_moge_resolution_level") layout.prop(self, "snapform_moge_resolution_level")
else: else:
layout.prop(self, "snapform_use_ray_pose")
draw_process_resolution(layout, self, context.object.data) draw_process_resolution(layout, self, context.object.data)
layout.prop(self, "snapform_create_camera") layout.prop(self, "snapform_create_camera")
layout.prop(self, "snapform_remove_background") layout.prop(self, "snapform_remove_background")

262
src/snapform/install_env.py
View File

@@ -1,5 +1,4 @@
import argparse import argparse
import filecmp
import json import json
import os import os
import platform import platform
@@ -28,47 +27,36 @@ PYPI_MIRRORS = (
"https://mirrors.aliyun.com/pypi/simple/", "https://mirrors.aliyun.com/pypi/simple/",
) )
PYPI_DEFAULT_INDEX = "https://pypi.org/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" FASTAPI_VERSION = "0.139.2"
UVICORN_VERSION = "0.51.0" UVICORN_VERSION = "0.51.0"
MANIFEST_REVISION = 5 ONNXRUNTIME_VERSION = "1.24.4"
TORCH_INDEXES = { NUMPY_VERSION = "2.4.2"
"cu128": "https://download.pytorch.org/whl/cu128", PILLOW_VERSION = "12.1.1"
"cpu": "https://download.pytorch.org/whl/cpu", OPENCV_VERSION = "4.13.0.92"
} TRIMESH_VERSION = "4.11.2"
HUGGINGFACE_HUB_VERSION = "1.4.1"
MANIFEST_REVISION = 7
_resolved_pypi_index = None _resolved_pypi_index = None
DEFAULT_MODEL_KEYS = (MOGE2_MODEL.key, BEN2_MODEL.key) DEFAULT_MODEL_KEYS = (MOGE2_MODEL.key, BEN2_MODEL.key)
RUNTIME_STAGE_COUNT = 3 RUNTIME_STAGE_COUNT = 2
ENVIRONMENT_SLICE = (0.02, 0.06) ENVIRONMENT_SLICE = (0.02, 0.06)
PYTORCH_SLICE = (0.06, 0.48) AI_COMPONENTS_SLICE = (0.10, 0.85)
XFORMERS_SLICE = (0.48, 0.62)
AI_COMPONENTS_SLICE = (0.62, 0.85)
MODELS_SLICE_START = 0.85 MODELS_SLICE_START = 0.85
MODELS_SLICE_END = 1.0 MODELS_SLICE_END = 1.0
RAMP_RATE = 0.02 RAMP_RATE = 0.02
CREEP_CAP = 0.60 CREEP_CAP = 0.60
INSTALL_ESTIMATES = { INSTALL_ESTIMATES = {
"PyTorch": 8 * 60, "ONNX Runtime": 3 * 60,
"AI components": 8 * 60,
} }
INSTALL_ESTIMATE_RANGES = { INSTALL_ESTIMATE_RANGES = {
"PyTorch": PYTORCH_SLICE, "ONNX Runtime": AI_COMPONENTS_SLICE,
"AI components": (XFORMERS_SLICE[0], AI_COMPONENTS_SLICE[1]),
} }
CHECK_CODE = ( CHECK_CODE = (
"import json; import torch; " "import json; import fastapi, uvicorn, onnxruntime; "
"from depth_anything_3.api import DepthAnything3; " "import cv2, numpy, PIL, trimesh; "
"from ben2 import BEN_Base; "
"from moge.model.v2 import MoGeModel; "
"import fastapi, uvicorn; "
"mps = getattr(torch.backends, 'mps', None); "
"print('SNAPFORM_DEVICES=' + json.dumps({" "print('SNAPFORM_DEVICES=' + json.dumps({"
"'torch': torch.__version__, " "'onnxruntime': onnxruntime.__version__, "
"'cuda': torch.cuda.is_available(), " "'providers': onnxruntime.get_available_providers()}))"
"'mps': bool(mps is not None and mps.is_available())}))"
) )
@@ -158,10 +146,6 @@ def estimated_remaining(label, progress, elapsed=0.0):
return max(expected_seconds * remaining_fraction - 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): def resolve_indexes(args):
"""???????????????????????????????????""" """???????????????????????????????????"""
global _resolved_pypi_index global _resolved_pypi_index
@@ -320,52 +304,13 @@ def venv_python(venv_dir):
return Path(venv_dir) / "bin" / "python" return Path(venv_dir) / "bin" / "python"
def install_stages(args, venv_dir, python, uv, vendor_dir=None): def install_stages(args, venv_dir, python, uv):
# macOS 的 PyTorch wheel 只在 PyPI 上，download.pytorch.org 没有 macosx 构建 """Build the lightweight ONNX-only runtime and model download stages."""
torch_index = ( runtime_package = (
None f"onnxruntime-directml=={ONNXRUNTIME_VERSION}"
if sys.platform == "darwin" if sys.platform == "win32"
else effective_torch_index(args.torch_variant) 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 = [ stages = [
( (
"Environment", "Environment",
@@ -374,13 +319,7 @@ def install_stages(args, venv_dir, python, uv, vendor_dir=None):
"plain", "plain",
), ),
( (
"PyTorch", "ONNX Runtime",
torch_command,
None,
"install",
),
(
"AI components",
[ [
uv, uv,
"pip", "pip",
@@ -388,38 +327,19 @@ def install_stages(args, venv_dir, python, uv, vendor_dir=None):
"-v", "-v",
"--python", "--python",
python, python,
da3_requirement,
f"addict=={ADDICT_VERSION}",
ben2_requirement,
utils3d_requirement,
moge_requirement,
f"fastapi=={FASTAPI_VERSION}", 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}", f"uvicorn=={UVICORN_VERSION}",
"setuptools",
"wheel",
], ],
None, None,
"install", "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 models = DOWNLOADABLE_MODELS
worker_script = Path(__file__).with_name("worker.py") worker_script = Path(__file__).with_name("worker.py")
models_dir = Path(args.models_dir) 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 = [] devices = []
machine = platform.machine().lower() machine = platform.machine().lower()
if sys.platform == "darwin" and machine in {"arm64", "aarch64"}: if sys.platform == "darwin" and machine in {"arm64", "aarch64"}:
devices.append("mps") devices.append("mps")
elif args.torch_variant.startswith("cu"): elif sys.platform == "win32":
devices.append("cuda") devices.append("directml")
devices.append("cpu") devices.append("cpu")
return devices return devices
@@ -592,82 +512,6 @@ def bundled_uv_archive():
return path if path.is_file() else None 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): def archive_is_valid(path):
if zipfile.is_zipfile(path): if zipfile.is_zipfile(path):
return True return True
@@ -800,14 +644,6 @@ def run_stage(
def install(args): 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) resolve_indexes(args)
venv_dir = Path(args.venv_dir) venv_dir = Path(args.venv_dir)
manifest_path = Path( manifest_path = Path(
@@ -823,13 +659,7 @@ def install(args):
uv = Path(resolved) if resolved else uv uv = Path(resolved) if resolved else uv
uv_missing = not uv.is_file() uv_missing = not uv.is_file()
python = venv_python(venv_dir) python = venv_python(venv_dir)
vendor_temp = tempfile.TemporaryDirectory(prefix="snapform_vendor_") stages = install_stages(args, venv_dir, python, uv)
vendor_dir = Path(vendor_temp.name)
copy_vendor_inputs(
bundled_vendor_directory(),
vendor_dir,
)
stages = install_stages(args, venv_dir, python, uv, vendor_dir)
env_stages = [stage for stage in stages if stage[3] != "download"] env_stages = [stage for stage in stages if stage[3] != "download"]
model_stages = [stage for stage in stages if stage[3] == "download"] model_stages = [stage for stage in stages if stage[3] == "download"]
model_span = ( model_span = (
@@ -838,23 +668,15 @@ def install(args):
else 0.0 else 0.0
) )
installed_packages = [] 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): for index, (label, command, environment, kind) in enumerate(env_stages):
if index == 0: if index == 0:
stage = 1 stage = 1
stage_start, stage_end = ENVIRONMENT_SLICE stage_start, stage_end = ENVIRONMENT_SLICE
elif index == 1: elif index == 1:
stage = 2 stage = 2
stage_start, stage_end = PYTORCH_SLICE stage_start, stage_end = AI_COMPONENTS_SLICE
else: else:
stage = 3 raise RuntimeError(f"Unexpected runtime stage: {label}")
component_index = index - 2
stage_start, stage_end = component_ranges[component_index]
if uv_missing and index == 0: if uv_missing and index == 0:
midpoint = stage_start + (stage_end - stage_start) * 0.5 midpoint = stage_start + (stage_end - stage_start) * 0.5
install_uv( install_uv(
@@ -880,9 +702,7 @@ def install(args):
) )
if isinstance(stage_installed, list): if isinstance(stage_installed, list):
installed_packages.extend(stage_installed) installed_packages.extend(stage_installed)
vendor_temp.cleanup() devices = installed_devices()
remove_legacy_vendor(venv_dir.parent)
devices = installed_devices(args)
venv_dir.mkdir(parents=True, exist_ok=True) venv_dir.mkdir(parents=True, exist_ok=True)
manifest_path.parent.mkdir(parents=True, exist_ok=True) manifest_path.parent.mkdir(parents=True, exist_ok=True)
ready_path = manifest_path ready_path = manifest_path
@@ -891,14 +711,13 @@ def install(args):
{ {
"manifest_revision": MANIFEST_REVISION, "manifest_revision": MANIFEST_REVISION,
"python": "3.12", "python": "3.12",
"torch": TORCH_VERSION, "onnxruntime": ONNXRUNTIME_VERSION,
"xformers": XFORMERS_VERSION,
"fastapi": FASTAPI_VERSION, "fastapi": FASTAPI_VERSION,
"uvicorn": UVICORN_VERSION, "uvicorn": UVICORN_VERSION,
"torch_variant": args.torch_variant, "runtime_variant": (
"directml" if sys.platform == "win32" else "default"
),
"devices": devices, "devices": devices,
"da3_commit": args.da3_commit,
"ben2_commit": args.ben2_commit,
"packages": sorted(set(installed_packages)), "packages": sorted(set(installed_packages)),
}, },
indent=2, indent=2,
@@ -952,9 +771,6 @@ def parse_args():
parser.add_argument("--venv-dir", required=True) parser.add_argument("--venv-dir", required=True)
parser.add_argument("--manifest", required=True) parser.add_argument("--manifest", required=True)
parser.add_argument("--models-dir", 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("--status", required=True)
parser.add_argument("--pypi-index", default=None) parser.add_argument("--pypi-index", default=None)
return parser.parse_args() return parser.parse_args()

131
src/snapform/models.py
View File

@@ -21,42 +21,33 @@ MODELS = {
"SMALL": ModelSpec( "SMALL": ModelSpec(
key="SMALL", key="SMALL",
label="DA3 Small", label="DA3 Small",
directory_name="da3_small", directory_name="DA3-SMALL",
repository="depth-anything/DA3-SMALL", repository="depth-anything/DA3-SMALL",
license_name="Apache 2.0", 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, supports_gs=False,
r2_directory="DA3Small", download_patterns=("model.onnx",),
ready_patterns=("model.onnx",),
r2_directory="DA3-SMALL",
r2_files=( r2_files=(
("config.json", 1202, "a486e29e82b7ab4a7d4cefc1ea4526cfe2ae438a572c8ca98917cfbcde7447d2"), ("model.onnx", 105484029, "7a9d81475b3e2334ace0d6d2c2a816b3c93081437e1866cc434ce1789125f631"),
("model.safetensors", 137248940, "364492e38a3a06d221ac75da7f6621ada3f2361cd24fde11ba79091e9f40efcf"),
), ),
), ),
"BASE": ModelSpec( "BASE": ModelSpec(
key="BASE", key="BASE",
label="DA3 Base", label="DA3 Base",
directory_name="da3_base", directory_name="DA3-BASE",
repository="depth-anything/DA3-BASE", repository="depth-anything/DA3-BASE",
license_name="Apache 2.0", 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, supports_gs=False,
r2_directory="DA3Base", download_patterns=("model.onnx",),
ready_patterns=("model.onnx",),
r2_directory="DA3-BASE",
r2_files=( r2_files=(
("config.json", 1205, "5e34115ebc17bd2d8d43033c5f72e9446ac8833fd61d3fa160b7e67e0bb5b7b5"), ("model.onnx", 412856697, "31f38fc4990539d7902ba267924b9731db1acb00d840c248a60b1865810ff9e6"),
("configuration.json", 74, "b684d553257b2ca379fd5dfdf96a6e4a633eec3b969c97119c9875465a7c07a8"),
("model.safetensors", 541518028, "e01067dc1659613083d9145a9a2547ccdbe6ccbbf83c4fe7b3e8a4e2bdae78b5"),
), ),
), ),
"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( "GIANT_1_1": ModelSpec(
key="GIANT_1_1", key="GIANT_1_1",
label="DA3 Giant 1.1", label="DA3 Giant 1.1",
@@ -82,52 +73,76 @@ MODELS = {
BEN2_MODEL = ModelSpec( BEN2_MODEL = ModelSpec(
key="BEN2_BASE", key="BEN2_BASE",
label="BEN2 Base", label="BEN2 Base",
directory_name="ben2_base", directory_name="BEN2-ONNX",
repository="PramaLLC/BEN2", repository="onnx-community/BEN2-ONNX",
license_name="MIT", license_name="Apache 2.0",
description="Fast high-quality background removal model", description="Fast high-quality FP16 ONNX background removal model",
supports_gs=False, supports_gs=False,
download_patterns=("config.json", "model.safetensors"), download_patterns=("onnx/model_fp16.onnx",),
r2_directory="BEN2Base", ready_patterns=("onnx/model_fp16.onnx",),
r2_files=( r2_directory="BEN2-ONNX",
("config.json", 124, "01ac876925577b52fa7c6af3f5b1d1420bc9b40016b6c25e2de89b38d8fff865"), r2_files=((
("model.safetensors", 380577976, "ea8b7907176a09667c86343dc7d00de6a6d871076cb90bb5f753618fd6fb3ebb"), "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_MODELS = {
"MOGE2_VITS_NORMAL": moge2_model( "MOGE2_VITS_NORMAL": ModelSpec(
"MOGE2_VITS_NORMAL", "MoGe-2 ViT-S Normal", "moge2_vits_normal", key="MOGE2_VITS_NORMAL",
"Ruicheng/moge-2-vits-normal", "Fastest MoGe-2 model", "MoGe2ViTSNormal", label="MoGe-2 ViT-S Normal",
140550416, "79a16621928c2bf0ed04659218c55c01075e950507f40bb3332fb4c873d3e1dc", 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": ModelSpec(
"MOGE2_VITB_NORMAL", "MoGe-2 ViT-B Normal", "moge2_vitb_normal", key="MOGE2_VITB_NORMAL",
"Ruicheng/moge-2-vitb-normal", "Balanced MoGe-2 model", "MoGe2ViTBNormal", label="MoGe-2 ViT-B Normal",
419110160, "16b8110e86d5dc5a849db120ca96ef3a223fd30b0c9146d1d81db504073da5f6", 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": ModelSpec(
"MOGE2_VITL_NORMAL", "MoGe-2 ViT-L Normal", "moge2_vitl_normal", key="MOGE2_VITL_NORMAL",
"Ruicheng/moge-2-vitl-normal", "Highest-quality MoGe-2 model", "MoGe2ViTLNormal", label="MoGe-2 ViT-L Normal",
1323815904, "280741fd09bc3f403ccff9967784c2a391b52d2c0742ae3efdb21d9f90cc1a01", 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 = { DOWNLOADABLE_MODELS = {
**MODELS, **MODELS,
BEN2_MODEL.key: BEN2_MODEL, BEN2_MODEL.key: BEN2_MODEL,

165
src/snapform/moge_onnx.py Normal file
View File

@@ -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)

55
src/snapform/onnx_runtime.py Normal file
View File

@@ -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)

5
src/snapform/panel.py
View File

@@ -186,7 +186,6 @@ def draw_da3_input_settings(layout, settings, model_property):
layout.prop(settings, "video_frames") layout.prop(settings, "video_frames")
layout.prop(settings, "process_res") layout.prop(settings, "process_res")
layout.prop(settings, "process_res_method") layout.prop(settings, "process_res_method")
layout.prop(settings, "use_ray_pose")
layout.prop(settings, "ref_view_strategy") layout.prop(settings, "ref_view_strategy")
@@ -198,7 +197,7 @@ class Da3OutputDebugPanel(bpy.types.Panel):
bl_region_type = "UI" bl_region_type = "UI"
bl_category = "SnapForm" bl_category = "SnapForm"
bl_options = {"DEFAULT_CLOSED"} bl_options = {"DEFAULT_CLOSED"}
bl_order = 0 bl_order = 1
def draw(self, context): def draw(self, context):
layout = self.layout layout = self.layout
@@ -248,7 +247,7 @@ class Moge2ApiDebugPanel(bpy.types.Panel):
bl_region_type = "UI" bl_region_type = "UI"
bl_category = "SnapForm" bl_category = "SnapForm"
bl_options = {"DEFAULT_CLOSED"} bl_options = {"DEFAULT_CLOSED"}
bl_order = 1 bl_order = 0
def draw(self, context): def draw(self, context):
layout = self.layout layout = self.layout

4
src/snapform/properties.py
View File

@@ -76,7 +76,7 @@ class SnapFormSettings(bpy.types.PropertyGroup):
name="Model", name="Model",
description="MoGe-2 model used for depth estimation", description="MoGe-2 model used for depth estimation",
items=moge2_model_enum_items, items=moge2_model_enum_items,
default=1, default=0,
) )
moge_export_format: bpy.props.EnumProperty( moge_export_format: bpy.props.EnumProperty(
name="Output Format", name="Output Format",
@@ -110,7 +110,7 @@ class SnapFormSettings(bpy.types.PropertyGroup):
) )
use_ray_pose: bpy.props.BoolProperty( use_ray_pose: bpy.props.BoolProperty(
name="Use Ray Pose", 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, default=False,
) )
ref_view_strategy: bpy.props.EnumProperty( ref_view_strategy: bpy.props.EnumProperty(

63
src/snapform/service.py
View File

@@ -74,9 +74,6 @@ class ServiceLock:
class ModelCache: class ModelCache:
def __init__(self): def __init__(self):
import torch
self.torch = torch
self.ben2_model = None self.ben2_model = None
self.ben2_key = None self.ben2_key = None
self.da3_model = None self.da3_model = None
@@ -85,75 +82,51 @@ class ModelCache:
self.moge_key = None self.moge_key = None
def get_ben2(self, model_dir, device): 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: if self.ben2_model is not None and self.ben2_key == key:
return self.ben2_model, 0.0 return self.ben2_model, 0.0
self.release_ben2() self.release_ben2()
try: return self._load_ben2(model_dir, requested_device, key)
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)
def _load_ben2(self, model_dir, device, key): def _load_ben2(self, model_dir, requested_device, key):
from ben2 import BEN_Base from ben2_onnx import create_session
started = time.perf_counter() started = time.perf_counter()
model = BEN_Base.from_pretrained(model_dir) model = create_session(model_dir, requested_device)
model = model.to(device=device)
model.eval()
self.ben2_model = model self.ben2_model = model
self.ben2_key = key self.ben2_key = key
elapsed_ms = (time.perf_counter() - started) * 1000.0 elapsed_ms = (time.perf_counter() - started) * 1000.0
return model, elapsed_ms return model, elapsed_ms
def get_da3(self, model_dir, device): 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: if self.moge_model is not None:
self.release_moge() self.release_moge()
if self.da3_model is not None and self.da3_key == key: if self.da3_model is not None and self.da3_key == key:
return self.da3_model return self.da3_model
self.release_da3() self.release_da3()
try: return self._load_da3(model_dir, requested_device, key)
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)
def _load_da3(self, model_dir, device, key): def _load_da3(self, model_dir, requested_device, key):
from depth_anything_3.api import DepthAnything3 from da3_onnx import create_session
model = DepthAnything3.from_pretrained(model_dir) model = create_session(model_dir, requested_device)
model = model.to(device=device)
model.eval()
self.da3_model = model self.da3_model = model
self.da3_key = key self.da3_key = key
return model return model
def get_moge(self, model_dir, device): def get_moge_onnx(self, model_dir, requested_device):
key = (str(Path(model_dir).resolve()), device.type) from moge_onnx import create_session
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.da3_model is not None: if self.da3_model is not None:
self.release_da3() self.release_da3()
if self.moge_model is not None and self.moge_key == key: if self.moge_model is not None and self.moge_key == key:
return self.moge_model return self.moge_model
self.release_moge() self.release_moge()
try: model = create_session(model_dir, requested_device)
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()
self.moge_model = model self.moge_model = model
self.moge_key = key self.moge_key = key
return model return model
@@ -201,8 +174,6 @@ class ModelCache:
def _release_cuda_cache(self): def _release_cuda_cache(self):
gc.collect() gc.collect()
if self.torch.cuda.is_available():
self.torch.cuda.empty_cache()
class ServiceState: class ServiceState:

378
src/snapform/worker.py
View File

@@ -6,11 +6,17 @@ import shutil
import tempfile import tempfile
import threading import threading
import time import time
import sys
from pathlib import Path from pathlib import Path
from urllib.parse import quote from urllib.parse import quote
from urllib.request import Request, urlopen 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"} IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
VIDEO_SUFFIXES = { VIDEO_SUFFIXES = {
".mp4", ".mp4",
@@ -434,6 +440,92 @@ def export_prediction_attributes(prediction, output_directory):
return output_path 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): def remove_background_faces(faces, foreground_mask):
"""Return only triangles whose vertices all belong to the foreground.""" """Return only triangles whose vertices all belong to the foreground."""
foreground_points = foreground_mask.reshape(-1) foreground_points = foreground_mask.reshape(-1)
@@ -1133,29 +1225,10 @@ def format_rate(rate):
return f"{size:.1f} {units[unit_index]}" 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( def infer_foreground(
input_path, input_path,
model_dir, model_dir,
device, device,
torch,
model_cache=None, model_cache=None,
cancel_check=None, cancel_check=None,
): ):
@@ -1165,14 +1238,13 @@ def infer_foreground(
check_cancelled(cancel_check) check_cancelled(cancel_check)
if input_path.suffix.lower() not in IMAGE_SUFFIXES: if input_path.suffix.lower() not in IMAGE_SUFFIXES:
raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image") raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
owns_model = model_cache is None from ben2_onnx import create_session, infer
if owns_model:
from ben2 import BEN_Base
owns_model = model_cache is None
requested_device = getattr(device, "type", device)
if owns_model:
load_started = time.perf_counter() load_started = time.perf_counter()
model = BEN_Base.from_pretrained(model_dir) model = create_session(model_dir, requested_device)
model = model.to(device=device)
model.eval()
load_ms = (time.perf_counter() - load_started) * 1000.0 load_ms = (time.perf_counter() - load_started) * 1000.0
else: else:
model, load_ms = model_cache.get_ben2(model_dir, device) model, load_ms = model_cache.get_ben2(model_dir, device)
@@ -1180,22 +1252,10 @@ def infer_foreground(
check_cancelled(cancel_check) check_cancelled(cancel_check)
with Image.open(input_path) as opened_image: with Image.open(input_path) as opened_image:
image = opened_image.convert("RGB") image = opened_image.convert("RGB")
if device.type == "cuda":
torch.cuda.synchronize()
inference_started = time.perf_counter() inference_started = time.perf_counter()
with torch.inference_mode(): foreground = infer(model, image)
foreground = model.inference(
image,
refine_foreground=False,
)
foreground = preserve_source_rgb(image, foreground) foreground = preserve_source_rgb(image, foreground)
if device.type == "cuda":
torch.cuda.synchronize()
inference_ms = (time.perf_counter() - inference_started) * 1000.0 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) check_cancelled(cancel_check)
return foreground, load_ms, inference_ms return foreground, load_ms, inference_ms
@@ -1216,7 +1276,6 @@ def infer_foreground_frames(
images, images,
model_dir, model_dir,
device, device,
torch,
output_directory, output_directory,
status_path, status_path,
start_progress, start_progress,
@@ -1238,7 +1297,6 @@ def infer_foreground_frames(
image_path, image_path,
model_dir, model_dir,
device, device,
torch,
model_cache=model_cache, model_cache=model_cache,
cancel_check=cancel_check, 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") write_status(args.status, 0.02, "Starting Environment")
check_cancelled(cancel_check) check_cancelled(cancel_check)
if model_cache is None:
import torch
else:
torch = model_cache.torch
output_directory = Path(args.output_dir) output_directory = Path(args.output_dir)
output_directory.mkdir(parents=True, exist_ok=True) output_directory.mkdir(parents=True, exist_ok=True)
device = select_device(args.device, torch) device = args.device
video_frames = getattr(args, "video_frames", 16) video_frames = getattr(args, "video_frames", 16)
start_frame = getattr(args, "start_frame", 1) 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, cancel_check=cancel_check,
) )
try: 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: if len(images) == 1:
write_status(args.status, 0.35, "Removing background") write_status(args.status, 0.35, "Removing background")
foreground, load_ms, inference_ms = infer_foreground( foreground, load_ms, inference_ms = infer_foreground(
images[0], images[0],
args.model_dir, args.model_dir,
device, device,
torch,
model_cache=model_cache, model_cache=model_cache,
cancel_check=cancel_check, cancel_check=cancel_check,
) )
@@ -1305,7 +1357,6 @@ def remove_background(args, model_cache=None, cancel_check=None):
images, images,
args.model_dir, args.model_dir,
device, device,
torch,
output_directory, output_directory,
args.status, args.status,
0.12, 0.12,
@@ -1517,41 +1568,33 @@ def export_moge_debug_mesh(
export_format, export_format,
): ):
import numpy as np import numpy as np
import utils3d import trimesh
from moge.utils.io import save_glb, save_ply
clean_mask = mask & ~utils3d.np.depth_map_edge(depth, rtol=0.04) clean_mask = mask & ~depth_edge_mask(depth, threshold=0.04)
uv = utils3d.np.uv_map(*depth.shape) faces = mesh_triangle_faces(
result = utils3d.np.build_mesh_from_map(
points, points,
image, depth,
uv, np.ones(depth.shape, dtype=np.float32),
normal, None,
mask=clean_mask, False,
tri=True, 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) 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) vertex_normals = vertex_normals * np.array([1.0, -1.0, -1.0], dtype=np.float32)
if export_format == "glb": mesh = trimesh.Trimesh(
save_glb( vertices=vertices,
output_directory / "mesh.glb", faces=faces,
vertices, vertex_colors=colors,
faces, vertex_normals=vertex_normals,
vertex_uvs, process=False,
(image * 255.0).astype(np.uint8), )
vertex_normals, mesh.remove_unreferenced_vertices()
) filename = "mesh.glb" if export_format == "glb" else "pointcloud.ply"
return mesh.export(Path(output_directory) / filename)
if export_format == "ply":
save_ply(
output_directory / "pointcloud.ply",
vertices,
np.zeros((0, 3), dtype=np.int32),
colors,
vertex_normals,
)
def export_moge_surface_archive( def export_moge_surface_archive(
@@ -1564,7 +1607,6 @@ def export_moge_surface_archive(
foreground_mask=None, foreground_mask=None,
): ):
import numpy as np import numpy as np
import utils3d
height, width = depth.shape height, width = depth.shape
finite_mask = np.isfinite(depth) & (depth > 0.0) finite_mask = np.isfinite(depth) & (depth > 0.0)
@@ -1572,34 +1614,28 @@ def export_moge_surface_archive(
if foreground_mask is not None: if foreground_mask is not None:
finite_mask &= foreground_mask finite_mask &= foreground_mask
depth_discontinuity = depth_discontinuity_ratio(depth) depth_discontinuity = depth_discontinuity_ratio(depth)
result = utils3d.np.build_mesh_from_map( faces = mesh_triangle_faces(
points, points,
utils3d.np.uv_map(height, width),
normal,
mask.astype(np.float32),
depth_discontinuity.astype(np.float32),
depth, depth,
mask=finite_mask, np.ones(depth.shape, dtype=np.float32),
tri=True, None,
False,
0.0,
vertex_mask=finite_mask,
) )
( vertices = points.reshape(-1, 3)
faces, u = np.linspace(0.0, 1.0, width, dtype=np.float32)
vertices, v = np.linspace(1.0, 0.0, height, dtype=np.float32)
vertex_uvs, uu, vv = np.meshgrid(u, v)
_vertex_normals, vertex_uvs = np.stack((uu, vv), axis=-1).reshape(-1, 2)
vertex_mask,
vertex_depth_discontinuity,
point_depth,
) = result
vertex_uvs = vertex_uvs * np.array([1.0, -1.0]) + np.array([0.0, 1.0])
pixel_intrinsics = intrinsics.copy() pixel_intrinsics = intrinsics.copy()
pixel_intrinsics[0] *= width pixel_intrinsics[0] *= width
pixel_intrinsics[1] *= height pixel_intrinsics[1] *= height
pixel_intrinsics[2] = (0.0, 0.0, 1.0) pixel_intrinsics[2] = (0.0, 0.0, 1.0)
point_depth = point_depth.astype(np.float32) point_depth = depth.reshape(-1).astype(np.float32)
face_mask = vertex_mask[faces].all(axis=1) face_mask = mask.reshape(-1)[faces].all(axis=1)
face_depth_discontinuity = vertex_depth_discontinuity[faces].max(axis=1) face_depth_discontinuity = depth_discontinuity.reshape(-1)[faces].max(axis=1)
masked_depth = point_depth[vertex_mask.astype(bool)] masked_depth = point_depth[mask.reshape(-1).astype(bool)]
reference_depth = np.median(masked_depth if len(masked_depth) else point_depth) reference_depth = np.median(masked_depth if len(masked_depth) else point_depth)
np.savez( np.savez(
output_path, output_path,
@@ -1619,14 +1655,13 @@ def export_moge_surface_archive(
def run_moge(args, model_cache=None, cancel_check=None): def run_moge(args, model_cache=None, cancel_check=None):
import numpy as np import numpy as np
from PIL import Image 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") 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 = Path(args.output_dir)
output_directory.mkdir(parents=True, exist_ok=True) output_directory.mkdir(parents=True, exist_ok=True)
images, temporary_directory = prepare_input( images, temporary_directory = prepare_input(
@@ -1637,15 +1672,15 @@ def run_moge(args, model_cache=None, cancel_check=None):
cancel_check=cancel_check, cancel_check=cancel_check,
) )
try: 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: if model_cache is None:
checkpoint = Path(args.model_dir) / "model.pt" model = create_session(args.model_dir, args.device)
model = MoGeModel.from_pretrained(checkpoint)
model = model.to(device=device)
model.eval()
else: 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 = [] depths = []
normals = [] normals = []
@@ -1662,21 +1697,17 @@ def run_moge(args, model_cache=None, cancel_check=None):
) )
with Image.open(image_path) as opened: with Image.open(image_path) as opened:
image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0 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) source_images.append(image)
with torch.inference_mode(): prediction = infer_moge_onnx(
prediction = model.infer( model,
image_tensor, image,
resolution_level=int(args.resolution_level), int(args.resolution_level),
force_projection=True, )
apply_mask=False, depth = prediction["depth"]
use_fp16=device.type == "cuda", normal = prediction["normal"]
) mask = prediction["mask"].astype(bool)
depth = prediction["depth"].detach().float().cpu().numpy() point = prediction["points"]
normal = prediction["normal"].detach().float().cpu().numpy() intrinsic = prediction["intrinsics"]
mask = prediction["mask"].detach().cpu().numpy().astype(bool)
point = prediction["points"].detach().float().cpu().numpy()
intrinsic = prediction["intrinsics"].detach().float().cpu().numpy()
depths.append(depth) depths.append(depth)
normals.append(normal) normals.append(normal)
masks.append(mask) masks.append(mask)
@@ -1727,8 +1758,7 @@ def run_moge(args, model_cache=None, cancel_check=None):
foreground, _, _ = infer_foreground( foreground, _, _ = infer_foreground(
images[0], images[0],
remove_background_model_dir, remove_background_model_dir,
device, args.device,
torch,
model_cache=model_cache, model_cache=model_cache,
cancel_check=cancel_check, cancel_check=cancel_check,
) )
@@ -1743,8 +1773,7 @@ def run_moge(args, model_cache=None, cancel_check=None):
infer_foreground_frames( infer_foreground_frames(
images, images,
remove_background_model_dir, remove_background_model_dir,
device, args.device,
torch,
output_directory, output_directory,
args.status, args.status,
0.87, 0.87,
@@ -1781,21 +1810,11 @@ def run_moge(args, model_cache=None, cancel_check=None):
shutil.rmtree(temporary_directory, ignore_errors=True) 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): def reconstruct(args, model_cache=None, cancel_check=None):
write_status(args.status, 0.02, "Starting Environment") write_status(args.status, 0.02, "Starting Environment")
check_cancelled(cancel_check) check_cancelled(cancel_check)
from da3_onnx import create_session, infer
if model_cache is None: from onnx_runtime import provider_label
import torch
from depth_anything_3.api import DepthAnything3
else:
torch = model_cache.torch
from depth_anything_3.utils.export import export
output_directory = Path(args.output_dir) output_directory = Path(args.output_dir)
output_directory.mkdir(parents=True, exist_ok=True) output_directory.mkdir(parents=True, exist_ok=True)
@@ -1810,7 +1829,6 @@ def reconstruct(args, model_cache=None, cancel_check=None):
try: try:
check_cancelled(cancel_check) check_cancelled(cancel_check)
device = select_device(args.device, torch)
remove_background_model_dir = getattr( remove_background_model_dir = getattr(
args, args,
"remove_background_model_dir", "remove_background_model_dir",
@@ -1823,8 +1841,7 @@ def reconstruct(args, model_cache=None, cancel_check=None):
foreground, load_ms, inference_ms = infer_foreground( foreground, load_ms, inference_ms = infer_foreground(
images[0], images[0],
remove_background_model_dir, remove_background_model_dir,
device, args.device,
torch,
model_cache=model_cache, model_cache=model_cache,
cancel_check=cancel_check, cancel_check=cancel_check,
) )
@@ -1841,8 +1858,7 @@ def reconstruct(args, model_cache=None, cancel_check=None):
load_ms, inference_ms = infer_foreground_frames( load_ms, inference_ms = infer_foreground_frames(
images, images,
remove_background_model_dir, remove_background_model_dir,
device, args.device,
torch,
output_directory, output_directory,
args.status, args.status,
0.12, 0.12,
@@ -1858,13 +1874,16 @@ def reconstruct(args, model_cache=None, cancel_check=None):
f"BEN2 inference {inference_ms:.0f} ms; " f"BEN2 inference {inference_ms:.0f} ms; "
f"model load {load_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: if model_cache is None:
model = DepthAnything3.from_pretrained(args.model_dir) model = create_session(args.model_dir, args.device)
model = model.to(device=device)
model.eval()
else: 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) check_cancelled(cancel_check)
write_status( write_status(
@@ -1872,27 +1891,14 @@ def reconstruct(args, model_cache=None, cancel_check=None):
0.25, 0.25,
f"Estimating depth for {len(images)} frame(s)", f"Estimating depth for {len(images)} frame(s)",
) )
export_kwargs = {} prediction = infer(
if args.export_format == "glb": model,
export_kwargs["glb"] = { [str(path) for path in images],
"conf_thresh_percentile": args.conf_thresh_percentile, process_res=args.process_res,
"num_max_points": args.num_max_points, process_res_method=args.process_res_method,
"show_cameras": bool(args.show_cameras), use_ray_pose=bool(args.use_ray_pose),
"export_depth_vis": False, ref_view_strategy=args.ref_view_strategy,
} )
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,
)
check_cancelled(cancel_check) check_cancelled(cancel_check)
write_status(args.status, 0.88, f"Exporting {args.export_format}") write_status(args.status, 0.88, f"Exporting {args.export_format}")
if args.export_format == "mesh": if args.export_format == "mesh":
@@ -1921,12 +1927,20 @@ def reconstruct(args, model_cache=None, cancel_check=None):
), ),
progress_callback=mesh_progress, progress_callback=mesh_progress,
) )
else: elif args.export_format == "mini_npz":
export( 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, prediction,
args.export_format, output_directory,
str(output_directory), args.conf_thresh_percentile,
**export_kwargs, 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"}: if args.export_format in {"glb", "gs_ply", "depth_vis", "mesh"}:
write_status( write_status(
@@ -1989,7 +2003,7 @@ def parser():
background.add_argument("--status", required=True) background.add_argument("--status", required=True)
background.add_argument( background.add_argument(
"--device", "--device",
choices=("auto", "cuda", "mps", "cpu"), choices=("auto", "cuda", "directml", "mps", "cpu"),
required=True, required=True,
) )
background.add_argument("--video-frames", type=int, default=16) 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("--start-frame", type=int, default=1)
reconstruction.add_argument( reconstruction.add_argument(
"--device", "--device",
choices=("auto", "cuda", "mps", "cpu"), choices=("auto", "cuda", "directml", "mps", "cpu"),
required=True, required=True,
) )
reconstruction.add_argument( reconstruction.add_argument(

194
sync_models.py Normal file
View File

@@ -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()

30
tests/test_addon.py
View File

@@ -280,6 +280,10 @@ class AddonTest(unittest.TestCase):
annotations["snapform_process_res_factor"]["subtype"], annotations["snapform_process_res_factor"]["subtype"],
"PERCENTAGE", "PERCENTAGE",
) )
self.assertEqual(
annotations["snapform_model"]["default"],
0,
)
def test_image_mesh_uses_o_prefixed_attributes_and_bundled_nodes(self): def test_image_mesh_uses_o_prefixed_attributes_and_bundled_nodes(self):
image_mesh = self.addon.image_mesh image_mesh = self.addon.image_mesh
@@ -644,8 +648,11 @@ class AddonTest(unittest.TestCase):
def test_conversion_dialog_exposes_requested_settings(self): def test_conversion_dialog_exposes_requested_settings(self):
operator = self.addon.ConvertImageToDepthPlane operator = self.addon.ConvertImageToDepthPlane
self.assertEqual(operator.__annotations__["snapform_model"]["default"], 1) self.assertEqual(
self.assertTrue(operator.__annotations__["snapform_use_ray_pose"]["default"]) operator.__annotations__["snapform_model"]["default"],
0,
)
self.assertFalse(operator.__annotations__["snapform_use_ray_pose"]["default"])
self.assertEqual( self.assertEqual(
operator.__annotations__["snapform_process_res_factor"]["default"], operator.__annotations__["snapform_process_res_factor"]["default"],
100.0, 100.0,
@@ -1094,7 +1101,6 @@ class AddonTest(unittest.TestCase):
for snapform_model, snapform_ray_pose in ( for snapform_model, snapform_ray_pose in (
("SMALL", True), ("SMALL", True),
("BASE", False), ("BASE", False),
("LARGE_1_1", True),
("GIANT_1_1", False), ("GIANT_1_1", False),
("NESTED_GIANT_LARGE_1_1", True), ("NESTED_GIANT_LARGE_1_1", True),
): ):
@@ -1947,8 +1953,8 @@ class AddonTest(unittest.TestCase):
self.addon.DebugPanel.bl_options, self.addon.DebugPanel.bl_options,
) )
panels = ( panels = (
self.addon.Da3OutputDebugPanel,
self.addon.Moge2ApiDebugPanel, self.addon.Moge2ApiDebugPanel,
self.addon.Da3OutputDebugPanel,
self.addon.Ben2RemovalDebugPanel, self.addon.Ben2RemovalDebugPanel,
) )
for panel in panels: for panel in panels:
@@ -1956,7 +1962,7 @@ class AddonTest(unittest.TestCase):
self.assertIn("DEFAULT_CLOSED", panel.bl_options) self.assertIn("DEFAULT_CLOSED", panel.bl_options)
self.assertEqual( self.assertEqual(
[panel.bl_label for panel in panels], [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): 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) for index, value in enumerate(command)
if value == "--allow-pattern" if value == "--allow-pattern"
] ]
self.assertEqual(patterns, ["config.json", "model.safetensors"]) self.assertEqual(patterns, ["onnx/model_fp16.onnx"])
label_index = command.index("--label") label_index = command.index("--label")
self.assertEqual(command[label_index + 1], "BEN2 Base") self.assertEqual(command[label_index + 1], "BEN2 Base")
@@ -2512,15 +2518,12 @@ class AddonTest(unittest.TestCase):
self.assertEqual(command[0], sys.executable) self.assertEqual(command[0], sys.executable)
self.assertIn("--uv", command) self.assertIn("--uv", command)
self.assertIn("C:/uv/uv.exe", command) self.assertIn("C:/uv/uv.exe", command)
self.assertEqual( self.assertNotIn("--torch-variant", command)
command[command.index("--torch-variant") + 1],
"cu128",
)
self.assertTrue( self.assertTrue(
any("install_env.py" in str(value) for value in command) 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") client = importlib.import_module("snapform.client")
with patch.object( with patch.object(
client, client,
@@ -2542,10 +2545,7 @@ class AddonTest(unittest.TestCase):
command = client.install_env_command(Path("status.json")) command = client.install_env_command(Path("status.json"))
self.assertNotIn("--pypi-index", command) self.assertNotIn("--pypi-index", command)
self.assertEqual( self.assertNotIn("--torch-variant", command)
command[command.index("--torch-variant") + 1],
"cpu",
)
self.assertIn("C:/uv/uv.exe", command) self.assertIn("C:/uv/uv.exe", command)
def test_completed_export_opens_job_directory(self): def test_completed_export_opens_job_directory(self):

53
tests/test_ben2_onnx.py Normal file
View File

@@ -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()

61
tests/test_da3_onnx.py Normal file
View File

@@ -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()

260
tests/test_install_env.py
View File

@@ -36,7 +36,7 @@ class InstallEnvTest(unittest.TestCase):
(vendor / "moge.tar.gz").write_bytes(b"moge") (vendor / "moge.tar.gz").write_bytes(b"moge")
return vendor return vendor
def test_installs_da3_from_archives_without_gsplat(self): def test_installs_onnx_runtime_without_training_frameworks(self):
bootstrap = load_install_env() bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory: with tempfile.TemporaryDirectory() as directory:
root = Path(directory) root = Path(directory)
@@ -71,77 +71,31 @@ class InstallEnvTest(unittest.TestCase):
commands += [ commands += [
call.args[4] for call in mocked_run_install.call_args_list 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( service_command = next(
command command
for command in commands for command in commands
if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command
) )
self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", service_command) self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", service_command)
self.assertIn("setuptools", service_command) self.assertIn(
self.assertNotIn("ninja", service_command) 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( check_command = bootstrap.verify_command(
bootstrap.venv_python(venv_dir) bootstrap.venv_python(venv_dir)
) )
self.assertTrue( self.assertFalse(any("torch" in str(value) for value in check_command))
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.assertTrue( self.assertTrue(
any("import fastapi, uvicorn" in str(value) for value in check_command) 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( self.assertFalse(
any("import gsplat" in str(value) for value in check_command) any("import gsplat" in str(value) for value in check_command)
) )
@@ -157,24 +111,24 @@ class InstallEnvTest(unittest.TestCase):
if any( if any(
repository in command repository in command
for repository in ( for repository in (
"Ruicheng/moge-2-vitb-normal", "Ruicheng/moge-2-vits-normal-onnx",
"PramaLLC/BEN2", "onnx-community/BEN2-ONNX",
) )
) )
] ]
self.assertEqual(len(default_downloads), 2) 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( ben2_download = next(
command command
for command in default_downloads for command in default_downloads
if "PramaLLC/BEN2" in command if "onnx-community/BEN2-ONNX" in command
) )
patterns = [ patterns = [
ben2_download[index + 1] ben2_download[index + 1]
for index, value in enumerate(ben2_download) for index, value in enumerate(ben2_download)
if value == "--allow-pattern" 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", ben2_download)
self.assertNotIn("--stage-label", ben2_download) self.assertNotIn("--stage-label", ben2_download)
self.assertEqual( self.assertEqual(
@@ -182,7 +136,7 @@ class InstallEnvTest(unittest.TestCase):
"BEN2 Base", "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() bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory: with tempfile.TemporaryDirectory() as directory:
root = Path(directory) root = Path(directory)
@@ -224,15 +178,15 @@ class InstallEnvTest(unittest.TestCase):
commands = [ commands = [
call.args[4] for call in mocked_run_install.call_args_list call.args[4] for call in mocked_run_install.call_args_list
] ]
torch_command = next( runtime_command = next(
command command
for command in commands for command in commands
if any( if any(
f"torch=={bootstrap.TORCH_VERSION}" in str(value) f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}" in str(value)
for value in command for value in command
) )
) )
self.assertNotIn("--index-url", torch_command) self.assertNotIn("--index-url", runtime_command)
self.assertFalse( self.assertFalse(
any( any(
"xformers" in str(value) "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() 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( with patch.object(
bootstrap.sys, bootstrap.sys,
"platform", "platform",
"darwin", "darwin",
), patch.object( ):
bootstrap.platform, stages = bootstrap.install_stages(
"machine", SimpleNamespace(
return_value="x86_64", models_dir="models",
), patch.object(bootstrap.shutil, "rmtree") as mocked_rmtree: status="status.json",
with self.assertRaisesRegex(RuntimeError, "Apple Silicon"): ),
bootstrap.install(args) ".venv",
"python",
mocked_rmtree.assert_not_called() "uv",
)
runtime = stages[1][1]
self.assertIn(
f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}",
runtime,
)
def test_run_stage_writes_stage_progress(self): def test_run_stage_writes_stage_progress(self):
bootstrap = load_install_env() bootstrap = load_install_env()
@@ -414,10 +365,10 @@ class InstallEnvTest(unittest.TestCase):
) )
self.assertEqual(ends[-1], 1.0) self.assertEqual(ends[-1], 1.0)
self.assertEqual(starts[0], bootstrap.ENVIRONMENT_SLICE[0]) self.assertEqual(starts[0], bootstrap.ENVIRONMENT_SLICE[0])
runtime_ranges = stage_ranges[:4] runtime_ranges = stage_ranges[:2]
self.assertEqual( self.assertEqual(
[item[0] for item in runtime_ranges], [item[0] for item in runtime_ranges],
[1, 2, 3, 3], [1, 2],
) )
self.assertTrue( self.assertTrue(
all(item[0] is None for item in stage_ranges[4:]) 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(12), "12s")
self.assertEqual(bootstrap.format_elapsed(83), "1m 23s") 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() bootstrap = load_install_env()
first = bootstrap.estimated_remaining( first = bootstrap.estimated_remaining(
"AI components", "ONNX Runtime",
bootstrap.XFORMERS_SLICE[0],
)
second = bootstrap.estimated_remaining(
"AI components",
bootstrap.AI_COMPONENTS_SLICE[0], 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.assertGreater(second, 0)
self.assertLess(second, first) self.assertLess(second, first)
@@ -727,7 +678,7 @@ class InstallEnvTest(unittest.TestCase):
ready["packages"], ready["packages"],
["numpy==1.26.4", "torch==2.9.1+cu128"], ["numpy==1.26.4", "torch==2.9.1+cu128"],
) )
self.assertEqual(ready["devices"], ["cuda", "cpu"]) self.assertEqual(ready["devices"], ["directml", "cpu"])
self.assertEqual( self.assertEqual(
ready["manifest_revision"], ready["manifest_revision"],
bootstrap.MANIFEST_REVISION, bootstrap.MANIFEST_REVISION,
@@ -814,7 +765,7 @@ class InstallEnvTest(unittest.TestCase):
def fake_run(command, environment=None): def fake_run(command, environment=None):
commands.append(command) 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) raise subprocess.CalledProcessError(1, command)
with tempfile.TemporaryDirectory() as directory: with tempfile.TemporaryDirectory() as directory:
@@ -852,7 +803,7 @@ class InstallEnvTest(unittest.TestCase):
self.assertIn("Environment is installed", status["message"]) self.assertIn("Environment is installed", status["message"])
self.assertTrue( self.assertTrue(
any( any(
"PramaLLC/BEN2" in map(str, command) "onnx-community/BEN2-ONNX" in map(str, command)
for command in commands for command in commands
) )
) )
@@ -983,25 +934,13 @@ class InstallEnvTest(unittest.TestCase):
for value in command for value in command
) )
) )
ai_command = next( runtime_command = next(
command command
for command in install_commands for command in install_commands
if any("depth-anything-3 @" in str(value) for value in command) if any("onnxruntime" 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
)
) )
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): def test_install_reuses_existing_uv(self):
bootstrap = load_install_env() bootstrap = load_install_env()
@@ -1047,81 +986,12 @@ class InstallEnvTest(unittest.TestCase):
self.assertEqual(uv_path.read_bytes(), b"binary") self.assertEqual(uv_path.read_bytes(), b"binary")
def test_copy_vendor_inputs_copies_missing_archives(self): def test_install_stages_use_only_runtime_packages(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):
bootstrap = load_install_env() bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory: with tempfile.TemporaryDirectory() as directory:
root = Path(directory) root = Path(directory)
venv_dir = root / ".venv" venv_dir = root / ".venv"
venv_dir.mkdir() 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 = root / "uv"
uv_path.write_bytes(b"uv") uv_path.write_bytes(b"uv")
args = SimpleNamespace( args = SimpleNamespace(
@@ -1141,16 +1011,14 @@ class InstallEnvTest(unittest.TestCase):
uv_path, uv_path,
) )
ai_command = next( runtime_command = next(
command command
for _label, command, _environment, kind in stages for _label, command, _environment, kind in stages
if kind == "install" 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() self.assertFalse(any(" @ file:" in str(value) for value in runtime_command))
ben2_local = (vendor / "ben2.tar.gz").as_uri() self.assertFalse(any("torch" in str(value) for value in runtime_command))
self.assertIn(f"depth-anything-3 @ {da3_local}", ai_command)
self.assertIn(f"ben2 @ {ben2_local}", ai_command)
def test_install_stages_download_only_default_models(self): def test_install_stages_download_only_default_models(self):
bootstrap = load_install_env() bootstrap = load_install_env()
@@ -1189,13 +1057,13 @@ class InstallEnvTest(unittest.TestCase):
if kind == "download" if kind == "download"
} }
self.assertIn( self.assertIn(
"moge2_vitb_normal", "moge-2-vits-normal-onnx",
" ".join(map(str, downloads["MoGe-2 ViT-B Normal"])), " ".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( self.assertEqual(
set(downloads), set(downloads),
{"MoGe-2 ViT-B Normal", "BEN2 Base"}, {"MoGe-2 ViT-S Normal", "BEN2 Base"},
) )
def test_install_does_not_verify_env(self): def test_install_does_not_verify_env(self):

64
tests/test_models.py
View File

@@ -30,12 +30,6 @@ class ModelTest(unittest.TestCase):
self.assertEqual(model.repository, "depth-anything/DA3-BASE") self.assertEqual(model.repository, "depth-anything/DA3-BASE")
self.assertEqual(model.license_name, "Apache 2.0") 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): def test_giant_models_expose_gaussian_head(self):
for model_key in ("GIANT_1_1", "NESTED_GIANT_LARGE_1_1"): for model_key in ("GIANT_1_1", "NESTED_GIANT_LARGE_1_1"):
with self.subTest(model_key=model_key): 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): def test_ben2_download_only_includes_required_hugging_face_files(self):
model = self.models.get_downloadable_model("BEN2_BASE") model = self.models.get_downloadable_model("BEN2_BASE")
self.assertEqual(model.repository, "PramaLLC/BEN2") self.assertEqual(model.repository, "onnx-community/BEN2-ONNX")
self.assertEqual(model.license_name, "MIT") self.assertEqual(model.license_name, "Apache 2.0")
self.assertEqual( self.assertEqual(model.download_patterns, ("onnx/model_fp16.onnx",))
model.download_patterns, self.assertEqual(model.ready_patterns, ("onnx/model_fp16.onnx",))
("config.json", "model.safetensors"), self.assertEqual(model.r2_files[0][0], "onnx/model_fp16.onnx")
)
self.assertNotIn("BEN2_BASE", self.models.MODELS) 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(): for model_key, model in self.models.MODELS.items():
if not model.enabled:
continue
with self.subTest(model_key=model_key): with self.subTest(model_key=model_key):
self.assertTrue( self.assertTrue(
model.directory_name.startswith("da3_"), model.directory_name.startswith("DA3-"),
model.directory_name, model.directory_name,
) )
self.assertEqual( self.assertEqual(
@@ -74,20 +69,37 @@ class ModelTest(unittest.TestCase):
model.directory_name, model.directory_name,
) )
def test_ben2_model_directory_uses_ben2_prefix(self): def test_ben2_model_directory_uses_repository_name(self):
self.assertEqual( self.assertEqual(
self.models.model_directory_name("BEN2_BASE"), self.models.model_directory_name("BEN2_BASE"),
"ben2_base", "BEN2-ONNX",
) )
def test_moge2_uses_single_checkpoint_and_mit_license(self): def test_enabled_da3_models_use_onnx_and_r2_checksums(self):
model = self.models.get_downloadable_model("MOGE2_VITB_NORMAL") 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") def test_moge2_models_use_official_onnx_models(self):
self.assertEqual(model.license_name, "MIT") for key, variant in (("MOGE2_VITS_NORMAL", "vits"),
self.assertEqual(model.download_patterns, ("model.pt",)) ("MOGE2_VITB_NORMAL", "vitb"),
self.assertEqual(model.ready_patterns, ("model.pt",)) ("MOGE2_VITL_NORMAL", "vitl")):
self.assertFalse(model.supports_gs) 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): def test_moge2_models_match_uploaded_r2_variants(self):
self.assertEqual( self.assertEqual(
@@ -96,5 +108,9 @@ class ModelTest(unittest.TestCase):
) )
self.assertEqual( self.assertEqual(
[model.r2_directory for model in self.models.MOGE2_MODELS.values()], [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",
],
) )

130
tests/test_moge_onnx.py Normal file
View File

@@ -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()

95
tests/test_pack.py
View File

@@ -1,7 +1,4 @@
import ast
import importlib.util import importlib.util
import io
import tarfile
import tempfile import tempfile
import tomllib import tomllib
import unittest import unittest
@@ -64,8 +61,6 @@ class PackTest(unittest.TestCase):
self.assertIn("service.py", names) self.assertIn("service.py", names)
self.assertNotIn("snapform/__init__.py", names) self.assertNotIn("snapform/__init__.py", names)
self.assertNotIn("importer.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( self.assertIn(
f"vendor/{pack.PLATFORMS[platform]}", f"vendor/{pack.PLATFORMS[platform]}",
names, names,
@@ -108,95 +103,5 @@ class PackTest(unittest.TestCase):
{"windows-x64", "macos-arm64", "linux-x64"}, {"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__": if __name__ == "__main__":
unittest.main() unittest.main()

60
tests/test_service.py
View File

@@ -40,7 +40,6 @@ class ServiceTest(unittest.TestCase):
def create_cache(self): def create_cache(self):
cache = self.service.ModelCache.__new__(self.service.ModelCache) cache = self.service.ModelCache.__new__(self.service.ModelCache)
cache.torch = SimpleNamespace(cuda=FakeCuda())
cache.ben2_model = None cache.ben2_model = None
cache.ben2_key = None cache.ben2_key = None
cache.da3_model = None cache.da3_model = None
@@ -52,65 +51,41 @@ class ServiceTest(unittest.TestCase):
def test_ben2_model_is_reused_for_matching_runtime_key(self): def test_ben2_model_is_reused_for_matching_runtime_key(self):
loads = [] loads = []
class FakeModel: ben2 = ModuleType("ben2_onnx")
@classmethod ben2.create_session = lambda model_dir, device: (
def from_pretrained(cls, model_dir): loads.append((model_dir, device)) or object()
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
cache = self.create_cache() cache = self.create_cache()
device = SimpleNamespace(type="cuda") 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) first, _load_ms = cache.get_ben2("model", device)
second, cached_load_ms = cache.get_ben2("model", device) second, cached_load_ms = cache.get_ben2("model", device)
self.assertIs(first, second) self.assertIs(first, second)
self.assertEqual(loads, ["model"]) self.assertEqual(loads, [("model", "cuda")])
self.assertEqual(cache.status()["ben2"], "model") self.assertEqual(cache.status()["ben2"], "model")
self.assertEqual(cached_load_ms, 0.0) self.assertEqual(cached_load_ms, 0.0)
def test_da3_model_is_reused_for_matching_runtime_key(self): def test_da3_model_is_reused_for_matching_runtime_key(self):
loads = [] loads = []
class FakeModel: da3 = ModuleType("da3_onnx")
@classmethod da3.create_session = lambda model_dir, device: (
def from_pretrained(cls, model_dir): loads.append((model_dir, device)) or object()
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
cache = self.create_cache() cache = self.create_cache()
device = SimpleNamespace(type="cuda") device = SimpleNamespace(type="cuda")
with patch.dict( with patch.dict(
sys.modules, sys.modules,
{ {
"depth_anything_3": package, "da3_onnx": da3,
"depth_anything_3.api": api,
}, },
): ):
first = cache.get_da3("model", device) first = cache.get_da3("model", device)
second = cache.get_da3("model", device) second = cache.get_da3("model", device)
self.assertIs(first, second) self.assertIs(first, second)
self.assertEqual(loads, ["model"]) self.assertEqual(loads, [("model", "cuda")])
self.assertEqual(cache.status()["da3"], "model") self.assertEqual(cache.status()["da3"], "model")
def test_depth_model_cache_keeps_only_da3_or_moge(self): def test_depth_model_cache_keeps_only_da3_or_moge(self):
@@ -135,13 +110,10 @@ class ServiceTest(unittest.TestCase):
moge = object() moge = object()
def load_moge(_model_dir, _device, key): fake_moge = ModuleType("moge_onnx")
cache.moge_model = moge fake_moge.create_session = lambda _model_dir, _device: moge
cache.moge_key = key with patch.dict(sys.modules, {"moge_onnx": fake_moge}):
return moge self.assertIs(cache.get_moge_onnx("moge", "cuda"), moge)
with patch.object(cache, "_load_moge", side_effect=load_moge):
self.assertIs(cache.get_moge("moge", device), moge)
self.assertIsNone(cache.da3_model) self.assertIsNone(cache.da3_model)
self.assertIsNotNone(cache.ben2_model) self.assertIsNotNone(cache.ben2_model)

278
tests/test_worker.py
View File

@@ -66,57 +66,6 @@ class WorkerTest(unittest.TestCase):
self.worker.check_cancelled(lambda: False) 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): def test_reconstruct_parser_accepts_background_model_directory(self):
args = self.worker.parser().parse_args( args = self.worker.parser().parse_args(
[ [
@@ -436,12 +385,6 @@ class WorkerTest(unittest.TestCase):
with self.assertRaisesRegex(RuntimeError, "Unsupported input type"): with self.assertRaisesRegex(RuntimeError, "Unsupported input type"):
self.worker.prepare_input(input_path, 16, directory) 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): def test_download_model_uses_actual_request_without_preflight(self):
calls = [] calls = []
fake_hugging_face = ModuleType("huggingface_hub") 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.device = lambda device_type: SimpleNamespace(type=device_type)
fake_torch.inference_mode = nullcontext fake_torch.inference_mode = nullcontext
fake_ben2 = ModuleType("ben2") fake_ben2 = ModuleType("ben2_onnx")
fake_ben2.BEN_Base = FakeModel 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 = ModuleType("PIL")
fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage()) fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())
@@ -614,7 +560,7 @@ class WorkerTest(unittest.TestCase):
sys.modules, sys.modules,
{ {
"torch": fake_torch, "torch": fake_torch,
"ben2": fake_ben2, "ben2_onnx": fake_ben2,
"PIL": fake_pil, "PIL": fake_pil,
}, },
): ):
@@ -635,7 +581,7 @@ class WorkerTest(unittest.TestCase):
self.assertEqual( self.assertEqual(
inference_options[0][1], inference_options[0][1],
{"refine_foreground": False}, {},
) )
self.assertEqual( self.assertEqual(
save_options[0], save_options[0],
@@ -690,8 +636,11 @@ class WorkerTest(unittest.TestCase):
) )
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type) fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
fake_torch.inference_mode = nullcontext fake_torch.inference_mode = nullcontext
fake_ben2 = ModuleType("ben2") fake_ben2 = ModuleType("ben2_onnx")
fake_ben2.BEN_Base = FakeModel 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 = ModuleType("PIL")
fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage()) fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())
@@ -699,7 +648,7 @@ class WorkerTest(unittest.TestCase):
sys.modules, sys.modules,
{ {
"torch": fake_torch, "torch": fake_torch,
"ben2": fake_ben2, "ben2_onnx": fake_ben2,
"PIL": fake_pil, "PIL": fake_pil,
}, },
): ):
@@ -734,7 +683,7 @@ class WorkerTest(unittest.TestCase):
self.assertEqual(len(inference_options), 2) self.assertEqual(len(inference_options), 2)
for _image, options in inference_options: for _image, options in inference_options:
self.assertEqual(options, {"refine_foreground": False}) self.assertEqual(options, {})
def test_remove_background_parser_accepts_video_frames(self): def test_remove_background_parser_accepts_video_frames(self):
args = self.worker.parser().parse_args( 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): def test_uses_da3_export_parameters_for_formats_and_debug_mesh(self):
inference_options = [] inference_options = []
export_calls = []
mesh_calls = [] mesh_calls = []
class FakeModel: prediction = SimpleNamespace(
@classmethod depth=np.full((1, 2, 2), 2.0, dtype=np.float32),
def from_pretrained(cls, model_dir): conf=np.ones((1, 2, 2), dtype=np.float32),
if model_dir != "model": extrinsics=np.eye(4, dtype=np.float32)[None, :3],
raise AssertionError(f"Unexpected model directory: {model_dir}") intrinsics=np.eye(3, dtype=np.float32)[None],
return cls() processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8),
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,
) )
fake_torch.inference_mode = nullcontext fake_da3 = ModuleType("da3_onnx")
fake_api = ModuleType("depth_anything_3.api") fake_da3.create_session = lambda _model_dir, _device: object()
fake_api.DepthAnything3 = FakeModel
fake_export_module = ModuleType("depth_anything_3.utils.export") def fake_infer(_session, images, **options):
fake_export_module.export = fake_export inference_options.append((images, options))
fake_utils = ModuleType("depth_anything_3.utils") return prediction
fake_utils.__path__ = []
fake_package = ModuleType("depth_anything_3") fake_da3.infer = fake_infer
fake_package.__path__ = [] fake_runtime = ModuleType("onnx_runtime")
fake_runtime.provider_label = lambda _session: "CPU"
def fake_mesh_export( def fake_mesh_export(
prediction, prediction,
@@ -893,6 +803,20 @@ class WorkerTest(unittest.TestCase):
output_path.touch() output_path.touch()
return output_path 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: with tempfile.TemporaryDirectory() as directory:
root = Path(directory) root = Path(directory)
input_path = root / "input.png" input_path = root / "input.png"
@@ -900,20 +824,28 @@ class WorkerTest(unittest.TestCase):
with patch.dict( with patch.dict(
sys.modules, sys.modules,
{ {
"torch": fake_torch, "da3_onnx": fake_da3,
"depth_anything_3": fake_package, "onnx_runtime": fake_runtime,
"depth_anything_3.api": fake_api,
"depth_anything_3.utils": fake_utils,
"depth_anything_3.utils.export": fake_export_module,
}, },
), patch.object( ), patch.object(
self.worker, self.worker,
"export_textured_mesh_glb", "export_textured_mesh_glb",
side_effect=fake_mesh_export, 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 = ( debug_formats = (
"glb", "glb",
"gs_ply",
"depth_vis", "depth_vis",
"mini_npz", "mini_npz",
"mesh", "mesh",
@@ -945,15 +877,11 @@ class WorkerTest(unittest.TestCase):
self.assertEqual(status["progress"], 1.0) self.assertEqual(status["progress"], 1.0)
self.assertIn(export_format, status["message"]) self.assertIn(export_format, status["message"])
result_directory = Path(args.output_dir) 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": if export_format == "depth_vis":
self.assertTrue((result_directory / "0000.jpg").is_file()) self.assertTrue((result_directory / "0000.jpg").is_file())
self.assertFalse((result_directory / "depth_vis").exists()) self.assertFalse((result_directory / "depth_vis").exists())
self.assertEqual(len(inference_options), 5) self.assertEqual(len(inference_options), 4)
self.assertEqual(len(export_calls), 4)
self.assertEqual(len(mesh_calls), 1) self.assertEqual(len(mesh_calls), 1)
for images, options in inference_options: for images, options in inference_options:
self.assertEqual(images, [str(input_path)]) self.assertEqual(images, [str(input_path)])
@@ -964,82 +892,21 @@ class WorkerTest(unittest.TestCase):
) )
self.assertTrue(options["use_ray_pose"]) self.assertTrue(options["use_ray_pose"])
self.assertEqual(options["ref_view_strategy"], "middle") 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][1], [input_path])
self.assertEqual(mesh_calls[0][2:], (35.0, True, 0.2)) 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): def test_reconstruct_copies_frames_for_multi_frame_input(self):
class FakeModel: fake_da3 = ModuleType("da3_onnx")
@classmethod fake_da3.create_session = lambda _model_dir, _device: object()
def from_pretrained(cls, model_dir): fake_da3.infer = lambda *_args, **_kwargs: SimpleNamespace(
return cls() depth=np.full((2, 2, 2), 2.0, dtype=np.float32),
conf=np.ones((2, 2, 2), dtype=np.float32),
def to(self, device): extrinsics=np.repeat(np.eye(4, dtype=np.float32)[None, :3], 2, axis=0),
return self intrinsics=np.repeat(np.eye(3, dtype=np.float32)[None], 2, axis=0),
processed_images=np.zeros((2, 2, 2, 3), dtype=np.uint8),
def 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_torch.inference_mode = nullcontext fake_runtime = ModuleType("onnx_runtime")
fake_api = ModuleType("depth_anything_3.api") fake_runtime.provider_label = lambda _session: "CPU"
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__ = []
with tempfile.TemporaryDirectory() as directory: with tempfile.TemporaryDirectory() as directory:
root = Path(directory) root = Path(directory)
@@ -1050,11 +917,8 @@ class WorkerTest(unittest.TestCase):
with patch.dict( with patch.dict(
sys.modules, sys.modules,
{ {
"torch": fake_torch, "da3_onnx": fake_da3,
"depth_anything_3": fake_package, "onnx_runtime": fake_runtime,
"depth_anything_3.api": fake_api,
"depth_anything_3.utils": fake_utils,
"depth_anything_3.utils.export": fake_export_module,
}, },
): ):
args = SimpleNamespace( args = SimpleNamespace(

27
vendor.lock.toml
View File

@@ -1,32 +1,5 @@
schema_version = 1 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] [uv]
version = "0.11.3" version = "0.11.3"