Update SnapForm to version 0.10.10 and add technical documentation

- Bump SnapForm version in uv.lock from 0.9.30 to 0.10.10.
- Introduce comprehensive technical documentation for SnapForm, detailing features, usage instructions, and development guidelines.
This commit is contained in:
2026-08-08 19:05:42 +08:00
parent d08d0e9c1a
commit aee460fe7f
22 changed files with 3905 additions and 1472 deletions

2
.gitignore vendored
View File

@@ -6,3 +6,5 @@ dist/
build/ build/
*.zip *.zip
.pytest_cache/ .pytest_cache/
sources/

355
README.md
View File

@@ -1,331 +1,64 @@
# SnapForm # SnapForm
`SnapForm` 是一套从图片快速构建空间结构的 Blender Extension。目前以 SnapForm 是一个 Blender 扩展,可以把一张照片、一段视频或一组序列帧快速变成有真实深度的立体面片。
[Depth Anything 3](https://github.com/ByteDance-Seed/Depth-Anything-3)
作为首个 AI 能力,在独立 Runtime 中运行 DA3,可将官方导出结果保存在独立 job 目录,
也可把 Image Empty 转换成带几何与材质置换的立体面片。Debug 面板还可运行
[BEN2](https://github.com/PramaLLC/BEN2) 背景剔除功能测试。
## 当前功能 ## 它能做什么
- 使用 Blender Extension manifest,可通过 `Install from Disk` 安装 - 一键把 Image Empty 转换为立体面片:让 AI 估算深度,或直接用原图作为深度
- 在独立 Python Runtime 的常驻 HTTP AI Host 中运行 Torch、DA3 与 BEN2,避免污染 Blender Python 进程 - 一键去除图片背景,生成透明前景图
- Blender 使用内置 `urllib.request` 访问仅监听 `127.0.0.1` 的 Host;请求带随机 Token - 支持视频与序列帧输入,逐帧处理,动画也能得到逐帧对应的立体面片
- 常驻 Host 复用 Torch、CUDA、BEN2 与最近使用的 DA3 模型,连续执行时不再重复加载 - 全部操作都在右键菜单里完成,不需要写任何代码
- 多个 Blender 实例自动连接同一个共享 Host;`Host` 面板显示状态、端口、客户端与模型缓存,并提供安装、启动、清理、重启和停止控制
- Runtime、模型和任务结果默认保存在 `~/.snapform`,可由多个 Blender 版本共享
- 支持 DA3 Small、Base、Large 1.1、Giant 1.1 和 Nested Giant Large 1.1
- 支持 BEN2 Base 单图背景剔除测试,输出透明 PNG 与推理耗时
- 支持单张 JPG、JPEG、PNG、WEBP 图片或常见视频文件输入
- 视频均匀抽帧,可设置参与重建的最大帧数
- `Host > Models` 默认折叠,展开后可逐项查看和下载所有可用模型
- Runtime 已就绪时不额外占用面板空间;模型列表每个模型只占一行
- `Debug` 默认折叠,内部的 DA3 与 BEN2 测试项也分别默认折叠
- 使用一个导出类型下拉菜单和一个 `Export` 按钮
- 只显示当前格式的专属参数;GLB 与 PLY 设置不会同时出现
- Debug 的模型选项按格式过滤;PLY 只提供两个带 Gaussian head 的 Giant 模型
- Debug 与 Convert to Depth Plane 的模型选项直接显示下载状态;未下载模型带
`(Missing)`,已下载模型保持原名
- 分别调用 DA3 的 `glb`、`gs_ply`、`depth_vis` 和 `mini_npz` exporter
- Image Empty 右键菜单可直接运行 DA3,根据 `mini_npz` 深度生成面片
- Image Empty 右键菜单可直接运行 BEN2,保留原图并换用新建的透明背景图片
- 使用浮点深度贴图、Simple Subdivision、Displace Modifier 和原图 Bump
构建立体面片
- 转换对话框可选择推理模式、处理分辨率和面片网格细分
- Debug 的 DA3 导出或 BEN2 背景剔除完成后只打开对应 job 文件夹,不向当前 Blender 文件添加任何数据
- 完整保留每次任务生成的所有官方文件
- 安装和模型下载使用后台子进程;推理提交给常驻 AI Host,状态栏实时显示阶段、百分比进度条与取消按钮
- 后台任务运行时,Install、Download 和 Export 按钮保持原位置并变灰,任务完成后恢复
## 使用 ## 安装
1. 使用 `uv run python pack.py` 构建 Extension ZIP。 1. 下载对应系统的扩展压缩包:
2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。 - Windows:`SnapForm.v0.10.10.windows-x64.zip`
3. 在 Extension Preferences 中确认 `uv` 路径和 Torch Backend。 - macOS(Apple 芯片):`SnapForm.v0.10.10.macos-arm64.zip`
4. 在 `Host` 点击 `Install Runtime`。首次安装会下载 Python、PyTorch、DA3、BEN2、FastAPI、Uvicorn, - macOS(Intel):`SnapForm.v0.10.10.macos-x64.zip`
并默认下载 `DA3 Small`、`DA3 Base` 和 `BEN2 Base` 模型。安装按 8 个阶段依次推进,每段进度 - Linux:`SnapForm.v0.10.10.linux-x64.zip`
从 0 到 100%:下载 uv(如需)、创建 Python Runtime、下载 PyTorch、安装 xFormers(macOS 跳过)、 2. 打开 Blender,进入 `Preferences > Extensions`
安装 AI 依赖、修补 DA3(跳过视频渲染),然后是三个模型的下载。进度文字统一为 `阶段/总数 阶段名: 状态`, 3. 点击 `Install from Disk`,选择下载的压缩包
模型下载阶段会实时显示下载速度,例如 `6/8 da3 small: downloading 14.2MB/s`。 4. 在扩展偏好设置中确认 `uv` 路径与 Torch 后端(保持默认即可);如果下载经常超时,
Runtime 就绪后即使模型下载中断也不影响 Host 使用,可稍后在 `Models` 面板单独补下。 可以在 `PyPI Index` 里填一个国内镜像,例如 `https://pypi.tuna.tsinghua.edu.cn/simple`
5. 在 3D View 按 `N`,打开 `SnapForm` 页签。 5. 打开 3D 视图右侧的 `SnapForm` 面板,点击 `Install Runtime`
6. 如需其他模型,展开 `Host > Models`,找到模型并点击该行的 `Download`。
7. 展开默认收起的 `Debug`,选择 `Export Format` 和本次执行使用的模型。
8. 选择一张图片或一个视频文件。视频可设置均匀抽取的帧数,然后点击 `Export`。
9. 测试背景剔除时,确认 `Host > Models` 中的 `BEN2 Base` 已就绪,然后在 `Debug` 的
`BEN2 Background Removal` 区域选择单张图片并点击 `Remove Background`。
完成后会打开 job 文件夹,其中的 `foreground.png` 带透明 Alpha;状态栏显示
BEN2 推理、模型加载与任务总耗时。当前测试固定关闭较慢的 foreground refinement。
## Image Empty 去背景 DA3、BEN2 源码包和 uv 已内置在扩展里,安装时直接从本地解包,不需要访问 GitHub。首次使用仍会自动
下载运行环境依赖(PyTorch 等)和 AI 模型,请保持网络畅通,下载需要一些时间。
1. 确认 `Host > Models` 中的 `BEN2 Base` 已就绪;首次安装 Runtime 时会默认下载。 ## 快速上手
2. 在 3D View 选中 Image Empty,右键选择 `Remove Background`。
3. 完成后以 Fake User 保留原 Image datablock,新建并 Pack 一张透明 PNG 到当前 `.blend`,然后只把
当前 Empty 的 `data` 换成新图;Empty 的 Transform、显示尺寸、Offset 及其他设置保持不变。
4. 新 Image 使用原图去掉格式后再追加 `.rb` 的名称。例如 `img.png` 生成 `img.rb`,
`img.png.001` 生成 `img.001.rb`。已有同名 datablock 时由 Blender 自动追加编号。
原图已 Pack 时会先写出临时 PNG 交给 BEN2,任务结束后删除临时输入;原 Image 仍保持 Pack ### 转换为立体面片
状态和原始内容。BEN2 的 `foreground.png` 也会保留在对应 job 目录中。
## Image Empty 转深度面片 1. 在 3D 视图中选中一个 Image Empty(空物体,显示方式为 Image)
2. 右键选择 `Convert to Depth Plane`
3. 在对话框中选择:
- `Depth from`:`Original` 直接用原图,`AI Generated` 让 AI 估算深度
- 需要去背景时勾选 `Remove Background`
- 其余选项保持默认即可
4. 点击确认,稍等片刻就得到带位移的立体面片
1. 在 3D View 选中 `Object Type = Empty`、`Display As = Image` 的对象。 ### 去除背景
2. 右键选择 `Convert to Depth Plane`,通过 `Depth from` 选择深度来源:
- `AI Generated`:显示 Model、Use Ray Pose 和 Resolution Scale。实际 Process Res 为图片长边像素数
乘以 Resolution Scale;范围为 `0.1–1.0`,默认 `1.0`。计算后的 Process Res 会直接显示在
`Resolution Scale (… px)` 标签中。
转换导出 `mini_npz`,使用前需要在
`Host > Models` 下载所选模型,并固定使用 `process_res_method="upper_bound_resize"`;设备和参考视图策略
继续使用场景设置。Model 下拉菜单只列出已经下载的 DA3 模型;若没有可用模型,对话框会引导到
`Host > Models` 下载。
- `Original`:直接把原图作为深度贴图,不显示估算参数。Remove Background 关闭时不启动
Runtime,也不要求源图片存在于磁盘。
`Original` 是默认深度来源。选择 `AI Generated` 时,其 Model、Use Ray Pose 与 Resolution Scale
参数会集中显示在独立边框中。两种模式都有默认折叠的 Options,内部依次为 Displace via、
Subdivision、Remove Background 和 Unshaded;Delete Original 保持在折叠项外。
Remove Background 默认关闭,与 Depth from 无关;开启后需要下载 `BEN2 Base`,并从原图生成、
Pack `.rb` 透明图片。它只替换生成面片材质的颜色图片,不修改源 Image Empty 的 Image datablock。
AI Generated 仍从原图生成 `.depth`;例如 `img.png` 会得到 `img.rb` 与 `img.depth`,不会生成
`img.rb.depth`。Original 仍用原图作为位移来源,不改变原有深度逻辑。
Displace via 可选 `Modifier`(默认)或 `Material`;
Unshaded 默认关闭;Subdivision 默认为 `5`,并默认删除原对象。
3. AI Generated 完成后会自动创建面片。面片几何先由 `mini_npz` 的相机内参与深度确定,再通过对象
Transform 匹配原 Image Empty 的可见尺寸、位置、旋转和 Offset。Original 模式直接使用源 Empty
的可见尺寸、Offset 与 Transform。开启 `Delete Original` 时,成功生成后
删除源 Empty。生成面片 Object 使用源 Empty 的原名;关闭 `Delete Original` 时不处理
同名冲突,由 Blender 自动为生成面片追加编号。
4. 短边分段数为 `2^Subdivision`,长边分段数根据图像宽高比取最接近的整数。例如高:宽为
`2:3`、Subdivision 为 `1` 时,直接生成 `2×3` 个面;Subdivision 为 `3` 时生成
`8×12` 个面。随后的 Simple Subdivision 修改器使用原生名称,Viewport Levels 固定为
`1`,Render Levels 固定为 `2`,不再随 Resolution Scale 或 Subdivision 动态变化。
所有任务使用统一的状态栏进度显示。变化中的阶段说明显示在右侧的紧凑进度条内部,取消按钮 选中 Image Empty 后,右键选择 `Remove Background`,完成后 Empty 会自动换成透明前景图。
位于进度条右侧。Image Depth 的阶段包括 Runtime 启动、输入准备、模型加载、深度估算、缓存
写出和面片构建。进度更新会强制刷新状态栏与窗口,无需移动鼠标即可实时显示。安装中断后残留的
半成品 Runtime 会在下次安装前自动清理,不会因为中断而无法再次安装。
导入过程精确保留 NPZ 内深度值之间的差值。例如深度 `2.0` 和 `5.5` 在 ### 视频与序列帧
`Displacement Strength = -1` 时,沿局部 Z 轴相差 `3.5 Blender Unit`。它不会把数值压缩到
`0..1`,也不再提供视觉归一化模式。
面片尺寸以有效深度的中位数 `Zref` 作为参考截面。先用 `mini_npz` 中的相机内参计算原始 SnapForm 支持视频(mp4、mkv、mov 等)和序列帧,会逐帧处理并生成逐帧对应的立体面片。处理范围默认
投影面积,再恢复源图纵横比: 跟随 Empty 上设置的 `Image Frames` 与 `Offset`,不需要额外配置。
```text ## 任务结果保存在哪里
RawWidth = depth_width × Zref × abs(Displacement Strength) / fx
RawHeight = depth_height × Zref × abs(Displacement Strength) / fy
Area = RawWidth × RawHeight
Aspect = source_image_width / source_image_height
Width = sqrt(Area × Aspect)
Height = sqrt(Area / Aspect)
```
这样即使 DA3 估计出的 `fx/fy` 不一致,也不会把正方形原图拉成长方形。主点 `cx/cy` 仍用于 任务结果保存在 `~/.snapform/jobs/` 下,每个任务一个文件夹。视频或序列帧生成的面片会引用这些文件夹
确定原始相机面片,随后对象 XY Transform 将它匹配到源 Image Empty。生成对象会记录 里的序列帧,请保留对应文件夹,不要随意删除。
`snapform_reference_depth`、`snapform_camera_plane_width` 和
`snapform_camera_plane_height`,便于检查实际采用的
参考深度和尺寸。
选择 `Modifier` 时,AI Generated 的 Displace Modifier 深度贴图使用 `Extension = Extend`, ## 常见问题
`Strength` 固定为 `-1`、`Midlevel` 固定为 `1`。Original 把原图指定给 Displace Texture,并设置
`Texture Coordinates = UV`、`Midlevel = 0`;其余修改器设置保留 Blender 默认值。
选择 `Material` 时,不创建 Displace Modifier;深度 Image Texture 的 Color 连接到 Displacement
节点的 Height,再连接 Material Output 的 Displacement,并把材质 Surface Displacement 设置为
`Displacement Only`。AI Generated 使用 `Midlevel = 1`、`Scale = -1`;Original 使用
`Midlevel = 0`、`Scale = 1`,与对应修改器值一致。无论选择哪种呈现方式,原始彩色贴图都会连接
Principled BSDF 的 `Base Color`,并通过
`Bump.Height → Bump.Normal → Principled.Normal` 提供表面细节。Bump 的基础 Distance 为
`0.001`,会乘以面片 Transform 的统一缩放倍率。原图 Color 还会经过 `Float Curve` 连接到
Principled BSDF 的 `Roughness`,曲线两端输出分别为 `0.5` 和 `0.1`;材质 IOR 保持
Principled BSDF 默认值 `1.5`。
开启 `Unshaded` 后,原图 Color 连接 Emission,原图 Alpha 作为 Mix Shader 系数,在
Transparent BSDF 与 Emission 之间混合后连接 Material Output;不创建 Principled BSDF、Bump 或
Float Curve。若 Displace via 选择 `Material`,深度贴图与 Displacement
节点仍会保留并连接 Material Output 的 Displacement。
这里的“原始距离”是 DA3 模型输出的尺度。普通 Small、Base、Large 和 Giant 模型的深度不应 - 首次使用需要下载模型,如果下载慢或失败,稍后在 `Host > Models` 里重试即可
被解释为现实世界的米;需要现实尺度时应选择带 metric depth 的 - 安装 Runtime 时提示下载超时:新版会自动探测并切换到镜像、失败时自动重试一次;如果仍然
`DA3 Nested Giant Large 1.1`。此外,当前结果仍是固定 XY 平面上的深度位移面,不是使用 超时,在扩展偏好设置的 `PyPI Index` 中填入镜像地址(如清华)后重新安装即可
相机内参逐像素反投影得到的完整 camera-space 网格;中位深度截面为现有 Displace Modifier - 找不到右键入口时,请确认选中了 Empty,且显示方式为 Image
工作流提供正确的整体高宽尺度。 - 动画面片不动:检查 Empty 的 `Image Frames` 是否大于 1,并确认播放场景帧
生成完成后会直接创建面片,不提供缓存重建入口或单独的导入参数。重新选择 ## 技术文档
`Convert to Depth Plane` 会按当前 Depth from 再次转换。
面片、深度 Image、Texture 与 Material 会在同一个 Blender Undo 步骤中创建;转换完成后执行
`Ctrl+Z` 会完整撤销本次转换,不会越过转换去回退更早的 Image Empty 操作。
开启 Remove Background 时,`.rb` Image 的创建、Pack 和深度面片生成也属于
同一个 Undo 步骤。
Image Empty 使用 Packed Image 时,SnapForm 会通过 Blender 写出临时 PNG 交给 DA3, 面向开发者的技术内容见 [docs/technical.md](docs/technical.md)。
不会 Unpack 或修改源 Image;任务结束后自动删除临时文件。普通图片仍直接读取源文件。DA3 的原始
`results.npz` 保留在对应 job 的 `exports/mini_npz/` 中,生成的浮点深度 Image 会 Pack 到
当前 `.blend`。Depth Image 与 Displace Texture 根据源 Image datablock 命名。例如 Empty Object
名为 `obj_img`、引用的 Image 名为 `img.png` 时,Depth Image 与 Texture 都命名为 `img.depth`;
生成的 Object、Mesh 与 Material 仍命名为 `obj_img`。Pack 使用 OpenEXR,因而不会把大于 `1.0`
的原始深度截断。已有同名 datablock 时由 Blender 自动追加 `.001` 等编号。
选择 PLY 后,模型菜单只会提供 `DA3 Giant 1.1` 和
`DA3 Nested Giant Large 1.1`。Small、Base 和 Large 没有 Gaussian head,无法由
DA3 生成 `gs_ply`。
每次任务会在 `~/.snapform/jobs/<任务目录>/` 中保留输出产物:
```text
<GLB 任务目录>/
└── scene.glb
<PLY 任务目录>/
└── gs_ply/*.ply
<Depth 任务目录>/
└── depth_vis/*.jpg
<Mini NPZ 任务目录>/
└── exports/mini_npz/results.npz
<BEN2 背景剔除任务目录>/
└── foreground.png
```
建议从 `DA3 Base` 开始。它使用 Apache 2.0 模型权重;`DA3 Large 1.1` 权重采用
CC BY-NC 4.0,仅限非商业用途。
## 输出格式说明
`Export Format` 下拉菜单分别把以下原始 `export_format` 传给 DA3:
- GLB:`glb`
- PLY:`gs_ply`,同时传入 `infer_gs=True`
- Depth Images:`depth_vis`
- Mini NPZ:`mini_npz`
界面使用 DA3 API 的参数名、选项和默认值。GLB 专属设置仅在选择 `GLB` 时显示,
`gs_views_interval` 仅在选择 `PLY (3DGS)` 时显示:
- `process_res=504`
- `process_res_method="upper_bound_resize"`
- `use_ray_pose=False`
- `ref_view_strategy="saddle_balanced"`
- `conf_thresh_percentile=40.0`
- `num_max_points=1_000_000`
- `show_cameras=True`
- `export_kwargs["gs_ply"]["gs_views_interval"]=1`
DA3 官方 `glb` 是彩色点云与可选相机线框;`gs_ply` 是带 Gaussian 属性的 PLY;
`depth_vis` 是原图与着色深度图的并排预览;`mini_npz` 保存深度、置信度和可用的相机参数。
GLB 任务会关闭其默认附带的深度预览,使四种 job 产物保持分离。
官方格式与限制可参阅
[DA3 Python API 文档](https://github.com/ByteDance-Seed/Depth-Anything-3/blob/main/docs/API.md)。
## Runtime
Extension 本身不包含 Torch。`Install Runtime` 使用 `uv` 创建独立 Python 3.12
环境,并安装。机器上即使没有安装 `uv` 或 Python 也没关系:安装程序由 Blender
自带的 Python 启动,缺 `uv` 时会自动下载官方独立二进制(约 30MB),再用它下载
Python 3.12 与全部依赖。`uv` 会缓存在扩展存储目录下,后续安装直接复用。
安装前会自动预检磁盘空间与网络连通性,未通过会直接给出英文报错,不会开始下载。
预检与安装支持环境变量切换镜像与代理:`PIP_INDEX_URL`(PyPI 镜像)、
`SNAPFORM_TORCH_INDEX`(PyTorch 索引)、`HF_ENDPOINT`(Hugging Face 镜像,
如 `https://hf-mirror.com`)、`HTTP(S)_PROXY`/`ALL_PROXY`(代理)。设置后预检会直接检查
这些地址。直连不通时预检会**自动回退到镜像源**(清华 PyPI、阿里云 PyTorch、hf-mirror),
无需手动配置,状态栏会提示 `using mirror sources`;镜像也不通才会报错。
单独下载模型时也会先检查 Hugging Face 连通性,直连不通自动切换到 `hf-mirror` 镜像。
Windows 上代理若只在系统设置里而未写入环境变量,安装程序会尝试自动从系统代理注入子进程。
- PyTorch 2.9.1
- TorchVision 0.24.1
- xFormers 0.0.33.post2(与 PyTorch 2.9.1 对应,仅 Windows/Linux;macOS 无 wheel 自动跳过)
- 固定到已验证 commit 的 Depth Anything 3
- 固定到已验证 commit 的 BEN2
- FastAPI 0.139.2 与 Uvicorn 0.51.0
- DA3 与 BEN2 通过 GitHub 官方源码包安装(无需 Git);SnapForm 直接调用 DA3
底层 export 接口,`gs_ply` 导出不会触发 `gs_video` 渲染,因此不需要 gsplat
与 C++ 编译工具链
- `addict` 2.4.0(补充 DA3 上游包元数据暂未声明的运行依赖)
默认 CUDA Backend 为 CUDA 12.8,也可以选择 CPU。CPU 可以用于验证,但实际推理
通常很慢。
CUDA 12.8 的 PyTorch wheel 自带 CUDA 运行库,无需安装 CUDA Toolkit;只需要
NVIDIA 驱动 ≥ 570(数据中心显卡可用更老驱动)。驱动不满足或没有 NVIDIA 显卡时,
设备会自动回退到 CPU(macOS 为 MPS),只有显式选择 CUDA 才会报错。
macOS(Apple Silicon)会自动从 PyPI 安装 CPU 版 PyTorch,推理设备默认使用 MPS
(Apple Metal)加速,也可以在设备选项里切换 CPU。PyTorch 自 2.3 起不再提供 Intel
Mac(x86_64)轮子,Intel Mac 暂不支持。
Extension 启用后只会查找已有的共享 AI Host,不会主动启动新进程。第一次执行 AI 操作或在侧边栏
手动点击 `Start` 时才会启动 Host。Host 只绑定回环地址,并让操作系统分配动态端口;第一个实例把
端口、进程和随机 Token 写入 `host.json`,其余 Blender 通过标准库
`urllib.request` 连接同一个 Host。`host.lock` 保证同一 Storage Root 同时只有一个 Host,Token
则阻止同机其他进程直接调用接口。Host 全局同时只执行一个 AI 任务。
每个 Blender 会注册独立客户端并每 5 秒发送心跳。正常关闭 Blender 时先注销自身;如果当前 Host
就是由这个 Blender 启动的,会直接终止它,无需等待。异常退出的客户端会在 10 秒后过期;Host
在最后一个客户端离开后仍会保留模型缓存等待 20 秒再退出,作为 Host 由其他进程启动等场景的兜底。
取消任务不会重启 Host,并会保留已经加载的模型;
如果取消发生在一次不可中断的 Torch 调用中,任务会在该调用返回后的最近检查点结束。
侧边栏 `Host` 面板以 `Running (Clients: n)` 表示空闲服务,以 `Running (Busy)` 表示正在执行任务,
并显示缓存模型和活动任务。空闲且只有当前 Blender 连接时,可以清理模型缓存、重启或停止 Host;
多客户端或任务执行期间不会允许破坏共享服务。
BEN2 在首次使用后保持加载;DA3 保留最近使用的一个模型。切换 DA3 模型时会释放旧模型;加载时如果
显存不足,会先释放另一类缓存模型并重试。
Runtime Schema 已升级,旧 Runtime 需要重新执行 `Install Runtime` 才能安装 HTTP Host 依赖。
Host 的启动和崩溃日志写入 `~/.snapform/ai-host.log`。
Runtime、模型和任务结果默认保存在用户主目录的 `~/.snapform`:
```text
~/.snapform/
├── runtime/
├── models/
├── jobs/
├── ai-host.log
├── host.json # Host 运行时存在
└── host.lock # Host 运行时持有
```
不同 Blender 版本安装相同 Extension 后会共用这个目录。也可以在 Extension
Preferences 的 `Storage Root` 中修改共享位置。不要把模型写入 Extension 安装目录,
因为该目录可能只读,并且升级时会被替换。不要同时从多个 Blender 实例执行
`Install Runtime`;推理 Host、任务队列和模型读取可以共享。
## 开发
```bash
uv run python -m unittest discover -s tests -v
uv run python pack.py
```
项目结构:
```text
.
├── pack.py
├── src/
│ └── snapform/
│ ├── blender_manifest.toml
│ ├── ai_host.py
│ ├── bootstrap_runtime.py
│ ├── models.py
│ ├── image_depth.py
│ ├── operators.py
│ ├── panel.py
│ ├── preferences.py
│ ├── properties.py
│ ├── runtime.py
│ └── worker.py
└── tests/
```
## 上游与许可证
- Depth Anything 3 代码采用 Apache 2.0。
- BEN2 代码与模型权重采用 MIT License。
- 本 Extension 代码采用 GPL-3.0-or-later。
- 模型权重许可证由所选 Hugging Face 模型仓库决定,Extension 下载前会在界面显示。

394
docs/technical.md Normal file
View File

@@ -0,0 +1,394 @@
# SnapForm 技术文档
本文档面向开发者和需要了解内部实现细节的用户,从原 README 迁移而来。普通用户请阅读项目根目录的
`README.md`。
`SnapForm` 是一套从图片快速构建空间结构的 Blender Extension。目前以
[Depth Anything 3](https://github.com/ByteDance-Seed/Depth-Anything-3)
作为首个 AI 能力,在独立 Runtime 中运行 DA3,可将官方导出结果保存在独立 job 目录,
也可把 Image Empty 转换成带几何与材质置换的立体面片。Debug 面板还可运行
[BEN2](https://github.com/PramaLLC/BEN2) 背景剔除功能测试。
## 当前功能
- 使用 Blender Extension manifest,可通过 `Install from Disk` 安装
- 在独立 Python Runtime 的常驻 HTTP AI Host 中运行 Torch、DA3 与 BEN2,避免污染 Blender Python 进程
- Blender 使用内置 `urllib.request` 访问仅监听 `127.0.0.1` 的 Host;请求带随机 Token
- 常驻 Host 复用 Torch、CUDA、BEN2 与最近使用的 DA3 模型,连续执行时不再重复加载
- 多个 Blender 实例自动连接同一个共享 Host;`Host` 面板显示状态、端口、客户端与模型缓存,并提供安装、启动、清理、重启和停止控制
- Runtime、模型和任务结果默认保存在 `~/.snapform`,可由多个 Blender 版本共享
- 支持 DA3 Small、Base、Large 1.1、Giant 1.1 和 Nested Giant Large 1.1
- 支持 BEN2 Base 单图背景剔除测试,输出透明 PNG 与推理耗时
- 支持单张 JPG、JPEG、PNG、WEBP 图片或常见视频文件输入
- 视频按顺序截取前若干帧,可设置参与重建的最大帧数
- `Host > Models` 默认折叠,展开后可逐项查看和下载所有可用模型
- Runtime 已就绪时不额外占用面板空间;模型列表每个模型只占一行
- `Debug` 默认折叠,内部的 DA3 与 BEN2 测试项也分别默认折叠
- 使用一个导出类型下拉菜单和一个 `Export` 按钮
- 只显示当前格式的专属参数;GLB 与 PLY 设置不会同时出现
- Debug 的模型选项按格式过滤;PLY 只提供两个带 Gaussian head 的 Giant 模型
- Debug 与 Convert to Depth Plane 的模型选项直接显示下载状态;未下载模型带
`(Missing)`,已下载模型保持原名
- 分别调用 DA3 的 `glb`、`gs_ply`、`depth_vis` 和 `mini_npz` exporter
- Image Empty 右键菜单可直接运行 DA3,根据 `mini_npz` 深度生成面片
- Image Empty 右键菜单可直接运行 BEN2,保留原图并换用新建的透明背景图片
- Image Empty 支持视频与序列帧输入:`Convert to Depth Plane` 逐帧估计深度,
`Remove Background` 逐帧抠图,入口保持不变
- 视频生成的序列帧以外部文件引用保留在 job 目录(Blender 无法整体 Pack 序列帧);
Displace 修改器与材质 Image Texture 均可按场景帧播放序列帧
- 使用浮点深度贴图、Simple Subdivision、Displace Modifier 和原图 Bump
构建立体面片
- 转换对话框可选择推理模式、处理分辨率和面片网格细分
- Debug 的 DA3 导出或 BEN2 背景剔除完成后只打开对应 job 文件夹,不向当前 Blender 文件添加任何数据
- 完整保留每次任务生成的所有官方文件
- 安装和模型下载使用后台子进程;推理提交给常驻 AI Host,状态栏实时显示阶段、百分比进度条与取消按钮
- 后台任务运行时,Install、Download 和 Export 按钮保持原位置并变灰,任务完成后恢复
## 使用
1. 使用 `uv run python pack.py` 分别构建 `windows-x64`、`macos-x64`、`macos-arm64`、
`linux-x64` 四个平台的 Extension ZIP(如 `SnapForm.v0.10.10.macos-arm64.zip`)。
每个包只含对应平台的 uv 二进制,并内置 DA3、BEN2 源码包。
2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。
3. 在 Extension Preferences 中确认 `uv` 路径和 Torch Backend。
4. 在 `Host` 点击 `Install Runtime`。首次安装会下载 Python、PyTorch、FastAPI、Uvicorn;
DA3、BEN2 源码包与 uv 已内置在扩展中,安装时从 `~/.snapform/sources/` 解包使用,
并默认下载 `DA3 Small`、`DA3 Base` 和 `BEN2 Base` 模型。安装按 8 个阶段依次推进,每段进度
从 0 到 100%:解包内置 uv、创建 Python Runtime、下载 PyTorch、安装 xFormers(macOS 跳过,
内置 DA3 源码包为官方原版并已移除 xformers 硬依赖,自动回退纯 PyTorch 实现)、安装 AI 依赖,
然后是三个模型的下载。进度文字统一为 `阶段/总数 阶段名: 状态`,
模型下载阶段会实时显示下载速度,例如 `6/8 da3 small: downloading 14.2MB/s`。
Runtime 就绪后即使模型下载中断也不影响 Host 使用,可稍后在 `Models` 面板单独补下。
安装不做启动网络预检:若网络无法稳定访问 pypi.org,可在扩展偏好设置的 `PyPI Index`
中指定镜像源(如清华、阿里),该地址会通过 `--pypi-index` 传给 bootstrap 并优先使用;
未指定时,pip 安装失败会自动重试一次清华镜像,模型下载由 worker 自行探测回退到 HF 镜像。
5. 在 3D View 按 `N`,打开 `SnapForm` 页签。
6. 如需其他模型,展开 `Host > Models`,找到模型并点击该行的 `Download`。
7. 展开默认收起的 `Debug`,选择 `Export Format` 和本次执行使用的模型。
8. 选择一张图片或一个视频文件。视频可设置从头截取的帧数,然后点击 `Export`。
9. 测试背景剔除时,确认 `Host > Models` 中的 `BEN2 Base` 已就绪,然后在 `Debug` 的
`BEN2 Background Removal` 区域选择单张图片并点击 `Remove Background`。
完成后会打开 job 文件夹,其中的 `foreground.png` 带透明 Alpha;状态栏显示
BEN2 推理、模型加载与任务总耗时。当前测试固定关闭较慢的 foreground refinement。
## Image Empty 去背景
1. 确认 `Host > Models` 中的 `BEN2 Base` 已就绪;首次安装 Runtime 时会默认下载。
2. 在 3D View 选中 Image Empty,右键选择 `Remove Background`。
3. 完成后以 Fake User 保留原 Image datablock,新建并 Pack 一张透明 PNG 到当前 `.blend`,然后只把
当前 Empty 的 `data` 换成新图;Empty 的 Transform、显示尺寸、Offset 及其他设置保持不变。
4. 新 Image 使用原图去掉格式后再追加 `.foreground` 的名称。例如 `img.png` 生成
`img.foreground`,`img.png.001` 生成 `img.001.foreground`。已有同名 datablock 时由
Blender 自动追加编号。
原图已 Pack 时会先写出临时 PNG 交给 BEN2,任务结束后删除临时输入;原 Image 仍保持 Pack
状态和原始内容。BEN2 的 `foreground.png` 也会保留在对应 job 目录中。
Image Empty 引用视频或序列帧时,Remove Background 直接执行,默认取 Empty 的 Image Frames
数值(未设置时为完整帧数),起始帧跟随 Empty 的 Offset;若 Offset 加 Frames 超出视频长度,
会把窗口整体收进视频内(保证实际处理 Frames 帧数),Offset 过小时从第 1 帧开始。输出
`foreground.0001.png` 等透明序列帧(帧号从 0001 开始,与 Blender 场景帧一一对应),并把
源帧以 `frame.0001.jpg` 等命名复制到 job 目录根下供材质颜色导入。替换 Empty 的 Image 为
加载这些帧的 SEQUENCE Image,从原 Empty 的 Start 位置开始播放(Start 原样保留),
Cyclic 沿用原 Empty 的设置。序列帧无法被 Blender 整体 Pack,因此保持外部引用 job 目录;
请保留对应 job 目录,否则序列会失效。
## Image Empty 转深度面片
1. 在 3D View 选中 `Object Type = Empty`、`Display As = Image` 的对象。
2. 右键选择 `Convert to Depth Plane`,通过 `Depth from` 选择深度来源:
- `AI Generated`:显示 Model、Use Ray Pose 和 Resolution Scale。实际 Process Res 为图片长边像素数
乘以 Resolution Scale;范围为 `0.1–1.0`,默认 `1.0`。计算后的 Process Res 会直接显示在
`Resolution Scale (… px)` 标签中。
转换导出 `mini_npz`,使用前需要在
`Host > Models` 下载所选模型,并固定使用 `process_res_method="upper_bound_resize"`;设备和参考视图策略
继续使用场景设置。Model 下拉菜单只列出已经下载的 DA3 模型;若没有可用模型,对话框会引导到
`Host > Models` 下载。
- `Original`:直接把原图作为深度贴图,不显示估算参数。单张图片不启动 Runtime;视频或
序列帧会启动轻量抽帧任务,把帧抽出来作为 `.depth` 序列(比直接引用视频稳定)。
`Original` 是默认深度来源。选择 `AI Generated` 时,其 Model、Use Ray Pose 与 Resolution Scale
参数会集中显示在独立边框中。两种模式都有默认折叠的 Options,内部依次为 Displace via、
Subdivision、Remove Background 和 Unshaded;Delete Original 保持在折叠项外。
Remove Background 默认关闭,与 Depth from 无关;开启后需要下载 `BEN2 Base`,并从原图生成、
Pack `.foreground` 透明图片。它只替换生成面片材质的颜色图片,不修改源 Image Empty 的
Image datablock。AI Generated 仍从原图生成 `.depth`;例如 `img.png` 会得到
`img.foreground` 与 `img.depth`,不会生成 `img.foreground.depth`。Original 仍用原图作为
位移来源,不改变原有深度逻辑。
Image Empty 引用视频或序列帧时,AI Generated 会按 Empty 的 Image Frames 数值逐帧估计
深度(未设置时为完整帧数),起始帧跟随 Empty 的 Offset,与 Remove Background 范围一致;
Offset 超出视频长度时同样把窗口收进视频内。材质节点与置换贴图的 Start 保持与原 Empty
一致;Original 模式(无 Remove Background)时,单帧会复制原图改名为 `.depth` 供位移使用,
视频或序列帧则抽帧保存为 job 目录根下的 `frame.0001.jpg` 等,并作为 `.depth` 序列使用。
材质颜色使用这些帧导入的 `.frame` 序列(Frames/Start/Offset 同步),开启 Remove
Background 时改用 `foreground.0001.png` 序列。AI Generated 的深度以 EXR 序列帧
(`depth.0001.exr` 起)保存在 job 目录根下;Original 模式(含 Remove Background)的位移
使用 `frame.*` 序列帧,保证逐帧对齐。
Displace via 可选 `Modifier`(默认)或 `Material`;
Unshaded 默认关闭;Subdivision 默认为 `5`,并默认删除原对象。
3. AI Generated 完成后会自动创建面片。面片几何先由 `mini_npz` 的相机内参与深度确定,再通过对象
Transform 匹配原 Image Empty 的可见尺寸、位置、旋转和 Offset。Original 模式直接使用源 Empty
的可见尺寸、Offset 与 Transform。开启 `Delete Original` 时,成功生成后
删除源 Empty。生成面片 Object 使用源 Empty 的原名;关闭 `Delete Original` 时不处理
同名冲突,由 Blender 自动为生成面片追加编号。
4. 短边分段数为 `2^Subdivision`,长边分段数根据图像宽高比取最接近的整数。例如高:宽为
`2:3`、Subdivision 为 `1` 时,直接生成 `2×3` 个面;Subdivision 为 `3` 时生成
`8×12` 个面。随后的 Simple Subdivision 修改器使用原生名称,Viewport Levels 固定为
`1`,Render Levels 固定为 `2`,不再随 Resolution Scale 或 Subdivision 动态变化。
所有任务使用统一的状态栏进度显示。变化中的阶段说明显示在右侧的紧凑进度条内部,取消按钮
位于进度条右侧。Image Depth 的阶段包括 Runtime 启动、输入准备、模型加载、深度估算、缓存
写出和面片构建。进度更新会强制刷新状态栏与窗口,无需移动鼠标即可实时显示。安装中断后残留的
半成品 Runtime 会在下次安装前自动清理,不会因为中断而无法再次安装。
导入过程精确保留 NPZ 内深度值之间的差值。例如深度 `2.0` 和 `5.5` 在
`Displacement Strength = -1` 时,沿局部 Z 轴相差 `3.5 Blender Unit`。它不会把数值压缩到
`0..1`,也不再提供视觉归一化模式。
面片尺寸以有效深度的中位数 `Zref` 作为参考截面。先用 `mini_npz` 中的相机内参计算原始
投影面积,再恢复源图纵横比:
```text
RawWidth = depth_width × Zref × abs(Displacement Strength) / fx
RawHeight = depth_height × Zref × abs(Displacement Strength) / fy
Area = RawWidth × RawHeight
Aspect = source_image_width / source_image_height
Width = sqrt(Area × Aspect)
Height = sqrt(Area / Aspect)
```
这样即使 DA3 估计出的 `fx/fy` 不一致,也不会把正方形原图拉成长方形。主点 `cx/cy` 仍用于
确定原始相机面片,随后对象 XY Transform 将它匹配到源 Image Empty。生成对象会记录
`snapform_reference_depth`、`snapform_camera_plane_width` 和
`snapform_camera_plane_height`,便于检查实际采用的
参考深度和尺寸。
选择 `Modifier` 时,AI Generated 的 Displace Modifier 深度贴图使用 `Extension = Extend`,
`Strength` 固定为 `-1`、`Midlevel` 固定为 `1`。Original 把原图指定给 Displace Texture,并设置
`Texture Coordinates = UV`、`Midlevel = 0`;其余修改器设置保留 Blender 默认值。
视频与序列帧生成的深度贴图是 SEQUENCE Image;Displace 修改器的 Texture 与材质的 Image Texture
节点都支持序列帧,通过 image_user 的 Frames / Offset / Start 设置把第 i 帧映射到场景
`Start+i` 帧(Start 沿用 Empty 的 Start),并启用 Auto Refresh,播放时随场景帧逐帧更新。
选择 `Material` 时,不创建 Displace Modifier;深度 Image Texture 的 Color 连接到 Displacement
节点的 Height,再连接 Material Output 的 Displacement,并把材质 Surface Displacement 设置为
`Displacement Only`。AI Generated 使用 `Midlevel = 1`、`Scale = -1`;Original 使用
`Midlevel = 0`、`Scale = 1`,与对应修改器值一致。无论选择哪种呈现方式,原始彩色贴图都会连接
Principled BSDF 的 `Base Color`,并通过
`Bump.Height → Bump.Normal → Principled.Normal` 提供表面细节。Bump 的基础 Distance 为
`0.001`,会乘以面片 Transform 的统一缩放倍率。原图 Color 还会经过 `Float Curve` 连接到
Principled BSDF 的 `Roughness`,曲线两端输出分别为 `0.5` 和 `0.1`;材质 IOR 保持
Principled BSDF 默认值 `1.5`。
开启 `Unshaded` 后,原图 Color 连接 Emission,原图 Alpha 作为 Mix Shader 系数,在
Transparent BSDF 与 Emission 之间混合后连接 Material Output;不创建 Principled BSDF、Bump 或
Float Curve。若 Displace via 选择 `Material`,深度贴图与 Displacement
节点仍会保留并连接 Material Output 的 Displacement。
这里的“原始距离”是 DA3 模型输出的尺度。普通 Small、Base、Large 和 Giant 模型的深度不应
被解释为现实世界的米;需要现实尺度时应选择带 metric depth 的
`DA3 Nested Giant Large 1.1`。此外,当前结果仍是固定 XY 平面上的深度位移面,不是使用
相机内参逐像素反投影得到的完整 camera-space 网格;中位深度截面为现有 Displace Modifier
工作流提供正确的整体高宽尺度。
生成完成后会直接创建面片,不提供缓存重建入口或单独的导入参数。重新选择
`Convert to Depth Plane` 会按当前 Depth from 再次转换。
面片、深度 Image、Texture 与 Material 会在同一个 Blender Undo 步骤中创建;转换完成后执行
`Ctrl+Z` 会完整撤销本次转换,不会越过转换去回退更早的 Image Empty 操作。
开启 Remove Background 时,`.foreground` Image 的创建、Pack 和深度面片生成也属于
同一个 Undo 步骤。
Image Empty 使用 Packed Image 时,SnapForm 会通过 Blender 写出临时 PNG 交给 DA3,
不会 Unpack 或修改源 Image;任务结束后自动删除临时文件。普通图片仍直接读取源文件。DA3 的原始
`results.npz` 保留在对应 job 文件夹根下,生成的浮点深度 Image 会 Pack 到
当前 `.blend`。Depth Image 与 Displace Texture 根据源 Image datablock 命名。例如 Empty Object
名为 `obj_img`、引用的 Image 名为 `img.png` 时,Depth Image 与 Texture 都命名为 `img.depth`;
生成的 Object、Mesh 与 Material 仍命名为 `obj_img`。Pack 使用 OpenEXR,因而不会把大于 `1.0`
的原始深度截断。已有同名 datablock 时由 Blender 自动追加 `.001` 等编号。
视频输入时 Blender 无法整体 Pack 序列帧,因此 `.depth` 与颜色序列改为引用 job 目录根下的
`depth.0001.exr`、`frame.0001.jpg`(`.frame`)或 `foreground.0001.png`,并在每帧保留原始
浮点深度;保存 `.blend` 后再次打开仍需对应 job 目录存在。
选择 PLY 后,模型菜单只会提供 `DA3 Giant 1.1` 和
`DA3 Nested Giant Large 1.1`。Small、Base 和 Large 没有 Gaussian head,无法由
DA3 生成 `gs_ply`。
每个任务在 `~/.snapform/jobs/` 下用独立文件夹包裹(命名格式为 `任务类型_时间号`,例如
`depth-plane_20260807-153000`),日志、状态与输出产物都放在该文件夹内:
```text
reconstruct_20260807-153000/
├── reconstruct_20260807-153000.log
├── reconstruct_20260807-153000.json
└── scene.glb
reconstruct_20260807-153000/
├── reconstruct_20260807-153000.log
├── reconstruct_20260807-153000.json
└── gs_ply/*.ply
reconstruct_20260807-153000/
├── reconstruct_20260807-153000.log
├── reconstruct_20260807-153000.json
└── depth_vis/*.jpg
reconstruct_20260807-153000/
├── reconstruct_20260807-153000.log
├── reconstruct_20260807-153000.json
└── results.npz
remove-background_20260807-153000/
├── remove-background_20260807-153000.log
├── remove-background_20260807-153000.json
├── foreground.png 单图结果
├── foreground.0001.png… 视频/序列帧结果
├── frame.0001.jpg… 视频/序列帧的源帧副本
└── depth.0001.exr… AI 深度序列帧
```
建议从 `DA3 Base` 开始。它使用 Apache 2.0 模型权重;`DA3 Large 1.1` 权重采用
CC BY-NC 4.0,仅限非商业用途。
## 输出格式说明
`Export Format` 下拉菜单分别把以下原始 `export_format` 传给 DA3:
- GLB:`glb`
- PLY:`gs_ply`,同时传入 `infer_gs=True`
- Depth Images:`depth_vis`
- Mini NPZ:`mini_npz`
界面使用 DA3 API 的参数名、选项和默认值。GLB 专属设置仅在选择 `GLB` 时显示,
`gs_views_interval` 仅在选择 `PLY (3DGS)` 时显示:
- `process_res=504`
- `process_res_method="upper_bound_resize"`
- `use_ray_pose=False`
- `ref_view_strategy="saddle_balanced"`
- `conf_thresh_percentile=40.0`
- `num_max_points=1_000_000`
- `show_cameras=True`
- `export_kwargs["gs_ply"]["gs_views_interval"]=1`
DA3 官方 `glb` 是彩色点云与可选相机线框;`gs_ply` 是带 Gaussian 属性的 PLY;
`depth_vis` 是原图与着色深度图的并排预览;`mini_npz` 保存深度、置信度和可用的相机参数。
GLB 任务会关闭其默认附带的深度预览,使四种 job 产物保持分离。
官方格式与限制可参阅
[DA3 Python API 文档](https://github.com/ByteDance-Seed/Depth-Anything-3/blob/main/docs/API.md)。
## Runtime
Extension 本身不包含 Torch。`Install Runtime` 使用 `uv` 创建独立 Python 3.12
环境,并安装。机器上即使没有安装 `uv` 或 Python 也没关系:安装程序由 Blender
自带的 Python 启动,缺 `uv` 时会自动下载官方独立二进制(约 30MB),再用它下载
Python 3.12 与全部依赖。`uv` 会缓存在扩展存储目录下,后续安装直接复用。
安装不做网络预检,直接开始;环境变量仍可切换镜像与代理:`PIP_INDEX_URL`(PyPI 镜像)、
`SNAPFORM_TORCH_INDEX`(PyTorch 索引)、`HF_ENDPOINT`(Hugging Face 镜像)、
`HTTP(S)_PROXY`/`ALL_PROXY`(代理)。未指定 PyPI 镜像时,
pip 安装失败会自动依次重试镜像(清华、阿里云);模型下载不预探测,官方 Hugging Face
失败后自动回退 ModelScope(`modelscope.cn`,仓库名与 Hugging Face 一致)。
GitHub 不再参与安装:DA3、BEN2、uv 已内置,缺少内置文件会直接报错提示重新安装,不再下载兜底。
单独下载模型时遵循同样的顺序:Hugging Face(可用 `HF_ENDPOINT` 指定其他镜像)失败后自动切换到
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 源码包随扩展打包,安装时复制到 `~/.snapform/sources/` 后由 uv 本地安装
(无需访问 GitHub);SnapForm 直接调用 DA3
底层 export 接口,`gs_ply` 导出不会触发 `gs_video` 渲染,因此不需要 gsplat
与 C++ 编译工具链
- `sources/` 下的 DA3/BEN2 源码包保持官方原版(固定 commit);`sources/patches/<包名>/`
存放 SnapForm 对源码的覆盖文件(移除 xformers 硬依赖、export 入口改懒加载、
射线位姿的 3×3 QR 分解在 MPS 上改走 CPU),
`pack.py` 打包时自动套入官方包,扩展 zip 内才是修补后的版本;更新上游只需替换官方包
并核对覆盖文件是否仍适用
- `addict` 2.4.0(补充 DA3 上游包元数据暂未声明的运行依赖)
默认 CUDA Backend 为 CUDA 12.8,也可以选择 CPU。CPU 可以用于验证,但实际推理
通常很慢。
CUDA 12.8 的 PyTorch wheel 自带 CUDA 运行库,无需安装 CUDA Toolkit;只需要
NVIDIA 驱动 ≥ 570(数据中心显卡可用更老驱动)。驱动不满足或没有 NVIDIA 显卡时,
设备会自动回退到 CPU(macOS 为 MPS),只有显式选择 CUDA 才会报错。
macOS(Apple Silicon)会自动从 PyPI 安装 CPU 版 PyTorch,推理设备默认使用 MPS
(Apple Metal)加速,也可以在设备选项里切换 CPU。PyTorch 自 2.3 起不再提供 Intel
Mac(x86_64)轮子,Intel Mac 暂不支持。DA3 射线位姿估计中的 3×3 QR 分解在 MPS
上尚未实现,SnapForm 的覆盖文件会单独把该分解放到 CPU 执行,其余推理仍走 MPS。
Extension 启用后只会查找已有的共享 AI Host,不会主动启动新进程。第一次执行 AI 操作或在侧边栏
手动点击 `Start` 时才会启动 Host。Host 只绑定回环地址,并让操作系统分配动态端口;第一个实例把
端口、进程和随机 Token 写入 `host.json`,其余 Blender 通过标准库
`urllib.request` 连接同一个 Host。`host.lock` 保证同一 Storage Root 同时只有一个 Host,Token
则阻止同机其他进程直接调用接口。Host 全局同时只执行一个 AI 任务。
每个 Blender 会注册独立客户端并每 5 秒发送心跳。正常关闭 Blender 时先注销自身;如果当前 Host
就是由这个 Blender 启动的,会直接终止它,无需等待。异常退出的客户端会在 10 秒后过期;Host
在最后一个客户端离开后仍会保留模型缓存等待 20 秒再退出,作为 Host 由其他进程启动等场景的兜底。
取消任务不会重启 Host,并会保留已经加载的模型;
如果取消发生在一次不可中断的 Torch 调用中,任务会在该调用返回后的最近检查点结束。
侧边栏 `Host` 面板以 `Running (Clients: n)` 表示空闲服务,以 `Running (Busy)` 表示正在执行任务,
并显示缓存模型和活动任务。空闲且只有当前 Blender 连接时,可以清理模型缓存、重启或停止 Host;
多客户端或任务执行期间不会允许破坏共享服务。
BEN2 在首次使用后保持加载;DA3 保留最近使用的一个模型。切换 DA3 模型时会释放旧模型;加载时如果
显存不足,会先释放另一类缓存模型并重试。
Runtime Schema 已升级,旧 Runtime 需要重新执行 `Install Runtime` 才能安装 HTTP Host 依赖。
Host 的启动和崩溃日志写入 `~/.snapform/ai-host.log`。
Runtime、模型和任务结果默认保存在用户主目录的 `~/.snapform`:
```text
~/.snapform/
├── runtime/
├── models/
├── jobs/
├── ai-host.log
├── host.json # Host 运行时存在
└── host.lock # Host 运行时持有
```
不同 Blender 版本安装相同 Extension 后会共用这个目录。也可以在 Extension
Preferences 的 `Storage Root` 中修改共享位置。不要把模型写入 Extension 安装目录,
因为该目录可能只读,并且升级时会被替换。不要同时从多个 Blender 实例执行
`Install Runtime`;推理 Host、任务队列和模型读取可以共享。
## 开发
```bash
uv run python -m unittest discover -s tests -v
uv run python pack.py
```
项目结构:
```text
.
├── pack.py
├── sources/ 官方 DA3/BEN2 源码包、patches 覆盖文件与各平台 uv 二进制
├── src/
│ └── snapform/
│ ├── blender_manifest.toml
│ ├── ai_host.py
│ ├── bootstrap_runtime.py
│ ├── models.py
│ ├── image_depth.py
│ ├── operators.py
│ ├── panel.py
│ ├── preferences.py
│ ├── properties.py
│ ├── runtime.py
│ └── worker.py
└── tests/
```
## 上游与许可证
- Depth Anything 3 代码采用 Apache 2.0。
- BEN2 代码与模型权重采用 MIT License。
- 本 Extension 代码采用 GPL-3.0-or-later。
- 模型权重许可证由所选 Hugging Face 模型仓库决定,Extension 下载前会在界面显示。

125
pack.py
View File

@@ -1,3 +1,5 @@
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
@@ -5,30 +7,135 @@ from zipfile import ZIP_DEFLATED, ZipFile
PROJECT_ROOT = Path(__file__).resolve().parent PROJECT_ROOT = Path(__file__).resolve().parent
SOURCE_DIR = PROJECT_ROOT / "src" / "snapform" SOURCE_DIR = PROJECT_ROOT / "src" / "snapform"
SOURCES_DIR = PROJECT_ROOT / "sources"
DIST_DIR = PROJECT_ROOT / "dist" DIST_DIR = PROJECT_ROOT / "dist"
PATCHES_DIR = SOURCES_DIR / "patches"
PLATFORMS = {
"windows-x64": "uv-x86_64-pc-windows-msvc.zip",
"macos-x64": "uv-x86_64-apple-darwin.tar.gz",
"macos-arm64": "uv-aarch64-apple-darwin.tar.gz",
"linux-x64": "uv-x86_64-unknown-linux-gnu.tar.gz",
}
BUNDLED_ARCHIVES = {
"sources/depth-anything-3.tar.gz": "depth-anything-3.tar.gz",
"sources/ben2.tar.gz": "ben2.tar.gz",
}
OVERLAYS = {
"depth-anything-3.tar.gz": "depth-anything-3",
}
def build_extension(): def bundled_source_bytes(filename):
manifest = tomllib.loads( """返回内置源码包内容;patches/<overlay>/ 目录中的文件会覆盖官方包对应路径。"""
(SOURCE_DIR / "blender_manifest.toml").read_text(encoding="utf-8") archive_path = SOURCES_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
output.addfile(info, overlay_file.open("rb"))
else:
with source.extractfile(member) as extracted:
output.addfile(member, extracted)
return buffer.getvalue()
def manifest_for_platform(platform):
"""Return the manifest text declaring only the given platform."""
text = (SOURCE_DIR / "blender_manifest.toml").read_text(encoding="utf-8")
lines = text.splitlines()
table_start = next(
(i for i, line in enumerate(lines) if line.startswith("[")),
len(lines),
) )
head = lines[:table_start]
tail = lines[table_start:]
head = [line for line in head if not line.strip().startswith("platforms =")]
# Insert into the top-level section (before any `[table]` header) so the
# field is part of the extension metadata, not a sub-key of a table.
# Place it after the `blender_version_min` line, alongside the other
# top-level keys (before `license`), keeping the file readable.
insert_at = next(
(i + 1 for i, line in enumerate(head) if line.strip().startswith("blender_version_min")),
len(head),
)
head.insert(insert_at, f'platforms = ["{platform}"]')
lines = head + tail
return "\n".join(lines) + "\n"
def bundled_source_files(platform):
entries = {
name: SOURCES_DIR / filename
for name, filename in BUNDLED_ARCHIVES.items()
}
entries[f"sources/{PLATFORMS[platform]}"] = (
SOURCES_DIR / PLATFORMS[platform]
)
return entries
def build_extension(platform):
if platform not in PLATFORMS:
raise ValueError(f"Unsupported platform: {platform}")
manifest = tomllib.loads(manifest_for_platform(platform))
output_path = DIST_DIR / ( output_path = DIST_DIR / (
f"{manifest['name']}.v{manifest['version']}.zip" f"{manifest['name']}.v{manifest['version']}.{platform}.zip"
) )
DIST_DIR.mkdir(exist_ok=True) DIST_DIR.mkdir(exist_ok=True)
sources = bundled_source_files(platform)
missing = [str(path) for path in sources.values() if not path.is_file()]
if missing:
raise RuntimeError(
"Missing bundled source files: " + ", ".join(missing)
)
with ZipFile(output_path, "w", ZIP_DEFLATED) as archive: with ZipFile(output_path, "w", ZIP_DEFLATED) as archive:
for source_path in sorted(SOURCE_DIR.rglob("*")): for source_path in sorted(SOURCE_DIR.rglob("*")):
if not source_path.is_file(): if not source_path.is_file():
continue continue
if "__pycache__" in source_path.parts: if "__pycache__" in source_path.parts:
continue continue
if source_path.name == "blender_manifest.toml":
continue
archive.write(source_path, source_path.relative_to(SOURCE_DIR)) archive.write(source_path, source_path.relative_to(SOURCE_DIR))
archive.writestr(
"blender_manifest.toml",
manifest_for_platform(platform),
)
for archive_name, source_path in sources.items():
archive.writestr(
archive_name,
bundled_source_bytes(source_path.name),
)
return output_path return output_path
def build_all():
return [build_extension(platform) for platform in PLATFORMS]
if __name__ == "__main__": if __name__ == "__main__":
extension_path = build_extension() extension_paths = build_all()
size_mib = extension_path.stat().st_size / (1024 * 1024) for extension_path in extension_paths:
print(f"Built {extension_path} ({size_mib:.2f} MiB)") size_mib = extension_path.stat().st_size / (1024 * 1024)
print(f"Built {extension_path} ({size_mib:.2f} MiB)")

View File

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

View File

@@ -278,6 +278,11 @@ class HostState:
model_cache=self.models, model_cache=self.models,
cancel_check=cancel_check, cancel_check=cancel_check,
) )
elif payload["command"] == "extract-frames":
worker.extract_frames(
args,
cancel_check=cancel_check,
)
else: else:
worker.reconstruct( worker.reconstruct(
args, args,
@@ -426,7 +431,11 @@ def create_app(token, instance_id="test"):
client_id = str(payload.get("client_id", "")) client_id = str(payload.get("client_id", ""))
command = payload.get("command") command = payload.get("command")
parameters = payload.get("parameters") parameters = payload.get("parameters")
if not job_id or command not in {"remove-background", "reconstruct"}: if not job_id or command not in {
"remove-background",
"reconstruct",
"extract-frames",
}:
raise HTTPException(status_code=422, detail="Invalid job request") raise HTTPException(status_code=422, detail="Invalid job request")
if not isinstance(parameters, dict) or not parameters.get("status"): if not isinstance(parameters, dict) or not parameters.get("status"):
raise HTTPException(status_code=422, detail="Missing job parameters") raise HTTPException(status_code=422, detail="Missing job parameters")

View File

@@ -1,7 +1,7 @@
schema_version = "1.0.0" schema_version = "1.0.0"
id = "snapform" id = "snapform"
version = "0.9.32" version = "0.10.10"
name = "SnapForm" name = "SnapForm"
tagline = "Build spatial structures from images, fast" tagline = "Build spatial structures from images, fast"
maintainer = "OmooLab" maintainer = "OmooLab"

View File

@@ -20,21 +20,13 @@ if str(SCRIPT_DIRECTORY) not in sys.path:
sys.path.insert(0, str(SCRIPT_DIRECTORY)) sys.path.insert(0, str(SCRIPT_DIRECTORY))
from models import BEN2_MODEL, MODELS from models import BEN2_MODEL, MODELS
from worker import HF_ENDPOINT, HF_MIRRORS, format_rate, reachable
UV_VERSION = "0.11.3"
PREFLIGHT_MIN_FREE_BYTES = 8 * 1024**3
PREFLIGHT_TIMEOUT = 8.0
MIRROR_TIMEOUT = 5.0
PYPI_MIRRORS = ( PYPI_MIRRORS = (
"https://pypi.tuna.tsinghua.edu.cn/simple", "https://pypi.tuna.tsinghua.edu.cn/simple",
"https://mirrors.aliyun.com/pypi/simple/", "https://mirrors.aliyun.com/pypi/simple/",
) )
TORCH_MIRRORS = { PYPI_DEFAULT_INDEX = "https://pypi.org/simple/"
"cu128": ("https://mirrors.aliyun.com/pytorch-wheels/cu128",),
"cpu": ("https://mirrors.aliyun.com/pytorch-wheels/cpu",),
}
TORCH_VERSION = "2.9.1" TORCH_VERSION = "2.9.1"
TORCHVISION_VERSION = "0.24.1" TORCHVISION_VERSION = "0.24.1"
XFORMERS_VERSION = "0.0.33.post2" XFORMERS_VERSION = "0.0.33.post2"
@@ -47,11 +39,7 @@ TORCH_INDEXES = {
"cpu": "https://download.pytorch.org/whl/cpu", "cpu": "https://download.pytorch.org/whl/cpu",
} }
_resolved_pypi_index = None _resolved_pypi_index = None
_resolved_hf_endpoint = None
_resolved_torch_index = None
DEFAULT_MODEL_KEYS = ("SMALL", "BASE", BEN2_MODEL.key) DEFAULT_MODEL_KEYS = ("SMALL", "BASE", BEN2_MODEL.key)
PREFLIGHT_END = 0.05
MIRROR_END = 0.07
RUNTIME_SLICE_START = 0.07 RUNTIME_SLICE_START = 0.07
RUNTIME_SLICE_END = 0.88 RUNTIME_SLICE_END = 0.88
VERIFY_SLICE_START = 0.88 VERIFY_SLICE_START = 0.88
@@ -101,8 +89,6 @@ def child_environment(environment=None):
process_environment[env_key] = value process_environment[env_key] = value
if _resolved_pypi_index: if _resolved_pypi_index:
process_environment["PIP_INDEX_URL"] = _resolved_pypi_index process_environment["PIP_INDEX_URL"] = _resolved_pypi_index
if _resolved_hf_endpoint:
process_environment["HF_ENDPOINT"] = _resolved_hf_endpoint
if environment: if environment:
process_environment.update(environment) process_environment.update(environment)
return process_environment return process_environment
@@ -131,127 +117,30 @@ def parse_package_count(line):
def effective_torch_index(backend): def effective_torch_index(backend):
if _resolved_torch_index:
return _resolved_torch_index
return os.environ.get("SNAPFORM_TORCH_INDEX") or TORCH_INDEXES.get(backend) return os.environ.get("SNAPFORM_TORCH_INDEX") or TORCH_INDEXES.get(backend)
def probe_hosts(urls, timeout): def resolve_indexes(args):
results = {} """???????????????????????????????????"""
global _resolved_pypi_index
def worker(url): _resolved_pypi_index = (
results[url] = reachable(url, timeout) getattr(args, "pypi_index", None)
or os.environ.get("PIP_INDEX_URL")
threads = [
threading.Thread(target=worker, args=(url,)) for url in urls
]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
return results
def resolve_endpoints(backend):
"""Resolve package indexes, falling back to mirrors when direct access fails; returns (resolved, failures, used_mirror)."""
preferred = {
"pypi": os.environ.get("PIP_INDEX_URL") or "https://pypi.org/simple/",
"hf": os.environ.get("HF_ENDPOINT") or HF_ENDPOINT,
"torch": os.environ.get("SNAPFORM_TORCH_INDEX")
or TORCH_INDEXES.get(backend),
"github": "https://github.com/",
}
if sys.platform == "darwin":
preferred["torch"] = None
mirror_map = {
"pypi": PYPI_MIRRORS,
"hf": HF_MIRRORS,
"torch": TORCH_MIRRORS.get(backend, ()),
"github": (),
}
results = probe_hosts(
[url for url in preferred.values() if url],
PREFLIGHT_TIMEOUT,
) )
resolved = {}
failures = []
for key, url in preferred.items():
if url is None:
resolved[key] = None
continue
if results[url]:
resolved[key] = url
continue
mirror_results = probe_hosts(mirror_map[key], MIRROR_TIMEOUT)
mirror = next(
(
candidate
for candidate in mirror_map[key]
if mirror_results.get(candidate)
),
None,
)
if mirror:
resolved[key] = mirror
else:
resolved[key] = url
failures.append(url)
used_mirror = any(
key in preferred
and preferred[key] is not None
and resolved[key] != preferred[key]
for key in resolved
)
return resolved, failures, used_mirror
def preflight(args): def _run_install_once(
"""Check disk space and network reachability; switch to mirrors when direct access fails and raise a clear English error when nothing is reachable."""
global _resolved_pypi_index, _resolved_hf_endpoint, _resolved_torch_index
problems = []
free_bytes = shutil.disk_usage(Path(args.runtime_dir).parent).free
if free_bytes < PREFLIGHT_MIN_FREE_BYTES:
needed_gb = PREFLIGHT_MIN_FREE_BYTES // (1024**3)
problems.append(
f"Not enough free disk space: {needed_gb} GB required, "
f"only {free_bytes / 1024**3:.1f} GB available."
)
resolved, failures, used_mirror = resolve_endpoints(args.backend)
_resolved_pypi_index = resolved["pypi"]
_resolved_hf_endpoint = resolved["hf"]
_resolved_torch_index = resolved["torch"]
unreachable = failures
if unreachable:
proxies = urllib.request.getproxies()
proxy_hint = (
"Proxy detected: " + ", ".join(
f"{key}={value}" for key, value in proxies.items()
)
if proxies
else "No proxy detected"
)
problems.append(
"Cannot reach: " + ", ".join(unreachable)
+ f". {proxy_hint}. Check network or proxy settings and retry."
)
if problems:
raise RuntimeError("Prerequisites check failed: " + " ".join(problems))
return resolved, used_mirror
def run_install(
status_path, status_path,
stage, stage,
stages, stages,
label, label,
command, command,
environment=None, environment,
start=0.0, start,
end=1.0, end,
clock=time.monotonic, clock,
): ):
"""Run a uv install command and smoothly map package-count events onto the stage progress range.""" """Run a uv install command once and map package-count events onto the stage progress range."""
print("Running:", subprocess.list2cmdline([str(value) for value in command])) print("Running:", subprocess.list2cmdline([str(value) for value in command]))
process_environment = child_environment(environment) process_environment = child_environment(environment)
process_environment["UV_LOG"] = "debug" process_environment["UV_LOG"] = "debug"
@@ -342,6 +231,42 @@ def run_install(
return installed return installed
def run_install(
status_path,
stage,
stages,
label,
command,
environment=None,
start=0.0,
end=1.0,
clock=time.monotonic,
):
"""Run a uv install command, retrying with mirrors when the default index fails."""
attempts = [environment]
if _resolved_pypi_index in (None, PYPI_DEFAULT_INDEX):
attempts.extend(
{"PIP_INDEX_URL": mirror} for mirror in PYPI_MIRRORS
)
for index, attempt_environment in enumerate(attempts):
try:
return _run_install_once(
status_path,
stage,
stages,
label,
command,
attempt_environment,
start,
end,
clock,
)
except subprocess.CalledProcessError:
if index + 1 == len(attempts):
raise
print(f"Retrying with PyPI mirror: {PYPI_MIRRORS[0]}")
def environment_python(runtime_dir): def environment_python(runtime_dir):
if sys.platform == "win32": if sys.platform == "win32":
return Path(runtime_dir) / "Scripts" / "python.exe" return Path(runtime_dir) / "Scripts" / "python.exe"
@@ -355,15 +280,22 @@ def install_stages(args, runtime_dir, python, uv):
if sys.platform == "darwin" if sys.platform == "darwin"
else effective_torch_index(args.backend) else effective_torch_index(args.backend)
) )
def github_archive_url(repository, commit): def archive_requirement(project_name, filename):
return f"https://github.com/{repository}/archive/{commit}.tar.gz" local = Path(args.runtime_dir).parent / "sources" / filename
if not local.is_file():
raise RuntimeError(
"Bundled source archive is missing; reinstall the extension: "
f"{local}"
)
return f"{project_name} @ {local.as_uri()}"
da3_requirement = ( da3_requirement = archive_requirement(
"depth-anything-3 @ " "depth-anything-3",
+ github_archive_url("ByteDance-Seed/Depth-Anything-3", args.da3_commit) "depth-anything-3.tar.gz",
) )
ben2_requirement = ( ben2_requirement = archive_requirement(
"ben2 @ " + github_archive_url("PramaLLC/BEN2", args.ben2_commit) "ben2",
"ben2.tar.gz",
) )
torch_command = [ torch_command = [
uv, uv,
@@ -488,12 +420,41 @@ def uv_download_target():
raise RuntimeError(f"Unsupported platform for uv download: {sys.platform}") raise RuntimeError(f"Unsupported platform for uv download: {sys.platform}")
def uv_download_url(): def bundled_sources_directory():
return Path(__file__).with_name("sources")
def bundled_uv_archive():
target, archive = uv_download_target() target, archive = uv_download_target()
return ( path = bundled_sources_directory() / f"uv-{target}.{archive}"
"https://github.com/astral-sh/uv/releases/download/" return path if path.is_file() else None
f"{UV_VERSION}/uv-{target}.{archive}"
)
def sync_bundled_sources(bundled_dir, sources_dir):
"""Copy bundled source archives into the storage sources directory."""
for name in ("depth-anything-3.tar.gz", "ben2.tar.gz"):
bundled = Path(bundled_dir) / name
if not bundled.is_file():
continue
target = Path(sources_dir) / name
if (
target.is_file()
and target.stat().st_size == bundled.stat().st_size
):
continue
Path(sources_dir).mkdir(parents=True, exist_ok=True)
shutil.copy2(bundled, target)
def archive_is_valid(path):
if zipfile.is_zipfile(path):
return True
try:
with tarfile.open(path) as archive:
archive.getmembers()
return True
except (OSError, tarfile.TarError):
return False
def extract_uv(archive_path, uv_path): def extract_uv(archive_path, uv_path):
@@ -511,8 +472,8 @@ def extract_uv(archive_path, uv_path):
os.chmod(uv_path, 0o755) os.chmod(uv_path, 0o755)
def download_uv( def install_uv(
uv_path, uv,
status_path, status_path,
stage=None, stage=None,
stages=None, stages=None,
@@ -520,50 +481,29 @@ def download_uv(
start=0.0, start=0.0,
end=1.0, end=1.0,
): ):
"""Download the standalone uv binary and map byte progress onto the stage progress range.""" """Extract the bundled uv binary; fail clearly when it is missing."""
stage_fields = { stage_fields = {
"stage": stage, "stage": stage,
"stages": stages, "stages": stages,
"stage_label": stage_label, "stage_label": stage_label,
} }
bundled_uv = bundled_uv_archive()
def write(progress, message): if bundled_uv is None:
write_status( raise RuntimeError(
status_path, "Bundled uv is missing; reinstall the extension matching your "
start + progress * (end - start), f"platform (searched {bundled_sources_directory()})"
message,
**stage_fields,
) )
if not archive_is_valid(bundled_uv):
write(0.0, "downloading") raise RuntimeError(f"Bundled uv archive is invalid: {bundled_uv}")
uv_path = Path(uv_path) print(f"Extracting bundled uv from {bundled_uv}")
uv_path.parent.mkdir(parents=True, exist_ok=True) write_status(
archive_path = uv_path.with_name(f"{uv_path.name}.part") status_path,
request = urllib.request.Request( start,
uv_download_url(), "extracting bundled uv",
headers={"User-Agent": "SnapForm-Runtime-Installer"}, **stage_fields,
) )
try: Path(uv).parent.mkdir(parents=True, exist_ok=True)
with urllib.request.urlopen(request, timeout=60) as response: extract_uv(bundled_uv, uv)
total = int(response.headers.get("Content-Length") or 0)
downloaded = 0
started = time.monotonic()
with archive_path.open("wb") as output:
while True:
chunk = response.read(65536)
if not chunk:
break
output.write(chunk)
downloaded += len(chunk)
progress = (
min(downloaded / total, 1.0) if total else 0.0
)
speed = downloaded / max(time.monotonic() - started, 1e-6)
write(progress, f"downloading {format_rate(speed)}")
write(1.0, "preparing")
extract_uv(archive_path, uv_path)
finally:
archive_path.unlink(missing_ok=True)
def run_stage( def run_stage(
@@ -636,10 +576,7 @@ def install(args):
"SnapForm requires Apple Silicon macOS; PyTorch stopped " "SnapForm requires Apple Silicon macOS; PyTorch stopped "
"shipping Intel Mac (x86_64) wheels after 2.2.2" "shipping Intel Mac (x86_64) wheels after 2.2.2"
) )
write_status(args.status, 0.0, "checking prerequisites") resolve_indexes(args)
_resolved, used_mirror = preflight(args)
if used_mirror:
write_status(args.status, MIRROR_END, "using mirror sources")
runtime_dir = Path(args.runtime_dir) runtime_dir = Path(args.runtime_dir)
runtime_dir.parent.mkdir(parents=True, exist_ok=True) runtime_dir.parent.mkdir(parents=True, exist_ok=True)
if not runtime_ready(runtime_dir): if not runtime_ready(runtime_dir):
@@ -650,6 +587,10 @@ 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 = environment_python(runtime_dir) python = environment_python(runtime_dir)
sync_bundled_sources(
bundled_sources_directory(),
Path(args.runtime_dir).parent / "sources",
)
stages = install_stages(args, runtime_dir, python, uv) stages = install_stages(args, runtime_dir, python, uv)
runtime_stages = [stage for stage in stages if stage[3] != "download"] runtime_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"]
@@ -677,7 +618,7 @@ def install(args):
install_span if kind == "install" else PLAIN_SLICE_SPAN install_span if kind == "install" else PLAIN_SLICE_SPAN
) )
if uv_missing and index == 1: if uv_missing and index == 1:
download_uv( install_uv(
uv, uv,
args.status, args.status,
stage=index, stage=index,
@@ -766,6 +707,7 @@ def parse_args():
parser.add_argument("--da3-commit", required=True) parser.add_argument("--da3-commit", required=True)
parser.add_argument("--ben2-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)
return parser.parse_args() return parser.parse_args()

View File

@@ -1,3 +1,5 @@
import re
import shutil
from pathlib import Path from pathlib import Path
import bpy import bpy
@@ -22,6 +24,17 @@ IMAGE_NAME_SUFFIXES = (
".bmp", ".bmp",
".tga", ".tga",
) )
VIDEO_NAME_SUFFIXES = (
".mp4",
".mov",
".avi",
".mkv",
".webm",
".m4v",
".mts",
".mpg",
".mpeg",
)
def image_depth_model_items(owner, context): def image_depth_model_items(owner, context):
@@ -47,6 +60,56 @@ def is_image_empty(candidate):
) )
def is_animated_image(image):
return getattr(image, "source", None) in {"MOVIE", "SEQUENCE"}
def video_total_frames(image):
"""Return the total animated frame count, or 1 for a still image."""
source = getattr(image, "source", None)
if source == "MOVIE":
return int(getattr(image, "frame_duration", 1) or 1)
if source == "SEQUENCE":
return len(collect_sequence_frames(image)) or 1
return 1
def image_user_settings(source_object):
"""Return the Empty's image user settings as a four-item tuple."""
image_user = getattr(source_object, "image_user", None)
if image_user is None:
return 1, 0, 0, False
return (
int(getattr(image_user, "frame_start", 1) or 1),
int(getattr(image_user, "frame_offset", 0) or 0),
int(getattr(image_user, "frame_duration", 0) or 0),
bool(getattr(image_user, "use_cyclic", False)),
)
def extract_frame_range(source_object):
"""Return the clamped (start_frame, frame_count) range extracted from the source."""
total_frames = video_total_frames(getattr(source_object, "data", None))
if total_frames <= 1:
return 1, total_frames
_frame_start, frame_offset, image_frames, _use_cyclic = (
image_user_settings(source_object)
)
frame_count = (
min(image_frames, total_frames)
if image_frames > 1
else total_frames
)
start_frame = min(1 + frame_offset, total_frames - frame_count + 1)
start_frame = max(start_frame, 1)
return start_frame, frame_count
def default_extract_frames(source_object):
"""Return the Empty's Image Frames value, or the full frame count when unset."""
return extract_frame_range(source_object)[1]
def image_path(image): def image_path(image):
filepath_from_user = getattr(image, "filepath_from_user", None) filepath_from_user = getattr(image, "filepath_from_user", None)
filepath = ( filepath = (
@@ -65,10 +128,22 @@ def image_is_packed(image):
return bool(getattr(image, "packed_files", ())) return bool(getattr(image, "packed_files", ()))
def remove_temporary_input(path):
"""Remove a temporary input file or frame directory."""
path = Path(path)
try:
if path.is_dir():
shutil.rmtree(path, ignore_errors=True)
else:
path.unlink(missing_ok=True)
except OSError:
pass
def image_base_name(image): def image_base_name(image):
name = image.name name = image.name
lower_name = name.lower() lower_name = name.lower()
for suffix in IMAGE_NAME_SUFFIXES: for suffix in IMAGE_NAME_SUFFIXES + VIDEO_NAME_SUFFIXES:
suffix_index = lower_name.rfind(suffix) suffix_index = lower_name.rfind(suffix)
if suffix_index <= 0: if suffix_index <= 0:
continue continue
@@ -85,8 +160,15 @@ def depth_data_name(source_object):
return f"{image_base_name(source_object.data)}.depth" return f"{image_base_name(source_object.data)}.depth"
def copy_depth_image(source_object):
"""Return a renamed copy of the source image for use as the depth map."""
depth_image = source_object.data.copy()
depth_image.name = depth_data_name(source_object)
return depth_image
def removed_background_data_name(source_object): def removed_background_data_name(source_object):
return f"{image_base_name(source_object.data)}.rb" return f"{image_base_name(source_object.data)}.foreground"
def scaled_process_resolution(image, scale): def scaled_process_resolution(image, scale):
@@ -117,10 +199,60 @@ def save_packed_image(image, scene):
return output_path return output_path
def collect_sequence_frames(image):
"""Return the numbered frame files of a SEQUENCE image, or [] when absent."""
source_path = image_path(image)
if not is_animated_image(image) or not source_path.is_file():
return []
match = re.match(
r"^(?P<prefix>.*?)(?P<digits>\d+)(?P<suffix>\.[^.]+)$",
source_path.name,
)
if match is None:
return []
prefix = match.group("prefix")
suffix = match.group("suffix").lower()
frames = []
for path in source_path.parent.iterdir():
frame_match = re.match(
r"^(?P<prefix>.*?)(?P<digits>\d+)(?P<suffix>\.[^.]+)$",
path.name,
)
if (
frame_match is not None
and frame_match.group("prefix") == prefix
and frame_match.group("suffix").lower() == suffix
):
frames.append((int(frame_match.group("digits")), path))
return [path for _digits, path in sorted(frames)]
def prepare_image_input(image, scene): def prepare_image_input(image, scene):
source_path = image_path(image) source_path = image_path(image)
if image_is_packed(image): if image_is_packed(image):
return save_packed_image(image, scene), True return save_packed_image(image, scene), True
if getattr(image, "source", None) == "MOVIE":
if not source_path.is_file():
raise RuntimeError(
"The image empty must reference an existing video file"
)
return source_path, False
sequence_frames = collect_sequence_frames(image)
if sequence_frames:
from .runtime import temporary_input_directory
temporary_directory = Path(temporary_input_directory())
try:
for index, frame_path in enumerate(sequence_frames):
suffix = frame_path.suffix.lower() or ".jpg"
shutil.copy2(
frame_path,
temporary_directory / f"frame.{index + 1:04d}{suffix}",
)
except Exception:
shutil.rmtree(temporary_directory, ignore_errors=True)
raise
return temporary_directory, True
if not source_path.is_file(): if not source_path.is_file():
raise RuntimeError( raise RuntimeError(
"The image empty must reference an existing image file" "The image empty must reference an existing image file"
@@ -128,12 +260,207 @@ def prepare_image_input(image, scene):
return source_path, False return source_path, False
def configure_image_user(
image_user,
frame_count=None,
frame_start=1,
frame_offset=0,
use_cyclic=None,
):
"""Map the image to scene frames and enable auto refresh."""
if image_user is None:
return
if hasattr(image_user, "frame_start"):
image_user.frame_start = frame_start
if hasattr(image_user, "frame_offset"):
image_user.frame_offset = frame_offset
if hasattr(image_user, "use_auto_refresh"):
image_user.use_auto_refresh = True
if use_cyclic is not None and hasattr(image_user, "use_cyclic"):
image_user.use_cyclic = use_cyclic
if frame_count is not None and hasattr(image_user, "frame_duration"):
image_user.frame_duration = frame_count
def pack_sequence_if_possible(image, expected_frames):
"""Pack a frame sequence when Blender supports it; otherwise keep external files."""
if expected_frames <= 0:
return False
try:
image.pack()
except (OSError, RuntimeError):
return False
if len(getattr(image, "packed_files", ())) >= expected_frames:
return True
try:
image.unpack(method="REMOVE")
except (OSError, RuntimeError):
pass
return False
def sorted_frame_paths(directory, prefix=None):
"""Return image files in a directory sorted by their trailing frame number."""
def frame_key(path):
match = re.search(r"(\d+)(?=\.[^.]+$)", path.name)
return (int(match.group(1)) if match is not None else 0, path.name)
return sorted(
(
path
for path in Path(directory).iterdir()
if path.is_file()
and path.suffix.lower() in IMAGE_NAME_SUFFIXES
and (prefix is None or path.stem.startswith(prefix))
),
key=frame_key,
)
def force_sequence_source(image):
"""Ensure a loaded image is a SEQUENCE datablock for its numbered frames."""
if getattr(image, "source", None) == "SEQUENCE":
return image
try:
image.source = "SEQUENCE"
except (OSError, RuntimeError, TypeError) as error:
raise RuntimeError(
"Blender cannot load the frame sequence"
) from error
if getattr(image, "source", None) != "SEQUENCE":
raise RuntimeError("Blender cannot load the frame sequence")
return image
def create_sequence_image(name, frame_directory, frame_prefix="frame."):
"""Load a numbered frame directory as a SEQUENCE image datablock."""
frame_paths = sorted_frame_paths(frame_directory, frame_prefix)
if not frame_paths:
raise RuntimeError(f"No image frames found in: {frame_directory}")
try:
image = bpy.data.images.load(
str(frame_paths[0]),
check_existing=False,
)
force_sequence_source(image)
except (OSError, RuntimeError) as error:
if "image" in locals():
bpy.data.images.remove(image, do_unlink=True)
raise RuntimeError("Unable to load the frame sequence") from error
image.name = name
pack_sequence_if_possible(image, len(frame_paths))
return image
def create_depth_sequence_image(name, depth_frames, scene, frame_directory):
"""Write one EXR per depth frame and return a SEQUENCE image datablock."""
import numpy as np
frame_directory = Path(frame_directory)
frame_directory.mkdir(parents=True, exist_ok=True)
height, width = depth_frames[0].shape
image = bpy.data.images.new(
name,
width=width,
height=height,
alpha=False,
float_buffer=True,
)
image.colorspace_settings.name = "Non-Color"
image_settings = scene.render.image_settings
original_file_format = image_settings.file_format
frame_paths = []
try:
image_settings.file_format = "OPEN_EXR"
for frame_index, depth_values in enumerate(depth_frames):
# NPZ arrays use a top-left origin, while Blender image pixels start below.
depth_pixels = np.flipud(depth_values)
rgba = np.empty((height, width, 4), dtype=np.float32)
rgba[..., 0] = depth_pixels
rgba[..., 1] = depth_pixels
rgba[..., 2] = depth_pixels
rgba[..., 3] = 1.0
image.pixels.foreach_set(rgba.ravel())
frame_path = (
frame_directory / f"depth.{frame_index + 1:04d}.exr"
)
image.save_render(str(frame_path), scene=scene)
if not frame_path.is_file():
raise RuntimeError(f"Blender did not write: {frame_path}")
frame_paths.append(frame_path)
finally:
image_settings.file_format = original_file_format
bpy.data.images.remove(image, do_unlink=True)
if not frame_paths:
raise RuntimeError("Unable to write the depth frame sequence")
try:
sequence_image = bpy.data.images.load(
str(frame_paths[0]),
check_existing=False,
)
force_sequence_source(sequence_image)
except (OSError, RuntimeError) as error:
if "sequence_image" in locals():
bpy.data.images.remove(sequence_image, do_unlink=True)
raise RuntimeError(
"Unable to load the depth frame sequence"
) from error
sequence_image.name = name
sequence_image.colorspace_settings.name = "Non-Color"
pack_sequence_if_possible(sequence_image, len(frame_paths))
return sequence_image
def create_removed_background_image( def create_removed_background_image(
source_object, source_object,
output_path, output_path,
assign_to_source=True, assign_to_source=True,
): ):
output_path = Path(output_path) output_path = Path(output_path)
if output_path.is_dir():
foreground_paths = sorted_frame_paths(output_path, "foreground.")
if not foreground_paths:
raise RuntimeError(
f"Background removal result has no foreground frames: "
f"{output_path}"
)
try:
result_image = bpy.data.images.load(
str(foreground_paths[0]),
check_existing=False,
)
force_sequence_source(result_image)
except (OSError, RuntimeError) as error:
if "result_image" in locals():
bpy.data.images.remove(result_image, do_unlink=True)
raise RuntimeError("Unable to load the background result") from error
source_image = source_object.data
original_fake_user = source_image.use_fake_user
try:
result_image.name = removed_background_data_name(source_object)
pack_sequence_if_possible(result_image, len(foreground_paths))
if assign_to_source:
image_user = getattr(source_object, "image_user", None)
(
original_frame_start,
_original_frame_offset,
_original_frame_duration,
original_use_cyclic,
) = image_user_settings(source_object)
source_image.use_fake_user = True
source_object.data = result_image
configure_image_user(
image_user,
len(foreground_paths),
frame_start=original_frame_start,
use_cyclic=original_use_cyclic,
)
except (OSError, RuntimeError) as error:
source_image.use_fake_user = original_fake_user
bpy.data.images.remove(result_image)
raise RuntimeError("Unable to load the background result") from error
return result_image
if not output_path.is_file(): if not output_path.is_file():
raise RuntimeError( raise RuntimeError(
f"Background removal result does not exist: {output_path}" f"Background removal result does not exist: {output_path}"
@@ -229,7 +556,7 @@ def apply_plane_transform(plane_object, source_object, camera_bounds):
return abs(float(plane_object.scale.z)) return abs(float(plane_object.scale.z))
def extract_depth_array(npz_path): def extract_depth_arrays(npz_path):
try: try:
import numpy as np import numpy as np
except ImportError as error: except ImportError as error:
@@ -256,13 +583,17 @@ def extract_depth_array(npz_path):
if depth.ndim >= 3 and depth.shape[-1] == 1: if depth.ndim >= 3 and depth.shape[-1] == 1:
depth = depth[..., 0] depth = depth[..., 0]
if depth.ndim == 3: if depth.ndim == 2:
depth = depth[0] depth = depth[None, ...]
if depth.ndim != 2: if depth.ndim != 3:
raise ValueError( raise ValueError(
f"Unsupported depth array shape: {tuple(depth.shape)}" f"Unsupported depth array shape: {tuple(depth.shape)}"
) )
return depth return [depth[index] for index in range(depth.shape[0])]
def extract_depth_array(npz_path):
return extract_depth_arrays(npz_path)[0]
def extract_intrinsics_array(npz_path): def extract_intrinsics_array(npz_path):
@@ -475,6 +806,13 @@ def create_material(
displacement_midlevel=0.0, displacement_midlevel=0.0,
displacement_scale=1.0, displacement_scale=1.0,
unshaded=False, unshaded=False,
frame_count=1,
frame_start=1,
color_frame_offset=0,
color_frame_duration=None,
depth_frame_offset=0,
depth_frame_duration=None,
use_cyclic=None,
): ):
material = bpy.data.materials.new(name) material = bpy.data.materials.new(name)
material.use_nodes = True material.use_nodes = True
@@ -491,6 +829,25 @@ def create_material(
color_texture.extension = "EXTEND" color_texture.extension = "EXTEND"
color_texture.location = (-640.0, 180.0) color_texture.location = (-640.0, 180.0)
output.location = (460.0, 180.0) output.location = (460.0, 180.0)
extracted_frames = frame_count if frame_count > 1 else None
color_frames = (
color_frame_duration
if color_frame_duration is not None
else extracted_frames
)
depth_frames = (
depth_frame_duration
if depth_frame_duration is not None
else extracted_frames
)
if frame_count > 1 or is_animated_image(color_image):
configure_image_user(
color_texture.image_user,
color_frames,
frame_start=frame_start,
frame_offset=color_frame_offset,
use_cyclic=use_cyclic,
)
links = material.node_tree.links links = material.node_tree.links
if unshaded: if unshaded:
@@ -558,6 +915,14 @@ def create_material(
depth_texture.image = displacement_image depth_texture.image = displacement_image
depth_texture.interpolation = "Linear" depth_texture.interpolation = "Linear"
depth_texture.extension = "EXTEND" depth_texture.extension = "EXTEND"
if frame_count > 1 or is_animated_image(displacement_image):
configure_image_user(
depth_texture.image_user,
depth_frames,
frame_start=frame_start,
frame_offset=depth_frame_offset,
use_cyclic=use_cyclic,
)
displacement.inputs["Midlevel"].default_value = ( displacement.inputs["Midlevel"].default_value = (
displacement_midlevel displacement_midlevel
) )
@@ -612,8 +977,35 @@ def build_image_depth_plane(
displacement_via="MODIFIER", displacement_via="MODIFIER",
unshaded=False, unshaded=False,
color_image=None, color_image=None,
depth_frames_directory=None,
): ):
source_name = source_object.name source_name = source_object.name
frame_count = 1
frames_image = None
if depth_frames_directory is not None:
frames_image = create_sequence_image(
depth_data_name(source_object),
depth_frames_directory,
)
frame_count = len(
sorted_frame_paths(depth_frames_directory, "frame.")
)
(
source_frame_start,
source_frame_offset,
source_frame_duration,
source_use_cyclic,
) = image_user_settings(source_object)
uses_extracted = frames_image is not None
texture_frame_count = (
frame_count if uses_extracted else source_frame_duration
)
texture_frame_offset = 0 if uses_extracted else source_frame_offset
depth_source_image = source_object.data
if frames_image is not None:
depth_source_image = frames_image
elif not is_animated_image(source_object.data):
depth_source_image = copy_depth_image(source_object)
mesh = create_plane_mesh( mesh = create_plane_mesh(
source_object, source_object,
subdivision_level=subdivision_level, subdivision_level=subdivision_level,
@@ -632,25 +1024,43 @@ def build_image_depth_plane(
depth_data_name(source_object), depth_data_name(source_object),
type="IMAGE", type="IMAGE",
) )
texture.image = source_object.data texture.image = depth_source_image
texture.extension = "EXTEND"
if uses_extracted or is_animated_image(source_object.data):
configure_image_user(
texture.image_user,
texture_frame_count,
frame_start=source_frame_start,
frame_offset=texture_frame_offset,
use_cyclic=source_use_cyclic,
)
displace = plane_object.modifiers.new("Displace", "DISPLACE") displace = plane_object.modifiers.new("Displace", "DISPLACE")
displace.texture = texture displace.texture = texture
displace.texture_coords = "UV" displace.texture_coords = "UV"
displace.mid_level = ORIGINAL_DISPLACEMENT_MIDLEVEL displace.mid_level = ORIGINAL_DISPLACEMENT_MIDLEVEL
depth_scale = abs(float(plane_object.scale.z)) depth_scale = abs(float(plane_object.scale.z))
displacement_image = (
depth_source_image if displacement_via == "MATERIAL" else None
)
material = create_material( material = create_material(
source_name, source_name,
color_image if color_image is not None else source_object.data, color_image if color_image is not None else source_object.data,
depth_scale, depth_scale,
displacement_image=( displacement_image=displacement_image,
source_object.data if displacement_via == "MATERIAL" else None
),
displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL, displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL,
displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH, displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH,
unshaded=unshaded, unshaded=unshaded,
frame_count=frame_count,
frame_start=source_frame_start,
color_frame_offset=texture_frame_offset,
color_frame_duration=texture_frame_count,
depth_frame_offset=texture_frame_offset,
depth_frame_duration=texture_frame_count,
use_cyclic=source_use_cyclic,
) )
mesh.materials.append(material) mesh.materials.append(material)
plane_object["snapform_frame_count"] = frame_count
return finish_depth_plane( return finish_depth_plane(
context, context,
source_object, source_object,
@@ -669,9 +1079,11 @@ def build_depth_plane(
unshaded=False, unshaded=False,
depth_name="", depth_name="",
color_image=None, color_image=None,
depth_frames_directory=None,
): ):
source_name = source_object.name source_name = source_object.name
depth = extract_depth_array(npz_path) depth_arrays = extract_depth_arrays(npz_path)
depth = depth_arrays[0]
intrinsics = extract_intrinsics_array(npz_path) intrinsics = extract_intrinsics_array(npz_path)
depth_values = prepare_depth_values(depth) depth_values = prepare_depth_values(depth)
bounds, reference_depth = camera_plane_bounds( bounds, reference_depth = camera_plane_bounds(
@@ -681,10 +1093,25 @@ def build_depth_plane(
DISPLACEMENT_STRENGTH, DISPLACEMENT_STRENGTH,
) )
depth_name = depth_name or depth_data_name(source_object) depth_name = depth_name or depth_data_name(source_object)
frame_count = len(depth_arrays)
(
source_frame_start,
source_frame_offset,
source_frame_duration,
source_use_cyclic,
) = image_user_settings(source_object)
depth_image = create_depth_image( depth_image = create_depth_image(
depth_name, depth_name,
depth_values, depth_values,
) )
if frame_count > 1 and depth_frames_directory is not None:
bpy.data.images.remove(depth_image, do_unlink=True)
depth_image = create_depth_sequence_image(
depth_name,
[prepare_depth_values(frame) for frame in depth_arrays],
context.scene,
depth_frames_directory,
)
mesh = create_plane_mesh(source_object, bounds, subdivision_level) mesh = create_plane_mesh(source_object, bounds, subdivision_level)
plane_object = bpy.data.objects.new( plane_object = bpy.data.objects.new(
source_name, source_name,
@@ -709,6 +1136,13 @@ def build_depth_plane(
) )
texture.image = depth_image texture.image = depth_image
texture.extension = "EXTEND" texture.extension = "EXTEND"
if frame_count > 1:
configure_image_user(
texture.image_user,
frame_count,
frame_start=source_frame_start,
use_cyclic=source_use_cyclic,
)
displace = plane_object.modifiers.new("Displace", "DISPLACE") displace = plane_object.modifiers.new("Displace", "DISPLACE")
displace.texture = texture displace.texture = texture
displace.texture_coords = "UV" displace.texture_coords = "UV"
@@ -726,12 +1160,22 @@ def build_depth_plane(
displacement_midlevel=DISPLACEMENT_MIDLEVEL, displacement_midlevel=DISPLACEMENT_MIDLEVEL,
displacement_scale=DISPLACEMENT_STRENGTH, displacement_scale=DISPLACEMENT_STRENGTH,
unshaded=unshaded, unshaded=unshaded,
frame_count=frame_count,
frame_start=source_frame_start,
color_frame_offset=0 if color_image is not None else source_frame_offset,
color_frame_duration=(
frame_count if color_image is not None else source_frame_duration
),
depth_frame_offset=0,
depth_frame_duration=frame_count,
use_cyclic=source_use_cyclic,
) )
mesh.materials.append(material) mesh.materials.append(material)
plane_object["snapform_reference_depth"] = reference_depth plane_object["snapform_reference_depth"] = reference_depth
plane_object["snapform_camera_plane_width"] = bounds[1] - bounds[0] plane_object["snapform_camera_plane_width"] = bounds[1] - bounds[0]
plane_object["snapform_camera_plane_height"] = bounds[3] - bounds[2] plane_object["snapform_camera_plane_height"] = bounds[3] - bounds[2]
plane_object["snapform_depth_scale"] = depth_scale plane_object["snapform_depth_scale"] = depth_scale
plane_object["snapform_frame_count"] = frame_count
return finish_depth_plane( return finish_depth_plane(
context, context,
@@ -758,6 +1202,12 @@ class BuildDepthPlane(bpy.types.Operator):
snapform_background_path: bpy.props.StringProperty( snapform_background_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"}, options={"HIDDEN", "SKIP_SAVE"},
) )
snapform_color_frames_dir: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_depth_frames_dir: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_subdivision_level: bpy.props.IntProperty( snapform_subdivision_level: bpy.props.IntProperty(
default=5, default=5,
min=0, min=0,
@@ -806,6 +1256,11 @@ class BuildDepthPlane(bpy.types.Operator):
self.snapform_background_path, self.snapform_background_path,
assign_to_source=False, assign_to_source=False,
) )
elif self.snapform_color_frames_dir:
color_image = create_sequence_image(
f"{image_base_name(source_object.data)}.frame",
self.snapform_color_frames_dir,
)
build_depth_plane( build_depth_plane(
context, context,
source_object, source_object,
@@ -816,6 +1271,9 @@ class BuildDepthPlane(bpy.types.Operator):
unshaded=self.snapform_unshaded, unshaded=self.snapform_unshaded,
depth_name=self.snapform_depth_data_name, depth_name=self.snapform_depth_data_name,
color_image=color_image, color_image=color_image,
depth_frames_directory=(
self.snapform_depth_frames_dir or None
),
) )
except (OSError, RuntimeError, ValueError) as error: except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error)) self.report({"ERROR"}, str(error))
@@ -834,6 +1292,12 @@ class BuildOriginalDepthPlane(bpy.types.Operator):
snapform_background_path: bpy.props.StringProperty( snapform_background_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"}, options={"HIDDEN", "SKIP_SAVE"},
) )
snapform_color_frames_dir: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_depth_frames_dir: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_subdivision_level: bpy.props.IntProperty( snapform_subdivision_level: bpy.props.IntProperty(
default=5, default=5,
min=0, min=0,
@@ -866,11 +1330,18 @@ class BuildOriginalDepthPlane(bpy.types.Operator):
) )
return {"CANCELLED"} return {"CANCELLED"}
try: try:
color_image = create_removed_background_image( color_image = None
source_object, if self.snapform_background_path:
self.snapform_background_path, color_image = create_removed_background_image(
assign_to_source=False, source_object,
) self.snapform_background_path,
assign_to_source=False,
)
elif self.snapform_color_frames_dir:
color_image = create_sequence_image(
f"{image_base_name(source_object.data)}.frame",
self.snapform_color_frames_dir,
)
build_image_depth_plane( build_image_depth_plane(
context, context,
source_object, source_object,
@@ -879,6 +1350,9 @@ class BuildOriginalDepthPlane(bpy.types.Operator):
displacement_via=self.snapform_displacement_via, displacement_via=self.snapform_displacement_via,
unshaded=self.snapform_unshaded, unshaded=self.snapform_unshaded,
color_image=color_image, color_image=color_image,
depth_frames_directory=(
self.snapform_depth_frames_dir or None
),
) )
except (OSError, RuntimeError, ValueError) as error: except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error)) self.report({"ERROR"}, str(error))
@@ -938,6 +1412,11 @@ class RemoveImageBackground(bpy.types.Operator):
from .operators import IMAGE_BACKGROUND_JOB from .operators import IMAGE_BACKGROUND_JOB
source_object = context.object source_object = context.object
total_frames = video_total_frames(source_object.data)
if total_frames > 1:
context.scene.snapform_settings.video_frames = (
default_extract_frames(source_object)
)
try: try:
source_path, delete_input = prepare_image_input( source_path, delete_input = prepare_image_input(
source_object.data, source_object.data,
@@ -956,7 +1435,7 @@ class RemoveImageBackground(bpy.types.Operator):
if "CANCELLED" not in result: if "CANCELLED" not in result:
return {"FINISHED"} return {"FINISHED"}
if delete_input: if delete_input:
source_path.unlink(missing_ok=True) remove_temporary_input(source_path)
return {"CANCELLED"} return {"CANCELLED"}
@@ -1089,13 +1568,45 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
layout.prop(self, "snapform_delete_original") layout.prop(self, "snapform_delete_original")
def execute(self, context): def execute(self, context):
from .operators import IMAGE_DEPTH_JOB, ORIGINAL_BACKGROUND_JOB from .operators import (
IMAGE_DEPTH_JOB,
ORIGINAL_BACKGROUND_JOB,
ORIGINAL_FRAMES_JOB,
)
source_object = context.object source_object = context.object
if video_total_frames(source_object.data) > 1:
context.scene.snapform_settings.video_frames = (
default_extract_frames(source_object)
)
if ( if (
self.snapform_depth_source == "ORIGINAL" self.snapform_depth_source == "ORIGINAL"
and not self.snapform_remove_background and not self.snapform_remove_background
): ):
if video_total_frames(source_object.data) > 1:
try:
source_path, delete_input = prepare_image_input(
source_object.data,
context.scene,
)
except RuntimeError as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
result = bpy.ops.snapform.run_job(
job_type=ORIGINAL_FRAMES_JOB,
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_subdivision_level=self.snapform_subdivision_level,
snapform_delete_source=self.snapform_delete_original,
snapform_delete_input=delete_input,
snapform_displacement_via=self.snapform_displacement_via,
snapform_unshaded=self.snapform_unshaded,
)
if "CANCELLED" not in result:
return {"FINISHED"}
if delete_input:
remove_temporary_input(source_path)
return {"CANCELLED"}
try: try:
build_image_depth_plane( build_image_depth_plane(
context, context,
@@ -1141,7 +1652,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
if "CANCELLED" not in result: if "CANCELLED" not in result:
return {"FINISHED"} return {"FINISHED"}
if delete_input: if delete_input:
source_path.unlink(missing_ok=True) remove_temporary_input(source_path)
return {"CANCELLED"} return {"CANCELLED"}
image = source_object.data image = source_object.data
@@ -1179,7 +1690,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
) )
if "CANCELLED" in result: if "CANCELLED" in result:
if delete_input: if delete_input:
source_path.unlink(missing_ok=True) remove_temporary_input(source_path)
return {"CANCELLED"} return {"CANCELLED"}
global _last_depth_plane_model global _last_depth_plane_model
_last_depth_plane_model = self.snapform_model _last_depth_plane_model = self.snapform_model

View File

@@ -1,3 +1,4 @@
import shutil
import subprocess import subprocess
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
@@ -18,8 +19,35 @@ IMAGE_DEPTH_JOB = "IMAGE_DEPTH"
BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND" BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND"
IMAGE_BACKGROUND_JOB = "IMAGE_BACKGROUND" IMAGE_BACKGROUND_JOB = "IMAGE_BACKGROUND"
ORIGINAL_BACKGROUND_JOB = "ORIGINAL_BACKGROUND" ORIGINAL_BACKGROUND_JOB = "ORIGINAL_BACKGROUND"
ORIGINAL_FRAMES_JOB = "ORIGINAL_FRAMES"
IMAGE_DEPTH_DEFAULT_PROCESS_RES = 1024 IMAGE_DEPTH_DEFAULT_PROCESS_RES = 1024
IMAGE_DEPTH_PROCESS_RES_METHOD = "upper_bound_resize" IMAGE_DEPTH_PROCESS_RES_METHOD = "upper_bound_resize"
BEN2_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
BEN2_VIDEO_SUFFIXES = {
".mp4",
".mov",
".avi",
".mkv",
".webm",
".m4v",
".mts",
".mpg",
".mpeg",
}
def is_ben2_input(input_path):
return input_path.is_dir() or (
input_path.suffix.lower()
in BEN2_IMAGE_SUFFIXES | BEN2_VIDEO_SUFFIXES
)
def empty_image_start_frame(source_object):
"""Return the first frame to extract, honoring the Empty's Offset within the clip length."""
from .image_depth import extract_frame_range
return extract_frame_range(source_object)[0]
def open_job_directory(result_directory): def open_job_directory(result_directory):
@@ -81,8 +109,12 @@ def close_active_job():
if _active_job.log_file is not None: if _active_job.log_file is not None:
_active_job.log_file.close() _active_job.log_file.close()
if _active_job.temporary_input_path is not None: if _active_job.temporary_input_path is not None:
temporary_path = _active_job.temporary_input_path
try: try:
_active_job.temporary_input_path.unlink(missing_ok=True) if temporary_path.is_dir():
shutil.rmtree(temporary_path, ignore_errors=True)
else:
temporary_path.unlink(missing_ok=True)
except OSError: except OSError:
pass pass
_active_job = None _active_job = None
@@ -102,6 +134,7 @@ class RunJob(bpy.types.Operator):
(BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""), (BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""),
(IMAGE_BACKGROUND_JOB, "Remove Image Empty Background", ""), (IMAGE_BACKGROUND_JOB, "Remove Image Empty Background", ""),
(ORIGINAL_BACKGROUND_JOB, "Build Original Background Plane", ""), (ORIGINAL_BACKGROUND_JOB, "Build Original Background Plane", ""),
(ORIGINAL_FRAMES_JOB, "Extract Original Depth Frames", ""),
), ),
options={"HIDDEN", "SKIP_SAVE"}, options={"HIDDEN", "SKIP_SAVE"},
) )
@@ -227,6 +260,7 @@ class RunJob(bpy.types.Operator):
IMAGE_DEPTH_JOB, IMAGE_DEPTH_JOB,
IMAGE_BACKGROUND_JOB, IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB, ORIGINAL_BACKGROUND_JOB,
ORIGINAL_FRAMES_JOB,
} }
and self.snapform_delete_input and self.snapform_delete_input
else None else None
@@ -267,6 +301,7 @@ class RunJob(bpy.types.Operator):
IMAGE_DEPTH_JOB, IMAGE_DEPTH_JOB,
IMAGE_BACKGROUND_JOB, IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB, ORIGINAL_BACKGROUND_JOB,
ORIGINAL_FRAMES_JOB,
}: }:
configured_input = self.snapform_input_path configured_input = self.snapform_input_path
elif self.job_type == BACKGROUND_REMOVAL_JOB: elif self.job_type == BACKGROUND_REMOVAL_JOB:
@@ -274,20 +309,24 @@ class RunJob(bpy.types.Operator):
else: else:
configured_input = settings.input_path configured_input = settings.input_path
input_path = Path(bpy.path.abspath(configured_input)) input_path = Path(bpy.path.abspath(configured_input))
if not input_path.is_file(): if not (input_path.is_file() or input_path.is_dir()):
raise ValueError("Select a valid image or video file") raise ValueError("Select a valid image, video, or frame directory")
start_frame = 1
if getattr(self, "snapform_source_object_name", ""):
source_object = bpy.data.objects.get(
self.snapform_source_object_name
)
if source_object is not None:
start_frame = empty_image_start_frame(source_object)
if self.job_type in { if self.job_type in {
BACKGROUND_REMOVAL_JOB, BACKGROUND_REMOVAL_JOB,
IMAGE_BACKGROUND_JOB, IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB, ORIGINAL_BACKGROUND_JOB,
}: }:
if input_path.suffix.lower() not in { if not is_ben2_input(input_path):
".jpg", raise ValueError(
".jpeg", "BEN2 requires an image, video, or directory of image frames"
".png", )
".webp",
}:
raise ValueError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
if not runtime.model_ready(BEN2_MODEL.key): if not runtime.model_ready(BEN2_MODEL.key):
raise RuntimeError(f"Download {BEN2_MODEL.label} first") raise RuntimeError(f"Download {BEN2_MODEL.label} first")
return ( return (
@@ -297,6 +336,18 @@ class RunJob(bpy.types.Operator):
paths["status"], paths["status"],
paths["result"], paths["result"],
BEN2_MODEL.key, BEN2_MODEL.key,
start_frame=start_frame,
),
paths,
)
if self.job_type == ORIGINAL_FRAMES_JOB:
return (
runtime.extract_frames_request(
settings,
input_path,
paths["status"],
paths["result"],
start_frame=start_frame,
), ),
paths, paths,
) )
@@ -342,6 +393,7 @@ class RunJob(bpy.types.Operator):
else None else None
), ),
use_ray_pose=use_ray_pose, use_ray_pose=use_ray_pose,
start_frame=start_frame,
remove_background=( remove_background=(
self.job_type == IMAGE_DEPTH_JOB self.job_type == IMAGE_DEPTH_JOB
and remove_background and remove_background
@@ -371,6 +423,8 @@ class RunJob(bpy.types.Operator):
return "Starting Image Empty background removal" return "Starting Image Empty background removal"
if self.job_type == ORIGINAL_BACKGROUND_JOB: if self.job_type == ORIGINAL_BACKGROUND_JOB:
return "Starting background removal" return "Starting background removal"
if self.job_type == ORIGINAL_FRAMES_JOB:
return "Extracting source frames"
if self.job_type == "DOWNLOAD_MODEL": if self.job_type == "DOWNLOAD_MODEL":
model = get_downloadable_model(self.model_key) model = get_downloadable_model(self.model_key)
return f"Starting {model.label} download" return f"Starting {model.label} download"
@@ -484,7 +538,7 @@ class RunJob(bpy.types.Operator):
self._redraw_ui(context, force=True) self._redraw_ui(context, force=True)
return {"CANCELLED"} return {"CANCELLED"}
status["message"] = "Depth plane is ready" status["message"] = "Depth plane is ready"
if job.kind == ORIGINAL_BACKGROUND_JOB: if job.kind in {ORIGINAL_BACKGROUND_JOB, ORIGINAL_FRAMES_JOB}:
try: try:
settings.status = "Building depth plane" settings.status = "Building depth plane"
settings.progress = 0.96 settings.progress = 0.96
@@ -515,24 +569,39 @@ class RunJob(bpy.types.Operator):
settings.status = status.get("message", "Task complete") settings.status = status.get("message", "Task complete")
settings.progress = 1.0 settings.progress = 1.0
try:
job.status_path.unlink(missing_ok=True)
except OSError:
pass
close_active_job() close_active_job()
self._redraw_ui(context, force=True) self._redraw_ui(context, force=True)
return {"FINISHED"} return {"FINISHED"}
def _build_depth_plane(self, context, job): def _build_depth_plane(self, context, job):
npz_path = ( result_directory = job.result_directory
job.result_directory npz_path = result_directory / "results.npz"
/ "exports" is_sequence = bool(list(result_directory.glob("frame.*")))
/ "mini_npz" foreground_paths = sorted(result_directory.glob("foreground.*.png"))
/ "results.npz" background_path = ""
) if job.remove_background:
background_path = (
str(result_directory)
if foreground_paths
else str(result_directory / "foreground.png")
)
result = bpy.ops.snapform.build_depth_plane( result = bpy.ops.snapform.build_depth_plane(
snapform_source_object_name=job.source_object_name, snapform_source_object_name=job.source_object_name,
snapform_npz_path=str(npz_path), snapform_npz_path=str(npz_path),
snapform_depth_data_name=job.depth_data_name, snapform_depth_data_name=job.depth_data_name,
snapform_background_path=( snapform_background_path=background_path,
str(job.result_directory / "foreground.png") snapform_color_frames_dir=(
if job.remove_background str(result_directory)
if is_sequence and not foreground_paths
else ""
),
snapform_depth_frames_dir=(
str(result_directory)
if is_sequence
else "" else ""
), ),
snapform_subdivision_level=job.subdivision_level, snapform_subdivision_level=job.subdivision_level,
@@ -547,10 +616,24 @@ class RunJob(bpy.types.Operator):
raise RuntimeError("Unable to create the depth plane Undo step") raise RuntimeError("Unable to create the depth plane Undo step")
def _build_original_depth_plane(self, job): def _build_original_depth_plane(self, job):
result_directory = job.result_directory
is_sequence = bool(list(result_directory.glob("frame.*")))
background_path = ""
color_frames_dir = ""
if job.kind == ORIGINAL_BACKGROUND_JOB:
background_path = str(
result_directory
if is_sequence
else result_directory / "foreground.png"
)
elif is_sequence:
color_frames_dir = str(result_directory)
result = bpy.ops.snapform.build_original_depth_plane( result = bpy.ops.snapform.build_original_depth_plane(
snapform_source_object_name=job.source_object_name, snapform_source_object_name=job.source_object_name,
snapform_background_path=str( snapform_background_path=background_path,
job.result_directory / "foreground.png" snapform_color_frames_dir=color_frames_dir,
snapform_depth_frames_dir=(
str(result_directory) if is_sequence else ""
), ),
snapform_subdivision_level=job.subdivision_level, snapform_subdivision_level=job.subdivision_level,
snapform_delete_source=job.delete_source, snapform_delete_source=job.delete_source,
@@ -564,7 +647,12 @@ class RunJob(bpy.types.Operator):
raise RuntimeError("Unable to create the depth plane Undo step") raise RuntimeError("Unable to create the depth plane Undo step")
def _apply_removed_background(self, job, push_undo=True): def _apply_removed_background(self, job, push_undo=True):
output_path = job.result_directory / "foreground.png" result_directory = job.result_directory
foreground_paths = sorted(result_directory.glob("foreground.*.png"))
output_path = (
result_directory if foreground_paths
else result_directory / "foreground.png"
)
result = bpy.ops.snapform.apply_removed_background( result = bpy.ops.snapform.apply_removed_background(
snapform_source_object_name=job.source_object_name, snapform_source_object_name=job.source_object_name,
snapform_output_path=str(output_path), snapform_output_path=str(output_path),

View File

@@ -35,6 +35,14 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
), ),
default="CPU" if sys.platform == "darwin" else "CU128", default="CPU" if sys.platform == "darwin" else "CU128",
) )
pypi_index_url: bpy.props.StringProperty(
name="PyPI Index",
description=(
"Package index for the isolated Runtime; leave empty to auto-detect "
"with mirror fallback. Example: https://pypi.tuna.tsinghua.edu.cn/simple"
),
default="",
)
def draw(self, _context): def draw(self, _context):
from .operators import active_job from .operators import active_job
@@ -44,6 +52,7 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
layout.prop(self, "storage_root") layout.prop(self, "storage_root")
layout.prop(self, "uv_executable") layout.prop(self, "uv_executable")
layout.prop(self, "torch_backend") layout.prop(self, "torch_backend")
layout.prop(self, "pypi_index_url")
layout.separator() layout.separator()
layout.label(text=f"Runtime: {runtime_directory()}") layout.label(text=f"Runtime: {runtime_directory()}")
if runtime_ready(): if runtime_ready():

View File

@@ -15,8 +15,8 @@ class SnapFormSettings(bpy.types.PropertyGroup):
subtype="FILE_PATH", subtype="FILE_PATH",
) )
video_frames: bpy.props.IntProperty( video_frames: bpy.props.IntProperty(
name="Video Frames", name="Extract Frames",
description="Maximum number of frames sampled evenly from a video", description="Maximum number of consecutive frames extracted from the start of a video or frame sequence",
default=16, default=16,
min=2, min=2,
soft_max=64, soft_max=64,

View File

@@ -24,6 +24,13 @@ RUNTIME_SCHEMA = 4
HOST_SCHEMA = 2 HOST_SCHEMA = 2
HOST_ADDRESS = "127.0.0.1" HOST_ADDRESS = "127.0.0.1"
DEFAULT_STORAGE_ROOT = Path.home() / ".snapform" DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"
JOB_FOLDER_NAMES = {
"INSTALL_RUNTIME": "install-runtime",
"DOWNLOAD_MODEL": "download-model",
"IMAGE_DEPTH": "depth-plane",
"IMAGE_BACKGROUND": "remove-background",
"ORIGINAL_FRAMES": "depth-plane",
}
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -122,15 +129,26 @@ def temporary_input_path():
return path / f"{uuid.uuid4().hex}.png" return path / f"{uuid.uuid4().hex}.png"
def temporary_input_directory():
path = jobs_directory() / "inputs" / f"{uuid.uuid4().hex}"
path.mkdir(parents=True, exist_ok=True)
return path
def job_folder_name(job_type):
"""Return the folder name segment for a job type."""
return JOB_FOLDER_NAMES.get(job_type, job_type.lower())
def new_job_paths(job_type): def new_job_paths(job_type):
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
token = uuid.uuid4().hex[:8] stem = f"{job_folder_name(job_type)}_{timestamp}"
stem = f"{timestamp}-{job_type.lower()}-{token}" directory = jobs_directory() / stem
directory = jobs_directory() directory.mkdir(parents=True, exist_ok=True)
return { return {
"log": directory / f"{stem}.log", "log": directory / f"{directory.name}.log",
"status": directory / f"{stem}.json", "status": directory / f"{directory.name}.json",
"result": directory / stem, "result": directory,
} }
@@ -155,8 +173,9 @@ def install_runtime_command(status_path):
executable = uv_executable() executable = uv_executable()
addon_preferences = preferences() addon_preferences = preferences()
backend = getattr(addon_preferences, "torch_backend", "CU128") backend = getattr(addon_preferences, "torch_backend", "CU128")
pypi_index = getattr(addon_preferences, "pypi_index_url", "") or ""
script = Path(__file__).with_name("bootstrap_runtime.py") script = Path(__file__).with_name("bootstrap_runtime.py")
return [ command = [
sys.executable, sys.executable,
str(script), str(script),
"--uv", "--uv",
@@ -174,6 +193,9 @@ def install_runtime_command(status_path):
"--status", "--status",
str(status_path), str(status_path),
] ]
if pypi_index:
command.extend(["--pypi-index", pypi_index])
return command
def host_log_path(): def host_log_path():
@@ -616,6 +638,8 @@ def remove_background_request(
status_path, status_path,
result_directory, result_directory,
model_key, model_key,
video_frames=None,
start_frame=1,
): ):
return HostJobRequest( return HostJobRequest(
{ {
@@ -627,6 +651,39 @@ def remove_background_request(
"output_dir": str(result_directory), "output_dir": str(result_directory),
"status": str(status_path), "status": str(status_path),
"device": settings.device.lower(), "device": settings.device.lower(),
"video_frames": (
settings.video_frames
if video_frames is None
else video_frames
),
"start_frame": start_frame,
},
}
)
def extract_frames_request(
settings,
input_path,
status_path,
result_directory,
video_frames=None,
start_frame=1,
):
return HostJobRequest(
{
"job_id": Path(status_path).stem,
"command": "extract-frames",
"parameters": {
"input": str(input_path),
"output_dir": str(result_directory),
"status": str(status_path),
"video_frames": (
settings.video_frames
if video_frames is None
else video_frames
),
"start_frame": start_frame,
}, },
} }
) )
@@ -643,6 +700,7 @@ def reconstruct_request(
process_res_method=None, process_res_method=None,
use_ray_pose=None, use_ray_pose=None,
remove_background=False, remove_background=False,
start_frame=1,
): ):
snapform_selected_key = model_key or selected_model_key(settings, export_format) snapform_selected_key = model_key or selected_model_key(settings, export_format)
snapform_selected_process_res = ( snapform_selected_process_res = (
@@ -673,6 +731,7 @@ def reconstruct_request(
"show_cameras": 1 if settings.show_cameras else 0, "show_cameras": 1 if settings.show_cameras else 0,
"gs_views_interval": settings.gs_views_interval, "gs_views_interval": settings.gs_views_interval,
"video_frames": settings.video_frames, "video_frames": settings.video_frames,
"start_frame": start_frame,
"device": settings.device.lower(), "device": settings.device.lower(),
"use_ray_pose": 1 if snapform_selected_ray_pose else 0, "use_ray_pose": 1 if snapform_selected_ray_pose else 0,
"ref_view_strategy": settings.ref_view_strategy, "ref_view_strategy": settings.ref_view_strategy,

View File

@@ -1,13 +1,14 @@
import argparse import argparse
import fnmatch
import json import json
import os import os
import shutil import shutil
import tempfile import tempfile
import threading import threading
import time import time
import urllib.error
import urllib.request
from pathlib import Path from pathlib import Path
from urllib.parse import quote, urlencode
from urllib.request import Request, urlopen
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"} IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
@@ -24,9 +25,10 @@ VIDEO_SUFFIXES = {
} }
SCENE_FILENAME = "scene.glb" SCENE_FILENAME = "scene.glb"
EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz") EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz")
HF_ENDPOINT = "https://huggingface.co/" HF_ENDPOINT = "https://huggingface.co"
HF_MIRRORS = ("https://hf-mirror.com",) MODELSCOPE_ENDPOINT = "https://modelscope.cn"
HF_PROBE_TIMEOUT = 5.0 MODELSCOPE_REVISION = "master"
MODELSCOPE_TIMEOUT = 60.0
class TaskCancelled(RuntimeError): class TaskCancelled(RuntimeError):
@@ -38,34 +40,6 @@ def check_cancelled(cancel_check):
raise TaskCancelled("Task cancelled") raise TaskCancelled("Task cancelled")
def reachable(url, timeout=HF_PROBE_TIMEOUT):
request = urllib.request.Request(
url,
method="HEAD",
headers={"User-Agent": "SnapForm-Runtime-Installer"},
)
try:
with urllib.request.urlopen(request, timeout=timeout):
return True
except urllib.error.HTTPError:
return True
except Exception:
return False
def resolve_hf_endpoint():
"""Probe Hugging Face reachability and fall back to a mirror when direct access fails; returns (endpoint, used_mirror)."""
configured = os.environ.get("HF_ENDPOINT")
if configured:
return configured, False
if reachable(HF_ENDPOINT):
return HF_ENDPOINT, False
for mirror in HF_MIRRORS:
if reachable(mirror):
return mirror, True
return HF_ENDPOINT, False
def write_status( def write_status(
path, path,
progress, progress,
@@ -106,56 +80,66 @@ def collect_images(directory):
) )
def first_frames(images, maximum_frames, start_frame=1):
"""Return up to maximum_frames image paths in order from start_frame."""
if maximum_frames <= 0:
return images
start_index = max(start_frame - 1, 0)
return images[start_index:start_index + maximum_frames]
def extract_video_frames( def extract_video_frames(
video_path, video_path,
output_directory, output_directory,
maximum_frames, maximum_frames,
cv2, cv2,
start_frame=1,
cancel_check=None, cancel_check=None,
): ):
output_directory = Path(output_directory)
output_directory.mkdir(parents=True, exist_ok=True)
capture = cv2.VideoCapture(str(video_path)) capture = cv2.VideoCapture(str(video_path))
if not capture.isOpened(): if not capture.isOpened():
raise RuntimeError(f"Unable to open video: {video_path}") raise RuntimeError(f"Unable to open video: {video_path}")
output_directory = Path(output_directory)
output_directory.mkdir(parents=True, exist_ok=True)
total_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
if total_frames <= 0:
capture.release()
raise RuntimeError("The video does not report a readable frame count")
import numpy as np
sample_count = min(maximum_frames, total_frames)
frame_indices = np.linspace(
0,
total_frames - 1,
sample_count,
dtype=np.int64,
)
image_paths = [] image_paths = []
skipped = 0
try: try:
for output_index, frame_index in enumerate(frame_indices): while len(image_paths) < maximum_frames:
check_cancelled(cancel_check) check_cancelled(cancel_check)
capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index))
success, frame = capture.read() success, frame = capture.read()
if not success: if not success:
raise RuntimeError( break
f"Unable to read video frame {int(frame_index)}" if skipped < start_frame - 1:
) skipped += 1
image_path = output_directory / f"frame-{output_index:04d}.jpg" continue
image_path = (
output_directory / f"frame.{len(image_paths) + 1:04d}.jpg"
)
if not cv2.imwrite(str(image_path), frame): if not cv2.imwrite(str(image_path), frame):
raise RuntimeError(f"Unable to write video frame: {image_path}") raise RuntimeError(f"Unable to write video frame: {image_path}")
image_paths.append(image_path) image_paths.append(image_path)
finally: finally:
capture.release() capture.release()
if not image_paths:
raise RuntimeError(f"Unable to read any video frame from: {video_path}")
return image_paths return image_paths
def prepare_input(input_path, video_frames, temporary_root, cancel_check=None): def prepare_input(
input_path,
video_frames,
temporary_root,
start_frame=1,
cancel_check=None,
):
check_cancelled(cancel_check) check_cancelled(cancel_check)
input_path = Path(input_path) input_path = Path(input_path)
if input_path.is_dir():
images = collect_images(input_path)
if not images:
raise RuntimeError(f"No image frames found in: {input_path}")
return first_frames(images, video_frames, start_frame), None
suffix = input_path.suffix.lower() suffix = input_path.suffix.lower()
if suffix in IMAGE_SUFFIXES: if suffix in IMAGE_SUFFIXES:
return [input_path], None return [input_path], None
@@ -174,6 +158,7 @@ def prepare_input(input_path, video_frames, temporary_root, cancel_check=None):
temporary_directory, temporary_directory,
video_frames, video_frames,
cv2, cv2,
start_frame=start_frame,
cancel_check=cancel_check, cancel_check=cancel_check,
) )
except Exception: except Exception:
@@ -182,6 +167,43 @@ def prepare_input(input_path, video_frames, temporary_root, cancel_check=None):
return images, temporary_directory return images, temporary_directory
def prepare_background_input(
input_path,
video_frames,
temporary_root,
start_frame=1,
cancel_check=None,
):
"""Resolve a BEN2 input to frame images; returns (images, temporary_directory)."""
input_path = Path(input_path)
if input_path.is_dir():
images = collect_images(input_path)
if not images:
raise RuntimeError(f"No image frames found in: {input_path}")
return first_frames(images, video_frames, start_frame), None
if input_path.suffix.lower() in IMAGE_SUFFIXES:
return [input_path], None
return prepare_input(
input_path,
video_frames,
temporary_root,
start_frame=start_frame,
cancel_check=cancel_check,
)
def copy_frames(images, output_directory):
"""Copy a multi-frame input as numbered frame files for the Blender sequence."""
if len(images) <= 1:
return
for index, image_path in enumerate(images):
suffix = Path(image_path).suffix.lower() or ".jpg"
shutil.copy2(
image_path,
Path(output_directory) / f"frame.{index + 1:04d}{suffix}",
)
class DownloadProgressReporter: class DownloadProgressReporter:
def __init__( def __init__(
self, self,
@@ -318,21 +340,8 @@ def download_model(args):
"stages": getattr(args, "stages", None), "stages": getattr(args, "stages", None),
"stage_label": getattr(args, "stage_label", None), "stage_label": getattr(args, "stage_label", None),
} }
write_status(
args.status,
start,
"checking download source",
**stage_kwargs,
)
endpoint, used_mirror = resolve_hf_endpoint()
if endpoint != os.environ.get("HF_ENDPOINT"):
os.environ["HF_ENDPOINT"] = endpoint
from huggingface_hub import snapshot_download from huggingface_hub import snapshot_download
message = f"downloading {label}"
if used_mirror:
message = f"{message} (mirror)"
write_status(args.status, start, message, **stage_kwargs)
progress_class = create_download_progress_class( progress_class = create_download_progress_class(
args.status, args.status,
args.repository, args.repository,
@@ -341,18 +350,213 @@ def download_model(args):
label=label, label=label,
**stage_kwargs, **stage_kwargs,
) )
snapshot_download( failures = []
repo_id=args.repository, for endpoint in download_endpoints():
local_dir=model_dir, if not hf_endpoint_reachable(endpoint):
tqdm_class=progress_class, failures.append((endpoint, "endpoint unreachable"))
allow_patterns=args.allow_pattern or None, continue
os.environ["HF_ENDPOINT"] = endpoint
message = f"downloading {label}"
if endpoint != HF_ENDPOINT:
message = f"{message} (mirror)"
write_status(args.status, start, message, **stage_kwargs)
try:
snapshot_download(
repo_id=args.repository,
local_dir=model_dir,
tqdm_class=progress_class,
allow_patterns=args.allow_pattern or None,
endpoint=endpoint,
)
_verify_model_complete(model_dir)
except Exception as error:
failures.append((endpoint, error))
continue
write_status(
args.status,
start + span,
"model ready",
**stage_kwargs,
)
return
# Hugging Face 官方源与已配置镜像均失败时,回退 ModelScope(仓库名与 Hugging Face 一致)
try:
write_status(
args.status,
start,
f"downloading {label} (mirror: modelscope.cn)",
**stage_kwargs,
)
download_from_modelscope(
args,
model_dir,
start=start,
span=span,
label=label,
stage_kwargs=stage_kwargs,
)
_verify_model_complete(model_dir)
except Exception as error:
failures.append((MODELSCOPE_ENDPOINT, error))
else:
write_status(
args.status,
start + span,
"model ready",
**stage_kwargs,
)
return
detail = "; ".join(
f"{source}: {error}" for source, error in failures
) )
write_status( raise RuntimeError(f"All download sources failed: {detail}")
def download_endpoints():
"""Hugging Face download candidates: configured mirror first, official endpoint otherwise."""
endpoint = os.environ.get("HF_ENDPOINT") or HF_ENDPOINT
return [endpoint.rstrip("/")]
def hf_endpoint_reachable(endpoint, timeout=10.0):
"""Quick connectivity probe so slow or blocked Hugging Face endpoints fail fast."""
try:
request = Request(
f"{endpoint.rstrip('/')}/api/models",
method="HEAD",
headers={"User-Agent": "SnapForm"},
)
with urlopen(request, timeout=timeout):
return True
except Exception:
return False
def _verify_model_complete(model_dir):
"""Refuse to report success when the required model files are missing."""
model_dir = Path(model_dir)
missing = []
if not (model_dir / "config.json").is_file():
missing.append("config.json")
if not any(model_dir.glob("*.safetensors")):
missing.append("*.safetensors")
if missing:
raise RuntimeError(
"Model download reported success but files are missing: "
+ ", ".join(missing)
)
def _hidden_repo_path(path):
return any(part.startswith(".") for part in Path(path).parts)
def _model_file_allowed(path, patterns):
if not patterns:
return True
return any(fnmatch.fnmatch(path, pattern) for pattern in patterns)
def modelscope_file_list(repository, revision=MODELSCOPE_REVISION):
"""List (path, size) file entries of a ModelScope repository."""
url = (
f"{MODELSCOPE_ENDPOINT}/api/v1/models/{repository}/repo/files?"
+ urlencode({"Revision": revision, "Recursive": "true"})
)
request = Request(url, headers={"User-Agent": "SnapForm"})
with urlopen(request, timeout=MODELSCOPE_TIMEOUT) as response:
payload = json.loads(response.read().decode("utf-8"))
entries = []
for entry in payload.get("Data", {}).get("Files", []):
path = entry.get("Path")
if path and entry.get("Type") == "blob" and not _hidden_repo_path(path):
entries.append((path, int(entry.get("Size") or 0)))
return entries
def _download_modelscope_file(
repository,
path,
destination,
progress=None,
retries=2,
):
"""Download one ModelScope file atomically with a per-chunk progress callback."""
url = (
f"{MODELSCOPE_ENDPOINT}/models/{repository}/resolve/"
f"{MODELSCOPE_REVISION}/{quote(path)}"
)
request = Request(url, headers={"User-Agent": "SnapForm"})
last_error = None
for _attempt in range(retries):
temporary = destination.with_name(f"{destination.name}.part")
try:
with urlopen(request, timeout=MODELSCOPE_TIMEOUT) as response:
with open(temporary, "wb") as output:
while True:
chunk = response.read(1024 * 1024)
if not chunk:
break
output.write(chunk)
if progress is not None:
progress(len(chunk))
temporary.replace(destination)
return
except Exception as error:
last_error = error
try:
temporary.unlink()
except OSError:
pass
raise last_error
def download_from_modelscope(
args,
model_dir,
start=0.05,
span=0.90,
label=None,
stage_kwargs=None,
):
"""Download a model from ModelScope using only standard library requests."""
repository = args.repository
files = modelscope_file_list(repository)
patterns = args.allow_pattern or None
selected = [
(path, size)
for path, size in files
if _model_file_allowed(path, patterns)
]
if not selected:
raise RuntimeError(
f"No matching files in ModelScope repository: {repository}"
)
reporter = DownloadProgressReporter(
args.status, args.status,
start + span, repository,
"model ready", start=start,
**stage_kwargs, span=span,
label=label or repository,
**(stage_kwargs or {}),
) )
total_size = sum(size for _path, size in selected)
downloaded = 0
def on_chunk(count):
nonlocal downloaded
downloaded += count
reporter.report(downloaded, total_size)
for path, _size in selected:
destination = Path(model_dir) / path
destination.parent.mkdir(parents=True, exist_ok=True)
_download_modelscope_file(
repository,
path,
destination,
progress=on_chunk,
)
def format_rate(rate): def format_rate(rate):
@@ -435,6 +639,49 @@ def infer_foreground(
return foreground, load_ms, inference_ms return foreground, load_ms, inference_ms
def infer_foreground_frames(
images,
model_dir,
device,
torch,
output_directory,
status_path,
start_progress,
progress_span,
model_cache=None,
cancel_check=None,
):
"""Run BEN2 on every frame and write foreground.0001.png etc."""
total_load_ms = 0.0
total_inference_ms = 0.0
for index, image_path in enumerate(images):
check_cancelled(cancel_check)
write_status(
status_path,
start_progress + progress_span * index / len(images),
f"Removing background ({index + 1}/{len(images)})",
)
foreground, load_ms, inference_ms = infer_foreground(
image_path,
model_dir,
device,
torch,
model_cache=model_cache,
cancel_check=cancel_check,
)
output_path = output_directory / f"foreground.{index + 1:04d}.png"
foreground.convert("RGBA").save(
output_path,
format="PNG",
compress_level=1,
)
if not output_path.is_file():
raise RuntimeError(f"BEN2 did not write: {output_path}")
total_load_ms += load_ms
total_inference_ms += inference_ms
return total_load_ms, total_inference_ms
def remove_background(args, model_cache=None, cancel_check=None): def remove_background(args, model_cache=None, cancel_check=None):
total_started = time.perf_counter() total_started = time.perf_counter()
write_status(args.status, 0.02, "Starting BEN2 Runtime") write_status(args.status, 0.02, "Starting BEN2 Runtime")
@@ -446,41 +693,90 @@ def remove_background(args, model_cache=None, cancel_check=None):
torch = model_cache.torch torch = model_cache.torch
output_directory = Path(args.output_dir) output_directory = Path(args.output_dir)
output_directory.mkdir(parents=True, exist_ok=False) output_directory.mkdir(parents=True, exist_ok=True)
device = select_device(args.device, torch) device = select_device(args.device, torch)
write_status(args.status, 0.12, f"Loading BEN2 on {device.type}") video_frames = getattr(args, "video_frames", 16)
write_status(args.status, 0.35, "Removing background") start_frame = getattr(args, "start_frame", 1)
foreground, load_ms, inference_ms = infer_foreground( images, temporary_directory = prepare_background_input(
args.input, args.input,
args.model_dir, video_frames,
device, output_directory.parent,
torch, start_frame=start_frame,
model_cache=model_cache,
cancel_check=cancel_check, cancel_check=cancel_check,
) )
try:
check_cancelled(cancel_check) write_status(args.status, 0.12, f"Loading BEN2 on {device.type}")
write_status(args.status, 0.92, "Writing transparent PNG") if len(images) == 1:
output_path = output_directory / "foreground.png" write_status(args.status, 0.35, "Removing background")
foreground.convert("RGBA").save( foreground, load_ms, inference_ms = infer_foreground(
output_path, images[0],
format="PNG", args.model_dir,
compress_level=1, device,
) torch,
if not output_path.is_file(): model_cache=model_cache,
raise RuntimeError("BEN2 did not write a foreground image") cancel_check=cancel_check,
)
check_cancelled(cancel_check)
write_status(args.status, 0.92, "Writing transparent PNG")
output_path = output_directory / "foreground.png"
foreground.convert("RGBA").save(
output_path,
format="PNG",
compress_level=1,
)
if not output_path.is_file():
raise RuntimeError("BEN2 did not write a foreground image")
else:
load_ms, inference_ms = infer_foreground_frames(
images,
args.model_dir,
device,
torch,
output_directory,
args.status,
0.12,
0.7,
model_cache=model_cache,
cancel_check=cancel_check,
)
copy_frames(images, output_directory)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
check_cancelled(cancel_check) check_cancelled(cancel_check)
total_ms = (time.perf_counter() - total_started) * 1000.0 total_ms = (time.perf_counter() - total_started) * 1000.0
frame_text = "" if len(images) == 1 else f" from {len(images)} frame(s)"
message = ( message = (
f"Background removed in {inference_ms:.0f} ms " f"Background removed{frame_text} in {inference_ms:.0f} ms "
f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)" f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)"
) )
print(message) print(message)
write_status(args.status, 1.0, message) write_status(args.status, 1.0, message)
def extract_frames(args, cancel_check=None):
"""Copy source frames into a sequence directory without loading any AI model."""
output_directory = Path(args.output_dir)
output_directory.mkdir(parents=True, exist_ok=True)
images, temporary_directory = prepare_input(
args.input,
getattr(args, "video_frames", 16),
output_directory.parent,
start_frame=getattr(args, "start_frame", 1),
cancel_check=cancel_check,
)
try:
copy_frames(images, output_directory)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
message = f"Extracted {len(images)} frame(s)"
print(message)
write_status(args.status, 1.0, message)
def exported_paths(output_directory, export_format): def exported_paths(output_directory, export_format):
output_directory = Path(output_directory) output_directory = Path(output_directory)
if export_format == "glb": if export_format == "glb":
@@ -491,7 +787,12 @@ def exported_paths(output_directory, export_format):
if export_format == "depth_vis": if export_format == "depth_vis":
return sorted((output_directory / "depth_vis").glob("*.jpg")) return sorted((output_directory / "depth_vis").glob("*.jpg"))
if export_format == "mini_npz": if export_format == "mini_npz":
result_path = output_directory / "exports" / "mini_npz" / "results.npz" result_path = output_directory / "results.npz"
if result_path.is_file():
return [result_path]
result_path = (
output_directory / "exports" / "mini_npz" / "results.npz"
)
return [result_path] if result_path.is_file() else [] return [result_path] if result_path.is_file() else []
return [] return []
@@ -543,12 +844,13 @@ def reconstruct(args, model_cache=None, cancel_check=None):
from depth_anything_3.utils.export import export 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=False) output_directory.mkdir(parents=True, exist_ok=True)
write_status(args.status, 0.08, "Preparing input") write_status(args.status, 0.08, "Preparing input")
images, temporary_directory = prepare_input( images, temporary_directory = prepare_input(
args.input, args.input,
args.video_frames, args.video_frames,
output_directory.parent, output_directory.parent,
start_frame=getattr(args, "start_frame", 1),
cancel_check=cancel_check, cancel_check=cancel_check,
) )
@@ -562,22 +864,36 @@ def reconstruct(args, model_cache=None, cancel_check=None):
) )
if remove_background_model_dir: if remove_background_model_dir:
write_status(args.status, 0.12, "Removing image background") write_status(args.status, 0.12, "Removing image background")
foreground, load_ms, inference_ms = infer_foreground( if len(images) == 1:
images[0], foreground, load_ms, inference_ms = infer_foreground(
remove_background_model_dir, images[0],
device, remove_background_model_dir,
torch, device,
model_cache=model_cache, torch,
cancel_check=cancel_check, model_cache=model_cache,
) cancel_check=cancel_check,
check_cancelled(cancel_check) )
foreground.convert("RGBA").save( check_cancelled(cancel_check)
output_directory / "foreground.png", foreground.convert("RGBA").save(
format="PNG", output_directory / "foreground.png",
compress_level=1, format="PNG",
) compress_level=1,
if not (output_directory / "foreground.png").is_file(): )
raise RuntimeError("BEN2 did not write a foreground image") if not (output_directory / "foreground.png").is_file():
raise RuntimeError("BEN2 did not write a foreground image")
else:
load_ms, inference_ms = infer_foreground_frames(
images,
remove_background_model_dir,
device,
torch,
output_directory,
args.status,
0.12,
0.05,
model_cache=model_cache,
cancel_check=cancel_check,
)
print( print(
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"
@@ -636,7 +952,21 @@ def reconstruct(args, model_cache=None, cancel_check=None):
raise RuntimeError( raise RuntimeError(
f"No {args.export_format} files were exported" f"No {args.export_format} files were exported"
) )
if args.export_format == "mini_npz":
nested_path = (
output_directory / "exports" / "mini_npz" / "results.npz"
)
result_path = output_directory / "results.npz"
if nested_path.is_file():
nested_path.replace(result_path)
result_paths = [result_path]
for directory in (nested_path.parent, nested_path.parent.parent):
try:
directory.rmdir()
except OSError:
pass
check_cancelled(cancel_check) check_cancelled(cancel_check)
copy_frames(images, output_directory)
write_status( write_status(
args.status, args.status,
1.0, 1.0,
@@ -673,6 +1003,15 @@ def parser():
choices=("auto", "cuda", "mps", "cpu"), choices=("auto", "cuda", "mps", "cpu"),
required=True, required=True,
) )
background.add_argument("--video-frames", type=int, default=16)
background.add_argument("--start-frame", type=int, default=1)
frames = commands.add_parser("extract-frames")
frames.add_argument("--input", required=True)
frames.add_argument("--output-dir", required=True)
frames.add_argument("--status", required=True)
frames.add_argument("--video-frames", type=int, default=16)
frames.add_argument("--start-frame", type=int, default=1)
reconstruction = commands.add_parser("reconstruct") reconstruction = commands.add_parser("reconstruct")
reconstruction.add_argument("--input", required=True) reconstruction.add_argument("--input", required=True)
@@ -709,6 +1048,7 @@ def parser():
required=True, required=True,
) )
reconstruction.add_argument("--video-frames", type=int, required=True) reconstruction.add_argument("--video-frames", type=int, required=True)
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", "mps", "cpu"),
@@ -735,6 +1075,8 @@ def main():
download_model(args) download_model(args)
elif args.command == "reconstruct": elif args.command == "reconstruct":
reconstruct(args) reconstruct(args)
elif args.command == "extract-frames":
extract_frames(args)
else: else:
remove_background(args) remove_background(args)
except Exception as error: except Exception as error:

File diff suppressed because it is too large Load Diff

View File

@@ -1,5 +1,4 @@
import importlib.util import importlib.util
import inspect
import sys import sys
import tempfile import tempfile
import unittest import unittest
@@ -112,21 +111,6 @@ class HostTest(unittest.TestCase):
self.assertEqual(loads, ["model"]) self.assertEqual(loads, ["model"])
self.assertEqual(cache.status()["da3"], "da3_model") self.assertEqual(cache.status()["da3"], "da3_model")
def test_http_host_has_health_jobs_auth_and_single_task_guard(self):
source = inspect.getsource(self.host.create_app)
self.assertIn('@app.get("/health")', source)
self.assertIn('@app.post("/jobs", status_code=202)', source)
self.assertIn("x_snapform_token", source)
self.assertIn("hmac.compare_digest", source)
self.assertIn("parameters[\"status\"]", source)
self.assertIn('@app.post("/clients", status_code=201)', source)
self.assertIn('@app.post("/clients/{client_id}/heartbeat")', source)
self.assertIn('@app.delete("/jobs/{job_id}", status_code=202)', source)
self.assertIn('@app.post("/models/clear")', source)
self.assertIn('@app.post("/shutdown", status_code=202)', source)
self.assertNotIn("os._exit", source)
def test_host_tracks_clients_and_rejects_shared_shutdown(self): def test_host_tracks_clients_and_rejects_shared_shutdown(self):
fake_models = SimpleNamespace( fake_models = SimpleNamespace(
status=lambda: {"ben2": "", "da3": ""}, status=lambda: {"ben2": "", "da3": ""},
@@ -218,6 +202,51 @@ class HostTest(unittest.TestCase):
self.host.remove_discovery(path, "owner") self.host.remove_discovery(path, "owner")
self.assertFalse(path.exists()) self.assertFalse(path.exists())
def test_host_dispatches_frame_extraction_without_models(self):
fake_models = SimpleNamespace(
status=lambda: {"ben2": "", "da3": ""},
close=lambda: None,
)
with patch.object(self.host, "ModelCache", return_value=fake_models):
state = self.host.HostState("token", "instance")
state.register_client("client", "Blender")
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
input_directory = root / "frames"
input_directory.mkdir()
(input_directory / "frame-0001.jpg").touch()
(input_directory / "frame-0002.jpg").touch()
result_directory = root / "result"
status_path = root / "status.json"
state.reserve(
"job",
"client",
"extract-frames",
str(status_path),
)
state.run_job(
{
"job_id": "job",
"command": "extract-frames",
"parameters": {
"input": str(input_directory),
"output_dir": str(result_directory),
"status": str(status_path),
"video_frames": 16,
"start_frame": 1,
},
}
)
status = self.host.json.loads(
status_path.read_text(encoding="utf-8")
)
self.assertTrue(
(result_directory / "frame.0001.jpg").is_file()
)
self.assertEqual(status["state"], "succeeded")
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()

View File

@@ -1,7 +1,9 @@
import importlib.util import importlib.util
import io
import json import json
import os import os
import subprocess import subprocess
import tarfile
import tempfile import tempfile
import time import time
import unittest import unittest
@@ -25,12 +27,24 @@ def load_bootstrap_runtime():
class BootstrapRuntimeTest(unittest.TestCase): class BootstrapRuntimeTest(unittest.TestCase):
def create_bundled_sources(self, root):
sources = root / "sources"
sources.mkdir()
(sources / "depth-anything-3.tar.gz").write_bytes(b"da3")
(sources / "ben2.tar.gz").write_bytes(b"ben2")
return sources
def test_installs_da3_from_archives_without_gsplat(self): def test_installs_da3_from_archives_without_gsplat(self):
bootstrap = load_bootstrap_runtime() bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory: with tempfile.TemporaryDirectory() as directory:
root = Path(directory) root = Path(directory)
runtime_dir = root / "runtime" runtime_dir = root / "runtime"
runtime_dir.mkdir() runtime_dir.mkdir()
self.create_bundled_sources(root)
da3_uri = (
root / "sources" / "depth-anything-3.tar.gz"
).as_uri()
ben2_uri = (root / "sources" / "ben2.tar.gz").as_uri()
uv_path = root / "uv" uv_path = root / "uv"
uv_path.write_bytes(b"uv") uv_path.write_bytes(b"uv")
args = SimpleNamespace( args = SimpleNamespace(
@@ -42,7 +56,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
ben2_commit="test-ben2-commit", ben2_commit="test-ben2-commit",
status=str(root / "status.json"), status=str(root / "status.json"),
) )
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap, bootstrap,
"run", "run",
) as mocked_run, patch.object( ) as mocked_run, patch.object(
@@ -64,13 +78,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
f"addict=={bootstrap.ADDICT_VERSION}", f"addict=={bootstrap.ADDICT_VERSION}",
da3_command, da3_command,
) )
self.assertTrue( self.assertIn(f"depth-anything-3 @ {da3_uri}", da3_command)
any(
"https://github.com/ByteDance-Seed/Depth-Anything-3/"
"archive/test-commit.tar.gz" in str(value)
for value in da3_command
)
)
self.assertFalse( self.assertFalse(
any("git+" in str(value) for value in da3_command) any("git+" in str(value) for value in da3_command)
) )
@@ -93,13 +101,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
for command in commands for command in commands
if any("ben2 @" in str(value) for value in command) if any("ben2 @" in str(value) for value in command)
) )
self.assertTrue( self.assertIn(f"ben2 @ {ben2_uri}", ben2_command)
any(
"https://github.com/PramaLLC/BEN2/archive/"
"test-ben2-commit.tar.gz" in str(value)
for value in ben2_command
)
)
self.assertFalse( self.assertFalse(
any("git+" in str(value) for value in ben2_command) any("git+" in str(value) for value in ben2_command)
) )
@@ -162,6 +164,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
root = Path(directory) root = Path(directory)
runtime_dir = root / "runtime" runtime_dir = root / "runtime"
runtime_dir.mkdir() runtime_dir.mkdir()
self.create_bundled_sources(root)
uv_path = root / "uv" uv_path = root / "uv"
uv_path.write_bytes(b"uv") uv_path.write_bytes(b"uv")
args = SimpleNamespace( args = SimpleNamespace(
@@ -181,7 +184,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
bootstrap.platform, bootstrap.platform,
"machine", "machine",
return_value="arm64", return_value="arm64",
), patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( ), patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap, bootstrap,
"run", "run",
), patch.object( ), patch.object(
@@ -329,6 +332,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
root = Path(directory) root = Path(directory)
runtime_dir = root / "runtime" runtime_dir = root / "runtime"
runtime_dir.mkdir() runtime_dir.mkdir()
self.create_bundled_sources(root)
uv_path = root / "uv" uv_path = root / "uv"
uv_path.write_bytes(b"uv") uv_path.write_bytes(b"uv")
args = SimpleNamespace( args = SimpleNamespace(
@@ -364,7 +368,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
) )
return [] return []
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap, bootstrap,
"run_stage", "run_stage",
side_effect=fake_run_stage, side_effect=fake_run_stage,
@@ -425,7 +429,8 @@ class BootstrapRuntimeTest(unittest.TestCase):
"3/3 packages", "3/3 packages",
], ],
) )
self.assertEqual(mocked_write.call_args_list[-1].args[1], 0.02) self.assertEqual(mocked_write.call_args_list[0].args[1], 0.02)
self.assertLess(mocked_write.call_args_list[-1].args[1], 0.03)
first = mocked_write.call_args_list[0] first = mocked_write.call_args_list[0]
self.assertEqual(first.kwargs["stage"], 2) self.assertEqual(first.kwargs["stage"], 2)
self.assertEqual(first.kwargs["stages"], 8) self.assertEqual(first.kwargs["stages"], 8)
@@ -525,146 +530,79 @@ class BootstrapRuntimeTest(unittest.TestCase):
self.assertEqual(first.kwargs["stages"], 8) self.assertEqual(first.kwargs["stages"], 8)
self.assertEqual(first.kwargs["stage_label"], "pytorch deps") self.assertEqual(first.kwargs["stage_label"], "pytorch deps")
def test_run_install_retries_once_with_mirror(self):
bootstrap = load_bootstrap_runtime()
previous = bootstrap._resolved_pypi_index
bootstrap._resolved_pypi_index = None
try:
with patch.object(
bootstrap,
"_run_install_once",
side_effect=[
subprocess.CalledProcessError(1, ["uv", "pip"]),
["numpy==1.26.4"],
],
) as mocked_install:
installed = bootstrap.run_install(
"status.json",
3,
8,
"ai deps",
["uv", "pip", "install", "x"],
)
finally:
bootstrap._resolved_pypi_index = previous
self.assertEqual(installed, ["numpy==1.26.4"])
self.assertEqual(mocked_install.call_count, 2)
self.assertEqual(
mocked_install.call_args_list[1].args[5],
{"PIP_INDEX_URL": bootstrap.PYPI_MIRRORS[0]},
)
def test_formats_elapsed_time(self): def test_formats_elapsed_time(self):
bootstrap = load_bootstrap_runtime() bootstrap = load_bootstrap_runtime()
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_preflight_rejects_low_disk_space(self): def test_resolve_indexes_uses_configured_pypi_index(self):
bootstrap = load_bootstrap_runtime() bootstrap = load_bootstrap_runtime()
args = SimpleNamespace( bootstrap.resolve_indexes(
backend="cu128", SimpleNamespace(pypi_index="https://pypi.tuna.tsinghua.edu.cn/simple")
runtime_dir=str(Path.cwd()),
) )
with patch.object(
bootstrap.shutil,
"disk_usage",
return_value=SimpleNamespace(free=1024**3),
), patch.object(bootstrap, "reachable", return_value=True):
with self.assertRaisesRegex(RuntimeError, "disk space"):
bootstrap.preflight(args)
def test_preflight_rejects_unreachable_network(self):
bootstrap = load_bootstrap_runtime()
args = SimpleNamespace(
backend="cu128",
runtime_dir=str(Path.cwd()),
)
with patch.object(
bootstrap.shutil,
"disk_usage",
return_value=SimpleNamespace(free=100 * 1024**3),
), patch.object(bootstrap, "reachable", return_value=False):
with self.assertRaisesRegex(RuntimeError, "Cannot reach"):
bootstrap.preflight(args)
def test_preflight_passes_when_checks_ok(self):
bootstrap = load_bootstrap_runtime()
args = SimpleNamespace(
backend="cu128",
runtime_dir=str(Path.cwd()),
)
with patch.object(
bootstrap.shutil,
"disk_usage",
return_value=SimpleNamespace(free=100 * 1024**3),
), patch.object(bootstrap, "reachable", return_value=True):
resolved, used_mirror = bootstrap.preflight(args)
self.assertFalse(used_mirror)
self.assertEqual( self.assertEqual(
resolved["pypi"], bootstrap._resolved_pypi_index,
"https://pypi.org/simple/",
)
def test_preflight_falls_back_to_mirrors(self):
bootstrap = load_bootstrap_runtime()
args = SimpleNamespace(
backend="cu128",
runtime_dir=str(Path.cwd()),
)
def fake_reachable(url, timeout=bootstrap.PREFLIGHT_TIMEOUT):
if "github.com" in url:
return True
if "tuna" in url or "hf-mirror" in url or "aliyun" in url:
return True
return False
with patch.object(
bootstrap.shutil,
"disk_usage",
return_value=SimpleNamespace(free=100 * 1024**3),
), patch.object(
bootstrap,
"reachable",
side_effect=fake_reachable,
):
resolved, used_mirror = bootstrap.preflight(args)
self.assertTrue(used_mirror)
self.assertEqual(
resolved["pypi"],
"https://pypi.tuna.tsinghua.edu.cn/simple", "https://pypi.tuna.tsinghua.edu.cn/simple",
) )
self.assertEqual(resolved["hf"], "https://hf-mirror.com")
self.assertEqual(
resolved["torch"],
"https://mirrors.aliyun.com/pytorch-wheels/cu128",
)
self.assertEqual(
bootstrap._resolved_torch_index,
"https://mirrors.aliyun.com/pytorch-wheels/cu128",
)
def test_preflight_uses_environment_overrides(self): def test_resolve_indexes_uses_environment_override(self):
bootstrap = load_bootstrap_runtime() bootstrap = load_bootstrap_runtime()
args = SimpleNamespace(
backend="cu128",
runtime_dir=str(Path.cwd()),
)
with patch.dict( with patch.dict(
os.environ, os.environ,
{ {"PIP_INDEX_URL": "https://pypi.example/simple"},
"PIP_INDEX_URL": "https://pypi.example/simple", ):
"HF_ENDPOINT": "https://hf.example", bootstrap.resolve_indexes(SimpleNamespace(pypi_index=None))
"SNAPFORM_TORCH_INDEX": "https://torch.example/whl/cu128",
},
), patch.object(
bootstrap.shutil,
"disk_usage",
return_value=SimpleNamespace(free=100 * 1024**3),
), patch.object(bootstrap, "reachable", return_value=True):
resolved, used_mirror = bootstrap.preflight(args)
self.assertFalse(used_mirror)
self.assertEqual(resolved["pypi"], "https://pypi.example/simple")
self.assertEqual(resolved["hf"], "https://hf.example")
self.assertEqual( self.assertEqual(
resolved["torch"], bootstrap._resolved_pypi_index,
"https://torch.example/whl/cu128", "https://pypi.example/simple",
) )
def test_child_environment_injects_resolved_mirrors(self): def test_child_environment_injects_resolved_pypi_index(self):
bootstrap = load_bootstrap_runtime() bootstrap = load_bootstrap_runtime()
previous_pypi = bootstrap._resolved_pypi_index previous_pypi = bootstrap._resolved_pypi_index
previous_hf = bootstrap._resolved_hf_endpoint
bootstrap._resolved_pypi_index = ( bootstrap._resolved_pypi_index = (
"https://pypi.tuna.tsinghua.edu.cn/simple" "https://pypi.tuna.tsinghua.edu.cn/simple"
) )
bootstrap._resolved_hf_endpoint = "https://hf-mirror.com"
try: try:
env = bootstrap.child_environment() env = bootstrap.child_environment()
finally: finally:
bootstrap._resolved_pypi_index = previous_pypi bootstrap._resolved_pypi_index = previous_pypi
bootstrap._resolved_hf_endpoint = previous_hf
self.assertEqual( self.assertEqual(
env["PIP_INDEX_URL"], env["PIP_INDEX_URL"],
"https://pypi.tuna.tsinghua.edu.cn/simple", "https://pypi.tuna.tsinghua.edu.cn/simple",
) )
self.assertEqual(env["HF_ENDPOINT"], "https://hf-mirror.com")
def test_child_environment_injects_registry_proxy(self): def test_child_environment_injects_registry_proxy(self):
bootstrap = load_bootstrap_runtime() bootstrap = load_bootstrap_runtime()
@@ -700,6 +638,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
root = Path(directory) root = Path(directory)
runtime_dir = root / "runtime" runtime_dir = root / "runtime"
runtime_dir.mkdir() runtime_dir.mkdir()
self.create_bundled_sources(root)
uv_path = root / "uv" uv_path = root / "uv"
uv_path.write_bytes(b"uv") uv_path.write_bytes(b"uv")
args = SimpleNamespace( args = SimpleNamespace(
@@ -711,7 +650,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
ben2_commit="test-ben2-commit", ben2_commit="test-ben2-commit",
status=str(root / "status.json"), status=str(root / "status.json"),
) )
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap, bootstrap,
"run", "run",
), patch.object( ), patch.object(
@@ -737,6 +676,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
root = Path(directory) root = Path(directory)
runtime_dir = root / "runtime" runtime_dir = root / "runtime"
runtime_dir.mkdir() runtime_dir.mkdir()
self.create_bundled_sources(root)
(runtime_dir / "partial-python.exe").write_text("junk") (runtime_dir / "partial-python.exe").write_text("junk")
args = SimpleNamespace( args = SimpleNamespace(
uv="uv", uv="uv",
@@ -747,7 +687,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
ben2_commit="test-ben2-commit", ben2_commit="test-ben2-commit",
status=str(root / "status.json"), status=str(root / "status.json"),
) )
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap, bootstrap,
"run", "run",
), patch.object( ), patch.object(
@@ -765,6 +705,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
root = Path(directory) root = Path(directory)
runtime_dir = root / "runtime" runtime_dir = root / "runtime"
runtime_dir.mkdir(parents=True) runtime_dir.mkdir(parents=True)
self.create_bundled_sources(root)
(runtime_dir / "Scripts").mkdir() (runtime_dir / "Scripts").mkdir()
(runtime_dir / "Scripts" / "python.exe").write_text("python") (runtime_dir / "Scripts" / "python.exe").write_text("python")
(runtime_dir / "snapform_ready.json").write_text( (runtime_dir / "snapform_ready.json").write_text(
@@ -781,7 +722,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
ben2_commit="test-ben2-commit", ben2_commit="test-ben2-commit",
status=str(root / "status.json"), status=str(root / "status.json"),
) )
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap, bootstrap,
"run", "run",
), patch.object( ), patch.object(
@@ -803,6 +744,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
root = Path(directory) root = Path(directory)
runtime_dir = root / "runtime" runtime_dir = root / "runtime"
runtime_dir.mkdir() runtime_dir.mkdir()
self.create_bundled_sources(root)
args = SimpleNamespace( args = SimpleNamespace(
uv="uv", uv="uv",
runtime_dir=str(runtime_dir), runtime_dir=str(runtime_dir),
@@ -812,7 +754,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
ben2_commit="test-ben2-commit", ben2_commit="test-ben2-commit",
status=str(root / "status.json"), status=str(root / "status.json"),
) )
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap, bootstrap,
"run", "run",
side_effect=fake_run, side_effect=fake_run,
@@ -826,12 +768,13 @@ class BootstrapRuntimeTest(unittest.TestCase):
self.assertEqual(status["progress"], 1.0) self.assertEqual(status["progress"], 1.0)
self.assertIn("Runtime is ready", status["message"]) self.assertIn("Runtime is ready", status["message"])
def test_install_downloads_missing_uv(self): def test_install_fails_when_bundled_uv_missing(self):
bootstrap = load_bootstrap_runtime() bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory: with tempfile.TemporaryDirectory() as directory:
root = Path(directory) root = Path(directory)
runtime_dir = root / "runtime" runtime_dir = root / "runtime"
runtime_dir.mkdir() runtime_dir.mkdir()
self.create_bundled_sources(root)
uv_path = root / "uv" / "uv.exe" uv_path = root / "uv" / "uv.exe"
args = SimpleNamespace( args = SimpleNamespace(
uv=str(uv_path), uv=str(uv_path),
@@ -842,7 +785,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
ben2_commit="test-ben2-commit", ben2_commit="test-ben2-commit",
status=str(root / "status.json"), status=str(root / "status.json"),
) )
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap, bootstrap,
"run", "run",
), patch.object( ), patch.object(
@@ -855,15 +798,101 @@ class BootstrapRuntimeTest(unittest.TestCase):
), patch.object( ), patch.object(
bootstrap.shutil, bootstrap.shutil,
"rmtree", "rmtree",
), patch.object(bootstrap, "download_uv") as mocked_download: ):
with self.assertRaisesRegex(
RuntimeError,
"Bundled uv is missing",
):
bootstrap.install(args)
def test_install_uses_bundled_sources_and_uv_without_network(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
runtime_dir = root / "runtime"
runtime_dir.mkdir()
bundled = root / "bundled"
bundled.mkdir()
(bundled / "depth-anything-3.tar.gz").write_bytes(b"da3")
(bundled / "ben2.tar.gz").write_bytes(b"ben2")
target, archive_name = bootstrap.uv_download_target()
archive = bundled / f"uv-{target}.{archive_name}"
if archive_name.endswith("zip"):
with zipfile.ZipFile(archive, "w") as bundle:
bundle.writestr("uv/uv.exe", b"uv-binary")
else:
with tarfile.open(archive, "w:gz") as tar:
content = b"uv-binary"
info = tarfile.TarInfo("uv")
info.size = len(content)
tar.addfile(info, io.BytesIO(content))
uv_path = root / "uv" / ("uv.exe" if os.name == "nt" else "uv")
args = SimpleNamespace(
uv=str(uv_path),
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
run_commands = []
install_commands = []
def fake_run(command, environment=None):
run_commands.append(command)
def fake_run_install(
status_path,
stage,
stages,
label,
command,
environment=None,
start=0.0,
end=1.0,
clock=None,
):
install_commands.append(command)
return []
with patch.object(
bootstrap,
"bundled_sources_directory",
return_value=bundled,
), patch.object(
bootstrap,
"resolve_indexes",
return_value=None,
), patch.object(
bootstrap,
"run",
side_effect=fake_run,
), patch.object(
bootstrap,
"run_install",
side_effect=fake_run_install,
):
bootstrap.install(args) bootstrap.install(args)
mocked_download.assert_called_once() self.assertEqual(uv_path.read_bytes(), b"uv-binary")
call_args = mocked_download.call_args all_commands = run_commands + install_commands
self.assertEqual(call_args.args[0], uv_path) self.assertFalse(
self.assertEqual(call_args.args[1], str(root / "status.json")) any(
self.assertEqual(call_args.kwargs["stage"], 1) "github.com" in str(value)
self.assertEqual(call_args.kwargs["stage_label"], "uv") for command in all_commands
for value in command
)
)
ai_command = next(
command
for command in install_commands
if any("depth-anything-3 @" in str(value) for value in command)
)
da3_uri = (root / "sources" / "depth-anything-3.tar.gz").as_uri()
ben2_uri = (root / "sources" / "ben2.tar.gz").as_uri()
self.assertIn(f"depth-anything-3 @ {da3_uri}", ai_command)
self.assertIn(f"ben2 @ {ben2_uri}", ai_command)
def test_install_reuses_existing_uv(self): def test_install_reuses_existing_uv(self):
bootstrap = load_bootstrap_runtime() bootstrap = load_bootstrap_runtime()
@@ -871,6 +900,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
root = Path(directory) root = Path(directory)
runtime_dir = root / "runtime" runtime_dir = root / "runtime"
runtime_dir.mkdir() runtime_dir.mkdir()
self.create_bundled_sources(root)
uv_path = root / "uv.exe" uv_path = root / "uv.exe"
uv_path.write_bytes(b"uv") uv_path.write_bytes(b"uv")
args = SimpleNamespace( args = SimpleNamespace(
@@ -882,17 +912,15 @@ class BootstrapRuntimeTest(unittest.TestCase):
ben2_commit="test-ben2-commit", ben2_commit="test-ben2-commit",
status=str(root / "status.json"), status=str(root / "status.json"),
) )
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
bootstrap, bootstrap,
"run", "run",
), patch.object( ), patch.object(
bootstrap, bootstrap,
"run_install", "run_install",
), patch.object(bootstrap, "download_uv") as mocked_download: ):
bootstrap.install(args) bootstrap.install(args)
mocked_download.assert_not_called()
def test_extract_uv_unpacks_standalone_binary(self): def test_extract_uv_unpacks_standalone_binary(self):
bootstrap = load_bootstrap_runtime() bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory: with tempfile.TemporaryDirectory() as directory:
@@ -906,6 +934,112 @@ class BootstrapRuntimeTest(unittest.TestCase):
self.assertEqual(uv_path.read_bytes(), b"binary") self.assertEqual(uv_path.read_bytes(), b"binary")
def test_sync_bundled_sources_copies_missing_archives(self):
bootstrap = load_bootstrap_runtime()
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")
sources = root / "sources"
bootstrap.sync_bundled_sources(bundled, sources)
self.assertEqual(
(sources / "depth-anything-3.tar.gz").read_bytes(),
b"da3",
)
self.assertEqual(
(sources / "ben2.tar.gz").read_bytes(),
b"ben2",
)
def test_install_stages_uses_local_source_archives(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
runtime_dir = root / "runtime"
runtime_dir.mkdir()
sources = root / "sources"
sources.mkdir()
(sources / "depth-anything-3.tar.gz").write_bytes(b"da3")
(sources / "ben2.tar.gz").write_bytes(b"ben2")
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
stages = bootstrap.install_stages(
args,
runtime_dir,
root / "python",
uv_path,
)
ai_command = next(
command
for _label, command, _environment, kind in stages
if kind == "install"
and any("depth-anything-3" in str(value) for value in command)
)
da3_local = (sources / "depth-anything-3.tar.gz").as_uri()
ben2_local = (sources / "ben2.tar.gz").as_uri()
self.assertIn(f"depth-anything-3 @ {da3_local}", ai_command)
self.assertIn(f"ben2 @ {ben2_local}", ai_command)
def test_install_uv_extracts_bundled_archive(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
uv_path = root / "uv"
bundled = root / "bundled"
bundled.mkdir(parents=True)
target, archive_name = bootstrap.uv_download_target()
archive = bundled / f"uv-{target}.{archive_name}"
if archive_name.endswith("zip"):
with zipfile.ZipFile(archive, "w") as bundle:
bundle.writestr("uv/uv.exe", b"uv-binary")
else:
with tarfile.open(archive, "w:gz") as tar:
content = b"uv-binary"
info = tarfile.TarInfo("uv")
info.size = len(content)
tar.addfile(info, io.BytesIO(content))
with patch.object(
bootstrap,
"bundled_sources_directory",
return_value=bundled,
):
bootstrap.install_uv(uv_path, str(root / "status.json"))
self.assertEqual(uv_path.read_bytes(), b"uv-binary")
def test_install_uv_raises_when_bundled_missing(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
uv_path = root / "uv"
bundled = root / "bundled"
bundled.mkdir()
with patch.object(
bootstrap,
"bundled_sources_directory",
return_value=bundled,
):
with self.assertRaisesRegex(
RuntimeError,
"Bundled uv is missing",
):
bootstrap.install_uv(uv_path, str(root / "status.json"))
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()

View File

@@ -143,49 +143,29 @@ class DownloadProgressTest(unittest.TestCase):
self.assertFalse(status_path.exists()) self.assertFalse(status_path.exists())
def test_resolve_hf_endpoint_keeps_direct_endpoint(self): def test_download_endpoints_keeps_direct_endpoint(self):
with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( with patch.dict(os.environ, {"HF_ENDPOINT": ""}):
self.worker, endpoints = self.worker.download_endpoints()
"reachable",
return_value=True,
):
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
self.assertEqual(endpoint, "https://huggingface.co/") self.assertEqual(
self.assertFalse(used_mirror) endpoints,
["https://huggingface.co"],
)
def test_resolve_hf_endpoint_falls_back_to_mirror(self): def test_download_endpoints_respects_environment_override(self):
with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example"}):
self.worker, endpoints = self.worker.download_endpoints()
"reachable",
side_effect=lambda url, timeout=None: "hf-mirror" in url,
):
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
self.assertEqual(endpoint, "https://hf-mirror.com") self.assertEqual(
self.assertTrue(used_mirror) endpoints,
["https://hf.example"],
)
def test_resolve_hf_endpoint_respects_environment_override(self): def test_download_endpoints_strips_trailing_slash(self):
with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example"}), patch.object( with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example/"}):
self.worker, endpoints = self.worker.download_endpoints()
"reachable",
) as mocked_reachable:
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
self.assertEqual(endpoint, "https://hf.example") self.assertEqual(endpoints, ["https://hf.example"])
self.assertFalse(used_mirror)
mocked_reachable.assert_not_called()
def test_resolve_hf_endpoint_keeps_direct_when_mirrors_unreachable(self):
with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker,
"reachable",
return_value=False,
):
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
self.assertEqual(endpoint, "https://huggingface.co/")
self.assertFalse(used_mirror)
if __name__ == "__main__": if __name__ == "__main__":

View File

@@ -1,5 +1,4 @@
import importlib.util import importlib.util
import inspect
import sys import sys
import tempfile import tempfile
import unittest import unittest
@@ -60,6 +59,22 @@ class ImageDepthTest(unittest.TestCase):
np.array(((1.0, 2.0), (3.0, 4.0)), dtype=np.float32), np.array(((1.0, 2.0), (3.0, 4.0)), dtype=np.float32),
) )
def test_extract_depth_arrays_returns_every_frame(self):
with tempfile.TemporaryDirectory() as directory:
npz_path = Path(directory) / "results.npz"
np.savez_compressed(
npz_path,
depth=np.arange(12, dtype=np.float32).reshape(3, 2, 2),
)
depth_frames = self.image_depth.extract_depth_arrays(npz_path)
self.assertEqual(len(depth_frames), 3)
np.testing.assert_array_equal(
depth_frames[1],
np.array(((4.0, 5.0), (6.0, 7.0)), dtype=np.float32),
)
def test_accepts_depths_key_with_single_channel(self): def test_accepts_depths_key_with_single_channel(self):
with tempfile.TemporaryDirectory() as directory: with tempfile.TemporaryDirectory() as directory:
npz_path = Path(directory) / "results.npz" npz_path = Path(directory) / "results.npz"
@@ -72,6 +87,261 @@ class ImageDepthTest(unittest.TestCase):
self.assertEqual(depth.shape, (2, 3)) self.assertEqual(depth.shape, (2, 3))
def test_is_animated_image_accepts_movie_and_sequence(self):
self.assertTrue(
self.image_depth.is_animated_image(
SimpleNamespace(source="MOVIE")
)
)
self.assertTrue(
self.image_depth.is_animated_image(
SimpleNamespace(source="SEQUENCE")
)
)
self.assertFalse(
self.image_depth.is_animated_image(
SimpleNamespace(source="FILE")
)
)
def test_video_total_frames_reads_movie_duration(self):
image = SimpleNamespace(source="MOVIE", frame_duration=97)
self.assertEqual(self.image_depth.video_total_frames(image), 97)
def test_video_total_frames_counts_sequence_files(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
(root / "foreground-0001.png").touch()
(root / "foreground-0002.png").touch()
image = SimpleNamespace(
source="SEQUENCE",
filepath=str(root / "foreground-0001.png"),
)
self.assertEqual(self.image_depth.video_total_frames(image), 2)
def test_video_total_frames_returns_one_for_still_image(self):
image = SimpleNamespace(source="FILE", frame_duration=1)
self.assertEqual(self.image_depth.video_total_frames(image), 1)
def test_default_extract_frames_uses_empty_image_frames(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="MOVIE", frame_duration=97),
image_user=SimpleNamespace(frame_duration=24),
)
self.assertEqual(self.image_depth.default_extract_frames(source_object), 24)
def test_default_extract_frames_caps_at_video_total(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="MOVIE", frame_duration=97),
image_user=SimpleNamespace(frame_duration=200),
)
self.assertEqual(self.image_depth.default_extract_frames(source_object), 97)
def test_default_extract_frames_falls_back_to_video_total(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="MOVIE", frame_duration=97),
)
self.assertEqual(self.image_depth.default_extract_frames(source_object), 97)
def test_extract_frame_range_honors_offset_within_clip(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="MOVIE", frame_duration=97),
image_user=SimpleNamespace(
frame_start=4,
frame_offset=70,
frame_duration=24,
),
)
self.assertEqual(
self.image_depth.extract_frame_range(source_object),
(71, 24),
)
def test_extract_frame_range_clamps_large_offset_to_clip_tail(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="MOVIE", frame_duration=97),
image_user=SimpleNamespace(
frame_start=1,
frame_offset=90,
frame_duration=24,
),
)
self.assertEqual(
self.image_depth.extract_frame_range(source_object),
(74, 24),
)
def test_extract_frame_range_clamps_negative_offset_to_start(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="MOVIE", frame_duration=97),
image_user=SimpleNamespace(
frame_start=1,
frame_offset=-5,
frame_duration=24,
),
)
self.assertEqual(
self.image_depth.extract_frame_range(source_object),
(1, 24),
)
def test_extract_frame_range_returns_full_clip_for_still(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="FILE", frame_duration=1),
image_user=SimpleNamespace(
frame_start=1,
frame_offset=50,
frame_duration=24,
),
)
self.assertEqual(
self.image_depth.extract_frame_range(source_object),
(1, 1),
)
def test_configure_image_user_preserves_cyclic_setting(self):
image_user = SimpleNamespace(
frame_start=0,
frame_offset=5,
use_auto_refresh=False,
use_cyclic=False,
frame_duration=1,
)
self.image_depth.configure_image_user(
image_user,
97,
frame_start=4,
frame_offset=70,
)
self.assertFalse(image_user.use_cyclic)
self.assertTrue(image_user.use_auto_refresh)
self.assertEqual(image_user.frame_start, 4)
self.assertEqual(image_user.frame_offset, 70)
self.assertEqual(image_user.frame_duration, 97)
self.image_depth.configure_image_user(
image_user,
use_cyclic=True,
)
self.assertTrue(image_user.use_cyclic)
def test_image_user_settings_reads_empty_values(self):
source_object = SimpleNamespace(
image_user=SimpleNamespace(
frame_start=4,
frame_offset=70,
frame_duration=12,
)
)
self.assertEqual(
self.image_depth.image_user_settings(source_object),
(4, 70, 12, False),
)
source_object.image_user.use_cyclic = True
self.assertEqual(
self.image_depth.image_user_settings(source_object),
(4, 70, 12, True),
)
def test_image_user_settings_defaults_without_user(self):
self.assertEqual(
self.image_depth.image_user_settings(
SimpleNamespace(image_user=None)
),
(1, 0, 0, False),
)
def test_image_base_name_strips_video_extension(self):
image = SimpleNamespace(name="微信视频2026-07-23_192756_414.mp4")
self.assertEqual(
self.image_depth.image_base_name(image),
"微信视频2026-07-23_192756_414",
)
def test_removed_background_name_strips_video_extension(self):
source_object = SimpleNamespace(data=SimpleNamespace(name="clip.mp4"))
self.assertEqual(
self.image_depth.removed_background_data_name(source_object),
"clip.foreground",
)
def test_force_sequence_source_keeps_loaded_sequence(self):
image = SimpleNamespace(source="SEQUENCE")
result = self.image_depth.force_sequence_source(image)
self.assertIs(result, image)
self.assertEqual(image.source, "SEQUENCE")
def test_force_sequence_source_switches_single_image(self):
image = SimpleNamespace(source="FILE")
result = self.image_depth.force_sequence_source(image)
self.assertIs(result, image)
self.assertEqual(image.source, "SEQUENCE")
def test_force_sequence_source_reports_unloadable_sequence(self):
class BlockedImage:
@property
def source(self):
return "FILE"
@source.setter
def source(self, value):
raise RuntimeError("blocked")
with self.assertRaisesRegex(RuntimeError, "cannot load"):
self.image_depth.force_sequence_source(BlockedImage())
def test_collect_sequence_frames_returns_numbered_files_in_order(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
(root / "frame-0001.jpg").touch()
(root / "frame-0002.jpg").touch()
(root / "frame-0003.jpg").touch()
(root / "notes.txt").touch()
image = SimpleNamespace(
source="SEQUENCE",
filepath=str(root / "frame-0002.jpg"),
)
frames = self.image_depth.collect_sequence_frames(image)
self.assertEqual(
[path.name for path in frames],
["frame-0001.jpg", "frame-0002.jpg", "frame-0003.jpg"],
)
def test_sorted_frame_paths_uses_numeric_frame_order(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
(root / "frame-2.jpg").touch()
(root / "frame-10.jpg").touch()
(root / "frame-1.jpg").touch()
frames = self.image_depth.sorted_frame_paths(root)
self.assertEqual(
[path.name for path in frames],
["frame-1.jpg", "frame-2.jpg", "frame-10.jpg"],
)
def test_accepts_single_depth_map_with_channel_axis(self): def test_accepts_single_depth_map_with_channel_axis(self):
with tempfile.TemporaryDirectory() as directory: with tempfile.TemporaryDirectory() as directory:
npz_path = Path(directory) / "results.npz" npz_path = Path(directory) / "results.npz"
@@ -205,16 +475,6 @@ class ImageDepthTest(unittest.TestCase):
(12, 8), (12, 8),
) )
def test_modifiers_use_native_names_and_fixed_levels(self):
source = inspect.getsource(self.image_depth.build_depth_plane)
self.assertIn('modifiers.new("Subdivision", "SUBSURF")', source)
self.assertIn('modifiers.new("Displace", "DISPLACE")', source)
self.assertNotIn("DA3 Subdivision", source)
self.assertNotIn("DA3 Displace", source)
self.assertIn("subdivision.levels = 1", source)
self.assertIn("subdivision.render_levels = 2", source)
def test_plane_fit_transform_matches_source_size_and_offset(self): def test_plane_fit_transform_matches_source_size_and_offset(self):
source = SimpleNamespace( source = SimpleNamespace(
data=SimpleNamespace(size=(800, 400)), data=SimpleNamespace(size=(800, 400)),
@@ -233,37 +493,6 @@ class ImageDepthTest(unittest.TestCase):
): ):
self.assertAlmostEqual(value, expected) self.assertAlmostEqual(value, expected)
def test_source_empty_is_only_removed_when_requested(self):
source = inspect.getsource(self.image_depth.finish_depth_plane)
self.assertIn("source_name = source_object.name", source)
self.assertIn("if delete_source:", source)
self.assertIn(
"bpy.data.objects.remove(source_object, do_unlink=True)",
source,
)
self.assertNotIn("source_object.name =", source)
self.assertIn(
"bpy.data.objects.remove(source_object, do_unlink=True)\n"
" plane_object.name = source_name",
source,
)
def test_original_depth_mode_uses_uv_and_zero_midlevel(self):
source = inspect.getsource(self.image_depth.build_image_depth_plane)
self.assertIn("texture.image = source_object.data", source)
self.assertIn('modifiers.new("Displace", "DISPLACE")', source)
self.assertIn("displace.texture = texture", source)
self.assertIn('displace.texture_coords = "UV"', source)
self.assertNotIn("displace.direction", source)
self.assertNotIn("displace.strength", source)
self.assertIn(
"displace.mid_level = ORIGINAL_DISPLACEMENT_MIDLEVEL",
source,
)
self.assertIn("plane_object.matrix_world =", source)
def test_process_resolution_scales_the_image_long_edge(self): def test_process_resolution_scales_the_image_long_edge(self):
image = SimpleNamespace(size=(1080, 1920)) image = SimpleNamespace(size=(1080, 1920))
@@ -272,11 +501,6 @@ class ImageDepthTest(unittest.TestCase):
960, 960,
) )
def test_mesh_uses_source_object_name(self):
source = inspect.getsource(self.image_depth.create_plane_mesh)
self.assertIn("bpy.data.meshes.new(source_object.name)", source)
def test_depth_data_uses_source_image_name(self): def test_depth_data_uses_source_image_name(self):
source_object = SimpleNamespace( source_object = SimpleNamespace(
name="obj_img", name="obj_img",
@@ -292,17 +516,20 @@ class ImageDepthTest(unittest.TestCase):
self.image_depth.depth_data_name(source_object), self.image_depth.depth_data_name(source_object),
"img.001.depth", "img.001.depth",
) )
estimate_source = inspect.getsource(
self.image_depth.build_depth_plane def test_copy_depth_image_renames_still_source(self):
) class FakeImage:
original_source = inspect.getsource( def __init__(self, name):
self.image_depth.build_image_depth_plane self.name = name
)
self.assertIn( def copy(self):
"depth_name = depth_name or depth_data_name(source_object)", return FakeImage("img.png.001")
estimate_source,
) source_object = SimpleNamespace(data=FakeImage("img.png"))
self.assertIn("depth_data_name(source_object)", original_source)
depth_image = self.image_depth.copy_depth_image(source_object)
self.assertEqual(depth_image.name, "img.depth")
def test_removed_background_name_uses_source_image_base_name(self): def test_removed_background_name_uses_source_image_base_name(self):
source_object = SimpleNamespace( source_object = SimpleNamespace(
@@ -311,213 +538,14 @@ class ImageDepthTest(unittest.TestCase):
self.assertEqual( self.assertEqual(
self.image_depth.removed_background_data_name(source_object), self.image_depth.removed_background_data_name(source_object),
"img.rb", "img.foreground",
) )
source_object.data.name = "img.png.001" source_object.data.name = "img.png.001"
self.assertEqual( self.assertEqual(
self.image_depth.removed_background_data_name(source_object), self.image_depth.removed_background_data_name(source_object),
"img.001.rb", "img.001.foreground",
) )
def test_depth_name_override_keeps_original_base_name(self):
source = inspect.getsource(self.image_depth.build_depth_plane)
self.assertIn(
"depth_name = depth_name or depth_data_name(source_object)",
source,
)
self.assertIn(
"color_image if color_image is not None else source_object.data",
source,
)
def test_original_depth_background_only_changes_material_color(self):
source = inspect.getsource(self.image_depth.build_image_depth_plane)
self.assertIn("texture.image = source_object.data", source)
self.assertIn(
"source_object.data if displacement_via == \"MATERIAL\" else None",
source,
)
self.assertIn(
"color_image if color_image is not None else source_object.data",
source,
)
def test_material_uses_source_object_name(self):
da3_source = inspect.getsource(self.image_depth.build_depth_plane)
image_source = inspect.getsource(
self.image_depth.build_image_depth_plane
)
self.assertIn(
"create_material(\n source_name,",
da3_source,
)
self.assertIn("create_material(\n source_name,", image_source)
self.assertNotIn("DA3DepthMaterial", da3_source)
def test_displace_texture_uses_extend(self):
source = inspect.getsource(
self.image_depth.build_depth_plane
)
self.assertIn(
"depth_name = depth_name or depth_data_name(source_object)",
source,
)
self.assertIn("create_depth_image(\n depth_name,", source)
self.assertIn("bpy.data.textures.new(\n depth_name,", source)
self.assertIn('texture.extension = "EXTEND"', source)
def test_material_supports_bump_and_optional_displacement(self):
source = inspect.getsource(self.image_depth.create_material)
self.assertIn('nodes.new("ShaderNodeBump")', source)
self.assertIn(
'color_texture.outputs["Color"], bump.inputs["Height"]',
source,
)
self.assertIn(
'bump.outputs["Normal"], shader.inputs["Normal"]',
source,
)
self.assertIn(
'bump.inputs["Distance"].default_value = '
"BUMP_BASE_DISTANCE * depth_scale",
source,
)
self.assertIn('nodes.new("ShaderNodeFloatCurve")', source)
self.assertIn(
'roughness_curve.inputs["Value"]',
source,
)
self.assertIn(
'shader.inputs["Roughness"]',
source,
)
self.assertIn(
'curve_points[0].location = (0.0, 0.5)',
source,
)
self.assertIn(
'curve_points[-1].location = (1.0, 0.1)',
source,
)
self.assertNotIn('shader.inputs["IOR"]', source)
self.assertIn(
'roughness_curve.location = (-300.0, 320.0)',
source,
)
self.assertIn('bump.location = (-300.0, -80.0)', source)
self.assertIn('nodes.new("ShaderNodeDisplacement")', source)
self.assertIn(
'depth_texture.outputs["Color"],\n'
' displacement.inputs["Height"]',
source,
)
self.assertIn(
'displacement.outputs["Displacement"],\n'
' output.inputs["Displacement"]',
source,
)
self.assertIn(
'displacement.inputs["Midlevel"].default_value',
source,
)
self.assertIn(
'displacement.inputs["Scale"].default_value',
source,
)
self.assertIn("set_displacement_only(material)", source)
def test_unshaded_material_uses_alpha_mixed_emission(self):
source = inspect.getsource(self.image_depth.create_material)
unshaded_start = source.index("if unshaded:")
shaded_start = source.index("else:", unshaded_start)
displacement_start = source.index(
"if displacement_image is not None:",
shaded_start,
)
unshaded_source = source[unshaded_start:shaded_start]
shaded_source = source[shaded_start:displacement_start]
self.assertIn('nodes.new("ShaderNodeEmission")', unshaded_source)
self.assertIn(
'nodes.new("ShaderNodeBsdfTransparent")',
unshaded_source,
)
self.assertIn('nodes.new("ShaderNodeMixShader")', unshaded_source)
self.assertIn(
'color_texture.outputs["Color"]',
unshaded_source,
)
self.assertIn('emission.inputs["Color"]', unshaded_source)
self.assertIn('color_texture.outputs["Alpha"]', unshaded_source)
self.assertIn('mix_shader.inputs["Fac"]', unshaded_source)
self.assertIn('transparent.outputs["BSDF"]', unshaded_source)
self.assertIn('mix_shader.inputs[1]', unshaded_source)
self.assertIn('emission.outputs["Emission"]', unshaded_source)
self.assertIn('mix_shader.inputs[2]', unshaded_source)
self.assertIn('mix_shader.outputs["Shader"]', unshaded_source)
self.assertNotIn("ShaderNodeBsdfPrincipled", unshaded_source)
self.assertNotIn("ShaderNodeBump", unshaded_source)
self.assertNotIn("ShaderNodeFloatCurve", unshaded_source)
self.assertIn("ShaderNodeBsdfPrincipled", shaded_source)
self.assertGreater(displacement_start, unshaded_start)
def test_displacement_values_follow_depth_source(self):
estimate_source = inspect.getsource(
self.image_depth.build_depth_plane
)
original_source = inspect.getsource(
self.image_depth.build_image_depth_plane
)
self.assertIn(
"displacement_midlevel=DISPLACEMENT_MIDLEVEL",
estimate_source,
)
self.assertIn(
"displacement_scale=DISPLACEMENT_STRENGTH",
estimate_source,
)
self.assertIn(
"displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL",
original_source,
)
self.assertIn(
"displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH",
original_source,
)
self.assertEqual(self.image_depth.DISPLACEMENT_MIDLEVEL, 1.0)
self.assertEqual(self.image_depth.DISPLACEMENT_STRENGTH, -1.0)
self.assertEqual(
self.image_depth.ORIGINAL_DISPLACEMENT_MIDLEVEL,
0.0,
)
self.assertEqual(
self.image_depth.ORIGINAL_DISPLACEMENT_STRENGTH,
1.0,
)
def test_material_displacement_mode_skips_displace_modifier(self):
estimate_source = inspect.getsource(
self.image_depth.build_depth_plane
)
original_source = inspect.getsource(
self.image_depth.build_image_depth_plane
)
for source in (estimate_source, original_source):
self.assertIn('if displacement_via == "MODIFIER":', source)
self.assertIn(
'if displacement_via == "MATERIAL" else None',
source,
)
self.assertIn("unshaded=unshaded", source)
def test_surface_displacement_is_set_to_displacement_only(self): def test_surface_displacement_is_set_to_displacement_only(self):
class Material: class Material:
surface_displacement_method = "BUMP" surface_displacement_method = "BUMP"
@@ -530,12 +558,6 @@ class ImageDepthTest(unittest.TestCase):
"DISPLACEMENT", "DISPLACEMENT",
) )
def test_plane_transform_scales_depth_axis(self):
source = inspect.getsource(self.image_depth.apply_plane_transform)
self.assertIn("(0.0, 0.0, scale_z, 0.0)", source)
self.assertIn("plane_object.scale.z", source)
def test_plane_bounds_keep_image_aspect_and_empty_offset(self): def test_plane_bounds_keep_image_aspect_and_empty_offset(self):
source = SimpleNamespace( source = SimpleNamespace(
data=SimpleNamespace(size=(800, 400)), data=SimpleNamespace(size=(800, 400)),

View File

@@ -1,4 +1,7 @@
import ast
import importlib.util import importlib.util
import io
import tarfile
import tempfile import tempfile
import unittest import unittest
from pathlib import Path from pathlib import Path
@@ -19,25 +22,148 @@ def load_pack():
class PackTest(unittest.TestCase): class PackTest(unittest.TestCase):
def test_builds_extension_root_layout(self): def test_builds_platform_specific_extension_zips(self):
pack = load_pack() pack = load_pack()
with tempfile.TemporaryDirectory() as directory: with tempfile.TemporaryDirectory() as directory:
original_dist = pack.DIST_DIR original_dist = pack.DIST_DIR
pack.DIST_DIR = Path(directory) pack.DIST_DIR = Path(directory)
try: try:
output_path = pack.build_extension() output_paths = pack.build_all()
finally: finally:
pack.DIST_DIR = original_dist pack.DIST_DIR = original_dist
self.assertEqual(output_path.name, "SnapForm.v0.9.32.zip") self.assertEqual(
with ZipFile(output_path) as archive: {path.name for path in output_paths},
names = set(archive.namelist()) {
"SnapForm.v0.10.10.windows-x64.zip",
"SnapForm.v0.10.10.macos-x64.zip",
"SnapForm.v0.10.10.macos-arm64.zip",
"SnapForm.v0.10.10.linux-x64.zip",
},
)
for output_path in output_paths:
platform = output_path.stem.split(".")[-1]
with ZipFile(output_path) as archive:
names = set(archive.namelist())
manifest = archive.read("blender_manifest.toml").decode(
"utf-8"
)
self.assertIn("__init__.py", names) self.assertIn("__init__.py", names)
self.assertIn("blender_manifest.toml", names) self.assertIn("blender_manifest.toml", names)
self.assertIn("image_depth.py", names) self.assertIn("image_depth.py", names)
self.assertNotIn("snapform_depth_plane.py", names) self.assertIn("worker.py", names)
self.assertIn("worker.py", names) self.assertIn("ai_host.py", names)
self.assertIn("ai_host.py", names) self.assertNotIn("snapform/__init__.py", names)
self.assertNotIn("importer.py", names) self.assertNotIn("importer.py", names)
self.assertNotIn("snapform/__init__.py", names) self.assertIn("sources/depth-anything-3.tar.gz", names)
self.assertIn("sources/ben2.tar.gz", names)
self.assertIn(
f"sources/{pack.PLATFORMS[platform]}",
names,
)
self.assertIn(f'platforms = ["{platform}"]', manifest)
self.assertNotIn(
"windows-x64"
if platform != "windows-x64"
else "macos-arm64",
manifest.split("platforms =", 1)[-1],
)
def test_rejects_unknown_platform(self):
pack = load_pack()
with self.assertRaisesRegex(ValueError, "Unsupported platform"):
pack.build_extension("windows-arm64")
def test_bundled_depth_anything_3_has_no_xformers_dependency(self):
pack = load_pack()
with tarfile.open(
fileobj=io.BytesIO(
pack.bundled_source_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_source_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.SOURCES_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_source_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.SOURCES_DIR / "depth-anything-3.tar.gz"
with tarfile.open(archive_path, "r:gz") as archive:
source = archive.extractfile(
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
"src/depth_anything_3/utils/ray_utils.py"
).read().decode("utf-8")
self.assertIn(
"Q_tilde, R_tilde = torch.linalg.qr(A_tilde)",
source,
)
self.assertNotIn("use_cpu", source)
if __name__ == "__main__":
unittest.main()

View File

@@ -1,5 +1,4 @@
import importlib.util import importlib.util
import inspect
import json import json
import os import os
import sys import sys
@@ -42,27 +41,11 @@ class WorkerTest(unittest.TestCase):
self.assertEqual([path.name for path in images], ["a.jpg", "b.PNG"]) self.assertEqual([path.name for path in images], ["a.jpg", "b.PNG"])
def test_reconstruct_reports_status_before_loading_torch(self):
source = inspect.getsource(self.worker.reconstruct)
self.assertLess(
source.index('"Starting AI runtime"'),
source.index("import torch"),
)
def test_ai_tasks_support_cooperative_cancellation(self): def test_ai_tasks_support_cooperative_cancellation(self):
with self.assertRaises(self.worker.TaskCancelled): with self.assertRaises(self.worker.TaskCancelled):
self.worker.check_cancelled(lambda: True) self.worker.check_cancelled(lambda: True)
self.worker.check_cancelled(lambda: False) self.worker.check_cancelled(lambda: False)
self.assertIn(
"cancel_check=cancel_check",
inspect.getsource(self.worker.reconstruct),
)
self.assertIn(
"cancel_check=cancel_check",
inspect.getsource(self.worker.remove_background),
)
def test_select_device_supports_mps_and_auto_fallback(self): def test_select_device_supports_mps_and_auto_fallback(self):
mps_torch = SimpleNamespace( mps_torch = SimpleNamespace(
@@ -115,13 +98,7 @@ class WorkerTest(unittest.TestCase):
with self.assertRaisesRegex(RuntimeError, "MPS was requested"): with self.assertRaisesRegex(RuntimeError, "MPS was requested"):
self.worker.select_device("mps", cpu_only_torch) self.worker.select_device("mps", cpu_only_torch)
def test_optional_background_removal_runs_before_da3_model_load(self): def test_reconstruct_parser_accepts_background_model_directory(self):
source = inspect.getsource(self.worker.reconstruct)
self.assertLess(
source.index("infer_foreground("),
source.index("DepthAnything3.from_pretrained"),
)
args = self.worker.parser().parse_args( args = self.worker.parser().parse_args(
[ [
"reconstruct", "reconstruct",
@@ -175,6 +152,199 @@ class WorkerTest(unittest.TestCase):
self.assertEqual(images, [image_path]) self.assertEqual(images, [image_path])
self.assertIsNone(temporary_directory) self.assertIsNone(temporary_directory)
def test_prepares_a_frame_directory(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
(root / "frame-0001.jpg").touch()
(root / "frame-0002.png").touch()
(root / "notes.txt").touch()
images, temporary_directory = self.worker.prepare_input(
root,
16,
directory,
)
self.assertEqual(
[path.name for path in images],
["frame-0001.jpg", "frame-0002.png"],
)
self.assertIsNone(temporary_directory)
def test_prepares_first_frames_of_a_frame_directory(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
for index in range(1, 6):
(root / f"frame-{index:04d}.jpg").touch()
images, temporary_directory = self.worker.prepare_input(
root,
3,
directory,
)
self.assertEqual(
[path.name for path in images],
["frame-0001.jpg", "frame-0002.jpg", "frame-0003.jpg"],
)
self.assertIsNone(temporary_directory)
def test_prepares_first_frames_after_offset_in_a_frame_directory(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
for index in range(1, 6):
(root / f"frame-{index:04d}.jpg").touch()
images, temporary_directory = self.worker.prepare_input(
root,
3,
directory,
start_frame=2,
)
self.assertEqual(
[path.name for path in images],
["frame-0002.jpg", "frame-0003.jpg", "frame-0004.jpg"],
)
self.assertIsNone(temporary_directory)
def test_extract_video_frames_reads_frames_sequentially(self):
class FakeCapture:
def __init__(self):
self.position = 0
def isOpened(self):
return True
def read(self):
if self.position >= 100:
return False, None
frame = self.position
self.position += 1
return True, frame
def release(self):
return None
class FakeCv2:
def __init__(self):
self.captures = []
def VideoCapture(self, _path):
capture = FakeCapture()
self.captures.append(capture)
return capture
def imwrite(self, path, frame):
Path(path).write_text(str(frame), encoding="utf-8")
return True
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
fake_cv2 = FakeCv2()
image_paths = self.worker.extract_video_frames(
root / "clip.mkv",
root / "out",
10,
fake_cv2,
)
self.assertEqual(
[path.name for path in image_paths],
[f"frame.{index:04d}.jpg" for index in range(1, 11)],
)
self.assertEqual(len(fake_cv2.captures), 1)
self.assertEqual(fake_cv2.captures[0].position, 10)
def test_extract_video_frames_starts_after_offset(self):
class FakeCapture:
def __init__(self):
self.position = 0
def isOpened(self):
return True
def read(self):
if self.position >= 100:
return False, None
frame = self.position
self.position += 1
return True, frame
def release(self):
return None
class FakeCv2:
def __init__(self):
self.captures = []
def VideoCapture(self, _path):
capture = FakeCapture()
self.captures.append(capture)
return capture
def imwrite(self, path, frame):
Path(path).write_text(str(frame), encoding="utf-8")
return True
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
fake_cv2 = FakeCv2()
image_paths = self.worker.extract_video_frames(
root / "clip.mkv",
root / "out",
10,
fake_cv2,
start_frame=71,
)
extracted = [
Path(path).read_text(encoding="utf-8")
for path in image_paths
]
self.assertEqual(extracted, [str(frame) for frame in range(70, 80)])
self.assertEqual(fake_cv2.captures[0].position, 80)
def test_extract_video_frames_tolerates_short_video_tail(self):
class ShortCapture:
def __init__(self):
self.position = 0
def isOpened(self):
return True
def read(self):
if self.position >= 99:
return False, None
frame = self.position
self.position += 1
return True, frame
def release(self):
return None
class FakeCv2:
def VideoCapture(self, _path):
return ShortCapture()
def imwrite(self, path, frame):
Path(path).write_text(str(frame), encoding="utf-8")
return True
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
image_paths = self.worker.extract_video_frames(
root / "clip.mkv",
root / "out",
10,
FakeCv2(),
)
self.assertEqual(
[path.name for path in image_paths],
[f"frame.{index:04d}.jpg" for index in range(1, 11)],
)
def test_npz_is_complete_requires_readable_archive(self): def test_npz_is_complete_requires_readable_archive(self):
with tempfile.TemporaryDirectory() as directory: with tempfile.TemporaryDirectory() as directory:
root = Path(directory) root = Path(directory)
@@ -235,10 +405,6 @@ class WorkerTest(unittest.TestCase):
self.worker, self.worker,
"create_download_progress_class", "create_download_progress_class",
return_value=object, return_value=object,
), patch.object(
self.worker,
"resolve_hf_endpoint",
return_value=(self.worker.HF_ENDPOINT, False),
): ):
args.model_dir = str(Path(directory) / "model") args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json") args.status = str(Path(directory) / "status.json")
@@ -248,13 +414,17 @@ class WorkerTest(unittest.TestCase):
calls[0]["allow_patterns"], calls[0]["allow_patterns"],
["config.json", "model.safetensors"], ["config.json", "model.safetensors"],
) )
self.assertEqual(calls[0]["endpoint"], self.worker.HF_ENDPOINT)
def test_download_model_switches_to_mirror_before_hf_import(self): def test_download_model_falls_back_to_modelscope_when_direct_fails(self):
calls = [] calls = []
fake_hugging_face = ModuleType("huggingface_hub") fake_hugging_face = ModuleType("huggingface_hub")
fake_hugging_face.snapshot_download = lambda **options: calls.append(
options def fake_snapshot_download(**options):
) calls.append(options)
raise OSError("Connection reset by peer")
fake_hugging_face.snapshot_download = fake_snapshot_download
args = SimpleNamespace( args = SimpleNamespace(
repository="PramaLLC/BEN2", repository="PramaLLC/BEN2",
model_dir="model", model_dir="model",
@@ -269,23 +439,151 @@ class WorkerTest(unittest.TestCase):
self.worker, self.worker,
"create_download_progress_class", "create_download_progress_class",
return_value=object, return_value=object,
), patch.object(
self.worker,
"resolve_hf_endpoint",
return_value=("https://hf-mirror.com", True),
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker, self.worker,
"write_status", "download_from_modelscope",
) as mocked_write: ) as mocked_modelscope:
def fake_download_from_modelscope(*args, **kwargs):
model_dir = Path(args[1])
model_dir.mkdir(parents=True, exist_ok=True)
(model_dir / "config.json").write_text("{}")
(model_dir / "model.safetensors").write_bytes(b"x")
mocked_modelscope.side_effect = fake_download_from_modelscope
args.model_dir = str(Path(directory) / "model") args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json") args.status = str(Path(directory) / "status.json")
self.worker.download_model(args) self.worker.download_model(args)
self.assertEqual(
os.environ["HF_ENDPOINT"], self.assertEqual(len(calls), 1)
"https://hf-mirror.com", self.assertEqual(calls[0]["endpoint"], self.worker.HF_ENDPOINT)
) self.assertEqual(mocked_modelscope.call_count, 1)
messages = [call.args[2] for call in mocked_write.call_args_list]
self.assertTrue(any("(mirror)" in message for message in messages)) def test_download_model_skips_unreachable_endpoint(self):
calls = []
fake_hugging_face = ModuleType("huggingface_hub")
fake_hugging_face.snapshot_download = lambda **options: calls.append(
options
)
args = SimpleNamespace(
repository="depth-anything/DA3-SMALL",
model_dir="model",
status="status.json",
allow_pattern=[],
)
with tempfile.TemporaryDirectory() as directory, patch.dict(
sys.modules,
{"huggingface_hub": fake_hugging_face},
), patch.object(
self.worker,
"create_download_progress_class",
return_value=object,
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker,
"hf_endpoint_reachable",
return_value=False,
), patch.object(
self.worker,
"download_from_modelscope",
) as mocked_modelscope, patch.object(
self.worker,
"_verify_model_complete",
return_value=None,
):
args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json")
self.worker.download_model(args)
self.assertEqual(calls, [])
self.assertEqual(mocked_modelscope.call_count, 1)
def test_download_model_refuses_success_when_files_missing(self):
fake_hugging_face = ModuleType("huggingface_hub")
fake_hugging_face.snapshot_download = lambda **options: None
args = SimpleNamespace(
repository="PramaLLC/BEN2",
model_dir="model",
status="status.json",
allow_pattern=[],
)
with tempfile.TemporaryDirectory() as directory, patch.dict(
sys.modules,
{"huggingface_hub": fake_hugging_face},
), patch.object(
self.worker,
"create_download_progress_class",
return_value=object,
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker,
"download_from_modelscope",
):
args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json")
with self.assertRaisesRegex(
RuntimeError,
"All download sources failed",
):
self.worker.download_model(args)
def test_download_model_reports_all_failed_sources(self):
fake_hugging_face = ModuleType("huggingface_hub")
def fake_snapshot_download(**options):
raise OSError("direct down")
fake_hugging_face.snapshot_download = fake_snapshot_download
args = SimpleNamespace(
repository="depth-anything/DA3-SMALL",
model_dir="model",
status="status.json",
allow_pattern=[],
)
with tempfile.TemporaryDirectory() as directory, patch.dict(
sys.modules,
{"huggingface_hub": fake_hugging_face},
), patch.object(
self.worker,
"create_download_progress_class",
return_value=object,
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker,
"download_from_modelscope",
side_effect=OSError("mirror down"),
):
args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json")
with self.assertRaisesRegex(RuntimeError, "modelscope.cn"):
self.worker.download_model(args)
def test_download_from_modelscope_downloads_selected_files(self):
args = SimpleNamespace(
repository="depth-anything/DA3-SMALL",
model_dir="model",
status="status.json",
allow_pattern=["config.json"],
)
downloaded = []
with tempfile.TemporaryDirectory() as directory, patch.object(
self.worker,
"modelscope_file_list",
return_value=[
("config.json", 100),
("model.safetensors", 200),
(".gitattributes", 10),
],
), patch.object(
self.worker,
"_download_modelscope_file",
side_effect=lambda *args, **kwargs: downloaded.append(args[1]),
), patch.object(self.worker, "write_status"):
args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json")
self.worker.download_from_modelscope(args, args.model_dir)
self.assertEqual(downloaded, ["config.json"])
def test_ben2_background_removal_writes_fast_transparent_png(self): def test_ben2_background_removal_writes_fast_transparent_png(self):
inference_options = [] inference_options = []
@@ -368,6 +666,166 @@ class WorkerTest(unittest.TestCase):
) )
self.assertIn("Background removed in", status["message"]) self.assertIn("Background removed in", status["message"])
def test_ben2_background_removal_processes_each_frame(self):
inference_options = []
class FakeImage:
def __enter__(self):
return self
def __exit__(self, *_args):
return None
def convert(self, mode):
self.mode = mode
return self
def save(self, path, **options):
Path(path).touch()
class FakeModel:
@classmethod
def from_pretrained(cls, model_dir):
return cls()
def to(model_self, device):
return model_self
def eval(self):
return None
def inference(self, image, **options):
inference_options.append((image, options))
return FakeImage()
fake_torch = ModuleType("torch")
fake_torch.cuda = SimpleNamespace(
is_available=lambda: False,
synchronize=lambda: None,
)
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
fake_torch.inference_mode = nullcontext
fake_ben2 = ModuleType("ben2")
fake_ben2.BEN_Base = FakeModel
fake_pil = ModuleType("PIL")
fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())
with tempfile.TemporaryDirectory() as directory, patch.dict(
sys.modules,
{
"torch": fake_torch,
"ben2": fake_ben2,
"PIL": fake_pil,
},
):
root = Path(directory)
input_directory = root / "frames"
input_directory.mkdir()
(input_directory / "frame-0001.jpg").touch()
(input_directory / "frame-0002.jpg").touch()
args = SimpleNamespace(
input=str(input_directory),
model_dir="model",
output_dir=str(root / "result"),
status=str(root / "status.json"),
device="cpu",
video_frames=8,
)
self.worker.remove_background(args)
result_directory = Path(args.output_dir)
self.assertTrue(
(result_directory / "foreground.0001.png").is_file()
)
self.assertTrue(
(result_directory / "foreground.0002.png").is_file()
)
self.assertTrue(
(result_directory / "frame.0001.jpg").is_file()
)
self.assertTrue(
(result_directory / "frame.0002.jpg").is_file()
)
self.assertEqual(len(inference_options), 2)
for _image, options in inference_options:
self.assertEqual(options, {"refine_foreground": False})
def test_remove_background_parser_accepts_video_frames(self):
args = self.worker.parser().parse_args(
[
"remove-background",
"--input",
"input.mp4",
"--model-dir",
"ben2",
"--output-dir",
"result",
"--status",
"status.json",
"--device",
"auto",
"--video-frames",
"24",
]
)
self.assertEqual(args.command, "remove-background")
self.assertEqual(args.video_frames, 24)
def test_extract_frames_parser_accepts_frame_range(self):
args = self.worker.parser().parse_args(
[
"extract-frames",
"--input",
"frames",
"--output-dir",
"result",
"--status",
"status.json",
"--video-frames",
"24",
"--start-frame",
"70",
]
)
self.assertEqual(args.command, "extract-frames")
self.assertEqual(args.video_frames, 24)
self.assertEqual(args.start_frame, 70)
def test_extract_frames_copies_selected_frames(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
input_directory = root / "frames"
input_directory.mkdir()
for index in range(1, 6):
(input_directory / f"frame-{index:04d}.jpg").touch()
args = SimpleNamespace(
input=str(input_directory),
output_dir=str(root / "result"),
status=str(root / "status.json"),
video_frames=3,
start_frame=2,
)
self.worker.extract_frames(args)
result_directory = Path(args.output_dir)
self.assertEqual(
sorted(
path.name
for path in result_directory.iterdir()
if path.is_file()
),
["frame.0001.jpg", "frame.0002.jpg", "frame.0003.jpg"],
)
status = json.loads(
Path(args.status).read_text(encoding="utf-8")
)
self.assertEqual(status["progress"], 1.0)
self.assertIn("3 frame", status["message"])
def test_uses_official_da3_export_parameters_for_all_four_formats(self): def test_uses_official_da3_export_parameters_for_all_four_formats(self):
inference_options = [] inference_options = []
export_calls = [] export_calls = []
@@ -500,6 +958,95 @@ class WorkerTest(unittest.TestCase):
infer_gs_values = [options["infer_gs"] for _, options in inference_options] infer_gs_values = [options["infer_gs"] for _, options in inference_options]
self.assertEqual(infer_gs_values, [False, True, False, False]) self.assertEqual(infer_gs_values, [False, True, False, False])
def test_reconstruct_copies_frames_for_multi_frame_input(self):
class FakeModel:
@classmethod
def from_pretrained(cls, model_dir):
return cls()
def to(self, device):
return self
def eval(self):
return None
def inference(self, images, **options):
return SimpleNamespace()
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_api = ModuleType("depth_anything_3.api")
fake_api.DepthAnything3 = FakeModel
fake_export_module = ModuleType("depth_anything_3.utils.export")
fake_export_module.export = fake_export
fake_utils = ModuleType("depth_anything_3.utils")
fake_utils.__path__ = []
fake_package = ModuleType("depth_anything_3")
fake_package.__path__ = []
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
input_directory = root / "frames"
input_directory.mkdir()
(input_directory / "frame-0001.jpg").touch()
(input_directory / "frame-0002.jpg").touch()
with patch.dict(
sys.modules,
{
"torch": fake_torch,
"depth_anything_3": fake_package,
"depth_anything_3.api": fake_api,
"depth_anything_3.utils": fake_utils,
"depth_anything_3.utils.export": fake_export_module,
},
):
args = SimpleNamespace(
input=str(input_directory),
video_frames=16,
output_dir=str(root / "result"),
export_format="mini_npz",
status=str(root / "status.json"),
device="cpu",
model_dir="model",
process_res=518,
process_res_method="upper_bound_resize",
use_ray_pose=1,
ref_view_strategy="first",
conf_thresh_percentile=40.0,
num_max_points=1000,
show_cameras=0,
gs_views_interval=1,
remove_background_model_dir="",
)
self.worker.reconstruct(args)
result_directory = Path(args.output_dir)
self.assertTrue(
(result_directory / "frame.0001.jpg").is_file()
)
self.assertTrue(
(result_directory / "frame.0002.jpg").is_file()
)
self.assertTrue(
(result_directory / "results.npz").is_file()
)
self.assertFalse(
(result_directory / "exports").exists()
)
status = json.loads(Path(args.status).read_text(encoding="utf-8"))
self.assertEqual(status["progress"], 1.0)
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()

2
uv.lock generated
View File

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